diff --git a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp
index fc29041bbf93..7e9f3bd7070d 100644
--- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp
+++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp
@@ -127,9 +127,8 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t
         if (radix_aux_indices.has_value() && radix_aux_logits.has_value())
         {
             // Caller-owned scratch with stable address (CUDA Graph safe;
-            // matches the heuristic_scratch convention noted above). All
-            // in-tree callers under CUDA Graph capture go through dsa.py's
-            // DSAtrtllmAttentionMetadata which always pre-allocates these.
+            // matches the heuristic_scratch convention noted above). The
+            // Python TopK module supplies these from its reusable buffer arena.
             auto const& ai = radix_aux_indices.value();
             auto const& al = radix_aux_logits.value();
             TORCH_CHECK(ai.is_cuda() && al.is_cuda(), "radix_aux_{indices,logits} must be CUDA tensors");
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py
index 80908e08b0ac..d08340286dbf 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py
@@ -20,7 +20,6 @@
     rotate_activation,
     split_prefill_chunks,
     transform_local_topk_and_prepare_pool_view,
-    warmup_heuristic_topk_decode,
 )
 from .metadata import DSAtrtllmAttentionMetadata, build_req_idx_per_token
 from .params import DSABackendForwardArgs, DSAMetadataParams, DSAParams
@@ -49,5 +48,4 @@
     "rotate_activation",
     "split_prefill_chunks",
     "transform_local_topk_and_prepare_pool_view",
-    "warmup_heuristic_topk_decode",
 ]
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
index 3e67594b1690..e32e806edb0e 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
@@ -5,10 +5,9 @@
 from __future__ import annotations
 
 import os
-import threading
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, List, Optional, Set, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -26,6 +25,7 @@
     maybe_execute_in_parallel,
 )
 from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding
+from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation
 from tensorrt_llm._torch.utils import Fp4QuantizedTensor, maybe_compile
 from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned
 from tensorrt_llm.deep_gemm import (
@@ -67,71 +67,9 @@
     hadamard_transform = None
     HAS_FAST_HADAMARD = False
 
-# Idempotency guard for warmup_heuristic_topk_decode — keyed by
-# (device_index, top_k, hint_size, num_cols). Prevents repeated allocations
-# and synchronizations when multiple Indexer modules invoke the warmup with
-# the same parameters during model construction.
-_HEURISTIC_TOPK_WARMUP_DONE: Set[Tuple[int, int, int, int]] = set()
-_HEURISTIC_TOPK_WARMUP_LOCK = threading.Lock()
 _DG_SCHEDULE_BLOCK_KV = 64
 
 
-def warmup_heuristic_topk_decode(
-    top_k: int = 2048, hint_size: int = 2048, num_cols: int = 4096
-) -> None:
-    """Pre-initialize cached hardware attributes in the C++ Scheme X dispatcher.
-
-    The dispatcher inside ``invokeIndexerTopKDecode`` lazily queries
-    ``cudaDeviceGetAttribute`` for ``MultiProcessorCount`` and
-    ``L2CacheSize`` on its first call. Those host-side queries must not
-    be issued during ``cudaStreamBeginCapture / EndCapture``: the values
-    captured there become frozen into the graph and cannot be refreshed
-    across replays on a different device.
-
-    This warmup issues one small heuristic decode call so the static
-    caches are populated before any CUDA Graph capture begins. Must be
-    called from the Indexer setup hook (``layer_idx == 0``) when
-    ``enable_heuristic_topk`` is true.
-
-    Repeated invocations with the same ``(device, top_k, hint_size,
-    num_cols)`` key are short-circuited so that constructing many Indexer
-    modules in the same process does not re-allocate scratch tensors or
-    issue redundant synchronizations.
-    """
-    key = (torch.cuda.current_device(), top_k, hint_size, num_cols)
-    with _HEURISTIC_TOPK_WARMUP_LOCK:
-        if key in _HEURISTIC_TOPK_WARMUP_DONE:
-            return
-        _HEURISTIC_TOPK_WARMUP_DONE.add(key)
-
-    device = torch.device("cuda")
-    logits = torch.zeros((1, num_cols), dtype=torch.float32, device=device)
-    seq_lens = torch.tensor([num_cols], dtype=torch.int32, device=device)
-    indices = torch.empty((1, top_k), dtype=torch.int32, device=device)
-    pre_idx = torch.zeros((1, hint_size), dtype=torch.int32, device=device)
-    scratch = torch.empty((top_k,), dtype=torch.float32, device=device)
-    # The default warmup geometry (num_cols=4096) falls below kSeqSmall=12288
-    # and routes to the Radix path with blocks_per_row=2 (num_rows=1 sweeps
-    # bp ∈ [2, maxByCols=2]). The cpp op rejects blocks_per_row > 1 without
-    # caller-owned radix aux scratch, so supply worst-case (kMaxBlocksPerRowDecode=10)
-    # buffers here. Cost is negligible (~80 KB) and the warmup is a one-shot.
-    _radix_max_bp = 10
-    radix_aux_indices = torch.empty((1, _radix_max_bp, top_k), dtype=torch.int32, device=device)
-    radix_aux_logits = torch.empty((1, _radix_max_bp, top_k), dtype=torch.float32, device=device)
-    torch.ops.trtllm.indexer_topk_decode(
-        logits,
-        seq_lens,
-        indices,
-        1,
-        top_k,
-        pre_idx=pre_idx,
-        heuristic_scratch=scratch,
-        radix_aux_indices=radix_aux_indices,
-        radix_aux_logits=radix_aux_logits,
-    )
-    torch.cuda.synchronize()
-
-
 def _pick_dsl_expand(
     next_n: int,
     num_sms: int,
@@ -699,11 +637,22 @@ def __init__(
         )
         self.mtp_index_share = sparse_params.mtp_index_share
 
-        if self._enable_heuristic_topk and layer_idx == 0:
-            # Populate static caches (sm_count, L2 cache size) inside the C++
-            # Scheme X dispatcher before any CUDA Graph capture so the host
-            # attribute queries do not end up frozen into a captured graph.
-            warmup_heuristic_topk_decode(top_k=self.index_topk)
+        if self.use_cute_dsl_topk:
+            decode_top_k_implementation = (
+                TopKImplementation.CUTE_DSL_GVR
+                if self._enable_heuristic_topk
+                else TopKImplementation.CUTE_DSL_RADIX
+            )
+        elif self._enable_heuristic_topk:
+            decode_top_k_implementation = TopKImplementation.CUDA_GVR
+        else:
+            decode_top_k_implementation = TopKImplementation.CUDA_RADIX
+        self.top_k = TopK(
+            self.index_topk,
+            prefill_implementation=TopKImplementation.CUDA_RADIX,
+            decode_implementation=decode_top_k_implementation,
+            compress_ratio=self.compress_ratio,
+        )
 
         # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM
         # (populated in cache_derived_state; maps to TF32 tensor cores on Ampere+)
@@ -1099,11 +1048,11 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata):
         """
         num_contexts = metadata.num_contexts
         num_generations = metadata.num_generations
+        gen_seq_lens = metadata.get_indexer_kv_lens(
+            metadata.kv_lens_cuda_runtime[num_contexts : num_contexts + num_generations]
+        )
+        metadata.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens
         if not metadata.use_expanded_buffers_for_mtp:
-            gen_seq_lens = metadata.get_indexer_kv_lens(
-                metadata.kv_lens_cuda_runtime[num_contexts : num_contexts + num_generations]
-            )
-            metadata.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens
             next_n_cap = metadata.kv_lens_cuda_2d.shape[1]
             metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_(
                 gen_seq_lens.unsqueeze(-1).expand(-1, next_n_cap)
@@ -1135,7 +1084,7 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata):
             )
 
     @staticmethod
-    def prepare(metadata: DSAtrtllmAttentionMetadata):
+    def prepare(metadata: DSAtrtllmAttentionMetadata) -> None:
         """
         Prepare indexer for the forward pass.
         This should be called during metadata.prepare() stage.
@@ -1379,7 +1328,6 @@ def sparse_attn_indexer(
         k_fp8: torch.Tensor,
         k_scale: torch.Tensor,
         weights: torch.Tensor,
-        use_custom_topk: bool = True,
         q_scale: Optional[torch.Tensor] = None,
         is_generation: Optional[bool] = None,
     ) -> torch.Tensor:
@@ -1406,6 +1354,10 @@ def sparse_attn_indexer(
         num_tokens = metadata.num_tokens
 
         num_gen_tokens = num_tokens - num_ctx_tokens
+        gvr_prior_indices = None
+        if self._enable_heuristic_topk:
+            local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx]
+            gvr_prior_indices = metadata.gvr_prior_indices[local_layer]
         if is_generation is None:
             has_prefill = num_contexts > 0
             has_decode = num_generations > 0
@@ -1431,8 +1383,6 @@ def sparse_attn_indexer(
             dtype=torch.int32,
             capture_graph=metadata.is_cuda_graph,
         )
-        if not use_custom_topk:
-            topk_indices_buffer[: hidden_states.shape[0]] = -1
 
         if has_prefill and not metadata.skip_indexer_for_ctx_reqs:
             # Use chunked prefill to reduce memory footprint
@@ -1517,36 +1467,15 @@ def sparse_attn_indexer(
                             chunk.cu_seqlen_ks[c0:c1],
                             chunk.cu_seqlen_ke[c0:c1],
                             tile_q_scale,
-                            clean_logits=not use_custom_topk,
+                            clean_logits=False,
+                        )
+                        self.top_k(
+                            logits,
+                            topk_indices_buffer[g0:g1, :],
+                            is_prefill=True,
+                            row_starts=chunk.cu_seqlen_ks[c0:c1],
+                            row_ends=chunk.cu_seqlen_ke[c0:c1],
                         )
-                        if use_custom_topk:
-                            torch.ops.trtllm.indexer_topk_prefill(
-                                logits,
-                                chunk.cu_seqlen_ks[c0:c1],
-                                chunk.cu_seqlen_ke[c0:c1],
-                                topk_indices_buffer[g0:g1, :],
-                                self.index_topk,
-                            )
-                        else:
-                            topk_indices = logits.topk(
-                                min(self.index_topk, logits.shape[-1]), dim=-1
-                            )[1]
-                            topk_indices -= chunk.cu_seqlen_ks[c0:c1][:, None]
-
-                            mask_lo = topk_indices >= 0
-                            mask_hi = (
-                                topk_indices
-                                - (chunk.cu_seqlen_ke[c0:c1] - chunk.cu_seqlen_ks[c0:c1])[:, None]
-                                < 0
-                            )
-                            mask = mask_lo & mask_hi
-
-                            # local indices per sequence
-                            topk_indices = topk_indices.masked_fill(~mask, -1)
-
-                            topk_indices_buffer[g0:g1, : topk_indices.shape[-1]] = topk_indices.to(
-                                dtype=torch.int32
-                            )
 
                     if apply_q_split:
                         q_sizes = [
@@ -1577,55 +1506,29 @@ def sparse_attn_indexer(
                     cu_seqlen_ks,
                     cu_seqlen_ke,
                     ctx_q_scale,
-                    clean_logits=not use_custom_topk,
+                    clean_logits=False,
+                )
+                self.top_k(
+                    logits,
+                    topk_indices_buffer[:num_ctx_tokens, :],
+                    is_prefill=True,
+                    row_starts=cu_seqlen_ks,
+                    row_ends=cu_seqlen_ke,
                 )
-                if use_custom_topk:
-                    torch.ops.trtllm.indexer_topk_prefill(
-                        logits,
-                        cu_seqlen_ks,
-                        cu_seqlen_ke,
-                        topk_indices_buffer[:num_ctx_tokens, :],
-                        self.index_topk,
-                    )
-                else:
-                    topk_indices = logits.topk(min(self.index_topk, logits.shape[-1]), dim=-1)[1]
-                    topk_indices -= cu_seqlen_ks[:, None]
-                    mask_lo = topk_indices >= 0
-                    mask_hi = topk_indices - (cu_seqlen_ke - cu_seqlen_ks)[:, None] < 0
-                    mask = mask_lo & mask_hi
-
-                    # local indices per sequence
-                    topk_indices = topk_indices.masked_fill(~mask, -1)
-                    topk_indices_buffer[:num_ctx_tokens, : topk_indices.shape[-1]] = (
-                        topk_indices.to(dtype=torch.int32)
-                    )
         elif has_prefill and metadata.skip_indexer_for_ctx_reqs:
             # Fill topk_indices_buffer with pre-defined dense topk indices
             topk_indices_buffer[:num_ctx_tokens, :] = metadata.topk_indices_buffer[
                 :num_ctx_tokens, :
             ]
 
-        # Prefill→decode GVR handoff: seed each finishing-prefill sequence's
-        # heuristic_prev_topk slot with its own last-context-token top-K, so
-        # the FIRST decode step of that sequence gets a warm-started preIdx
-        # (~60-75% set-overlap with the eventual decode top-K on this
-        # workload) instead of the all-zero / all-(-1) cold start that the
-        # default `heuristic_prev_topk.zero_()` initialization leaves behind.
-        # Without this, GVR P2 secant on decode step 0 runs from a benign
-        # but uninformative seed (kernel +1 offset on zeros → all indices
-        # point at compressed-token position 1), wasting iterations.
-        # Slot convention (mirrors the existing decode write-back at the
-        # bottom of the decode block): new gens from finishing prefill
-        # append after currently-active gens, i.e., slots
-        # [num_generations : num_generations + num_contexts].
-        if self._enable_heuristic_topk and has_prefill and not metadata.skip_indexer_for_ctx_reqs:
-            local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx]
-            ctx_seq_lens = metadata.seq_lens[:num_contexts]
-            # Per-sequence last context-token offset (exclusive cumsum minus 1).
-            last_ctx_idx = (torch.cumsum(ctx_seq_lens, dim=0) - 1).to(dtype=torch.long)
-            metadata.heuristic_prev_topk[
-                local_layer, num_generations : num_generations + num_contexts
-            ].copy_(topk_indices_buffer[last_ctx_idx, :])
+        # Dense skip outputs remain valid priors if a later step enters GVR.
+        if has_prefill:
+            self.top_k.update_gvr_prior_from_prefill(
+                topk_indices_buffer[:num_ctx_tokens],
+                metadata.seq_lens[:num_contexts],
+                gvr_prior_indices,
+                request_offset=num_generations,
+            )
 
         reuse_topk = (
             self.mtp_index_share
@@ -1802,101 +1705,31 @@ def sparse_attn_indexer(
                     decode_q_scale,
                 )
 
-            if use_custom_topk:
-                # Kernel expects kv_lens (total cache length), not seq_lens (new tokens)
-                # This is because rowEnd = seq_len - next_n + offset + 1
-                gen_kv_lens_cuda = metadata.kv_lens_cuda_runtime[
-                    num_contexts : num_contexts + num_generations
-                ]
-
-                pre_idx = None
-                heuristic_scratch = None
-                if self._enable_heuristic_topk:
-                    local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx]
-                    # Pass prev_topk directly; the +1 temporal offset is
-                    # handled inside the C++ kernel (preIdxOffset += 1).
-                    pre_idx = metadata.heuristic_prev_topk[local_layer, :num_generations]
-                    if not metadata.use_cute_dsl_topk:
-                        heuristic_scratch = metadata.heuristic_scratch_values[:num_gen_tokens]
-
-                if self.use_cute_dsl_topk and self._enable_heuristic_topk:
-                    torch.ops.trtllm.cute_dsl_gvr_topk_decode(
-                        logits_decode,
-                        pre_idx,
-                        gen_kv_lens_cuda,
-                        topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :],
-                        self.index_topk,
-                        next_n=next_n,
-                        compress_ratio=self.compress_ratio,
-                        max_seq_len=indexer_max_seq_len,
-                        order_row=metadata.kv_lens_row_reorder,
-                    )
-                elif self.use_cute_dsl_topk and (self.compress_ratio == 1 or next_n == 1):
-                    torch.ops.trtllm.cute_dsl_indexer_topk_decode(
-                        logits_decode,
-                        context_lens if self.compress_ratio > 1 else gen_kv_lens_cuda,
-                        topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :],
-                        self.index_topk,
-                        next_n,
-                    )
-                else:
-                    torch.ops.trtllm.indexer_topk_decode(
-                        logits_decode,
-                        gen_kv_lens_cuda,
-                        topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :],
-                        next_n,
-                        self.index_topk,
-                        pre_idx=pre_idx,
-                        heuristic_scratch=heuristic_scratch,
-                        compress_ratio=self.compress_ratio,
-                        radix_aux_indices=metadata.radix_aux_indices,
-                        radix_aux_logits=metadata.radix_aux_logits,
-                    )
-            else:
-                # padded
-                positions = (
-                    torch.arange(logits_decode.shape[-1], device=q_decode.device)
-                    .unsqueeze(0)
-                    .expand(num_gen_tokens, -1)
-                )
-                row_indices = torch.arange(num_gen_tokens, device=q_decode.device) // next_n
-                next_n_offset = torch.arange(num_gen_tokens, device=q_decode.device) % next_n
-                index_end_pos = (context_lens[row_indices] - next_n + next_n_offset).unsqueeze(1)
-                # index_end_pos: [B * N, 1]
-                mask = positions <= index_end_pos
-                # mask: [B * N, L]
-                logits_decode = logits_decode.masked_fill(~mask, float("-inf"))
-                topk_indices_decode = logits_decode.topk(
-                    min(self.index_topk, logits_decode.shape[-1]), dim=-1
-                )[1].to(torch.int32)  # [B * N, K]
-                # ensure we don't set indices for the top k
-                # that is out of range(masked already)
-                # this will happen if context length is shorter than K
-                mask_decode = topk_indices_decode <= index_end_pos
-
-                # local indices per sequence
-                topk_indices_decode = topk_indices_decode.masked_fill(~mask_decode, -1)
-                # Store in buffer
-                topk_indices_buffer[
-                    token_offset : token_offset + num_gen_tokens, : topk_indices_decode.shape[-1]
-                ] = topk_indices_decode.to(dtype=torch.int32)
-
-            if self._enable_heuristic_topk:
-                local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx]
-                decode_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
-                last_mtp_topk = decode_topk[next_n - 1 :: next_n]
-                prev_topk_dst = metadata.heuristic_prev_topk[local_layer, :num_generations]
-                if do_multi_stream() and self.aux_stream is not None:
-                    # Fork the write-back onto the aux stream to overlap with this
-                    # layer's core sparse attention; joined by maybe_join_prev_topk_copy().
-                    self.prev_topk_copy_events[0].record()
-                    with torch.cuda.stream(self.aux_stream):
-                        self.prev_topk_copy_events[0].wait()
-                        prev_topk_dst.copy_(last_mtp_topk)
-                        self.prev_topk_copy_events[1].record()
-                    self._prev_topk_copy_pending = True
-                else:
-                    prev_topk_dst.copy_(last_mtp_topk)
+            # Native/GVR use logical lengths; radix uses score-column lengths.
+            gen_kv_lens_cuda = metadata.kv_lens_cuda_runtime[
+                num_contexts : num_contexts + num_generations
+            ]
+            scan_lengths = metadata.gen_indexer_kv_lens_cuda_runtime
+            assert scan_lengths is not None
+
+            gvr_ext_kwargs = (
+                {
+                    "gvr_prior_indices": gvr_prior_indices[:num_generations],
+                    "gvr_row_order": metadata.kv_lens_row_reorder,
+                }
+                if gvr_prior_indices is not None
+                else None
+            )
+            self.top_k(
+                logits_decode,
+                topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :],
+                is_prefill=False,
+                sequence_lengths=gen_kv_lens_cuda,
+                scan_lengths=scan_lengths,
+                next_n=next_n,
+                max_seq_len=indexer_max_seq_len,
+                gvr_ext_kwargs=gvr_ext_kwargs,
+            )
 
         elif has_decode and metadata.skip_indexer_for_gen_reqs:
             # Fill topk_indices_buffer with pre-defined dense topk indices
@@ -1904,6 +1737,23 @@ def sparse_attn_indexer(
                 metadata.topk_indices_buffer[num_ctx_tokens:num_tokens, :]
             )
 
+        # Keep the GVR prior current for computed, reused, and dense-skip TopK.
+        if gvr_prior_indices is not None and has_decode:
+            next_n = num_gen_tokens // num_generations
+            decode_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
+            last_mtp_topk = decode_topk[next_n - 1 :: next_n]
+            prev_topk_dst = gvr_prior_indices[:num_generations]
+            if do_multi_stream() and self.aux_stream is not None:
+                # Overlap the GVR prior write-back with this layer's sparse attention.
+                self.prev_topk_copy_events[0].record()
+                with torch.cuda.stream(self.aux_stream):
+                    self.prev_topk_copy_events[0].wait()
+                    prev_topk_dst.copy_(last_mtp_topk)
+                    self.prev_topk_copy_events[1].record()
+                self._prev_topk_copy_pending = True
+            else:
+                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:
@@ -1947,7 +1797,7 @@ def _mtp_last_accepted_rows(
         return gen_topk[base + offset]
 
     def maybe_join_prev_topk_copy(self) -> None:
-        """Join the aux-stream heuristic prev_topk write-back, if forked."""
+        """Join the aux-stream GVR prior write-back, if forked."""
         if self._prev_topk_copy_pending:
             self.prev_topk_copy_events[1].wait()
             self._prev_topk_copy_pending = False
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
index 998366944daf..2854c71b6f96 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
@@ -31,14 +31,8 @@
 
 ModelConfig = tensorrt_llm.bindings.ModelConfig
 
-# dtype of the indexer MQA-logits that feed the top-k. All paged_mqa_logits
-# paths produce fp32 today (DSL fp8/fp4 default output_dtype=fp32; DeepGEMM
-# fp8 hardcodes kFloat; DeepGEMM fp4 defaults logits_dtype=kFloat32 and is not
-# overridden here), and the decode forward feeds logits to the top-k without a
-# cast. dtype is a top-k compile-key dimension, so the warmup pre-compiles for
-# exactly this value. If a paged_mqa_logits caller ever emits a different dtype
-# (e.g. overriding the DeepGEMM fp4 logits_dtype to bf16), update this constant
-# or the warmup silently compiles the wrong variant.
+# Indexer MQA-logits are currently always fp32. The dtype is part of the
+# CuTe DSL Top-K compile key, so warmup must use the runtime dtype.
 _INDEXER_LOGITS_DTYPE = torch.float32
 
 if TYPE_CHECKING:
@@ -58,8 +52,6 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata):
     """Attention metadata for DSA (Dense Sparse Attention) with indexer state."""
 
     sparse_metadata_params: Optional[DSAMetadataParams] = None
-    # Store reference to indexer for preparation stage
-    indexer: Optional["Indexer"] = None
     # Chunked prefill metadata for indexer (prefill-only, no CUDA graph needed)
     indexer_prefill_chunks: Optional[List[IndexerPrefillChunkMetadata]] = None
     # Max chunk size for two-level chunking:
@@ -160,6 +152,9 @@ def __post_init__(self):
         self.use_cute_dsl_topk = (
             sparse_metadata_params.use_cute_dsl_topk and IS_CUTLASS_DSL_AVAILABLE
         )
+        self.enable_gvr_topk = (
+            sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100
+        )
         self.kv_lens_row_reorder = None
         capture_graph = self.is_cuda_graph
         # Plain DSA has no compression and uses the default [1]. DeepSeek-V4's
@@ -313,38 +308,17 @@ def get_indexer_max_seq_len(self) -> int:
         return max(1, self.kv_cache_manager.max_seq_len // self._indexer_compress_ratio)
 
     def warmup_cute_dsl_radix_topk(self, next_n: int) -> None:
-        """Pre-compile the radix-filter CuTe DSL decode top-k during warmup.
-
-        Eager decode iters (mixed prefill+decode batch, or ``cuda_graph``
-        disabled) whose ``num_rows`` lands in a ``cluster_size`` band that
-        graph capture did not exercise otherwise pay a first-touch JIT stall
-        on a live request. ``num_cols`` is fixed at ``indexer_max_seq_len``,
-        so only the ``cluster_size`` dimension needs sweeping; delegate to the
-        custom-op warmup helper, which owns the band enumeration.
-
-        ``next_n`` (a compile-key dimension) is supplied by the caller from
-        the engine's static spec-decode config.
-
-        No-op unless decode actually routes to
-        ``cute_dsl_indexer_topk_decode``: heuristic top-k uses the GVR kernel
-        and plain (no cute_dsl_topk) decode uses the C++ op. Called once from
-        ``ModelEngine.warmup``.
-        """
-        if not self.use_cute_dsl_topk or self.enable_heuristic_topk:
+        """Pre-compile CuTe DSL radix variants not covered by engine warmup."""
+        sparse_params = self.sparse_metadata_params
+        if not self.use_cute_dsl_topk or (
+            sparse_params.enable_heuristic_topk and get_sm_version() >= 100
+        ):
             return
         if self.kv_cache_manager is None:
             return
-        top_k = getattr(self.sparse_metadata_params, "index_topk", None)
+        top_k = self.sparse_mla_topk
         if not top_k:
             return
-        # The radix-filter DSL kernel does not support a compressed indexer
-        # combined with multi-row MTP: decode dispatches to it only when
-        # compress_ratio == 1 or next_n == 1. The compress_ratio > 1 &&
-        # next_n > 1 case routes to the C++ op (or GVR when heuristic top-k is
-        # on), so there is nothing to pre-compile here.
-        # TODO: extending the radix-filter path to compress_ratio > 1 &&
-        # next_n > 1 is straightforward; once the dispatch above is relaxed to
-        # use it there, drop this guard so the case is pre-compiled too.
         if self._indexer_compress_ratio > 1 and next_n > 1:
             return
         try:
@@ -353,6 +327,7 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None:
             )
         except ImportError:
             return
+
         warmup_cute_dsl_radix_topk_decode(
             top_k=int(top_k),
             num_cols=int(self.get_indexer_max_seq_len()),
@@ -480,11 +455,11 @@ def on_update_kv_lens(self):
         self._compute_kv_lens_row_reorder()
         self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True)
 
-    def _compute_kv_lens_row_reorder(self):
-        """Prepare longest-job-first row order for GVR top-k."""
+    def _compute_kv_lens_row_reorder(self) -> None:
+        """Prepare the longest-job-first GVR row order once per forward step."""
         next_n = 1 + self.max_draft_tokens
         if (
-            self.enable_heuristic_topk
+            self.enable_gvr_topk
             and self.use_cute_dsl_topk
             and self.num_generations * next_n >= 2 * self.num_sms
         ):
@@ -560,37 +535,6 @@ def create_buffers_for_mla_rope_append(self, capture_graph=False):
             pin_memory=prefer_pinned(),
         )
 
-    def _create_radix_aux_buffers(self, capture_graph=False):
-        # Persistent scratch for Radix-split-work indexer path (blocks_per_row > 1).
-        # Mirrors the fix the Heuristic path applied: per-call th::empty inside
-        # indexer_topk_decode produces stale pointers under CUDA Graph replay when
-        # the caching allocator is perturbed by chunked prefill at high CONC.
-        # Sized to the worst case kMaxBlocksPerRowDecode=10 from
-        # cpp/tensorrt_llm/kernels/indexerTopK.cu, times the max number of
-        # generation rows (num_seqs * (1 + max_draft_tokens)); the cpp op aborts
-        # if this is smaller than num_rows*blocks_per_row*index_topk. Allocated
-        # unconditionally: even with enable_heuristic_topk=True the dispatcher can
-        # fall back to Radix when canUseHeuristic returns False (small numColumns,
-        # etc.). MUST be re-created whenever max_draft_tokens changes (see
-        # update_spec_dec_param) or it is left too small once MTP raises the
-        # generation-row count.
-        _radix_max_blocks_per_row = 10
-        _radix_max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens)
-        self.radix_aux_indices = self.get_empty(
-            self.cuda_graph_buffers,
-            (_radix_max_gen_tokens, _radix_max_blocks_per_row, self.num_sparse_topk),
-            cache_name="radix_aux_indices",
-            dtype=torch.int32,
-            capture_graph=capture_graph,
-        )
-        self.radix_aux_logits = self.get_empty(
-            self.cuda_graph_buffers,
-            (_radix_max_gen_tokens, _radix_max_blocks_per_row, self.num_sparse_topk),
-            cache_name="radix_aux_logits",
-            dtype=torch.float32,
-            capture_graph=capture_graph,
-        )
-
     def create_buffers_for_indexer(self, capture_graph=False):
         sparse_metadata_params = self.sparse_metadata_params
         if not isinstance(sparse_metadata_params, DSAMetadataParams):
@@ -695,7 +639,8 @@ def create_buffers_for_indexer(self, capture_graph=False):
                 pin_memory=prefer_pinned(),
             )
         # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation.
-        # These buffers will be allocated dynamically in Indexer.prepare() based on actual total_kv_len to save memory.
+        # Allocate these buffers dynamically in Indexer.prepare()
+        # based on the actual total_kv_len to save memory.
         if self.enable_context_mla_with_cached_kv:
             self.slot_mapping_fp8_fullkv = None
             self.slot_mapping_scale_fullkv = None
@@ -786,37 +731,19 @@ def create_buffers_for_indexer(self, capture_graph=False):
                 device="cpu",
                 pin_memory=prefer_pinned(),
             )
-        # Per-layer persistent buffers for heuristic TopK pre_idx.
-        # Indexed by [local_layer_idx, generation_position, :].
-        # The graph captures reads/writes on these stable-address buffers;
-        # each replay's write becomes the next replay's read (feedback loop).
-        self.enable_heuristic_topk = (
-            sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100
-        )
-        if self.enable_heuristic_topk:
-            num_local_layers = self.kv_cache_manager.num_local_layers
-            self.heuristic_prev_topk = self.get_empty(
+        if self.enable_gvr_topk:
+            self.gvr_prior_indices = self.get_empty(
                 self.cuda_graph_buffers,
-                (num_local_layers, self.max_num_sequences, self.num_sparse_topk),
-                cache_name="heuristic_prev_topk",
+                (
+                    self.kv_cache_manager.num_local_layers,
+                    self.max_num_sequences,
+                    self.num_sparse_topk,
+                ),
+                cache_name="gvr_prior_indices",
                 dtype=torch.int32,
                 capture_graph=capture_graph,
             )
-            # Zero-initialize so the first decode step's pre_idx (kernel
-            # adds +1 offset) points to index 1 — a valid but benign candidate.
-            # Without this, uninitialized memory produces random hint indices.
-            self.heuristic_prev_topk.zero_()
-            # The C++ top-k path needs a stable scratch address.
-            if not self.use_cute_dsl_topk:
-                max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens)
-                self.heuristic_scratch_values = self.get_empty(
-                    self.cuda_graph_buffers,
-                    (max_gen_tokens, self.num_sparse_topk),
-                    cache_name="heuristic_scratch_values",
-                    dtype=torch.float32,
-                    capture_graph=capture_graph,
-                )
-            # GVR row order also needs a stable address for CUDA graphs.
+            self.gvr_prior_indices.zero_()
             if self.use_cute_dsl_topk:
                 self.kv_lens_row_reorder_buffer = self.get_empty(
                     self.cuda_graph_buffers,
@@ -825,12 +752,6 @@ def create_buffers_for_indexer(self, capture_graph=False):
                     dtype=torch.int32,
                     capture_graph=capture_graph,
                 )
-
-        # Persistent scratch for the Radix-split-work indexer path. Re-created
-        # in update_spec_dec_param when max_draft_tokens changes so it stays
-        # large enough for the MTP generation-row count.
-        self._create_radix_aux_buffers(capture_graph=capture_graph)
-
         # Create expanded buffers for MTP support
         self.create_expanded_buffers(capture_graph=capture_graph)
 
@@ -938,23 +859,6 @@ def update_spec_dec_param(
         init_shape = self.kv_lens_expanded_host.shape[0]
         if self.max_num_sequences * (1 + self.max_draft_tokens) != init_shape:
             self.create_expanded_buffers(capture_graph=capture_graph)
-            # Resize heuristic scratch buffer for new max_draft_tokens.
-            if self.enable_heuristic_topk and not self.use_cute_dsl_topk:
-                max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens)
-                self.heuristic_scratch_values = self.get_empty(
-                    self.cuda_graph_buffers,
-                    (max_gen_tokens, self.num_sparse_topk),
-                    cache_name="heuristic_scratch_values",
-                    dtype=torch.float32,
-                    capture_graph=capture_graph,
-                )
-            # The Radix-split-work scratch (radix_aux_*) is sized the same way
-            # (num_seqs * (1 + max_draft_tokens) rows) and is allocated
-            # unconditionally, so it must be resized here too -- otherwise the
-            # cpp indexer_topk_decode op aborts once MTP raises max_draft_tokens
-            # ("radix_aux_* must hold at least num_rows*blocks_per_row*index_topk
-            # elements").
-            self._create_radix_aux_buffers(capture_graph=capture_graph)
 
     def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor:
         """Refresh INDEX_KEY offsets and return their physical pool slots."""
diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py
index 06ce761b28e1..8462a4bd8cc7 100644
--- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py
+++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py
@@ -36,6 +36,7 @@
 from tensorrt_llm.logger import logger
 
 from ...distributed import allgather
+from ...modules.top_k import TopK, TopKImplementation
 from ...pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
 from ...pyexecutor.llm_request import LlmRequestState
 from ...pyexecutor.resource_manager import KVCacheCompressionManager
@@ -632,12 +633,13 @@ def _select_kept_ordinals(self, request_count: int) -> None:
         """Select top-k tokens and settle score ties into kept-ordinal rows."""
         rows = request_count * self._selection_rows_per_request
         # The trailing 1 is next_n: decode scores one query token per request.
-        torch.ops.trtllm.cute_dsl_indexer_topk_decode(
+        self._selection_top_k(
             self._selection_scores_rows[:rows],
-            self._selection_row_lengths[:rows],
             self._provisional_rows[:rows],
-            self.budget,
-            1,
+            is_prefill=False,
+            sequence_lengths=self._selection_row_lengths[:rows],
+            scan_lengths=self._selection_row_lengths[:rows],
+            next_n=1,
         )
         settle_ties(
             self._selection_scores_rows,
@@ -799,6 +801,10 @@ def _allocate_metadata_buffers(
 
     def _allocate_selection_buffers(self, device: torch.device, *, tp_size: int) -> None:
         """Allocate fixed manager-lifetime TopK inputs and outputs."""
+        self._selection_top_k = TopK(
+            self.budget,
+            decode_implementation=TopKImplementation.CUTE_DSL_RADIX,
+        )
         request_capacity = self._request_capacity
         selection_width = self._selection_width_capacity
         union = self.eviction_mode == "union"
diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py
new file mode 100644
index 000000000000..704ce11283bc
--- /dev/null
+++ b/tensorrt_llm/_torch/modules/top_k.py
@@ -0,0 +1,364 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Reusable index-selection Top-K module for sparse inference paths."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+import torch
+import torch.nn as nn
+
+from ..memory_buffer_utils import get_memory_buffers
+
+
+class TopKImplementation(str, Enum):
+    """Top-K implementations grouped by backend and algorithm."""
+
+    TORCH = "torch"
+    CUDA_RADIX = "cuda_radix"
+    CUTE_DSL_RADIX = "cute_dsl_radix"
+    CUDA_GVR = "cuda_gvr"
+    CUTE_DSL_GVR = "cute_dsl_gvr"
+
+
+_GVR_IMPLEMENTATIONS = {
+    TopKImplementation.CUDA_GVR,
+    TopKImplementation.CUTE_DSL_GVR,
+}
+_MAX_RADIX_BLOCKS_PER_ROW = 10
+
+
+class TopK(nn.Module):
+    """Select Top-K indices for sparse prefill and decode paths.
+
+    GVR decode state is owned by the caller so it can be shared with the
+    request metadata and retain a stable address across CUDA Graph replays.
+    """
+
+    _memory_buffers = get_memory_buffers()
+
+    def __init__(
+        self,
+        top_k: int,
+        *,
+        prefill_implementation: TopKImplementation | None = None,
+        decode_implementation: TopKImplementation | None = None,
+        compress_ratio: int = 1,
+    ) -> None:
+        super().__init__()
+        self.top_k = top_k
+        self.prefill_implementation = TopKImplementation(
+            prefill_implementation or TopKImplementation.CUDA_RADIX
+        )
+        self.decode_implementation = TopKImplementation(
+            decode_implementation or TopKImplementation.CUDA_RADIX
+        )
+        self.compress_ratio = compress_ratio
+
+    def forward(
+        self,
+        scores: torch.Tensor,
+        output_indices: torch.Tensor,
+        *,
+        is_prefill: bool,
+        row_starts: torch.Tensor | None = None,
+        row_ends: torch.Tensor | None = None,
+        sequence_lengths: torch.Tensor | None = None,
+        scan_lengths: torch.Tensor | None = None,
+        next_n: int = 1,
+        max_seq_len: int | None = None,
+        gvr_ext_kwargs: dict[str, torch.Tensor | None] | None = None,
+    ) -> torch.Tensor:
+        """Write prefill or decode Top-K indices into ``output_indices``.
+
+        Args:
+            scores: Top-K input scores with shape ``[num_rows, num_columns]``.
+            output_indices: Int32 output with shape ``[num_rows, top_k]``.
+            is_prefill: Whether to run the prefill implementation.
+            row_starts: Per-row inclusive starts for prefill.
+            row_ends: Per-row exclusive ends for prefill.
+            sequence_lengths: Per-request logical KV lengths for decode.
+            scan_lengths: Per-request score-column lengths for decode.
+            next_n: Number of decode rows per request.
+            max_seq_len: Maximum decode score width used for GVR kernel tuning.
+            gvr_ext_kwargs: GVR-only keyword arguments. ``gvr_prior_indices``
+                is the required caller-owned int32 previous selection with
+                shape ``[num_requests, top_k]`` on ``scores.device``.
+                ``gvr_row_order`` is an optional int32 request ordering with
+                shape ``[num_requests]`` on the same device.
+
+        Returns:
+            ``output_indices`` after the selected implementation writes it.
+        """
+        if is_prefill:
+            assert row_starts is not None and row_ends is not None
+            return self._forward_prefill(scores, row_starts, row_ends, output_indices)
+
+        assert sequence_lengths is not None and scan_lengths is not None
+        return self._forward_decode(
+            scores,
+            sequence_lengths,
+            scan_lengths,
+            output_indices,
+            next_n,
+            max_seq_len,
+            gvr_ext_kwargs,
+        )
+
+    def _forward_prefill(
+        self,
+        scores: torch.Tensor,
+        row_starts: torch.Tensor,
+        row_ends: torch.Tensor,
+        output_indices: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.prefill_implementation == TopKImplementation.TORCH:
+            return self._forward_prefill_torch(
+                scores,
+                row_starts,
+                row_ends,
+                output_indices,
+            )
+        if self.prefill_implementation != TopKImplementation.CUDA_RADIX:
+            raise NotImplementedError(
+                f"{self.prefill_implementation.value} does not support prefill Top-K"
+            )
+        torch.ops.trtllm.indexer_topk_prefill(
+            scores,
+            row_starts,
+            row_ends,
+            output_indices,
+            self.top_k,
+        )
+        return output_indices
+
+    def _forward_decode(
+        self,
+        scores: torch.Tensor,
+        sequence_lengths: torch.Tensor,
+        scan_lengths: torch.Tensor,
+        output_indices: torch.Tensor,
+        next_n: int,
+        max_seq_len: int | None,
+        gvr_ext_kwargs: dict[str, torch.Tensor | None] | None,
+    ) -> torch.Tensor:
+        if self.decode_implementation == TopKImplementation.TORCH:
+            return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n)
+
+        if self.decode_implementation in _GVR_IMPLEMENTATIONS:
+            return self._forward_decode_gvr(
+                scores,
+                sequence_lengths,
+                output_indices,
+                next_n,
+                max_seq_len=max_seq_len,
+                **(gvr_ext_kwargs or {}),
+            )
+
+        return self._forward_decode_radix(
+            scores,
+            sequence_lengths,
+            scan_lengths,
+            output_indices,
+            next_n,
+        )
+
+    def _forward_decode_radix(
+        self,
+        scores: torch.Tensor,
+        sequence_lengths: torch.Tensor,
+        scan_lengths: torch.Tensor,
+        output_indices: torch.Tensor,
+        next_n: int,
+    ) -> torch.Tensor:
+        use_cute_dsl = self.decode_implementation == TopKImplementation.CUTE_DSL_RADIX and not (
+            self.compress_ratio > 1 and next_n > 1
+        )
+        if use_cute_dsl:
+            torch.ops.trtllm.cute_dsl_indexer_topk_decode(
+                scores,
+                scan_lengths,
+                output_indices,
+                self.top_k,
+                next_n,
+            )
+            return output_indices
+
+        radix_indices, radix_values = self._get_radix_workspace(scores)
+        torch.ops.trtllm.indexer_topk_decode(
+            scores,
+            sequence_lengths,
+            output_indices,
+            next_n,
+            self.top_k,
+            pre_idx=None,
+            heuristic_scratch=None,
+            compress_ratio=self.compress_ratio,
+            radix_aux_indices=radix_indices,
+            radix_aux_logits=radix_values,
+        )
+        return output_indices
+
+    def _get_workspace(
+        self,
+        scores: torch.Tensor,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        buffer_name: str,
+    ) -> torch.Tensor:
+        device_buffer_name = f"{buffer_name}_{scores.device}"
+        if scores.is_cuda:
+            with torch.cuda.device(scores.device):
+                capture_graph = torch.cuda.is_current_stream_capturing()
+                return self._memory_buffers.get_buffer(
+                    shape,
+                    dtype=dtype,
+                    buffer_name=device_buffer_name,
+                    reserve_buffer=capture_graph,
+                )
+        return self._memory_buffers.get_buffer(
+            shape,
+            dtype=dtype,
+            buffer_name=device_buffer_name,
+            reserve_buffer=False,
+        )
+
+    def _get_radix_workspace(
+        self, scores: torch.Tensor
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        if scores.dtype != torch.float32:
+            # The C++ bf16/fp16 entry has no split-work tier or aux-buffer
+            # arguments and rejects widths that would require split work.
+            return None, None
+
+        shape = (scores.shape[0], _MAX_RADIX_BLOCKS_PER_ROW, self.top_k)
+        radix_indices = self._get_workspace(
+            scores,
+            shape,
+            torch.int32,
+            "top_k_radix_indices_workspace",
+        )
+        radix_values = self._get_workspace(
+            scores,
+            shape,
+            torch.float32,
+            "top_k_radix_values_workspace",
+        )
+        return radix_indices, radix_values
+
+    def _forward_decode_gvr(
+        self,
+        scores: torch.Tensor,
+        sequence_lengths: torch.Tensor,
+        output_indices: torch.Tensor,
+        next_n: int,
+        max_seq_len: int | None,
+        gvr_prior_indices: torch.Tensor | None = None,
+        gvr_row_order: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        assert gvr_prior_indices is not None
+        if self.decode_implementation == TopKImplementation.CUDA_GVR:
+            workspace = self._get_workspace(
+                scores,
+                (scores.shape[0], self.top_k),
+                scores.dtype,
+                "top_k_cuda_gvr_workspace",
+            )
+            radix_indices, radix_values = self._get_radix_workspace(scores)
+            torch.ops.trtllm.indexer_topk_decode(
+                scores,
+                sequence_lengths,
+                output_indices,
+                next_n,
+                self.top_k,
+                pre_idx=gvr_prior_indices,
+                heuristic_scratch=workspace,
+                compress_ratio=self.compress_ratio,
+                radix_aux_indices=radix_indices,
+                radix_aux_logits=radix_values,
+            )
+        else:
+            assert max_seq_len is not None
+            torch.ops.trtllm.cute_dsl_gvr_topk_decode(
+                scores,
+                gvr_prior_indices,
+                sequence_lengths,
+                output_indices,
+                self.top_k,
+                next_n=next_n,
+                compress_ratio=self.compress_ratio,
+                max_seq_len=max_seq_len,
+                order_row=gvr_row_order,
+            )
+        return output_indices
+
+    def update_gvr_prior_from_prefill(
+        self,
+        output_indices: torch.Tensor,
+        request_lengths: torch.Tensor,
+        gvr_prior_indices: torch.Tensor | None,
+        *,
+        request_offset: int = 0,
+    ) -> None:
+        """Update GVR prior indices from each prefill request's last row.
+
+        Args:
+            output_indices: Int32 prefill selections with shape
+                ``[num_prefill_rows, top_k]``.
+            request_lengths: Per-request prefill row counts.
+            gvr_prior_indices: Int32 caller-owned state on
+                ``output_indices.device`` with shape ``[capacity, top_k]``.
+                The slice starting at ``request_offset`` is updated in place.
+            request_offset: First request row to update in the prior state.
+        """
+        if self.decode_implementation not in _GVR_IMPLEMENTATIONS:
+            return
+        assert gvr_prior_indices is not None
+        last_rows = (torch.cumsum(request_lengths, dim=0) - 1).to(dtype=torch.long)
+        num_requests = request_lengths.shape[0]
+        gvr_prior_indices[request_offset : request_offset + num_requests].copy_(
+            output_indices[last_rows]
+        )
+
+    def _forward_prefill_torch(
+        self,
+        scores: torch.Tensor,
+        row_starts: torch.Tensor,
+        row_ends: torch.Tensor,
+        output_indices: torch.Tensor,
+    ) -> torch.Tensor:
+        output_indices.fill_(-1)
+        selected_count = min(self.top_k, scores.shape[1])
+        if selected_count == 0:
+            return output_indices
+        columns = torch.arange(scores.shape[1], device=scores.device).unsqueeze(0)
+        valid = (columns >= row_starts.unsqueeze(1)) & (columns < row_ends.unsqueeze(1))
+        selected = scores.masked_fill(~valid, float("-inf")).topk(selected_count, dim=-1).indices
+        selected_valid = torch.gather(valid, 1, selected)
+        selected = selected - row_starts.unsqueeze(1)
+        selected = selected.masked_fill(~selected_valid, -1)
+        output_indices[:, :selected_count].copy_(selected.to(torch.int32))
+        return output_indices
+
+    def _forward_decode_torch(
+        self,
+        scores: torch.Tensor,
+        scan_lengths: torch.Tensor,
+        output_indices: torch.Tensor,
+        next_n: int,
+    ) -> torch.Tensor:
+        output_indices.fill_(-1)
+        selected_count = min(self.top_k, scores.shape[1])
+        if selected_count == 0:
+            return output_indices
+        positions = torch.arange(scores.shape[1], device=scores.device).unsqueeze(0)
+        row_indices = torch.arange(scores.shape[0], device=scores.device) // next_n
+        next_n_offsets = torch.arange(scores.shape[0], device=scores.device) % next_n
+        row_ends = scan_lengths[row_indices] - next_n + next_n_offsets + 1
+        valid = positions < row_ends.unsqueeze(1)
+        selected = scores.masked_fill(~valid, float("-inf")).topk(selected_count, dim=-1).indices
+        selected_valid = torch.gather(valid, 1, selected)
+        selected = selected.masked_fill(~selected_valid, -1)
+        output_indices[:, :selected_count].copy_(selected.to(torch.int32))
+        return output_indices
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 76552749651f..afbfefea9993 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -54,6 +54,8 @@
     transform_local_topk_and_prepare_pool_view,
 )
 from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata
+from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation
 from tensorrt_llm._torch.pyexecutor._util import get_kv_cache_manager_cls
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role
 from tensorrt_llm._torch.speculative.interface import (
@@ -83,6 +85,15 @@ def has_deep_gemm():
         return False
 
 
+def _set_torch_top_k(indexer: Indexer) -> None:
+    indexer.top_k = TopK(
+        indexer.index_topk,
+        prefill_implementation=TopKImplementation.TORCH,
+        decode_implementation=TopKImplementation.TORCH,
+        compress_ratio=indexer.compress_ratio,
+    )
+
+
 def test_metadata_cache_geometry_comes_from_sparse_metadata_params():
     sparse_config = DeepSeekV4SparseAttentionConfig(
         compress_ratios=[1, 4, 128],
@@ -112,6 +123,161 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params():
     assert metadata._tokens_per_block == 64
 
 
+@pytest.mark.parametrize(
+    "enable_heuristic,use_cute_dsl,sm_version,compress_ratio,next_n,should_warmup",
+    [
+        (True, True, 100, 1, 1, False),
+        (False, True, 100, 1, 1, True),
+        (True, True, 90, 1, 1, True),
+        (False, True, 100, 4, 2, False),
+        (False, False, 100, 1, 1, False),
+    ],
+)
+def test_metadata_warmup_cute_dsl_radix_topk_dispatch(
+    enable_heuristic,
+    use_cute_dsl,
+    sm_version,
+    compress_ratio,
+    next_n,
+    should_warmup,
+):
+    metadata = SimpleNamespace(
+        sparse_metadata_params=SimpleNamespace(
+            enable_heuristic_topk=enable_heuristic,
+        ),
+        use_cute_dsl_topk=use_cute_dsl,
+        num_sparse_topk=512,
+        sparse_mla_topk=384,
+        kv_cache_manager=SimpleNamespace(),
+        _indexer_compress_ratio=compress_ratio,
+        get_indexer_max_seq_len=Mock(return_value=32768),
+        num_sms=148,
+    )
+
+    with (
+        patch(
+            "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.get_sm_version",
+            return_value=sm_version,
+        ),
+        patch(
+            "tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops.warmup_cute_dsl_radix_topk_decode"
+        ) as cute_dsl_radix,
+    ):
+        DSAtrtllmAttentionMetadata.warmup_cute_dsl_radix_topk(metadata, next_n)
+
+    if should_warmup:
+        cute_dsl_radix.assert_called_once_with(
+            top_k=384,
+            num_cols=32768,
+            next_n=next_n,
+            dtype=torch.float32,
+            num_sms=148,
+        )
+    else:
+        cute_dsl_radix.assert_not_called()
+
+
+def test_kv_lens_row_reorder_threshold():
+    """Prepare row order only when CuTe DSL GVR has enough decode rows."""
+    num_sms = 16
+    next_n = 2
+
+    def make_mock(num_generations, kv_lens_list):
+        kv_lens_cuda = torch.tensor(kv_lens_list, dtype=torch.int32, device="cuda")
+        row_order_buffer = torch.zeros(64, dtype=torch.int32, device="cuda")
+        return SimpleNamespace(
+            enable_gvr_topk=True,
+            use_cute_dsl_topk=True,
+            num_generations=num_generations,
+            num_sms=num_sms,
+            max_draft_tokens=next_n - 1,
+            num_contexts=0,
+            num_seqs=num_generations,
+            kv_lens_cuda=kv_lens_cuda,
+            kv_lens_row_reorder_buffer=row_order_buffer,
+            kv_lens_row_reorder=None,
+        )
+
+    kv_lens = [4, 1, 8, 2, 16, 3, 12, 6, 7, 9, 5, 11, 13, 10, 14, 15]
+
+    metadata_below = make_mock(1, [1000])
+    DSAtrtllmAttentionMetadata._compute_kv_lens_row_reorder(metadata_below)
+    assert metadata_below.kv_lens_row_reorder is None
+
+    metadata_at = make_mock(num_sms, kv_lens)
+    DSAtrtllmAttentionMetadata._compute_kv_lens_row_reorder(metadata_at)
+    row_order = metadata_at.kv_lens_row_reorder.cpu().tolist()
+    assert [kv_lens[i] for i in row_order] == sorted(kv_lens, reverse=True)
+
+
+@skip_pre_hopper
+def test_gvr_prior_writeback_uses_aux_stream():
+    batch_size = 2
+    index_topk = 4
+    cache_manager, sparse_config = create_dsa_cache_manager(
+        batch_size=batch_size,
+        head_dim=128,
+        tokens_per_block=64,
+        max_seq_len=64,
+        num_layers=1,
+        index_topk=index_topk,
+    )
+    try:
+        request_ids = list(range(batch_size))
+        kv_lens = torch.full((batch_size,), index_topk, dtype=torch.int32)
+        cache_manager.add_dummy_requests(
+            request_ids,
+            kv_lens.tolist(),
+            is_gen=False,
+            prepare_resource=True,
+        )
+        metadata = _create_mock_metadata(
+            request_ids,
+            batch_size,
+            num_contexts=0,
+            num_generations=batch_size,
+            seq_lens=torch.ones(batch_size, dtype=torch.int32),
+            kv_lens=kv_lens,
+            num_cached_tokens=[index_topk - 1] * batch_size,
+            cache_manager=cache_manager,
+            num_ctx_tokens=0,
+            num_tokens=batch_size,
+            index_topk=index_topk,
+            enable_indexer_skip=True,
+        )
+        indexer = create_indexer(sparse_config)
+        indexer._enable_heuristic_topk = True
+        indexer.aux_stream = torch.cuda.Stream()
+        metadata.gvr_prior_indices = torch.zeros(
+            (cache_manager.num_local_layers, batch_size, index_topk),
+            device="cuda",
+            dtype=torch.int32,
+        )
+        hidden_states = torch.empty((batch_size, 1), device="cuda")
+        unused = torch.empty((batch_size, 1), device="cuda")
+
+        with with_multi_stream(True):
+            topk_indices = indexer.sparse_attn_indexer(
+                metadata,
+                hidden_states,
+                unused,
+                unused,
+                unused,
+                unused,
+            )
+            assert indexer._prev_topk_copy_pending
+            indexer.maybe_join_prev_topk_copy()
+
+        local_layer = cache_manager.layer_offsets[indexer.layer_idx]
+        torch.testing.assert_close(
+            metadata.gvr_prior_indices[local_layer, :batch_size],
+            topk_indices,
+        )
+        assert not indexer._prev_topk_copy_pending
+    finally:
+        cache_manager.shutdown()
+
+
 def test_shared_topk_lifecycle():
     sparse_config = DeepSeekSparseAttentionConfig(
         index_n_heads=1,
@@ -138,11 +304,11 @@ def test_shared_topk_lifecycle():
     metadata.kv_cache_manager = SimpleNamespace(max_blocks_per_seq=2)
     metadata.enable_context_mla_with_cached_kv = False
     metadata.enable_indexer_skip = False
+    metadata.enable_gvr_topk = False
     metadata.get_empty = Mock(
         side_effect=lambda _, shape, **kwargs: torch.empty(tuple(shape), dtype=kwargs["dtype"])
     )
     metadata._create_kv_lens_2d_buffer = Mock()
-    metadata._create_radix_aux_buffers = Mock()
     metadata.create_expanded_buffers = Mock()
 
     with patch(
@@ -229,6 +395,46 @@ def test_indexer_post_load_weights_caches_fused_weight():
     assert not hasattr(indexer, "_weights_transformed")
 
 
+@skip_pre_hopper
+@pytest.mark.parametrize(
+    "use_cute_dsl,enable_heuristic,expected_decode",
+    [
+        (False, False, TopKImplementation.CUDA_RADIX),
+        (True, False, TopKImplementation.CUTE_DSL_RADIX),
+        (False, True, TopKImplementation.CUDA_GVR),
+        (True, True, TopKImplementation.CUTE_DSL_GVR),
+    ],
+)
+def test_indexer_configures_one_top_k_module(
+    use_cute_dsl,
+    enable_heuristic,
+    expected_decode,
+):
+    sparse_config = DeepSeekSparseAttentionConfig(
+        index_head_dim=128,
+        index_n_heads=32,
+        index_topk=128,
+        use_cute_dsl_topk=use_cute_dsl,
+        enable_heuristic_topk=enable_heuristic,
+    )
+
+    with (
+        patch(
+            "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.IS_CUTLASS_DSL_AVAILABLE",
+            True,
+        ),
+        patch(
+            "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.get_sm_version",
+            return_value=100,
+        ),
+    ):
+        indexer = create_indexer(sparse_config)
+
+    assert isinstance(indexer.top_k, TopK)
+    assert indexer.top_k.prefill_implementation == TopKImplementation.CUDA_RADIX
+    assert indexer.top_k.decode_implementation == expected_decode
+
+
 def _ceil_to_ue8m0(x: torch.Tensor):
     """Round tensor values up to the nearest power of two (UE8M0 format)."""
     return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
@@ -718,6 +924,7 @@ def __init__(self):
             self.num_contexts = num_contexts
             self.num_generations = num_generations
             self._num_seqs = num_contexts + num_generations
+            self.max_num_sequences = batch_size
             self.max_draft_tokens = max_draft_tokens
             self.num_sparse_topk = index_topk
             self.enable_indexer_skip = enable_indexer_skip
@@ -2660,7 +2867,7 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, compr
 @pytest.mark.parametrize("seq_len_range", [(2048, 8192), (512, 1024)])
 def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_len_range):
     """
-    Test that use_custom_topk=True and use_custom_topk=False produce identical results
+    Test that the production and Torch Top-K modules produce identical results
     in the decode phase of sparse_attn_indexer.
 
     This test validates:
@@ -2806,7 +3013,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l
 
     try:
         topk_indices_custom = indexer.sparse_attn_indexer(
-            metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True
+            metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights
         )
     except Exception as e:
         pytest.skip(f"Custom topk not available: {e}")
@@ -2830,8 +3037,9 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l
 
     Indexer.prepare(metadata_fallback)
     indexer._update_k_cache(k_fp8, k_scale, metadata_fallback)
+    _set_torch_top_k(indexer)
     topk_indices_fallback = indexer.sparse_attn_indexer(
-        metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False
+        metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights
     )
 
     # Test with indexer skip enabled
@@ -2858,7 +3066,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l
 
         try:
             topk_indices_skip = indexer.sparse_attn_indexer(
-                metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True
+                metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights
             )
         except Exception as e:
             raise RuntimeError(f"Error when testing indexer skip: {e}")
@@ -2975,14 +3183,6 @@ def make_inputs(n_tokens):
             max_draft_tokens=md,
         )
         Indexer.prepare(meta0)
-        # indexer_topk_decode needs caller-owned radix aux buffers for small gen batches.
-        _radix_bp = 10
-        meta0.radix_aux_indices = torch.zeros(
-            (step0_tokens, _radix_bp, index_topk), device="cuda", dtype=torch.int32
-        )
-        meta0.radix_aux_logits = torch.zeros(
-            (step0_tokens, _radix_bp, index_topk), device="cuda", dtype=torch.float32
-        )
     else:  # context: first gen round -- step 0 runs the context/prefill path.
         step0_tokens = kv_lens.sum().item()
         meta0 = _create_mock_metadata(
@@ -3009,9 +3209,7 @@ def make_inputs(n_tokens):
     h0, q0, k0_fp8, k0_scale, w0 = make_inputs(step0_tokens)
     indexer._update_k_cache(k0_fp8, k0_scale, meta0)
     try:
-        topk0 = indexer.sparse_attn_indexer(
-            meta0, h0, q0, k0_fp8, k0_scale, w0, use_custom_topk=True
-        )
+        topk0 = indexer.sparse_attn_indexer(meta0, h0, q0, k0_fp8, k0_scale, w0)
     except Exception as e:
         pytest.skip(f"Custom topk not available: {e}")
 
@@ -3049,7 +3247,7 @@ def make_inputs(n_tokens):
         meta.indexer_skip_topk = True
         hs, qs, ks_fp8, ks_scale, ws = make_inputs(batch_size)
         indexer._update_k_cache(ks_fp8, ks_scale, meta)
-        topk = indexer.sparse_attn_indexer(meta, hs, qs, ks_fp8, ks_scale, ws, use_custom_topk=True)
+        topk = indexer.sparse_attn_indexer(meta, hs, qs, ks_fp8, ks_scale, ws)
         assert torch.equal(topk, stash[:batch_size, :]), (
             f"{step0_mode} draft reuse step {step} should copy the stash 1:1 (next_n=1)"
         )
@@ -3062,7 +3260,7 @@ def make_inputs(n_tokens):
 @pytest.mark.parametrize("chunk_size", [1024, 2048])
 def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chunk_size):
     """
-    Test chunked prefill: use_custom_topk=True vs use_custom_topk=False
+    Test chunked prefill: production Top-K vs Torch Top-K
     with metadata.indexer_prefill_chunks != None.
 
     This test validates:
@@ -3135,7 +3333,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun
 
     try:
         topk_indices_custom = indexer.sparse_attn_indexer(
-            metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True
+            metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights
         )
     except Exception as e:
         pytest.skip(f"Custom topk not available: {e}")
@@ -3158,8 +3356,9 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun
 
     Indexer.prepare(metadata_fallback)
     indexer._update_k_cache(k_fp8, k_scale, metadata_fallback)
+    _set_torch_top_k(indexer)
     topk_indices_fallback = indexer.sparse_attn_indexer(
-        metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False
+        metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights
     )
 
     # Validation
@@ -3179,7 +3378,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun
 @pytest.mark.parametrize("seq_len_range", [(1, 512)])
 def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, seq_len_range):
     """
-    Test single-pass prefill: use_custom_topk=True vs use_custom_topk=False
+    Test single-pass prefill: production Top-K vs Torch Top-K
     with metadata.indexer_prefill_chunks == None (else branch).
     """
     torch.manual_seed(42)
@@ -3243,7 +3442,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk,
 
     try:
         topk_indices_custom = indexer.sparse_attn_indexer(
-            metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True
+            metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights
         )
     except Exception as e:
         pytest.skip(f"Custom topk not available: {e}")
@@ -3269,8 +3468,9 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk,
     # Force single-pass path by setting indexer_prefill_chunks to None
     metadata_fallback.indexer_prefill_chunks = None
 
+    _set_torch_top_k(indexer)
     topk_indices_fallback = indexer.sparse_attn_indexer(
-        metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False
+        metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights
     )
 
     # Test with indexer skip enabled
@@ -3292,14 +3492,32 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk,
     Indexer.prepare(metadata_skip)
     indexer._update_k_cache(k_fp8, k_scale, metadata_skip)
     metadata_skip.indexer_prefill_chunks = None
+    indexer.top_k = TopK(
+        index_topk,
+        prefill_implementation=TopKImplementation.CUDA_RADIX,
+        decode_implementation=TopKImplementation.CUTE_DSL_GVR,
+    )
+    indexer._enable_heuristic_topk = True
+    metadata_skip.gvr_prior_indices = torch.zeros(
+        (cache_manager.num_local_layers, batch_size, index_topk),
+        device="cuda",
+        dtype=torch.int32,
+    )
 
     try:
         topk_indices_skip = indexer.sparse_attn_indexer(
-            metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True
+            metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights
         )
     except Exception as e:
         raise RuntimeError(f"Indexer skip not available: {e}")
 
+    last_rows = torch.cumsum(metadata_skip.seq_lens[:batch_size], dim=0) - 1
+    local_layer = cache_manager.layer_offsets[layer_idx]
+    torch.testing.assert_close(
+        metadata_skip.gvr_prior_indices[local_layer, :batch_size],
+        topk_indices_skip[last_rows],
+    )
+
     # Validation
     ## Custom vs fallback
     num_exact_matches, total_similarity, _ = validate_topk_indices(
@@ -3399,12 +3617,13 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip):
 
     # Test custom kernel
     topk_custom = indexer.sparse_attn_indexer(
-        metadata, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True
-    )
+        metadata, hidden_states, q_fp8, k_fp8, k_scale, weights
+    ).clone()
 
     # Test fallback
+    _set_torch_top_k(indexer)
     topk_fallback = indexer.sparse_attn_indexer(
-        metadata, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False
+        metadata, hidden_states, q_fp8, k_fp8, k_scale, weights
     )
 
     # Test with indexer skip enabled
@@ -3428,7 +3647,7 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip):
         Indexer.prepare(metadata_skip)
         indexer._update_k_cache(k_fp8, k_scale, metadata_skip)
         topk_indices_skip = indexer.sparse_attn_indexer(
-            metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True
+            metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights
         )
 
     # Validate: custom and fallback should match
@@ -3783,55 +4002,3 @@ 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"
     )
-
-
-def test_kv_lens_row_reorder_threshold():
-    """_compute_kv_lens_row_reorder engages iff num_generations * next_n >= 2 * num_sms,
-    and produces a descending argsort of gen_kv_lens when active."""
-    num_sms = 16  # small synthetic value; threshold = 2 * 16 = 32 rows
-    next_n = 2  # max_draft_tokens=1 → next_n = 1 + 1 = 2
-
-    def make_mock(num_generations, kv_lens_list):
-        kv_cuda = torch.tensor(kv_lens_list, dtype=torch.int32, device="cuda")
-        buf = torch.zeros(64, dtype=torch.int32, device="cuda")
-        ns = SimpleNamespace(
-            enable_heuristic_topk=True,
-            use_cute_dsl_topk=True,
-            num_generations=num_generations,
-            num_sms=num_sms,
-            max_draft_tokens=next_n - 1,
-            num_contexts=0,
-            num_seqs=num_generations,
-            kv_lens_cuda=kv_cuda,
-            kv_lens_row_reorder_buffer=buf,
-            kv_lens_row_reorder=None,
-        )
-        ns._compute_kv_lens_row_reorder = (
-            lambda: DSAtrtllmAttentionMetadata._compute_kv_lens_row_reorder(ns)
-        )
-        return ns
-
-    # Fixed unsorted sequence for deterministic sort verification (len == num_sms)
-    kv_vals = [4, 1, 8, 2, 16, 3, 12, 6, 7, 9, 5, 11, 13, 10, 14, 15]
-
-    # Below threshold: 1 * 2 = 2 < 32 → None
-    md_below = make_mock(1, [1000])
-    md_below._compute_kv_lens_row_reorder()
-    assert md_below.kv_lens_row_reorder is None
-
-    # At threshold: num_sms * 2 = 32 → engages, verify descending argsort
-    md_at = make_mock(num_sms, kv_vals)
-    md_at._compute_kv_lens_row_reorder()
-    assert md_at.kv_lens_row_reorder is not None
-    reorder = md_at.kv_lens_row_reorder.cpu().tolist()
-    assert [kv_vals[i] for i in reorder] == sorted(kv_vals, reverse=True), (
-        "order_row must be a descending argsort of gen_kv_lens"
-    )
-
-    # Above threshold: (num_sms + 1) * 2 = 34 > 32 → also engages with correct sort
-    kv_vals2 = kv_vals + [100]
-    md_above = make_mock(num_sms + 1, kv_vals2)
-    md_above._compute_kv_lens_row_reorder()
-    assert md_above.kv_lens_row_reorder is not None
-    reorder2 = md_above.kv_lens_row_reorder.cpu().tolist()
-    assert [kv_vals2[i] for i in reorder2] == sorted(kv_vals2, reverse=True)
diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py
new file mode 100644
index 000000000000..9682d7981ab5
--- /dev/null
+++ b/tests/unittest/_torch/modules/test_top_k.py
@@ -0,0 +1,330 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the reusable sparse index-selection Top-K module."""
+
+from contextlib import nullcontext
+from unittest.mock import Mock, call
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation
+
+
+def test_prefill_torch_masks_dirty_scores_and_pads_output() -> None:
+    scores = torch.tensor(
+        [
+            [1000.0, 9.0, 8.0, 7.0, 1000.0, 1000.0],
+            [1000.0, 1000.0, 3.0, 1000.0, 1000.0, 1000.0],
+            [1000.0, 1000.0, 1000.0, 1000.0, 1000.0, 1000.0],
+        ]
+    )
+    row_starts = torch.tensor([1, 2, 4], dtype=torch.int32)
+    row_ends = torch.tensor([4, 3, 4], dtype=torch.int32)
+    output = torch.full((3, 4), 77, dtype=torch.int32)
+
+    result = TopK(4, prefill_implementation=TopKImplementation.TORCH)(
+        scores,
+        output,
+        is_prefill=True,
+        row_starts=row_starts,
+        row_ends=row_ends,
+    )
+
+    assert result is output
+    assert output.tolist() == [[0, 1, 2, -1], [0, -1, -1, -1], [-1, -1, -1, -1]]
+
+
+def test_decode_torch_uses_scan_lengths() -> None:
+    scores = torch.tensor(
+        [
+            [1.0, 2.0, 3.0, 1000.0, 1000.0, 1000.0],
+            [1.0, 2.0, 3.0, 4.0, 1000.0, 1000.0],
+        ]
+    )
+    logical_lengths = torch.tensor([16], dtype=torch.int32)
+    scan_lengths = torch.tensor([3], dtype=torch.int32)
+    output = torch.full((2, 4), 77, dtype=torch.int32)
+
+    result = TopK(
+        4,
+        decode_implementation=TopKImplementation.TORCH,
+        compress_ratio=4,
+    )(
+        scores,
+        output,
+        is_prefill=False,
+        sequence_lengths=logical_lengths,
+        scan_lengths=scan_lengths,
+        next_n=2,
+    )
+
+    assert result is output
+    assert output.tolist() == [[1, 0, -1, -1], [2, 1, 0, -1]]
+
+
+def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None:
+    cute_dsl = Mock()
+    trtllm = Mock()
+    monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_indexer_topk_decode", cute_dsl)
+    monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", trtllm)
+
+    top_k = TopK(
+        2,
+        decode_implementation=TopKImplementation.CUTE_DSL_RADIX,
+        compress_ratio=4,
+    )
+    logical_lengths = torch.tensor([16], dtype=torch.int32)
+    scan_lengths = torch.tensor([4], dtype=torch.int32)
+
+    scores = torch.randn(1, 4)
+    output = torch.empty(1, 2, dtype=torch.int32)
+    top_k(
+        scores,
+        output,
+        is_prefill=False,
+        sequence_lengths=logical_lengths,
+        scan_lengths=scan_lengths,
+        next_n=1,
+    )
+    cute_dsl.assert_called_once_with(scores, scan_lengths, output, 2, 1)
+    trtllm.assert_not_called()
+
+    cute_dsl.reset_mock()
+    scores = torch.randn(2, 4)
+    output = torch.empty(2, 2, dtype=torch.int32)
+    radix_indices = torch.empty(2, 10, 2, dtype=torch.int32)
+    radix_values = torch.empty(2, 10, 2)
+    buffers = Mock()
+    buffers.get_buffer.side_effect = [radix_indices, radix_values]
+    monkeypatch.setattr(TopK, "_memory_buffers", buffers)
+    top_k(
+        scores,
+        output,
+        is_prefill=False,
+        sequence_lengths=logical_lengths,
+        scan_lengths=scan_lengths,
+        next_n=2,
+    )
+    cute_dsl.assert_not_called()
+    assert buffers.get_buffer.call_args_list == [
+        call(
+            (2, 10, 2),
+            dtype=torch.int32,
+            buffer_name="top_k_radix_indices_workspace_cpu",
+            reserve_buffer=False,
+        ),
+        call(
+            (2, 10, 2),
+            dtype=torch.float32,
+            buffer_name="top_k_radix_values_workspace_cpu",
+            reserve_buffer=False,
+        ),
+    ]
+    trtllm.assert_called_once_with(
+        scores,
+        logical_lengths,
+        output,
+        2,
+        2,
+        pre_idx=None,
+        heuristic_scratch=None,
+        compress_ratio=4,
+        radix_aux_indices=radix_indices,
+        radix_aux_logits=radix_values,
+    )
+
+
+def test_gvr_uses_caller_prior_state(monkeypatch) -> None:
+    gvr = Mock()
+    monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr)
+    top_k = TopK(
+        2,
+        decode_implementation=TopKImplementation.CUTE_DSL_GVR,
+        compress_ratio=4,
+    )
+    scores = torch.randn(1, 8)
+    logical_lengths = torch.tensor([32], dtype=torch.int32)
+    scan_lengths = torch.tensor([8], dtype=torch.int32)
+    output = torch.empty(1, 2, dtype=torch.int32)
+    prior_indices = torch.zeros(1, 2, dtype=torch.int32)
+    gvr.side_effect = lambda *args, **kwargs: output.copy_(torch.tensor([[5, 3]]))
+
+    top_k(
+        scores,
+        output,
+        is_prefill=False,
+        sequence_lengths=logical_lengths,
+        scan_lengths=scan_lengths,
+        next_n=1,
+        max_seq_len=16,
+        gvr_ext_kwargs={"gvr_prior_indices": prior_indices},
+    )
+
+    args, kwargs = gvr.call_args
+    assert args[0] is scores
+    assert args[1] is prior_indices
+    assert args[2] is logical_lengths
+    assert args[3] is output
+    assert args[4] == 2
+    assert kwargs == {
+        "next_n": 1,
+        "compress_ratio": 4,
+        "max_seq_len": 16,
+        "order_row": None,
+    }
+    assert prior_indices.tolist() == [[0, 0]]
+
+
+def test_gvr_uses_caller_prepared_row_order(monkeypatch) -> None:
+    gvr = Mock(side_effect=lambda *args, **kwargs: args[3].zero_())
+    monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr)
+    top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR)
+    next_n = 2
+    lengths = torch.tensor([4, 1, 8, 2], dtype=torch.int32)
+    row_order = torch.tensor([2, 0, 3, 1], dtype=torch.int32)
+
+    top_k(
+        torch.randn(lengths.shape[0] * next_n, 8),
+        torch.empty(lengths.shape[0] * next_n, 2, dtype=torch.int32),
+        is_prefill=False,
+        sequence_lengths=lengths,
+        scan_lengths=lengths,
+        next_n=next_n,
+        max_seq_len=8,
+        gvr_ext_kwargs={
+            "gvr_prior_indices": torch.zeros(lengths.shape[0], 2, dtype=torch.int32),
+            "gvr_row_order": row_order,
+        },
+    )
+
+    assert gvr.call_args.kwargs["order_row"] is row_order
+
+
+def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None:
+    top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR)
+    prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32)
+    prior_indices = torch.zeros(3, 2, dtype=torch.int32)
+
+    top_k.update_gvr_prior_from_prefill(
+        prefill_indices,
+        torch.tensor([2, 1], dtype=torch.int32),
+        prior_indices,
+        request_offset=1,
+    )
+
+    assert prior_indices.tolist() == [[0, 0], [2, 3], [4, 5]]
+
+
+def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None:
+    decode = Mock()
+    monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode)
+    top_k = TopK(1)
+    scores = torch.randn(1, 8)
+    lengths = torch.tensor([8], dtype=torch.int32)
+    output = torch.empty((1, 1), dtype=torch.int32)
+    radix_indices = torch.empty(1, 10, 1, dtype=torch.int32)
+    radix_values = torch.empty(1, 10, 1)
+    buffers = Mock()
+    buffers.get_buffer.side_effect = [radix_indices, radix_values]
+    monkeypatch.setattr(TopK, "_memory_buffers", buffers)
+
+    result = top_k(
+        scores,
+        output,
+        is_prefill=False,
+        sequence_lengths=lengths,
+        scan_lengths=lengths,
+    )
+
+    assert top_k.prefill_implementation == TopKImplementation.CUDA_RADIX
+    assert top_k.decode_implementation == TopKImplementation.CUDA_RADIX
+    assert result is output
+    assert buffers.get_buffer.call_count == 2
+    decode.assert_called_once_with(
+        scores,
+        lengths,
+        output,
+        1,
+        1,
+        pre_idx=None,
+        heuristic_scratch=None,
+        compress_ratio=1,
+        radix_aux_indices=radix_indices,
+        radix_aux_logits=radix_values,
+    )
+
+
+def test_cuda_gvr_reserves_workspace_during_capture(monkeypatch) -> None:
+    decode = Mock(side_effect=lambda *args, **kwargs: args[2].copy_(torch.tensor([[3, 1]])))
+    monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True))
+    device_context = Mock(side_effect=lambda _: nullcontext())
+    monkeypatch.setattr(torch.cuda, "device", device_context)
+
+    top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR)
+    scores = Mock(
+        shape=(1, 8),
+        dtype=torch.float32,
+        is_cuda=True,
+        device=torch.device("cuda", 3),
+    )
+    lengths = torch.tensor([8], dtype=torch.int32)
+    output = torch.empty(1, 2, dtype=torch.int32)
+    radix_indices = torch.empty(1, 10, 2, dtype=torch.int32)
+    radix_values = torch.empty(1, 10, 2)
+    workspace = torch.empty(1, 2)
+    prior_indices = torch.zeros(1, 2, dtype=torch.int32)
+    buffers = Mock()
+    buffers.get_buffer.side_effect = [workspace, radix_indices, radix_values]
+    monkeypatch.setattr(TopK, "_memory_buffers", buffers)
+
+    top_k(
+        scores,
+        output,
+        is_prefill=False,
+        sequence_lengths=lengths,
+        scan_lengths=lengths,
+        gvr_ext_kwargs={"gvr_prior_indices": prior_indices},
+    )
+
+    assert buffers.get_buffer.call_args_list == [
+        call(
+            (scores.shape[0], 2),
+            dtype=scores.dtype,
+            buffer_name="top_k_cuda_gvr_workspace_cuda:3",
+            reserve_buffer=True,
+        ),
+        call(
+            (scores.shape[0], 10, 2),
+            dtype=torch.int32,
+            buffer_name="top_k_radix_indices_workspace_cuda:3",
+            reserve_buffer=True,
+        ),
+        call(
+            (scores.shape[0], 10, 2),
+            dtype=torch.float32,
+            buffer_name="top_k_radix_values_workspace_cuda:3",
+            reserve_buffer=True,
+        ),
+    ]
+    assert device_context.call_args_list == [call(scores.device)] * 3
+    runtime_call = decode.call_args_list[-1]
+    assert runtime_call.kwargs["pre_idx"] is prior_indices
+    assert runtime_call.kwargs["heuristic_scratch"].data_ptr() == workspace.data_ptr()
+    assert runtime_call.kwargs["radix_aux_indices"] is radix_indices
+    assert runtime_call.kwargs["radix_aux_logits"] is radix_values
+    assert prior_indices.tolist() == [[0, 0]]
+
+
+def test_unsupported_prefill_implementation_raises() -> None:
+    top_k = TopK(1, prefill_implementation=TopKImplementation.CUTE_DSL_RADIX)
+
+    with pytest.raises(NotImplementedError, match="does not support prefill Top-K"):
+        top_k(
+            torch.ones(1, 1),
+            torch.empty(1, 1, dtype=torch.int32),
+            is_prefill=True,
+            row_starts=torch.zeros(1, dtype=torch.int32),
+            row_ends=torch.ones(1, dtype=torch.int32),
+        )