feat(top_k_per_row): expose parametric k via runtime arg

[valarLip]

Summary

The radix backend (standalone_stable_radix_11bits) supports arbitrary k at runtime, but the public C++/Python wrappers top_k_per_row_prefill and top_k_per_row_decode hardcoded kTopK = 2048. This PR plumbs k through end-to-end so callers (e.g. DeepSeek-V4 indexer with k = index_topk = 1024) can use the same kernel without forcing 2048.

Backward compatible: the default k = 2048 matches previous behavior; existing callers do not need changes.

Changes

C++ surface (header / pybind / kernel)

• csrc/include/topk_per_row.h — both wrapper signatures get int64_t k = 2048 default.
• csrc/include/rocm_ops.hpp — pybind exposes py::arg(""k"") = 2048 for both ops.
• csrc/kernels/topk_per_row_kernels.cu
  • invokeComputeTopkLastDimWorkspaceSize<T, Phase>(numRows, stride0) gains int k_param = 2048; the internal constexpr int k = 2048 becomes runtime int k = k_param and is threaded into the radix call.
  • Explicit <float> instantiation updated.
  • top_k_per_row_prefill / top_k_per_row_decode bodies replace static constexpr int kTopK = 2048 with static_cast<int>(k) and pass it into both the workspace-size query and standalone_stable_radix_11bits.

Python surface

• aiter/ops/topk.py — top_k_per_row_prefill and top_k_per_row_decode (the radix-backed variants) gain k: int = 2048 so the @compile_ops-generated torch op schema accepts the new kwarg.
• The _fast ASM-kernel variants (top_k_per_row_prefill_fast, top_k_per_row_decode_fast) are intentionally not modified — their precompiled .co blobs hardcode k=2048 and cannot honor a runtime k.

Tests

• op_tests/test_topk_per_row.py
  • Default --top_k widened from [2048] to [512, 1024, 2048].
  • run_top_k_per_row_{prefill,decode} thread k through to the kernel call.
  • _fast decode is skipped when k != 2048; the dispatcher also asserts k == 2048 for _fast to surface misuse loudly.

Local validation

Ran python op_tests/test_topk_per_row.py on MI355X (gfx950) with the new defaults:

Prefill (9 cases, k ∈ {512, 1024, 2048} × ctx ∈ {1024, 4096, 16384}): all all_close = True.

Decode (18 cases, bs ∈ {4, 16} × ctx ∈ {1024, 4096, 16384} × k ∈ {512, 1024, 2048}, next_n=1): all all_close = True.

Sample decode results:

batch_size	context_len	top_k	all_close	us
4	1024	512	True	7.50
4	1024	1024	True	2.11
4	1024	2048	True	2.12
16	4096	512	True	11.38
16	4096	1024	True	10.44
16	4096	2048	True	10.26
16	16384	512	True	17.28
16	16384	1024	True	18.07
16	16384	2048	True	19.33

Test plan

• Unit test: python op_tests/test_topk_per_row.py parity vs torch.topk for k ∈ {512, 1024, 2048}.
• Downstream consumer (DeepSeek-V4 indexer in ATOM): GSM8K-100 num_fewshot=3 = 0.97 / 0.97 with k=1024 prefill + decode paths.
• CI: confirm the matrix in op_tests/test_topk_per_row.py runs on the target gfx variants.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests
ci:atom	ATOM benchmark (DeepSeek-R1 + GPT-OSS)
ci:vllm	vLLM benchmark
ci:all	All of the above

Add labels via the sidebar or gh pr edit 3012 --add-label <label>

[Copilot]

Pull request overview

This PR exposes a runtime-configurable k for the radix-backed top_k_per_row_prefill / top_k_per_row_decode ops (previously hardcoded to 2048), wiring the parameter through C++ wrappers, pybind, Python op schema, and the operator test script.

Changes:

• Add k (default 2048) to the C++ public APIs and pybind bindings for top_k_per_row_{prefill,decode}.
• Thread runtime k into the radix workspace sizing and kernel invocation in topk_per_row_kernels.cu.
• Extend op_tests/test_topk_per_row.py to test multiple k values and guard _fast decode to k==2048.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
csrc/include/topk_per_row.h	Adds defaulted k parameter to the public C++ function declarations.
csrc/include/rocm_ops.hpp	Exposes k to Python via pybind with default 2048.
csrc/kernels/topk_per_row_kernels.cu	Threads runtime k into workspace sizing + radix top-k calls.
aiter/ops/topk.py	Updates @compile_ops function signatures to accept k with default 2048.
op_tests/test_topk_per_row.py	Passes k through to ops; expands default --top_k list; skips _fast decode when k!=2048.

Comments suppressed due to low confidence (1)

csrc/kernels/topk_per_row_kernels.cu:2444

• invokeComputeTopkLastDimWorkspaceSize signature changed to include k_param. There are other translation units (e.g. csrc/kernels/topk_plain_kernels.cu) that forward-declare/extern template this function with the old 2-arg signature and call it without passing k, which will cause an unresolved symbol or incorrect workspace sizing once this lands. Update those declarations/calls to match the new signature (and pass the runtime k where appropriate).

template <typename T, aiter::Phase phase = aiter::Phase::Prefill>
int64_t invokeComputeTopkLastDimWorkspaceSize(int32_t numRows,
                                              int32_t stride0,
                                              int k_param = 2048)
{
    using IdxT = int32_t;

    size_t buf_size = 0;

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.