[DSA] Trim top-k v2 output modes and tighten its PDL waits (#35041)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh

@@ -28,6 +28,11 @@ namespace sglang {
 namespace impl = device::topk;
 using impl::TopKProblem;
 
+enum class TopKMode {
+  INDICES,     ///< raw selected indices into `out`; `page_table` unused
+  PAGE_TABLE,  ///< page-table-transformed indices into `out`
+};
+
 using Register2 = impl::TopKRegister<2>;  // <= 8192, register-resident, 1 read
 using Register4 = impl::TopKRegister<4>;  // <= 16384, register-resident, 1 read
 using Streaming = impl::TopKStreaming;
@@ -64,7 +69,6 @@ struct TopKLaunchParams {
   const int32_t* __restrict__ seq_lens;
   const int32_t* __restrict__ page_table;
   int32_t* __restrict__ page_indices;
-  int32_t* __restrict__ raw_indices;      // optional raw (pre-transform) indices output; nullptr if unused
   const PlanItem* __restrict__ metadata;  // [0]=GlobalMetadata, [1+i]=PlanItem
   int64_t score_stride;
   int64_t page_table_stride;
@@ -89,7 +93,6 @@ struct TopKLaunchParams {
     return TopKProblem{
         .in = scores + batch_id * score_stride,
         .out = page_indices + batch_id * k,
-        .raw_out = raw_indices != nullptr ? raw_indices + batch_id * k : nullptr,
         .page_table = page_table + batch_id * page_table_stride,
         .topk = topk,
         .seq_len = seq_len,
@@ -133,12 +136,17 @@ SGL_DEVICE void for_each_item(uint32_t topk, const F& f) {
   }
 }
 
-template <bool kPDL>
+template <bool kPDL, TopKMode kMode>
 SGL_DEVICE void trivial_transform(const TopKProblem& problem) {
   device::PDLWaitPrimary<kPDL>();
   device::PDLTriggerSecondary<kPDL>();
   for_each_item(problem.topk, [&](uint32_t tx, uint32_t) {
-    problem.transform_output(tx, tx < problem.seq_len ? static_cast<int32_t>(tx) : -1);
+    const auto idx = tx < problem.seq_len ? static_cast<int32_t>(tx) : -1;
+    if constexpr (kMode == TopKMode::INDICES) {
+      problem.emit(tx, idx);
+    } else {
+      problem.transform_output(tx, idx);
+    }
   });
 }
 
@@ -160,17 +168,25 @@ SGL_DEVICE void problem_transform(TopKProblem& problem, int32_t* output_ptr) {
  * - Level 2: max_seq_len <= cluster_floor  -> trivial + register<4> + streaming
  * - Level 3: max_seq_len > cluster_floor   -> + epilogue process of cluster path
  */
-template <bool kPDL, int kLevel>
+template <bool kPDL, int kLevel, TopKMode kMode>
 TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams params) {
   device::enable_smem_spilling();
   auto problem = params.problem(blockIdx.x);
   constexpr uint32_t kU32Max = std::numeric_limits<uint32_t>::max();
-  __shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
-  if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
-  __shared__ int32_t topk_indices[kMaxTopK];
-  problem.out = topk_indices;
-
   constexpr bool kHandleCluster = (kLevel == 3);
+  // Only the cluster path consumes the cluster kernel's output, so only it waits
+  // on that kernel (kPDLFinal). Every other path waits at most on the indexer
+  // (kPDLEarly) and must not be held on an SM slot until the long-running
+  // persistent pool retires -- that would serialize the short items behind it.
+  constexpr bool kPDLEarly = kPDL && !kHandleCluster;
+  constexpr bool kPDLFinal = kPDL && kHandleCluster;
+  __shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
+  if (problem.seq_len <= problem.topk) return trivial_transform<kPDLEarly, kMode>(problem);
+
+  constexpr bool kNeedStaging = kMode != TopKMode::INDICES;
+  __shared__ int32_t s_topk_indices[kNeedStaging ? kMaxTopK : 1];
+  if constexpr (kNeedStaging) problem.out = s_topk_indices;
+
   // non-trivial path: dispatch based on level and seq_len
   const auto cluster_threshold = kHandleCluster ? params.cluster_threshold() : kU32Max;
   if constexpr (kLevel == 0) {
@@ -181,34 +197,35 @@ TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams param
     Register4::forward<kPDL>(problem, &smem);  // max_seq_len <= 16384 guarantees seq <= 16384
   } else {
     static_assert(kLevel == 2 || kLevel == 3, "we only support level = 0,1,2,3 now");
-    // if using cluster, we can delay the PDL wait
-    constexpr bool kPDLEarly = kPDL && !kHandleCluster;
-    constexpr bool kPDLFinal = kPDL && kHandleCluster;
     if (problem.seq_len <= kReg4MaxSeqLen) {
       Register4::forward<kPDLEarly>(problem, &smem);
     } else if (problem.seq_len <= cluster_threshold) {
       Streaming::forward<kPDLEarly>(problem, &smem);
-    } else {  // cluster path do nothing here
+    } else {
+      // Cluster path: the pool already selected into our output row; the only
+      // work left is the epilogue, so this is the one path that waits for it.
       problem.out = params.get_output_ptr(blockIdx.x);
-    }
       device::PDLWaitPrimary<kPDLFinal>();
     }
+  }
 
-  // page-table transform pass (gathers kept out of the hot scatter loop),
-  // then trigger the dependent kernel only after the full output is written.
   device::PDLTriggerSecondary<kPDL>();
+  if constexpr (kNeedStaging) {
     __syncthreads();
     problem_transform(problem, params.get_output_ptr(blockIdx.x));
+  }
 }
 
-template <bool kPDL>
+template <bool kPDL, TopKMode kMode>
 CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLaunchParams params) {
   device::enable_smem_spilling();
   auto problem = params.problem(blockIdx.x);
   __shared__ impl::MaxSmem<Streaming::Smem, Cluster::Smem> smem;
-  if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
+  if (problem.seq_len <= problem.topk) return trivial_transform<kPDL, kMode>(problem);
-  __shared__ int32_t topk_indices[kMaxTopK];
+
-  problem.out = topk_indices;
+  constexpr bool kNeedStaging = kMode != TopKMode::INDICES;
+  __shared__ int32_t s_topk_indices[kNeedStaging ? kMaxTopK : 1];
+  if constexpr (kNeedStaging) problem.out = s_topk_indices;
 
   // randomly elect one worker rank to avoid workload imbalance
   const auto worker_rank = blockIdx.x % kClusterSize;
@@ -217,30 +234,34 @@ CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLau
   if (problem.seq_len <= kReg4MaxSeqLen) {
     if (blockIdx.y != worker_rank) return;
     Register4::forward<kPDL>(problem, &smem);
-    device::PDLWaitPrimary<kPDL>();
     __syncthreads();
   } else if (problem.seq_len <= params.cluster_floor) {
     if (blockIdx.y != worker_rank) return;
     Streaming::forward<kPDL>(problem, &smem);
-    device::PDLWaitPrimary<kPDL>();
     __syncthreads();
   } else {
     auto cluster = cooperative_groups::this_cluster();
-    problem.out = cluster.map_shared_rank(topk_indices, worker_rank);
+    if constexpr (kNeedStaging) {
-    Cluster::forward<kPDL>(problem, &smem);  // write to peer's output shared memory
+      problem.out = cluster.map_shared_rank(s_topk_indices, worker_rank);
-    device::PDLWaitPrimary<kPDL>();
+    }
+    Cluster::forward<kPDL>(problem, &smem);
+    if constexpr (kNeedStaging) {
       cluster.sync();
       if (blockIdx.y != worker_rank) return;
     }
+  }
 
+  device::PDLTriggerSecondary<kPDL>();
+  if constexpr (kNeedStaging) {
     // Only the elected worker reaches here, and it mapped `topk_indices` to
     // itself, so `problem.out` is this block's own buffer. Stating that keeps the
     // shared::cluster address out of the load problem_transform issues -- which is
     // load-bearing, not an optimization: without it cicc segfaults on CUDA 13.1+
     // for sm_90a (issue #32830, previously worked around by copying `problem` in
     // #32910). Verified: dropping this line reproduces the crash on 13.1/13.2/13.3.
-  __builtin_assume(problem.out == topk_indices);
+    __builtin_assume(problem.out == s_topk_indices);
     problem_transform(problem, params.get_output_ptr(blockIdx.x));
+  }
 }
 
 // --- Plan: choose cluster_threshold from the seq_len distribution -----------
@@ -360,11 +381,10 @@ struct TopKKernel {
   static void transform(
       const tvm::ffi::TensorView scores,
       const tvm::ffi::TensorView seq_lens,
-      const tvm::ffi::TensorView page_table,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> page_table,
       const tvm::ffi::TensorView page_indices,
       const uint32_t page_size,
-      const tvm::ffi::TensorView metadata,
+      const tvm::ffi::TensorView metadata) {
-      const tvm::ffi::Optional<tvm::ffi::TensorView> raw_indices) {
     using namespace host;
     auto B = SymbolicSize{"batch_size"};
     auto Bp1 = SymbolicSize{"batch_size_plus_1"};
@@ -384,11 +404,19 @@ struct TopKKernel {
         .with_dtype<int32_t>()
         .with_device(device_)
         .verify(seq_lens);
+    // Absent means "no page transform": `page_indices` then receives the raw
+    // selected indices and nothing dereferences a page table.
+    const int32_t* page_table_ptr = nullptr;
+    int64_t page_table_stride = 0;
+    if (page_table.has_value()) {
       TensorMatcher({B, -1})  // page_table
           .with_strides({P, 1})
           .with_dtype<int32_t>()
           .with_device(device_)
-        .verify(page_table);
+          .verify(page_table.value());
+      page_table_ptr = static_cast<const int32_t*>(page_table.value().data_ptr());
+      page_table_stride = P.unwrap();
+    }
     TensorMatcher({B, K})  // page_indices
         .with_dtype<int32_t>()
         .with_device(device_)
@@ -398,12 +426,6 @@ struct TopKKernel {
         .with_device(device_)
         .verify(metadata);
 
-    int32_t* raw_indices_ptr = nullptr;
-    if (raw_indices.has_value()) {
-      TensorMatcher({B, K}).with_dtype<int32_t>().with_device(device_).verify(raw_indices.value());
-      raw_indices_ptr = static_cast<int32_t*>(raw_indices.value().data_ptr());
-    }
-
     RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2");
     RuntimeCheck(S.unwrap() % 4 == 0, "score_stride must be a multiple of 4 (16-byte vectorized load)");
     RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
@@ -425,12 +447,11 @@ struct TopKKernel {
     const auto params = TopKLaunchParams{
         .scores = static_cast<const float*>(scores.data_ptr()),
         .seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
-        .page_table = static_cast<const int32_t*>(page_table.data_ptr()),
+        .page_table = page_table_ptr,
         .page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
-        .raw_indices = raw_indices_ptr,
         .metadata = static_cast<const PlanItem*>(metadata.data_ptr()),
         .score_stride = S.unwrap(),
-        .page_table_stride = P.unwrap(),
+        .page_table_stride = page_table_stride,
         .topk = topk,
         .page_bits = page_bits,
         .cluster_floor = (batch_size <= kSmallBatchLowFloor) ? kClusterFloorSmall : kClusterFloor,
@@ -438,11 +459,21 @@ struct TopKKernel {
 
     const bool use_cluster = (max_seq_len > params.cluster_floor) && (batch_size <= kClusterMaxBatch);
     constexpr bool kUsePDL = true;
+    const auto mode = page_table.has_value() ? TopKMode::PAGE_TABLE : TopKMode::INDICES;
+    const auto dispatch = [&]<typename F>(F&& f) {
+      switch (mode) {
+        case TopKMode::INDICES:
+          return f.template operator()<TopKMode::INDICES>();
+        default:
+          return f.template operator()<TopKMode::PAGE_TABLE>();
+      }
+    };
+    dispatch([&]<TopKMode kMode>() {
       if (use_cluster) {
         if (batch_size <= kNumPersistentClusters) {
           LaunchKernel({batch_size, kClusterSize}, kBlockSize, device)
               .config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
-            .launch(topk_small_batch_kernel<kUsePDL>, params);
+              .launch(topk_small_batch_kernel<kUsePDL, kMode>, params);
         } else {
           const uint32_t num_clusters = std::min(batch_size, kNumPersistentClusters);
           LaunchKernel({num_clusters, kClusterSize}, kBlockSize, device)
@@ -450,21 +481,22 @@ struct TopKKernel {
               .launch(topk_persistent_cluster_kernel<kUsePDL>, params);
           LaunchKernel(batch_size, kBlockSize, device)
               .config({.use_pdl = kUsePDL})
-            .launch(topk_main_kernel<kUsePDL, /*kLevel=*/3>, params);
+              .launch(topk_main_kernel<kUsePDL, /*kLevel=*/3, kMode>, params);
         }
       } else if (max_seq_len <= kReg2MaxSeqLen) {
         LaunchKernel(batch_size, kBlockSize, device)
             .config({.use_pdl = kUsePDL})
-          .launch(topk_main_kernel<kUsePDL, /*kLevel=*/0>, params);
+            .launch(topk_main_kernel<kUsePDL, /*kLevel=*/0, kMode>, params);
       } else if (max_seq_len <= kReg4MaxSeqLen) {
         LaunchKernel(batch_size, kBlockSize, device)
             .config({.use_pdl = kUsePDL})
-          .launch(topk_main_kernel<kUsePDL, /*kLevel=*/1>, params);
+            .launch(topk_main_kernel<kUsePDL, /*kLevel=*/1, kMode>, params);
       } else {
         LaunchKernel(batch_size, kBlockSize, device)
             .config({.use_pdl = kUsePDL})
-          .launch(topk_main_kernel<kUsePDL, /*kLevel=*/2>, params);
+            .launch(topk_main_kernel<kUsePDL, /*kLevel=*/2, kMode>, params);
       }
+    });
   }
 };
 

python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh

@@ -168,25 +168,19 @@ SGL_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, uint3
 
 /// One batch element's worth of work. `emit(pos, raw_idx)` writes the selected raw
 /// index to output slot `pos`; `transform_output` then applies the page-table
-/// transform in a separate pass (and records the raw index in `raw_out` if set).
+/// transform in a separate pass.
 struct TopKProblem {
   const float* __restrict__ in;
   int32_t* __restrict__ out;  // page_indices [topk]
-  int32_t* __restrict__ raw_out;  // optional raw (pre-transform) indices [topk]; nullptr if unused
   const int32_t* __restrict__ page_table;
   uint32_t topk;
   uint32_t seq_len;
   uint32_t page_bits;
 
-  // Write the raw selected index; the page-table transform is applied afterwards
-  // by transform_output() in a separate, pipelined pass. Keeping the per-element
-  // page_table gather off the atomic-serialized scatter loop is measurably faster
-  // for both short and long context.
   SGL_DEVICE void emit(uint32_t pos, uint32_t raw_idx) const {
     out[pos] = static_cast<int32_t>(raw_idx);
   }
   SGL_DEVICE void transform_output(uint32_t t, int32_t raw) const {
-    if (raw_out != nullptr) raw_out[t] = raw;
     out[t] = raw < 0 ? -1 : page_to_indices(page_table, raw, page_bits);
   }
 };

python/sglang/kernels/ops/attention/dsv4/topk.py

@@ -89,13 +89,21 @@ def plan_topk_v2(seq_lens: torch.Tensor, static_threshold: int = 0) -> torch.Ten
 def topk_transform_512_v2(
     scores: torch.Tensor,
     seq_lens: torch.Tensor,
-    page_tables: torch.Tensor,
+    page_tables: Optional[torch.Tensor],
     out_page_indices: torch.Tensor,
     page_size: int,
     metadata: torch.Tensor,
-    out_raw_indices: Optional[torch.Tensor] = None,
 ) -> None:
-    """Fused top-k + page-table transform (DeepSeek-V4 top-k v2 kernel).
+    """Fused top-k + optional page-table transform (DeepSeek-V4 top-k v2 kernel).
+
+    Two output modes, chosen by whether ``page_tables`` is given and resolved to
+    a device-side template parameter, so an unused page-table gather is compiled
+    out rather than skipped at runtime:
+
+    * ``page_tables=None`` -- ``out_page_indices`` receives the raw selected
+      indices and no page table is read.
+    * ``page_tables`` given -- ``out_page_indices`` receives the page-table
+      transform of them.
 
     IMPORTANT: every entry of ``seq_lens`` must be NON-NEGATIVE, and
     ``metadata`` must come from :func:`plan_topk_v2` over the same ``seq_lens``
@@ -114,5 +122,4 @@ def topk_transform_512_v2(
         out_page_indices,
         page_size,
         metadata,
-        out_raw_indices,
     )

python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py

@@ -289,8 +289,6 @@ def _topk_transform_v2_paged(
     assert 0 < topk <= 2048, f"v2 top-k supports 0 < topk <= 2048, got {topk=}"
 
     page_table = attn_metadata.real_page_table
-    assert page_table.dtype == torch.int32
-    lengths_i32 = lengths.to(torch.int32)
 
     # The plan is preprocessed once per forward (DSAMetadata.topk_v2_plan,
     # refreshed in-place under CUDA graph) and reused across layers. A missing or
@@ -302,8 +300,8 @@ def _topk_transform_v2_paged(
     ), "topk_v2_plan must be preprocessed per forward (see DSAMetadata.topk_v2_plan)"
 
     page_size = attn_metadata.page_size
-    out = logits.new_full((num_rows, topk), -1, dtype=torch.int32)
+    out = logits.new_empty((num_rows, topk), dtype=torch.int32)
-    topk_transform_512_v2(logits, lengths_i32, page_table, out, page_size, plan)
+    topk_transform_512_v2(logits, lengths, page_table, out, page_size, plan)
     return out
 
 

test/registered/kernels/ops/attention/test_topk_v2.py

@@ -141,28 +141,41 @@ def _reference(scores, seq_lens, k):
     return ref
 
 
+def _plan(seq_lens):
+    """Plan, then break stream adjacency with the transform launch.
+
+    The transform kernel prefetches the plan metadata BEFORE its PDL wait, which
+    is only legal while the plan kernel is not the immediately preceding kernel
+    in the stream -- in production the plan is built during per-forward metadata
+    prep, a whole model forward earlier. Launching the two back to back would
+    read the plan through a programmatic dependency that guarantees no memory
+    visibility, so mirror the production separation instead.
+    """
+    metadata = plan_topk_v2(seq_lens)
+    torch.cuda.synchronize()
+    return metadata
+
+
 def _run(scores, seq_lens, page_table, inv_cpu, k):
     batch = scores.shape[0]
+    metadata = _plan(seq_lens)
     out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
-    metadata = plan_topk_v2(seq_lens)
     topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata)
     torch.cuda.synchronize()
     out_cpu = out.cpu().tolist()
     return [_invert(out_cpu[i], inv_cpu[i]) for i in range(batch)]
 
 
-def _run_raw(scores, seq_lens, page_table, k):
+def _run_raw(scores, seq_lens, k):
-    """Run the kernel and return its optional raw (pre-transform) top-k index
+    """Run with no page table and return the selected indices per row, dropping
-    output per row, dropping -1 padding -- the selected positions themselves,
+    -1 padding -- the selected positions themselves, NOT a page transform."""
-    NOT the page-table transform of them."""
     batch = scores.shape[0]
+    metadata = _plan(seq_lens)
     out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
-    raw = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
+    topk_transform_512_v2(scores, seq_lens, None, out, PAGE_SIZE, metadata)
-    metadata = plan_topk_v2(seq_lens)
-    topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata, raw)
     torch.cuda.synchronize()
-    raw_cpu = raw.cpu().tolist()
+    out_cpu = out.cpu().tolist()
-    return [[v for v in raw_cpu[i] if v != -1] for i in range(batch)]
+    return [[v for v in out_cpu[i] if v != -1] for i in range(batch)]
 
 
 @pytest.mark.parametrize("page_mode", ["identity", "perm"])
@@ -229,68 +242,25 @@ def test_topk_v2_ragged(batch: int, shape: str, k: int, per_row_pt: bool) -> Non
     _assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, lengths.cpu(), k)
 
 
-@pytest.mark.parametrize("page_mode", ["identity", "perm"])
-@pytest.mark.parametrize(
-    "batch,seq",
-    [
-        (8, 256),  # trivial
-        (8, 4096),  # register
-        (4, 131072),  # fused small-batch cluster
-        (64, 131072),  # persistent cluster + main<3> epilogue
-        (256, 131072),  # non-cluster streaming
-    ],
-)
-@torch.inference_mode()
-def test_topk_v2_raw_indices(batch: int, seq: int, page_mode: str) -> None:
-    """The optional raw-index output must be the pre-transform position of each
-    transformed output slot (out[j] == page_to_indices(raw[j])), and -1 aligns."""
-    k = 512
-    torch.manual_seed(batch * 131 + seq)
-    device = "cuda"
-    width = (seq + 3) & ~3
-    scores = torch.randn(batch, width, dtype=torch.float32, device=device)[:, :seq]
-    seq_lens = torch.full((batch,), seq, dtype=torch.int32, device=device)
-    num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE
-    page_table, inv_cpu = _make_page_table(batch, num_pages, page_mode, device)
-    out = torch.full((batch, k), -1, dtype=torch.int32, device=device)
-    raw = torch.full((batch, k), -1, dtype=torch.int32, device=device)
-
-    metadata = plan_topk_v2(seq_lens)
-    topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata, raw)
-    torch.cuda.synchronize()
-
-    out_cpu, raw_cpu = out.cpu().tolist(), raw.cpu().tolist()
-    for i in range(batch):
-        for j in range(k):
-            o, r = out_cpu[i][j], raw_cpu[i][j]
-            if o == -1:
-                assert r == -1, f"b={i} j={j}: out=-1 but raw={r}"
-            else:
-                inv = (int(inv_cpu[i][o >> PAGE_BITS]) << PAGE_BITS) | (o & PAGE_MASK)
-                assert r == inv, f"b={i} j={j}: raw={r} != inverse(out)={inv}"
-
-
 @pytest.mark.parametrize("k", [512, 1024, 2048])
 @pytest.mark.parametrize("batch,seq", FIXED_CONFIGS)
 @torch.inference_mode()
 def test_topk_v2_output_indices(batch: int, seq: int, k: int) -> None:
-    """Validate the raw (pre-transform) index output DIRECTLY against torch.topk.
+    """Validate the selected indices DIRECTLY against torch.topk.
 
-    Unlike ``test_topk_v2`` -- which checks the page-transformed output and inverts
+    Runs the no-page-table mode, so the output is the selected positions
-    it through the page table -- this exercises the selected indices themselves, so
+    themselves. Unlike ``test_topk_v2`` -- which checks the page-transformed
-    it isolates the top-k selection from the page-table transform. A permuted page
+    output and inverts it through the page table -- this isolates the top-k
-    table is used so raw != out, catching any bug that leaks transformed page
+    selection from the transform, and it is the only coverage of that mode.
-    indices into the raw buffer. Covers every dispatch template/boundary.
+    Covers every dispatch template/boundary.
     """
     torch.manual_seed(batch * 100003 + seq * 7 + k + 1)
     device = "cuda"
     width = (seq + 3) & ~3
     scores = torch.randn(batch, width, dtype=torch.float32, device=device)[:, :seq]
     seq_lens = torch.full((batch,), seq, dtype=torch.int32, device=device)
-    num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE
-    page_table, _ = _make_page_table(batch, num_pages, "perm", device)
 
-    our_raw = _run_raw(scores, seq_lens, page_table, k)
+    our_raw = _run_raw(scores, seq_lens, k)
     ref_raw = _reference(scores, seq_lens, k)
     _assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, seq_lens.cpu(), k)