[ROCm] Make DSA indexer top-k exact with cooperative selection (#37591)

3rdparty/amd/wheel/sgl-kernel/rocm_hipify.py

@@ -18,7 +18,7 @@ sources = [
     "csrc/common_extension_rocm.cc",
     "csrc/elementwise/activation.cu",
     "csrc/elementwise/pos_enc.cu",
-    "csrc/elementwise/topk.cu",
+    # topk.hip is maintained as native HIP instead of being generated from topk.cu.
     "csrc/grammar/apply_token_bitmask_inplace_cuda.cu",
     "csrc/kvcacheio/transfer.cu",
     "csrc/memory/weak_ref_tensor.cpp",

python/sglang/kernels/aot/csrc/elementwise/topk.hip

@@ -0,0 +1,325 @@
+/* Copyright 2025 SGLang Team. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// ROCm implementation of the DSA indexer top-k ops. Same three entry points, same
+// semantics and same tensor contracts as csrc/elementwise/topk.cu, which the ROCm
+// build takes this file instead of; see include/hip/dsa_topk_coop.cuh for what the
+// kernel does differently and why.
+
+#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
+#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
+#include <c10/util/Exception.h>
+#include <hip/hip_runtime.h>
+#include <torch/all.h>
+
+#include <cstdint>
+#include <cstdlib>
+#include <optional>
+
+#include "hip/dsa_topk_coop.cuh"
+
+namespace {
+
+using namespace sgl::dsa_topk;
+
+// The k every caller of these ops uses (DeepSeek-V3.2 and GLM-5.2 index_topk). It is
+// a template parameter so the emit bounds and the -1 padding stay compile-time.
+constexpr uint32_t kTopK = 2048;
+
+// 12-bit bins (16 KB) plus 4096 ties (32 KB) is 48 KB of LDS, inside the 64 KB static
+// ceiling that gfx942 imposes. The bin count is what keeps the tie set small enough to
+// fit: see the precision discussion in dsa_topk_coop.cuh.
+constexpr uint32_t kHistBits = 12;
+constexpr uint32_t kTieCap = 4096;
+constexpr uint32_t kBlock = 1024;
+
+using Ws = CoopMbWorkspace<kHistBits, kTieCap>;
+
+// Blocks per row, 0 meaning "use the one-block kernel".
+//
+// Splitting a row buys two things, both of which only matter when one block per row
+// leaves the device idle: it divides the LDS histogram contention that clustered
+// logits create, and it gives a row more than one block's worth of memory parallelism.
+// It costs three extra launches and a second read of the row in scatter, so once the
+// batch alone fills the machine it is a straight loss.
+//
+// The two gates below are both needed because the terms move independently. The cost
+// is close to fixed, while the benefit grows with row length, since that is what sets
+// how much contention there is to divide. Measured on MI355X under graph replay (us,
+// split against one-block, clustered / diffuse scores, page_size=1 transform):
+//   len  32768  b1   57.7/30.7 vs  57.8/26.1     b64   63.4/33.5 vs  58.2/26.5
+//   len  49152  b1   68.9/31.0 vs  77.2/31.6     b64   78.2/36.2 vs  78.5/32.3
+//   len  65536  b1   81.5/31.1 vs  98.3/37.2     b64   95.1/41.3 vs  99.7/39.7
+//   len 100500  b1  102.6/32.1 vs 135.5/46.6     b64  123.7/49.5 vs 143.0/54.7
+// Below 64K the win turns into a loss, so short rows stay on the one-block kernel.
+// From batch 128 up one block per row already fills the machine, hence requiring at
+// least four blocks per row rather than merely more than one; on captured GLM-5.2
+// decode logits at 20 rows the split is worth 2.1-2.8x over one-block.
+//
+// `row_len_hint` is the score buffer's row stride, not the true sequence length, which
+// lives on the device. That is deliberate: the decision has to be identical on capture
+// and on replay for a graph-captured decode, and the stride is the only length-like
+// quantity the host can see. The cost of that is a wide buffer holding short rows,
+// which is the 32768/b1 line above -- 18% on a case the kernel already wins by 2x.
+int row_split(int batch, int64_t row_len_hint) {
+  // Escape hatch for benchmarking the two paths against each other; 0 forces the
+  // one-block kernel. An empty value means unset, not zero.
+  if (const char* e = std::getenv("SGL_DSA_TOPK_ROW_SPLIT"); e != nullptr && *e != '\0') {
+    const int v = std::atoi(e);
+    if (v >= 0) {
+      return v;
+    }
+  }
+  int dev = 0;
+  (void)hipGetDevice(&dev);
+  static thread_local int cached_dev = -1;
+  static thread_local int cached_cu = 0;
+  if (cached_dev != dev) {
+    hipDeviceProp_t p{};
+    (void)hipGetDeviceProperties(&p, dev);
+    cached_dev = dev;
+    cached_cu = p.multiProcessorCount;
+  }
+
+  constexpr int64_t kMinRowLen = 65536;
+  constexpr int kMinSplit = 4;
+  // Holds the slice at >= 2048 elements given kMinRowLen, below which a block cannot
+  // amortise the global atomics of its own histogram flush.
+  constexpr int kMaxSplit = 32;
+
+  if (row_len_hint < kMinRowLen) {
+    return 0;
+  }
+  const int g = std::min(cached_cu / std::max(batch, 1), kMaxSplit);
+  return g < kMinSplit ? 0 : g;
+}
+
+void launch(
+    const float* input,
+    int32_t* out_idx,
+    const int32_t* row_starts,
+    const int32_t* lengths,
+    const OutMap& map,
+    int batch,
+    int64_t stride,
+    at::Device device,
+    hipStream_t stream) {
+  const int g = row_split(batch, stride);
+
+  if (g == 0) {
+    CoopParams<kTopK> p{};
+    p.input = input;
+    p.out_idx = out_idx;
+    p.row_starts = row_starts;
+    p.lengths = lengths;
+    p.map = map;
+    p.stride = stride;
+    coop_topk_kernel<kTopK, kHistBits, kTieCap, kBlock><<<batch, kBlock, 0, stream>>>(p);
+    return;
+  }
+
+  // The row-split path's scratch. Taken from the caching allocator rather than a
+  // process-level buffer grown with hipMalloc: the address of such a buffer is baked
+  // into every HIP graph captured while it was current, both as a kernel argument and
+  // as the destination of the memset below, so freeing it leaves those graphs writing
+  // to address space the allocator no longer owns.
+  const size_t ws_bytes = Ws::bytes(static_cast<size_t>(batch));
+  at::Tensor ws = at::empty({static_cast<int64_t>(ws_bytes)}, at::TensorOptions().dtype(at::kByte).device(device));
+
+  CoopMbParams<kTopK> p{};
+  p.input = input;
+  p.out_idx = out_idx;
+  p.row_starts = row_starts;
+  p.lengths = lengths;
+  p.map = map;
+  p.ws = ws.data_ptr();
+  p.stride = stride;
+  p.batch = static_cast<uint32_t>(batch);
+
+  // Histograms, state and counters must start at zero; the workspace groups them into
+  // one contiguous prefix so this is a single launch regardless of batch.
+  C10_HIP_CHECK(hipMemsetAsync(p.ws, 0, Ws::zero_bytes(static_cast<size_t>(batch)), stream));
+
+  const dim3 grid(static_cast<unsigned>(batch), static_cast<unsigned>(g), 1);
+  coop_mb_hist0<kTopK, kHistBits, kTieCap, kBlock><<<grid, kBlock, 0, stream>>>(p, static_cast<uint32_t>(g));
+  coop_mb_hist1<kTopK, kHistBits, kTieCap, kBlock><<<grid, kBlock, 0, stream>>>(p, static_cast<uint32_t>(g));
+  coop_mb_scatter<kTopK, kHistBits, kTieCap, kBlock><<<grid, kBlock, 0, stream>>>(p, static_cast<uint32_t>(g));
+  coop_mb_refine<kTopK, kHistBits, kTieCap, kBlock><<<batch, kBlock, 0, stream>>>(p);
+}
+
+#define CHECK_HIP(x) TORCH_CHECK(x.is_cuda(), #x " must be a HIP tensor")
+#define CHECK_SAME_DEVICE(x, ref) TORCH_CHECK(x.device() == ref.device(), #x " must be on the same device as " #ref)
+
+struct Common {
+  const float* input;
+  const int32_t* lengths;
+  const int32_t* row_starts;
+  int64_t batch;
+  int64_t stride;
+};
+
+Common
+check_common(const at::Tensor& score, const at::Tensor& lengths, const std::optional<at::Tensor>& row_starts_opt) {
+  CHECK_HIP(score);
+  CHECK_HIP(lengths);
+  CHECK_SAME_DEVICE(lengths, score);
+  TORCH_CHECK(score.dim() == 2 && score.scalar_type() == at::kFloat, "score must be a float32 [B, L] tensor");
+  TORCH_CHECK(score.stride(1) == 1, "score must be contiguous along the last dim");
+  const auto B = score.size(0);
+  TORCH_CHECK(
+      lengths.dim() == 1 && lengths.is_contiguous() && lengths.scalar_type() == at::kInt,
+      "lengths must be a contiguous int32 [B] tensor");
+  TORCH_CHECK(lengths.size(0) == B, "lengths must have one entry per score row");
+
+  const int32_t* row_starts = nullptr;
+  if (row_starts_opt.has_value()) {
+    const auto& rs = row_starts_opt.value();
+    CHECK_HIP(rs);
+    CHECK_SAME_DEVICE(rs, score);
+    TORCH_CHECK(
+        rs.dim() == 1 && rs.is_contiguous() && rs.scalar_type() == at::kInt,
+        "row_starts must be a contiguous int32 [B] tensor");
+    TORCH_CHECK(rs.size(0) == B, "row_starts must have one entry per score row");
+    row_starts = rs.data_ptr<int32_t>();
+  }
+
+  return Common{score.data_ptr<float>(), lengths.data_ptr<int32_t>(), row_starts, B, score.stride(0)};
+}
+
+void check_output(const at::Tensor& out, int64_t B) {
+  TORCH_CHECK(
+      out.dim() == 2 && out.is_contiguous() && out.scalar_type() == at::kInt,
+      "output must be a contiguous int32 [B, topk] tensor");
+  TORCH_CHECK(out.size(0) == B, "output must have one row per score row");
+  TORCH_CHECK(out.size(1) == static_cast<int64_t>(kTopK), "these ops are instantiated for topk=", kTopK, " only");
+}
+
+}  // namespace
+
+void fast_topk_interface(
+    const at::Tensor& score, at::Tensor& indices, const at::Tensor& lengths, std::optional<at::Tensor> row_starts_opt) {
+  const auto c = check_common(score, lengths, row_starts_opt);
+  const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(score));
+  CHECK_HIP(indices);
+  CHECK_SAME_DEVICE(indices, score);
+  check_output(indices, c.batch);
+  if (c.batch == 0) {
+    return;
+  }
+
+  launch(
+      c.input,
+      indices.data_ptr<int32_t>(),
+      c.row_starts,
+      c.lengths,
+      OutMap{},
+      static_cast<int>(c.batch),
+      c.stride,
+      score.device(),
+      c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream());
+  C10_HIP_CHECK(hipGetLastError());
+}
+
+void fast_topk_transform_interface(
+    const at::Tensor& score,
+    const at::Tensor& lengths,
+    at::Tensor& dst_page_table,
+    const at::Tensor& src_page_table,
+    const at::Tensor& cu_seqlens_q,
+    std::optional<at::Tensor> row_starts_opt) {
+  const auto c = check_common(score, lengths, row_starts_opt);
+  const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(score));
+  CHECK_HIP(dst_page_table);
+  CHECK_HIP(src_page_table);
+  CHECK_HIP(cu_seqlens_q);
+  CHECK_SAME_DEVICE(dst_page_table, score);
+  CHECK_SAME_DEVICE(src_page_table, score);
+  CHECK_SAME_DEVICE(cu_seqlens_q, score);
+  check_output(dst_page_table, c.batch);
+  TORCH_CHECK(
+      src_page_table.dim() == 2 && src_page_table.stride(1) == 1 && src_page_table.scalar_type() == at::kInt,
+      "src_page_table must be an int32 [prefill_bs, num_slots] tensor, contiguous along the last dim");
+  TORCH_CHECK(
+      cu_seqlens_q.dim() == 1 && cu_seqlens_q.is_contiguous() && cu_seqlens_q.scalar_type() == at::kInt,
+      "cu_seqlens_q must be a contiguous int32 tensor");
+  const auto prefill_bs = cu_seqlens_q.size(0) - 1;
+  TORCH_CHECK(src_page_table.size(0) == prefill_bs, "src_page_table must have one row per sequence");
+  TORCH_CHECK(prefill_bs <= c.batch, "cu_seqlens_q describes more sequences than there are score rows");
+  if (c.batch == 0) {
+    return;
+  }
+
+  OutMap map{};
+  map.page_table = src_page_table.data_ptr<int32_t>();
+  map.pt_stride = src_page_table.stride(0);
+  // Decode is one query token per sequence, so the page-table row is the logits row
+  // and the per-block search over cu_seqlens_q is dead work. Extend, draft-extend and
+  // target-verify all expand several rows per sequence and need the map.
+  const bool is_decode = !row_starts_opt.has_value() && prefill_bs == c.batch;
+  if (!is_decode) {
+    map.cu_seqlens_q = cu_seqlens_q.data_ptr<int32_t>();
+    map.prefill_bs = static_cast<uint32_t>(prefill_bs);
+  }
+
+  launch(
+      c.input,
+      dst_page_table.data_ptr<int32_t>(),
+      c.row_starts,
+      c.lengths,
+      map,
+      static_cast<int>(c.batch),
+      c.stride,
+      score.device(),
+      c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream());
+  C10_HIP_CHECK(hipGetLastError());
+}
+
+void fast_topk_transform_ragged_interface(
+    const at::Tensor& score,
+    const at::Tensor& lengths,
+    at::Tensor& topk_indices_ragged,
+    const at::Tensor& topk_indices_offset,
+    std::optional<at::Tensor> row_starts_opt) {
+  const auto c = check_common(score, lengths, row_starts_opt);
+  const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(score));
+  CHECK_HIP(topk_indices_ragged);
+  CHECK_HIP(topk_indices_offset);
+  CHECK_SAME_DEVICE(topk_indices_ragged, score);
+  CHECK_SAME_DEVICE(topk_indices_offset, score);
+  check_output(topk_indices_ragged, c.batch);
+  TORCH_CHECK(
+      topk_indices_offset.dim() == 1 && topk_indices_offset.is_contiguous() &&
+          topk_indices_offset.scalar_type() == at::kInt,
+      "topk_indices_offset must be a contiguous int32 [B] tensor");
+  TORCH_CHECK(topk_indices_offset.size(0) == c.batch, "topk_indices_offset must have one entry per score row");
+  if (c.batch == 0) {
+    return;
+  }
+
+  OutMap map{};
+  map.offsets = topk_indices_offset.data_ptr<int32_t>();
+
+  launch(
+      c.input,
+      topk_indices_ragged.data_ptr<int32_t>(),
+      c.row_starts,
+      c.lengths,
+      map,
+      static_cast<int>(c.batch),
+      c.stride,
+      score.device(),
+      c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream());
+  C10_HIP_CHECK(hipGetLastError());
+}

python/sglang/kernels/aot/setup_rocm.py

@@ -48,7 +48,8 @@ sources = [
     "csrc/elementwise/activation.cu",
     "csrc/elementwise/deepseek_v4_topk.cu",
     "csrc/elementwise/dsv4_norm_rope.cu",
-    "csrc/elementwise/topk.cu",
+    # Native HIP implementation of the same three ops exposed by topk.cu.
+    "csrc/elementwise/topk.hip",
     "csrc/grammar/apply_token_bitmask_inplace_cuda.cu",
     "csrc/moe/moe_align_kernel.cu",
     "csrc/moe/moe_topk_softmax_kernels.cu",

python/sglang/kernels/aot/tests/test_topk.py

@@ -292,5 +292,167 @@ def test_deepseek_v4_topk_transform(bs: int, c4_len: int) -> None:
     )
 
 
+def _make_scores(kind: str, bs: int, width: int, seed: int) -> torch.Tensor:
+    """Score distributions that stress the coarse stage of a histogram top-k.
+
+    Everything above uses ``torch.randn``, which spreads over enough exponents that
+    even a narrow coarse key separates it -- 127 populated buckets out of 256 on an
+    8-bit key. Real DSA indexer logits are far more concentrated than that. Captured
+    from a GLM-5.2 decode at 134,849 tokens of context, six consecutive indexer calls
+    populate 4 to 126 buckets, with the largest holding 6% to 88% of the row.
+
+    The three below bracket that regime, and each is calibrated to what it provokes on
+    an 8-bit fp16 coarse key: `banded` populates 4 buckets with 47% in the largest,
+    `narrow` collapses to a single bucket, and `subnormal` to two.
+    """
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    if kind == "diffuse":
+        return torch.randn(bs, width, generator=g, device="cuda", dtype=torch.float32)
+    if kind == "banded":
+        # The value range of the worst real capture, [54, 88].
+        return 54.0 + 34.0 * torch.rand(
+            bs, width, generator=g, device="cuda", dtype=torch.float32
+        )
+    if kind == "narrow":
+        # A row whose spread is small next to its magnitude, which is what makes a
+        # truncating coarse key run out of buckets.
+        return 70.0 + torch.randn(
+            bs, width, generator=g, device="cuda", dtype=torch.float32
+        )
+    if kind == "subnormal":
+        # Same shape as diffuse, scaled below fp16's smallest normal: a coarse key that
+        # rounds through fp16 cannot separate this row at all, an fp32 one is unaffected.
+        return 1e-16 * torch.randn(
+            bs, width, generator=g, device="cuda", dtype=torch.float32
+        )
+    raise ValueError(kind)
+
+
+def assert_exact(
+    score: torch.Tensor, indices: torch.Tensor, seq_len: int, k: int
+) -> None:
+    """The selected scores must be the top-k scores, as a multiset.
+
+    Stricter than ``assert_equal`` on purpose. Comparing index sets has to forgive
+    tie-breaking, and that forgiveness is what lets a kernel selecting from a silently
+    truncated candidate set pass: the indices it returns are all in range and all
+    distinct, they are simply not the largest.
+    """
+    for i in range(score.shape[0]):
+        want = torch.sort(
+            torch.topk(score[i, :seq_len], k).values, descending=True
+        ).values
+        got = torch.sort(score[i, :seq_len][indices[i].long()], descending=True).values
+        assert torch.equal(got, want), (
+            f"row {i}: {int((got != want).sum())}/{k} selected scores are not the top-{k}"
+        )
+
+
+def assert_exact_rows(
+    score: torch.Tensor,
+    indices: torch.Tensor,
+    lengths: torch.Tensor,
+    k: int,
+    row_starts: Optional[torch.Tensor] = None,
+) -> None:
+    """Exact value-multiset check for variable windows and relative indices."""
+    for i in range(score.shape[0]):
+        length = int(lengths[i])
+        start = 0 if row_starts is None else int(row_starts[i])
+        row_indices = indices[i].long()
+        assert torch.all((row_indices >= 0) & (row_indices < length))
+        row = score[i, start : start + length]
+        want = torch.sort(torch.topk(row, k).values).values
+        got = torch.sort(row[row_indices]).values
+        assert torch.equal(got, want), (
+            f"row {i}: {int((got != want).sum())}/{k} selected scores are not the top-{k}"
+        )
+
+
+@pytest.mark.skipif(
+    torch.version.hip is None or torch.cuda.device_count() < 2,
+    reason="requires a multi-GPU ROCm runner",
+)
+@torch.inference_mode()
+def test_topk_uses_score_device_and_rejects_mixed_devices() -> None:
+    current_device = torch.cuda.current_device()
+    score_device = (current_device + 1) % torch.cuda.device_count()
+    score = torch.randn(1, 16384, dtype=torch.float32, device=score_device)
+    lengths = torch.full((1,), 16384, dtype=torch.int32, device=score_device)
+
+    indices = fast_topk_v2(score, lengths, 2048)
+    assert_exact(score, indices, 16384, 2048)
+    assert torch.cuda.current_device() == current_device
+
+    with pytest.raises(RuntimeError, match="same device"):
+        fast_topk_v2(score, lengths.to(f"cuda:{current_device}"), 2048)
+
+
+@pytest.mark.skipif(
+    torch.version.hip is None,
+    reason="the CUDA kernel in csrc/elementwise/topk.cu shares this limitation; only "
+    "the ROCm one (csrc/elementwise/topk.hip) is exact on these distributions",
+)
+@pytest.mark.parametrize("kind", ["diffuse", "banded", "narrow", "subnormal"])
+@pytest.mark.parametrize("bs", [1, 4, 64])
+@pytest.mark.parametrize("seq_len", [16384, 65536, 100500])
+@torch.inference_mode()
+def test_topk_is_exact_for_indexer_distributions(
+    kind: str, bs: int, seq_len: int
+) -> None:
+    k = 2048
+    score = _make_scores(kind, bs, MAX_SEQ_LEN, seed=seq_len + bs)
+    lengths = torch.full((bs,), seq_len, dtype=torch.int32, device="cuda")
+    assert_exact(score, fast_topk_v2(score, lengths, k), seq_len, k)
+
+
+@pytest.mark.skipif(
+    torch.version.hip is None,
+    reason="the cooperative top-k implementation is only built on ROCm",
+)
+@pytest.mark.parametrize("kind", ["banded", "narrow", "subnormal"])
+@pytest.mark.parametrize("force_one_block", [False, True])
+@torch.inference_mode()
+def test_topk_variants_are_exact_on_variable_windows(
+    kind: str, force_one_block: bool, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Cover both dispatches and all three APIs on adversarial score distributions."""
+    if force_one_block:
+        monkeypatch.setenv("SGL_DSA_TOPK_ROW_SPLIT", "0")
+    else:
+        monkeypatch.delenv("SGL_DSA_TOPK_ROW_SPLIT", raising=False)
+
+    bs, k = 4, 2048
+    score = _make_scores(kind, bs, MAX_SEQ_LEN, seed=1200 + force_one_block)
+    lengths = torch.tensor(
+        [65536, 70000, 90000, 100500], dtype=torch.int32, device="cuda"
+    )
+    row_starts = torch.tensor([0, 127, 511, 0], dtype=torch.int32, device="cuda")
+
+    raw = fast_topk_v2(score, lengths, k, row_starts=row_starts)
+    assert_exact_rows(score, raw, lengths, k, row_starts)
+
+    ragged_offsets = torch.tensor(
+        [17, 200000, 400000, 600000], dtype=torch.int32, device="cuda"
+    )
+    ragged = fast_topk_transform_ragged_fused(
+        score, lengths, ragged_offsets, k, row_starts=row_starts
+    )
+    assert_exact_rows(score, ragged - ragged_offsets[:, None], lengths, k, row_starts)
+
+    logical = torch.arange(MAX_SEQ_LEN, dtype=torch.int32, device="cuda")
+    multipliers = torch.tensor([1, 3, 5, 7], dtype=torch.int32, device="cuda")
+    shifts = torch.tensor([19, 43, 71, 101], dtype=torch.int32, device="cuda")
+    page_table = (
+        logical[None, :] * multipliers[:, None] + shifts[:, None]
+    ) % MAX_SEQ_LEN
+    cu_seqlens_q = torch.arange(bs + 1, dtype=torch.int32, device="cuda")
+    mapped = fast_topk_transform_fused(score, lengths, page_table, cu_seqlens_q, k)
+    inverse_page_table = torch.empty_like(page_table)
+    inverse_page_table.scatter_(1, page_table.long(), logical[None, :].expand(bs, -1))
+    mapped_raw = inverse_page_table.gather(1, mapped.long())
+    assert_exact_rows(score, mapped_raw, lengths, k)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))