[DSV4] Support raw-index output in TopK v2 (#33672)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>

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

@@ -183,6 +183,20 @@ def _run_raw(scores, seq_lens, k):
     return [[v for v in out_cpu[i] if v != -1] for i in range(batch)]
 
 
+def _run_dual(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)
+    raw = torch.full_like(out, -1)
+    topk_transform_paged_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata, raw)
+    torch.cuda.synchronize()
+    out_cpu = out.cpu().tolist()
+    raw_cpu = raw.cpu().tolist()
+    transformed_raw = [_invert(out_cpu[i], inv_cpu[i]) for i in range(batch)]
+    direct_raw = [[v for v in raw_cpu[i] if v != -1] for i in range(batch)]
+    return transformed_raw, direct_raw
+
+
 @pytest.mark.parametrize("page_mode", ["identity", "perm"])
 @pytest.mark.parametrize("k", [512, 1024, 2048])
 @pytest.mark.parametrize("batch,seq", FIXED_CONFIGS)
@@ -270,6 +284,27 @@ def test_topk_v2_output_indices(batch: int, seq: int, k: int) -> None:
     _assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, seq_lens.cpu(), k)
 
 
+@pytest.mark.parametrize(
+    "batch,seq", [(8, 256), (8, 8192), (4, 32768), (2, 131072), (31, 131072)]
+)
+@torch.inference_mode()
+def test_topk_v2_dual_output(batch: int, seq: int) -> None:
+    """The dual mode returns the same selection before and after page transform."""
+    k = 512
+    torch.manual_seed(batch * 100003 + seq * 7 + k + 2)
+    device = "cuda"
+    scores = torch.randn(batch, seq, dtype=torch.float32, device=device)
+    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, "perm", device)
+
+    transformed_raw, direct_raw = _run_dual(scores, seq_lens, page_table, inv_cpu, k)
+    for row in range(batch):
+        assert sorted(transformed_raw[row]) == sorted(direct_raw[row])
+    ref_raw = _reference(scores, seq_lens, k)
+    _assert_topk_close(scores.cpu(), ref_raw, direct_raw, batch, seq_lens.cpu(), k)
+
+
 # --- ragged entry point ------------------------------------------------------
 # Rows select inside `[row_start, row_start + seq_len)` of their score row and
 # emit `position + offset`. The window start is an arbitrary token offset, so