feat(topk): Add top_k_varlen with GVR and radix backends for sparse-attention KV selection - #3901
Merged
Merged
Conversation
dhiraj113
requested review from
aleozlx,
bkryu,
cyx-6,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 9, 2026 23:05
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a unified ChangesTop-K decode
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant top_k_decode
participant GvrTopKLBPrepareKernel
participant GvrTopKLBKernel
participant GvrTopKKernel
Caller->>top_k_decode: logits, seq_lens, top_k, pre_idx
top_k_decode->>GvrTopKLBPrepareKernel: prepare request ordering
GvrTopKLBPrepareKernel-->>top_k_decode: order_row and counters
top_k_decode->>GvrTopKLBKernel: launch load-balanced decode
GvrTopKLBKernel->>GvrTopKKernel: process long and short rows
GvrTopKKernel-->>Caller: indices and optional values
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces GVR Top-K CuTe DSL kernels and APIs tailored for NVIDIA Blackwell (sm_100+) GPUs, including a load-balanced decode variant and block prefix sum kernels. The review feedback highlights critical issues with the JIT compilation caching functions (_compile_gvr, _compile_lb_prepare, and _compile_lb). Specifically, the cache keys must include the active CUDA device ID to prevent context mismatches in multi-GPU environments. Furthermore, static dimensions like batch_size, num_rows, and N should be replaced with symbolic integers (cute.sym_int()) in the cache keys to avoid severe recompilation latency spikes during dynamic batching and decoding. Finally, an optimization is suggested in block_scan.py to refactor a thread-dependent loop into a statically uniform loop to enable compiler unrolling.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+70
to
+86
| @functools.cache | ||
| def _compile_gvr( | ||
| cute_dtype, | ||
| top_k: int, | ||
| next_n: int, | ||
| enable_unroll_4: bool, | ||
| enable_phase3_unroll: bool, | ||
| use_constant_hint: bool, | ||
| min_blocks_per_mp: int, | ||
| use_256bit_load: bool, | ||
| num_threads_per_block: int, | ||
| enable_warp_parallel_reduce: bool, | ||
| compress_ratio: int, | ||
| return_output_values: bool, | ||
| cluster_size: int, | ||
| seqlen_sorted: bool, | ||
| ): |
Contributor
There was a problem hiding this comment.
The JIT compilation cache _compile_gvr is decorated with @functools.cache but does not include the active CUDA device ID in its cache key. In multi-GPU environments (such as Tensor Parallelism or pipeline parallelism), calling this function on a different GPU will hit the cache and return the compiled function handle from the first GPU's CUDA context. This will cause runtime failures like CUDA_ERROR_INVALID_HANDLE or CUDA_ERROR_INVALID_CONTEXT. Adding device_id as a parameter ensures context-safe caching across multiple GPUs.
@functools.cache
def _compile_gvr(
device_id: int,
cute_dtype,
top_k: int,
next_n: int,
enable_unroll_4: bool,
enable_phase3_unroll: bool,
use_constant_hint: bool,
min_blocks_per_mp: int,
use_256bit_load: bool,
num_threads_per_block: int,
enable_warp_parallel_reduce: bool,
compress_ratio: int,
return_output_values: bool,
cluster_size: int,
seqlen_sorted: bool,
):
Comment on lines
+366
to
+381
| compiled = _compile_gvr( | ||
| cute_dtype, | ||
| top_k, | ||
| next_n, | ||
| enable_unroll_4, | ||
| enable_phase3_unroll, | ||
| use_constant_hint, | ||
| min_blocks_per_mp, | ||
| use_256bit_load, | ||
| num_threads_per_block, | ||
| enable_warp_parallel_reduce, | ||
| compress_ratio, | ||
| return_output_values, | ||
| cluster_size, | ||
| seqlen_sorted, | ||
| ) |
Contributor
There was a problem hiding this comment.
Pass the active CUDA device ID to _compile_gvr to ensure context-safe caching across multiple GPUs.
device_id = torch.cuda.current_device()
compiled = _compile_gvr(
device_id,
cute_dtype,
top_k,
next_n,
enable_unroll_4,
enable_phase3_unroll,
use_constant_hint,
min_blocks_per_mp,
use_256bit_load,
num_threads_per_block,
enable_warp_parallel_reduce,
compress_ratio,
return_output_values,
cluster_size,
seqlen_sorted,
)
Comment on lines
+423
to
+429
| @functools.cache | ||
| def _compile_lb_prepare( | ||
| num_threads: int, | ||
| batch_size: int, | ||
| long_threshold: int, | ||
| compress_ratio: int, | ||
| ): |
Contributor
There was a problem hiding this comment.
The JIT compilation cache _compile_lb_prepare takes batch_size as a static integer in its cache key. This means any change in the runtime batch size (common in dynamic batching) will trigger a new JIT compilation, causing severe latency spikes. We should use cute.sym_int() to compile a shape-generic kernel instead. Additionally, we should include device_id in the cache key to prevent multi-GPU context mismatch issues.
Suggested change
| @functools.cache | |
| def _compile_lb_prepare( | |
| num_threads: int, | |
| batch_size: int, | |
| long_threshold: int, | |
| compress_ratio: int, | |
| ): | |
| @functools.cache | |
| def _compile_lb_prepare( | |
| device_id: int, | |
| num_threads: int, | |
| long_threshold: int, | |
| compress_ratio: int, | |
| ): |
Comment on lines
+430
to
+451
| prep = GvrTopKLBPrepareKernel( | ||
| long_threshold=long_threshold, | ||
| compress_ratio=compress_ratio, | ||
| num_threads=num_threads, | ||
| ) | ||
| fake_seq = cute.runtime.make_fake_compact_tensor( | ||
| cutlass.Int32, (batch_size,), stride_order=(0,) | ||
| ) | ||
| fake_order = cute.runtime.make_fake_compact_tensor( | ||
| cutlass.Int32, (num_threads,), stride_order=(0,) | ||
| ) | ||
| fake_ctr = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (2,), stride_order=(0,)) | ||
| fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) | ||
| return cute.compile( | ||
| prep, | ||
| fake_seq, | ||
| fake_order, | ||
| fake_ctr, | ||
| cutlass.Int32(0), | ||
| stream=fake_stream, | ||
| options="--enable-tvm-ffi", | ||
| ) |
Contributor
There was a problem hiding this comment.
Refactor the body of _compile_lb_prepare to use cute.sym_int() for the dynamic batch_size dimension.
prep = GvrTopKLBPrepareKernel(
long_threshold=long_threshold,
compress_ratio=compress_ratio,
num_threads=num_threads,
)
n_batch = cute.sym_int()
fake_seq = cute.runtime.make_fake_compact_tensor(
cutlass.Int32, (n_batch,), stride_order=(0,)
)
fake_order = cute.runtime.make_fake_compact_tensor(
cutlass.Int32, (num_threads,), stride_order=(0,)
)
fake_ctr = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (2,), stride_order=(0,))
fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
return cute.compile(
prep,
fake_seq,
fake_order,
fake_ctr,
cutlass.Int32(0),
stream=fake_stream,
options="--enable-tvm-ffi",
)| if counters is None: | ||
| counters = torch.zeros(2, dtype=torch.int32, device=seq_lens.device) | ||
|
|
||
| compiled = _compile_lb_prepare(max_batch_size, batch_size, long_threshold, compress_ratio) |
Contributor
There was a problem hiding this comment.
Pass the active CUDA device ID to _compile_lb_prepare and remove the dynamic batch_size from the cache key.
Suggested change
| compiled = _compile_lb_prepare(max_batch_size, batch_size, long_threshold, compress_ratio) | |
| device_id = torch.cuda.current_device() | |
| compiled = _compile_lb_prepare(device_id, max_batch_size, long_threshold, compress_ratio) |
Comment on lines
+454
to
+466
| @functools.cache | ||
| def _compile_lb( | ||
| cute_dtype, | ||
| top_k: int, | ||
| next_n: int, | ||
| num_rows: int, | ||
| N: int, | ||
| compress_ratio: int, | ||
| max_batch_size: int, | ||
| num_threads: int, | ||
| cluster_size: int, | ||
| return_output_values: bool, | ||
| ): |
Contributor
There was a problem hiding this comment.
The JIT compilation cache _compile_lb takes num_rows and N as static integers in its cache key. Since N (the sequence length / KV cache length) changes on every single decode step, this will trigger a full JIT compilation on every single token generation step, causing severe latency spikes. We should use cute.sym_int() to compile a shape-generic kernel instead. Additionally, we should include device_id in the cache key to prevent multi-GPU context mismatch issues.
@functools.cache
def _compile_lb(
device_id: int,
cute_dtype,
top_k: int,
next_n: int,
compress_ratio: int,
max_batch_size: int,
num_threads: int,
cluster_size: int,
return_output_values: bool,
):
Comment on lines
+467
to
+511
| kernel = GvrTopKLBKernel( | ||
| dtype=cute_dtype, | ||
| top_k=top_k, | ||
| next_n=next_n, | ||
| num_threads=num_threads, | ||
| compress_ratio=compress_ratio, | ||
| return_output_values=return_output_values, | ||
| cluster_size=cluster_size, | ||
| max_batch_size=max_batch_size, | ||
| ) | ||
| n_groups = num_rows // next_n | ||
| fake_logits = cute.runtime.make_fake_compact_tensor( | ||
| cute_dtype, (num_rows, N), stride_order=(1, 0), assumed_align=16 | ||
| ) | ||
| fake_pre_idx = cute.runtime.make_fake_compact_tensor( | ||
| cutlass.Int32, (n_groups, top_k), stride_order=(1, 0), assumed_align=16 | ||
| ) | ||
| fake_seq = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (n_groups,), stride_order=(0,)) | ||
| fake_out_v = ( | ||
| cute.runtime.make_fake_compact_tensor( | ||
| cute_dtype, (num_rows, top_k), stride_order=(1, 0), assumed_align=16 | ||
| ) | ||
| if return_output_values | ||
| else None | ||
| ) | ||
| fake_out_i = cute.runtime.make_fake_compact_tensor( | ||
| cutlass.Int32, (num_rows, top_k), stride_order=(1, 0), assumed_align=16 | ||
| ) | ||
| fake_order = cute.runtime.make_fake_compact_tensor( | ||
| cutlass.Int32, (max_batch_size,), stride_order=(0,) | ||
| ) | ||
| fake_ctr = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (2,), stride_order=(0,)) | ||
| fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) | ||
| return cute.compile( | ||
| kernel, | ||
| fake_logits, | ||
| fake_pre_idx, | ||
| fake_seq, | ||
| fake_out_v, | ||
| fake_out_i, | ||
| fake_order, | ||
| fake_ctr, | ||
| stream=fake_stream, | ||
| options="--enable-tvm-ffi", | ||
| ) |
Contributor
There was a problem hiding this comment.
Refactor the body of _compile_lb to use cute.sym_int() for the dynamic num_rows and N dimensions.
kernel = GvrTopKLBKernel(
dtype=cute_dtype,
top_k=top_k,
next_n=next_n,
num_threads=num_threads,
compress_ratio=compress_ratio,
return_output_values=return_output_values,
cluster_size=cluster_size,
max_batch_size=max_batch_size,
)
n_rows = cute.sym_int()
n_cols = cute.sym_int()
n_groups = n_rows // next_n
fake_logits = cute.runtime.make_fake_compact_tensor(
cute_dtype, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16
)
fake_pre_idx = cute.runtime.make_fake_compact_tensor(
cutlass.Int32, (n_groups, top_k), stride_order=(1, 0), assumed_align=16
)
fake_seq = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (n_groups,), stride_order=(0,))
fake_out_v = (
cute.runtime.make_fake_compact_tensor(
cute_dtype, (n_rows, top_k), stride_order=(1, 0), assumed_align=16
)
if return_output_values
else None
)
fake_out_i = cute.runtime.make_fake_compact_tensor(
cutlass.Int32, (n_rows, top_k), stride_order=(1, 0), assumed_align=16
)
fake_order = cute.runtime.make_fake_compact_tensor(
cutlass.Int32, (max_batch_size,), stride_order=(0,)
)
fake_ctr = cute.runtime.make_fake_compact_tensor(
cutlass.Int32, (2,), stride_order=(0,)
)
fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
return cute.compile(
kernel,
fake_logits,
fake_pre_idx,
fake_seq,
fake_out_v,
fake_out_i,
fake_order,
fake_ctr,
stream=fake_stream,
options="--enable-tvm-ffi",
)
Comment on lines
+660
to
+671
| compiled = _compile_lb( | ||
| cute_dtype, | ||
| top_k, | ||
| next_n, | ||
| num_rows, | ||
| N, | ||
| compress_ratio, | ||
| max_batch_size, | ||
| num_threads, | ||
| cluster_size, | ||
| return_output_values, | ||
| ) |
Contributor
There was a problem hiding this comment.
Pass the active CUDA device ID to _compile_lb and remove the dynamic num_rows and N from the cache key.
Suggested change
| compiled = _compile_lb( | |
| cute_dtype, | |
| top_k, | |
| next_n, | |
| num_rows, | |
| N, | |
| compress_ratio, | |
| max_batch_size, | |
| num_threads, | |
| cluster_size, | |
| return_output_values, | |
| ) | |
| device_id = torch.cuda.current_device() | |
| compiled = _compile_lb( | |
| device_id, | |
| cute_dtype, | |
| top_k, | |
| next_n, | |
| compress_ratio, | |
| max_batch_size, | |
| num_threads, | |
| cluster_size, | |
| return_output_values, | |
| ) |
Comment on lines
+152
to
+164
| for i in range(tidx, num_bins, num_threads_per_block): | ||
| val = input[i] | ||
| val, total_sum = block_prefix_sum_kernel( | ||
| val, | ||
| s_warp_sums, | ||
| tidx, | ||
| num_threads_per_block, | ||
| num_warps, | ||
| barrier_id=0, | ||
| need_total_sum=True, | ||
| ) | ||
| output[i] = val + previous_sum | ||
| previous_sum = previous_sum + total_sum |
Contributor
There was a problem hiding this comment.
The current loop for i in range(tidx, num_bins, num_threads_per_block) uses a thread-dependent start index tidx. Because tidx is a runtime value, this prevents the compiler from unrolling or optimizing the loop at compile-time. Since num_bins is guaranteed to be a multiple of num_threads_per_block, we can refactor this into a statically uniform loop over the number of iterations. This allows the compiler to fully unroll and optimize the loop structure.
Suggested change
| for i in range(tidx, num_bins, num_threads_per_block): | |
| val = input[i] | |
| val, total_sum = block_prefix_sum_kernel( | |
| val, | |
| s_warp_sums, | |
| tidx, | |
| num_threads_per_block, | |
| num_warps, | |
| barrier_id=0, | |
| need_total_sum=True, | |
| ) | |
| output[i] = val + previous_sum | |
| previous_sum = previous_sum + total_sum | |
| num_iters = num_bins // num_threads_per_block | |
| for step in range(num_iters): | |
| i = step * num_threads_per_block + tidx | |
| val = input[i] | |
| val, total_sum = block_prefix_sum_kernel( | |
| val, | |
| s_warp_sums, | |
| tidx, | |
| num_threads_per_block, | |
| num_warps, | |
| barrier_id=0, | |
| need_total_sum=True, | |
| ) | |
| output[i] = val + previous_sum | |
| previous_sum = previous_sum + total_sum |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
flashinfer/topk_blackwell.py (1)
173-196: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider exposing a
trace=argument on the tensor-in/tensor-out APIs.Per repo guidance,
@flashinfer_apipublic APIs should provide atrace=argument sofi_trace()can auto-generate benchmark JSON — but only when inputs and outputs are all tensors.gvr_topk_sort_prepare(tensor→tensor) andgvr_topk_lb_prepare(tensors→tensors) qualify cleanly; the decode APIs return anOptional[torch.Tensor](Nonewhenreturn_output_values=False), so per the learning they may legitimately be left withouttrace=. Please addtrace=where the I/O is fully tensor-typed.As per coding guidelines: "Every public API decorated with
@flashinfer_apishould also provide atrace=argument sofi_trace()can auto-generate benchmark-definition JSON." Based on learnings, only passtrace=when the API's inputs and outputs are tensors.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/topk_blackwell.py` around lines 173 - 196, Add a trace= argument to the `@flashinfer_api-decorated` tensor-in/tensor-out functions gvr_topk_sort_prepare and gvr_topk_lb_prepare, wiring it through according to existing tracing conventions so fi_trace() can generate benchmark JSON. Do not add trace= to gvr_topk_decode or other APIs whose outputs can be None.Sources: Coding guidelines, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/topk_blackwell.py`:
- Around line 298-312: Add an upfront assertion in the top-k entry point,
alongside the existing input validations, requiring top_k to be one of 512,
1024, or 2048 and reporting the received value. Use the function’s top_k
parameter and existing validation block so unsupported values fail before
cute.compile or kernel launch.
- Around line 562-581: Add the same can_use_gvr_topk hardware/CuTe-DSL
availability guard used by gvr_topk_decode to both gvr_topk_lb_prepare and
gvr_topk_lb_decode before referencing their kernels; raise the documented clear
RuntimeError when unsupported, preventing bare NameError failures.
- Around line 33-35: Make flashinfer/topk_blackwell.py import-safe when
nvidia-cutlass-dsl is unavailable: defer the cutlass and cutlass.cute imports
into the Blackwell-specific functions or guard them with an optional import,
while preserving the existing torch import and ensuring those code paths fail
clearly only when invoked.
---
Nitpick comments:
In `@flashinfer/topk_blackwell.py`:
- Around line 173-196: Add a trace= argument to the `@flashinfer_api-decorated`
tensor-in/tensor-out functions gvr_topk_sort_prepare and gvr_topk_lb_prepare,
wiring it through according to existing tracing conventions so fi_trace() can
generate benchmark JSON. Do not add trace= to gvr_topk_decode or other APIs
whose outputs can be None.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 2f25815e-2bce-4f2c-85b3-e2948b917921
📥 Commits
Reviewing files that changed from the base of the PR and between ce52279 and 8edd577753d9294e6272b11632cf70bce4bebe02.
📒 Files selected for processing (9)
flashinfer/__init__.pyflashinfer/cute_dsl/top_k/__init__.pyflashinfer/cute_dsl/top_k/block_scan.pyflashinfer/cute_dsl/top_k/gvr_topk_decode.pyflashinfer/cute_dsl/top_k/gvr_topk_decode_lb.pyflashinfer/cute_dsl/top_k/pdl_utils.pyflashinfer/topk_blackwell.pytests/topk/__init__.pytests/topk/test_gvr_topk.py
Comment on lines
+298
to
+312
| assert logits.is_cuda, "logits must be on CUDA" | ||
| assert logits.dim() == 2, f"logits must be 2D, got shape {logits.shape}" | ||
| assert pre_idx.dim() == 2 and pre_idx.dtype == torch.int32 | ||
| assert seq_lens.dim() == 1 and seq_lens.dtype == torch.int32 | ||
| if seqlen_sorted: | ||
| assert ( | ||
| order_row is not None | ||
| and order_row.dtype == torch.int32 | ||
| and order_row.is_cuda | ||
| and order_row.shape == seq_lens.shape | ||
| ), ( | ||
| "seqlen_sorted=True requires order_row: int32[batch_size] on CUDA " | ||
| f"(expected shape {tuple(seq_lens.shape)}, got " | ||
| f"{tuple(order_row.shape) if order_row is not None else None})" | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Validate top_k against the supported set.
The docstring states top_k must be in {512, 1024, 2048}, but there is no runtime check. An out-of-set value falls through to cute.compile/kernel launch and fails with an opaque error. A cheap up-front assert improves the failure mode.
🛡️ Suggested validation
assert seq_lens.dim() == 1 and seq_lens.dtype == torch.int32
+ assert top_k in (512, 1024, 2048), f"top_k must be one of {{512, 1024, 2048}}; got {top_k}"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| assert logits.is_cuda, "logits must be on CUDA" | |
| assert logits.dim() == 2, f"logits must be 2D, got shape {logits.shape}" | |
| assert pre_idx.dim() == 2 and pre_idx.dtype == torch.int32 | |
| assert seq_lens.dim() == 1 and seq_lens.dtype == torch.int32 | |
| if seqlen_sorted: | |
| assert ( | |
| order_row is not None | |
| and order_row.dtype == torch.int32 | |
| and order_row.is_cuda | |
| and order_row.shape == seq_lens.shape | |
| ), ( | |
| "seqlen_sorted=True requires order_row: int32[batch_size] on CUDA " | |
| f"(expected shape {tuple(seq_lens.shape)}, got " | |
| f"{tuple(order_row.shape) if order_row is not None else None})" | |
| ) | |
| assert logits.is_cuda, "logits must be on CUDA" | |
| assert logits.dim() == 2, f"logits must be 2D, got shape {logits.shape}" | |
| assert pre_idx.dim() == 2 and pre_idx.dtype == torch.int32 | |
| assert seq_lens.dim() == 1 and seq_lens.dtype == torch.int32 | |
| assert top_k in (512, 1024, 2048), f"top_k must be one of {{512, 1024, 2048}}; got {top_k}" | |
| if seqlen_sorted: | |
| assert ( | |
| order_row is not None | |
| and order_row.dtype == torch.int32 | |
| and order_row.is_cuda | |
| and order_row.shape == seq_lens.shape | |
| ), ( | |
| "seqlen_sorted=True requires order_row: int32[batch_size] on CUDA " | |
| f"(expected shape {tuple(seq_lens.shape)}, got " | |
| f"{tuple(order_row.shape) if order_row is not None else None})" | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/topk_blackwell.py` around lines 298 - 312, Add an upfront
assertion in the top-k entry point, alongside the existing input validations,
requiring top_k to be one of 512, 1024, or 2048 and reporting the received
value. Use the function’s top_k parameter and existing validation block so
unsupported values fail before cute.compile or kernel launch.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (1)
tests/topk/test_gvr_topk.py (1)
259-259: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: replace ambiguous
×(U+00D7) withx.Ruff flags the multiplication sign in these docstrings/strings (also lines 23 and 28) as ambiguous (RUF001/RUF002).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/topk/test_gvr_topk.py` at line 259, Replace the ambiguous multiplication sign “×” with the ASCII letter “x” in the affected docstrings and strings, including the stress-test description and the occurrences near lines 23 and 28, to satisfy Ruff RUF001/RUF002.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@benchmarks/routines/sampling.py`:
- Around line 2170-2290: Remove the later duplicate testTopKDecode function
definition, keeping the earlier implementation as the sole definition so it is
not shadowed or allowed to diverge.
- Around line 1952-1958: Guard the gvr branch in run_backend so it is only
selected when top_k is one of the supported values 512, 1024, or 2048; otherwise
skip gvr or raise a clear validation error before dispatch. Also update
testTopKDecode’s top_k help text or validation to document these gvr
requirements, including the legacy cuda expansion to radix and gvr.
In `@flashinfer/cute_dsl/top_k/config.py`:
- Line 46: Run ruff format on flashinfer/cute_dsl/top_k/config.py and commit the
resulting formatting changes; update the docstrings in the affected
configuration definitions to replace the Unicode multiplication sign “×” with
ASCII “x” so RUF002 passes.
- Around line 25-26: Remove the unused field and Optional imports from the
config module’s import section, retaining only imports referenced by the file to
satisfy ruff F401 checks.
- Line 118: Fix the Ruff SIM300 violation in the thread-count selection
expression by rewriting the comparison in standard operand order, changing the
`num_rows <= num_sms` check within the `num_threads_per_block` assignment to the
equivalent non-Yoda form while preserving the existing logic.
In `@flashinfer/topk_blackwell.py`:
- Around line 267-268: Add a TraceTemplate for the public top_k_decode function
and pass it via the trace= argument to its `@flashinfer_api` decorator, matching
the trace template pattern used by related APIs in topk_blackwell.py.
---
Nitpick comments:
In `@tests/topk/test_gvr_topk.py`:
- Line 259: Replace the ambiguous multiplication sign “×” with the ASCII letter
“x” in the affected docstrings and strings, including the stress-test
description and the occurrences near lines 23 and 28, to satisfy Ruff
RUF001/RUF002.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 0f6387a8-d64a-49eb-8c9d-698772187371
📥 Commits
Reviewing files that changed from the base of the PR and between 8edd577753d9294e6272b11632cf70bce4bebe02 and a0d66f121841d6e0e3acbd970d83a51cbb2cfacf.
📒 Files selected for processing (7)
benchmarks/routines/flashinfer_benchmark_utils.pybenchmarks/routines/sampling.pyflashinfer/__init__.pyflashinfer/cute_dsl/top_k/__init__.pyflashinfer/cute_dsl/top_k/config.pyflashinfer/topk_blackwell.pytests/topk/test_gvr_topk.py
Comment on lines
+1952
to
+1958
| def run_backend(backend, logits): | ||
| if backend == "radix": | ||
| return flashinfer.top_k_decode(logits, seq_lens, top_k, pre_idx=None) | ||
| elif backend == "gvr": | ||
| return flashinfer.top_k_decode(logits, seq_lens, top_k, pre_idx=pre_idx) | ||
| else: | ||
| raise ValueError(f"Unsupported backend: {backend}") |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm GVR top_k restriction in the backend implementation.
fd -t f topk_blackwell.py | xargs rg -nP -C3 '512|1024|2048|top_k' | rg -n 'assert|Support|raise|512|1024|2048' | head -50
ast-grep run --pattern '_run_gvr($$$)' --lang python $(fd -t f topk_blackwell.py) 2>/dev/nullRepository: flashinfer-ai/flashinfer
Length of output: 996
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant argument parsing and backend selection logic.
printf '\n== parse_sampling_args / top_k defaults ==\n'
ast-grep outline benchmarks/routines/sampling.py --view expanded | sed -n '1,220p'
printf '\n== targeted matches ==\n'
rg -n -C 3 'top_k|backends|cuda|radix|gvr|parse_sampling_args' benchmarks/routines/sampling.py flashinfer/topk_blackwell.py
printf '\n== surrounding code in benchmarks/routines/sampling.py ==\n'
sed -n '1880,1995p' benchmarks/routines/sampling.py
printf '\n== surrounding code in flashinfer/topk_blackwell.py ==\n'
sed -n '130,190p' flashinfer/topk_blackwell.pyRepository: flashinfer-ai/flashinfer
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the top_k_decode benchmark and the backend filtering helper in detail.
printf '\n== sampling.py top_k_decode section ==\n'
sed -n '1900,1995p' benchmarks/routines/sampling.py
printf '\n== filtering helper definition ==\n'
rg -n -C 4 'def filter_backends_by_compute_capability|def .*filter_backends_by_compute_capability' -S benchmarks flashinfer
printf '\n== helper body if found ==\n'
helper_file=$(rg -l 'def filter_backends_by_compute_capability' benchmarks flashinfer | head -n 1)
if [ -n "${helper_file:-}" ]; then
line=$(rg -n 'def filter_backends_by_compute_capability' "$helper_file" | head -n 1 | cut -d: -f1)
start=$((line-20))
end=$((line+80))
sed -n "${start},${end}p" "$helper_file"
fiRepository: flashinfer-ai/flashinfer
Length of output: 6881
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("benchmarks/routines/sampling.py")
lines = p.read_text().splitlines()
for i in range(1890, 1998):
if i <= len(lines):
print(f"{i}:{lines[i-1]}")
PY
printf '\n== helper search ==\n'
rg -n -C 5 'filter_backends_by_compute_capability' benchmarks flashinferRepository: flashinfer-ai/flashinfer
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect any runtime validation for GVR top_k in the implementation.
rg -n -C 4 'top_k.*512|512.*top_k|1024|2048|ValueError|assert .*top_k|supported.*top_k|requires.*512|requires.*1024|requires.*2048' flashinfer/topk_blackwell.py
# Show the public API docs and the GVR call path around the backend switch.
sed -n '280,380p' flashinfer/topk_blackwell.py
# Check whether the benchmark docs mention the special top_k requirement.
sed -n '1900,1960p' benchmarks/routines/sampling.pyRepository: flashinfer-ai/flashinfer
Length of output: 8737
Guard gvr for unsupported top_k values
testTopKDecode defaults --top_k to 50, but flashinfer.top_k_decode only supports GVR for {512, 1024, 2048}. Since legacy cuda expands to ["radix", "gvr"], the default Blackwell run will raise in the gvr path unless it’s skipped or the help text calls out the required values.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@benchmarks/routines/sampling.py` around lines 1952 - 1958, Guard the gvr
branch in run_backend so it is only selected when top_k is one of the supported
values 512, 1024, or 2048; otherwise skip gvr or raise a clear validation error
before dispatch. Also update testTopKDecode’s top_k help text or validation to
document these gvr requirements, including the legacy cuda expansion to radix
and gvr.
| num_threads_per_block : int | ||
| CTA block size. Valid values: ``512`` or ``1024``. | ||
| enable_unroll_4 : bool | ||
| Unroll Phase-2/3 inner scan loop 4× for LSU-pipelining ILP. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Run ruff format and resolve RUF002. The ruff-format hook reported files were modified, so CI is red. Additionally, the × (U+00D7) in the docstrings (Lines 46, 61, 67) triggers RUF002; replace with x if you want the check clean.
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 46-46: Docstring contains ambiguous × (MULTIPLICATION SIGN). Did you mean x (LATIN SMALL LETTER X)?
(RUF002)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/cute_dsl/top_k/config.py` at line 46, Run ruff format on
flashinfer/cute_dsl/top_k/config.py and commit the resulting formatting changes;
update the docstrings in the affected configuration definitions to replace the
Unicode multiplication sign “×” with ASCII “x” so RUF002 passes.
Sources: Linters/SAST tools, Pipeline failures
Comment on lines
+267
to
+268
| @flashinfer_api | ||
| def top_k_decode( |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- flashinfer/topk_blackwell.py outline ---'
ast-grep outline flashinfer/topk_blackwell.py --view expanded || true
echo '--- relevant section around top_k_decode ---'
nl -ba flashinfer/topk_blackwell.py | sed -n '220,340p'
echo '--- search for flashinfer_api usages with trace= ---'
rg -n "`@flashinfer_api`|trace=" flashinfer -g '*.py' | sed -n '1,240p'Repository: flashinfer-ai/flashinfer
Length of output: 737
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- top_k_decode slice ---'
python3 - <<'PY'
from pathlib import Path
path = Path('flashinfer/topk_blackwell.py')
for i, line in enumerate(path.read_text().splitlines(), 1):
if 240 <= i <= 320:
print(f"{i:4}: {line}")
PY
echo '--- flashinfer_api decorator usages with trace= ---'
python3 - <<'PY'
from pathlib import Path
import re
for path in Path('flashinfer').rglob('*.py'):
text = path.read_text()
if '`@flashinfer_api`' in text:
print(f'FILE: {path}')
for m in re.finditer(r'`@flashinfer_api`(?:\([^\n]*\))?', text):
start = text.rfind('\n', 0, m.start()) + 1
end = text.find('\n', m.end())
print(' ', text[start:end])
PY
echo '--- explicit trace= search ---'
rg -n "trace=" flashinfer -g '*.py' | sed -n '1,240p'Repository: flashinfer-ai/flashinfer
Length of output: 30254
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- search for top_k_decode trace/template references ---'
rg -n "top_k_decode|top_k_.*trace|TraceTemplate" flashinfer -g '*.py' | sed -n '1,240p'
echo '--- flashinfer/topk.py outline ---'
ast-grep outline flashinfer/topk.py --view expanded || true
echo '--- topk.py relevant sections around public APIs ---'
python3 - <<'PY'
from pathlib import Path
path = Path('flashinfer/topk.py')
for i, line in enumerate(path.read_text().splitlines(), 1):
if 620 <= i <= 840:
print(f"{i:4}: {line}")
PYRepository: flashinfer-ai/flashinfer
Length of output: 29461
Add a TraceTemplate for top_k_decode flashinfer/topk_blackwell.py still uses a bare @flashinfer_api, so this public tensor API does not opt into fi_trace. Add a matching trace= template and pass it to the decorator.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/topk_blackwell.py` around lines 267 - 268, Add a TraceTemplate for
the public top_k_decode function and pass it via the trace= argument to its
`@flashinfer_api` decorator, matching the trace template pattern used by related
APIs in topk_blackwell.py.
Sources: Coding guidelines, Learnings
dhiraj113
force-pushed
the
dhiraj113/gvr_topk
branch
from
July 10, 2026 17:41
a0d66f1 to
822e312
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
flashinfer/cute_dsl/top_k/block_scan.py (1)
152-164: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a test for the multi-tile path.
test_block_prefix_sumonly covers thenum_bins == num_threads_per_blockbranch, so thenum_bins > num_threads_per_blockloop atrange(tidx, num_bins, num_threads_per_block)still lacks coverage.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/cute_dsl/top_k/block_scan.py` around lines 152 - 164, Add coverage for the multi-tile branch in test_block_prefix_sum by adding a case where num_bins exceeds num_threads_per_block, exercising multiple iterations of the range(tidx, num_bins, num_threads_per_block) loop and asserting the complete prefix-sum output against expected values.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/cute_dsl/top_k/gvr_topk_decode.py`:
- Line 735: Rewrite the comparison in the top-k decode loop from the Yoda-style
form in the surrounding condition to the equivalent conventional form, placing
the variable expression on the left and the constant expression on the right;
preserve the existing semantics and behavior.
- Around line 970-973: Retain the explicit `nv == vlo or nv == vhi` comparisons
in the midpoint fallback within the top-k decode logic, and add a targeted
SIM109 suppression for this condition; do not replace it with tuple-membership
because traced `cutlass.Float32` values must lower safely under `@cute.jit`.
In `@tests/topk/test_gvr_topk.py`:
- Around line 23-28: Replace every ambiguous multiplication sign “×” in the
test-matrix comments and the `test_large_batch` docstring with ASCII “x”,
preserving the surrounding text and formatting so Ruff RUF001/RUF002 passes.
- Line 136: Reverse the comparison in the top-k validation condition so the
variable is on the left: change the Yoda-style `N < top_k` check to `top_k > N`
in the relevant test logic.
---
Nitpick comments:
In `@flashinfer/cute_dsl/top_k/block_scan.py`:
- Around line 152-164: Add coverage for the multi-tile branch in
test_block_prefix_sum by adding a case where num_bins exceeds
num_threads_per_block, exercising multiple iterations of the range(tidx,
num_bins, num_threads_per_block) loop and asserting the complete prefix-sum
output against expected values.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: ad1d020c-d9c9-42d0-ad6e-5f6c804fa309
📥 Commits
Reviewing files that changed from the base of the PR and between a0d66f121841d6e0e3acbd970d83a51cbb2cfacf and 822e312251aabb5a7df038a5af43497717417ff7.
📒 Files selected for processing (12)
benchmarks/routines/flashinfer_benchmark_utils.pybenchmarks/routines/sampling.pyflashinfer/__init__.pyflashinfer/cute_dsl/top_k/__init__.pyflashinfer/cute_dsl/top_k/block_scan.pyflashinfer/cute_dsl/top_k/config.pyflashinfer/cute_dsl/top_k/gvr_topk_decode.pyflashinfer/cute_dsl/top_k/gvr_topk_decode_lb.pyflashinfer/cute_dsl/top_k/pdl_utils.pyflashinfer/topk_blackwell.pytests/topk/__init__.pytests/topk/test_gvr_topk.py
🚧 Files skipped from review as they are similar to previous changes (6)
- flashinfer/init.py
- benchmarks/routines/flashinfer_benchmark_utils.py
- flashinfer/cute_dsl/top_k/pdl_utils.py
- flashinfer/cute_dsl/top_k/gvr_topk_decode_lb.py
- benchmarks/routines/sampling.py
- flashinfer/topk_blackwell.py
| # max k: floor((N - i - vec_w) / step_elem) | ||
| # N_iters = max_k + 1 | ||
| big_iters = cutlass.Int32(0) | ||
| if N > i + cutlass.Int32(vec_w - 1): |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix pipeline-blocking Ruff SIM300 (Yoda condition).
pre-commit fails on this line. The rewrite is a semantically identical comparison flip.
Proposed fix
- if N > i + cutlass.Int32(vec_w - 1):
+ if i + cutlass.Int32(vec_w - 1) < N:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if N > i + cutlass.Int32(vec_w - 1): | |
| if i + cutlass.Int32(vec_w - 1) < N: |
🧰 Tools
🪛 GitHub Actions: pre-commit / 0_pre-commit.txt
[error] 735-735: Ruff (SIM300) Yoda condition detected. Rewrite as i + cutlass.Int32(vec_w - 1) < N.
🪛 GitHub Actions: pre-commit / pre-commit
[error] 735-735: ruff-check (SIM300): Yoda condition detected. Rewrite as i + cutlass.Int32(vec_w - 1) < N.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/cute_dsl/top_k/gvr_topk_decode.py` at line 735, Rewrite the
comparison in the top-k decode loop from the Yoda-style form in the surrounding
condition to the equivalent conventional form, placing the variable expression
on the left and the constant expression on the right; preserve the existing
semantics and behavior.
Source: Pipeline failures
Comment on lines
+970
to
+973
| if nv == vlo or nv == vhi: | ||
| # Bracket exhausted — try midpoint, else give up. | ||
| nv = (vlo + vhi) * cutlass.Float32(0.5) | ||
| if nv == vlo or nv == vhi: |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Does nvidia cutlass DSL cute.jit support Python in membership operator on dynamic Float32 values
💡 Result:
No, NVIDIA CUTLASS CuTe DSL does not support the Python in membership operator for dynamic Float32 values within @cute.jit functions [1][2]. CuTe DSL functions are JIT-compiled by tracing Python code and converting supported operations into an intermediate representation (IR) suitable for GPU execution [3][4]. While the DSL supports common Python control flow constructs (such as if/else and for loops) and certain operators by lowering them to IR [2][4], it only supports a restricted subset of Python’s features [2]. The documentation explicitly notes that CuTe DSL imposes limitations on what can be treated as a dynamic value [2]. Specifically, while the DSL can perform operations on primitive types like Float32 and Int32, it does not implement support for Python's membership operator (in) for dynamic values [2]. Composite types such as lists, tuples, and dictionaries can be used for meta-programming or as static containers to store dynamic values, but they cannot be dynamically modified or queried with membership operators during kernel execution [2]. Attempting to use unsupported Python syntax or operators inside a @cute.jit function will generally result in a runtime error during the JIT compilation or execution process [1][2].
Citations:
- 1: https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/limitations.html
- 2: https://docs.nvidia.com/cutlass/4.4.1/media/docs/pythonDSL/limitations.html
- 3: https://docs.nvidia.com/cutlass/4.5.2/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.html
- 4: https://docs.nvidia.com/cutlass/media/docs/pythonDSL/cute_dsl_general/dsl_control_flow.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant section and nearby definitions in the target file.
FILE="flashinfer/cute_dsl/top_k/gvr_topk_decode.py"
wc -l "$FILE"
sed -n '940,990p' "$FILE" | cat -nRepository: flashinfer-ai/flashinfer
Length of output: 2761
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search for any existing SIM109 suppressions or similar patterns in the repo.
rg -n "SIM109|noqa: SIM109|membership operator|tuple-membership|cute\.jit" flashinfer -SRepository: flashinfer-ai/flashinfer
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect CUTLASS CuTe DSL docs in the repo if vendored references exist.
rg -n "limitations|membership|in operator|Float32|cute\.jit" -S .Repository: flashinfer-ai/flashinfer
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Look for existing SIM109 suppressions or nearby lint-avoidance patterns in cute DSL code.
rg -n "# noqa: SIM109|SIM109|Ruff|noqa" flashinfer/cute_dsl flashinfer/gdn_kernels flashinfer/fused_moe flashinfer/mamba -S --max-count 80Repository: flashinfer-ai/flashinfer
Length of output: 7080
Keep the explicit comparisons and suppress SIM109 here. Tuple-membership on traced cutlass.Float32 values doesn’t lower cleanly under @cute.jit, so if nv in (vlo, vhi): is not a safe autofix.
🧰 Tools
🪛 GitHub Actions: pre-commit / 0_pre-commit.txt
[error] 970-970: Ruff (SIM109) Use nv in (vlo, vhi) instead of multiple equality comparisons.
[error] 973-973: Ruff (SIM109) Use nv in (vlo, vhi) instead of multiple equality comparisons.
🪛 GitHub Actions: pre-commit / pre-commit
[error] 970-970: ruff-check (SIM109): Use nv in (vlo, vhi) instead of multiple equality comparisons.
[error] 973-973: ruff-check (SIM109): Use nv in (vlo, vhi) instead of multiple equality comparisons.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/cute_dsl/top_k/gvr_topk_decode.py` around lines 970 - 973, Retain
the explicit `nv == vlo or nv == vhi` comparisons in the midpoint fallback
within the top-k decode logic, and add a targeted SIM109 suppression for this
condition; do not replace it with tuple-membership because traced
`cutlass.Float32` values must lower safely under `@cute.jit`.
Source: Pipeline failures
Comment on lines
+23
to
+28
| test_basic_decode — dtype × top_k × N × batch; works on all GPUs | ||
| test_return_values — return_values=True correctness | ||
| test_next_n — next_n=2 (V3.2 speculative-decode stride) | ||
| test_compress_ratio — compress_ratio=4 (DSv4 KV compression) | ||
| test_preallocated_outputs — pre-allocated out_indices / out_values | ||
| test_large_batch — stress: large batch × long rows |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Replace ambiguous × (U+00D7) with ASCII x. The multiplication signs in the test-matrix string (Lines 23, 28) trigger ruff RUF001, and the × in the test_large_batch docstring (Line 259) trips RUF002, keeping the ruff check red.
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 23-23: String contains ambiguous × (MULTIPLICATION SIGN). Did you mean x (LATIN SMALL LETTER X)?
(RUF001)
[warning] 23-23: String contains ambiguous × (MULTIPLICATION SIGN). Did you mean x (LATIN SMALL LETTER X)?
(RUF001)
[warning] 23-23: String contains ambiguous × (MULTIPLICATION SIGN). Did you mean x (LATIN SMALL LETTER X)?
(RUF001)
[warning] 28-28: String contains ambiguous × (MULTIPLICATION SIGN). Did you mean x (LATIN SMALL LETTER X)?
(RUF001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/topk/test_gvr_topk.py` around lines 23 - 28, Replace every ambiguous
multiplication sign “×” in the test-matrix comments and the `test_large_batch`
docstring with ASCII “x”, preserving the surrounding text and formatting so Ruff
RUF001/RUF002 passes.
Source: Linters/SAST tools
dhiraj113
force-pushed
the
dhiraj113/gvr_topk
branch
from
July 20, 2026 21:30
822e312 to
5e1383f
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (5)
flashinfer/cute_dsl/top_k/config.py (3)
25-26: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUnused imports still present (CI failing).
fieldandOptionalare imported but never used, still tripping ruff F401.🔧 Proposed fix
-from dataclasses import dataclass, field -from typing import Optional +from dataclasses import dataclass🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/cute_dsl/top_k/config.py` around lines 25 - 26, Remove the unused field and Optional imports from the import section of config.py, leaving dataclass imported for the existing implementation.Source: Pipeline failures
118-118: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winYoda condition still present (CI failing).
Ruff SIM300 fails the pipeline on this comparison.
🔧 Proposed fix
- num_threads_per_block = 1024 if (num_rows <= num_sms and N >= n_thresh_t) else 512 + num_threads_per_block = 1024 if (num_rows <= num_sms and n_thresh_t <= N) else 512🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/cute_dsl/top_k/config.py` at line 118, Update the comparison in the num_threads_per_block assignment to use standard operand ordering rather than a Yoda condition, while preserving the existing conditional behavior and values.Source: Pipeline failures
46-46: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAmbiguous
×in docstring still present.Triggers RUF002; replace with ASCII
xfor a clean lint.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/cute_dsl/top_k/config.py` at line 46, Update the Phase-2/3 inner scan loop docstring in the visible configuration documentation to replace the ambiguous multiplication symbol “×” with ASCII “x”, preserving the rest of the wording and behavior.Source: Linters/SAST tools
tests/topk/test_gvr_topk.py (2)
23-28: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAmbiguous
×characters still present.The test-matrix strings (Lines 23, 28) trip RUF001, and the
test_large_batchdocstring (Line 259) trips RUF002.Also applies to: 257-259
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/topk/test_gvr_topk.py` around lines 23 - 28, Replace the ambiguous multiplication signs in the test-matrix descriptions with ASCII wording or symbols that satisfy RUF001, including the entries for test_basic_decode and test_large_batch. Update the test_large_batch docstring similarly to remove the ambiguous character and preserve its intended stress-test description.Source: Linters/SAST tools
136-136: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winYoda condition still present (CI failing).
Ruff SIM300 fails the pipeline on this comparison.
🔧 Proposed fix
- if N < top_k: + if top_k > N: pytest.skip("N < top_k")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/topk/test_gvr_topk.py` at line 136, Update the comparison in the top-k test from the Yoda-style form to the conventional variable-first form so Ruff SIM300 passes, preserving the existing condition and behavior.Source: Pipeline failures
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/cute_dsl/top_k/config.py`:
- Line 1: Run ruff formatting across all three affected files:
flashinfer/cute_dsl/top_k/config.py (including the multi-line raise statements
in GvrTopKLBConfig.__post_init__), flashinfer/cute_dsl/top_k/block_scan.py, and
tests/topk/test_gvr_topk.py. A single repository-wide ruff format or pre-commit
run is sufficient.
In `@flashinfer/cute_dsl/top_k/gvr_topk_decode.py`:
- Line 1993: Update the comparison in the top-k decode logic to use the
conventional operand order, placing the literal-derived top_k value on the left
and N on the right, while preserving the existing Int32 conversion and `@cute.jit`
behavior.
---
Duplicate comments:
In `@flashinfer/cute_dsl/top_k/config.py`:
- Around line 25-26: Remove the unused field and Optional imports from the
import section of config.py, leaving dataclass imported for the existing
implementation.
- Line 118: Update the comparison in the num_threads_per_block assignment to use
standard operand ordering rather than a Yoda condition, while preserving the
existing conditional behavior and values.
- Line 46: Update the Phase-2/3 inner scan loop docstring in the visible
configuration documentation to replace the ambiguous multiplication symbol “×”
with ASCII “x”, preserving the rest of the wording and behavior.
In `@tests/topk/test_gvr_topk.py`:
- Around line 23-28: Replace the ambiguous multiplication signs in the
test-matrix descriptions with ASCII wording or symbols that satisfy RUF001,
including the entries for test_basic_decode and test_large_batch. Update the
test_large_batch docstring similarly to remove the ambiguous character and
preserve its intended stress-test description.
- Line 136: Update the comparison in the top-k test from the Yoda-style form to
the conventional variable-first form so Ruff SIM300 passes, preserving the
existing condition and behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: cd52314a-f001-4ed1-842e-a0a5e69867b8
📥 Commits
Reviewing files that changed from the base of the PR and between 822e312251aabb5a7df038a5af43497717417ff7 and 5e1383f7501b17e4fe3a8fa00d2d4902a8c65b8b.
📒 Files selected for processing (12)
benchmarks/routines/flashinfer_benchmark_utils.pybenchmarks/routines/sampling.pyflashinfer/__init__.pyflashinfer/cute_dsl/top_k/__init__.pyflashinfer/cute_dsl/top_k/block_scan.pyflashinfer/cute_dsl/top_k/config.pyflashinfer/cute_dsl/top_k/gvr_topk_decode.pyflashinfer/cute_dsl/top_k/gvr_topk_decode_lb.pyflashinfer/cute_dsl/top_k/pdl_utils.pyflashinfer/topk_blackwell.pytests/topk/__init__.pytests/topk/test_gvr_topk.py
🚧 Files skipped from review as they are similar to previous changes (6)
- flashinfer/init.py
- benchmarks/routines/flashinfer_benchmark_utils.py
- flashinfer/topk_blackwell.py
- flashinfer/cute_dsl/top_k/gvr_topk_decode_lb.py
- flashinfer/cute_dsl/top_k/pdl_utils.py
- benchmarks/routines/sampling.py
| @@ -0,0 +1,195 @@ | |||
| # Copyright (c) 2026, the FlashInfer team. | |||
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
ruff-format hook fails identically across three files — run formatting once for the whole PR. All three files trip the same pre-commit ruff-format hook because the diff wasn't formatted before commit; a single ruff format pass (or pre-commit run --all-files) fixes all three.
flashinfer/cute_dsl/top_k/config.py#L1-L1: runruff formaton this file (also collapses the multi-lineraise ValueError(...)calls inGvrTopKLBConfig.__post_init__, e.g. lines 188-195, into ruff's preferred single-line form where they fit).flashinfer/cute_dsl/top_k/block_scan.py#L1-L1: runruff formaton this file.tests/topk/test_gvr_topk.py#L1-L1: runruff formaton this file.
🧰 Tools
🪛 GitHub Actions: pre-commit / 0_pre-commit.txt
[error] 1-1: pre-commit hook ruff-format failed: files were modified by this hook. Run pre-commit run --all-files locally to apply formatting.
📍 Affects 3 files
flashinfer/cute_dsl/top_k/config.py#L1-L1(this comment)flashinfer/cute_dsl/top_k/block_scan.py#L1-L1tests/topk/test_gvr_topk.py#L1-L1
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/cute_dsl/top_k/config.py` at line 1, Run ruff formatting across
all three affected files: flashinfer/cute_dsl/top_k/config.py (including the
multi-line raise statements in GvrTopKLBConfig.__post_init__),
flashinfer/cute_dsl/top_k/block_scan.py, and tests/topk/test_gvr_topk.py. A
single repository-wide ruff format or pre-commit run is sufficient.
Source: Pipeline failures
| smem_input = None | ||
|
|
||
| # ---- Degenerate path: N <= top_k → copy input as-is ---- | ||
| if N <= cutlass.Int32(top_k): |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix pipeline-blocking Ruff SIM300 (Yoda condition).
pre-commit fails here. Unlike the SIM109 case above, this is a plain Int32 comparison and the flip lowers identically under @cute.jit.
Proposed fix
- if N <= cutlass.Int32(top_k):
+ if cutlass.Int32(top_k) >= N:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if N <= cutlass.Int32(top_k): | |
| if cutlass.Int32(top_k) >= N: |
🧰 Tools
🪛 GitHub Actions: pre-commit / 0_pre-commit.txt
[error] 1993-1993: pre-commit hook ruff-check (ruff) failed: SIM300 Yoda condition detected (if N <= cutlass.Int32(top_k):). Rewrite as cutlass.Int32(top_k) >= N.
🪛 GitHub Actions: pre-commit / pre-commit
[error] 1993-1993: ruff check (SIM300): Yoda condition detected; rewrite as cutlass.Int32(top_k) >= N
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/cute_dsl/top_k/gvr_topk_decode.py` at line 1993, Update the
comparison in the top-k decode logic to use the conventional operand order,
placing the literal-derived top_k value on the left and N on the right, while
preserving the existing Int32 conversion and `@cute.jit` behavior.
Source: Pipeline failures
Collaborator
Author
There was a problem hiding this comment.
is this still an issue
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="flashinfer/cute_dsl/top_k/gvr_topk_decode.py"
printf '%s\n' '--- Current comparison near the reviewed line ---'
sed -n '1988,1998p' "$file"
printf '%s\n' '--- Targeted Ruff SIM300 check ---'
if command -v ruff >/dev/null 2>&1; then
ruff check --select SIM300 "$file"
elif command -v uv >/dev/null 2>&1; then
uv run ruff check --select SIM300 "$file"
else
echo "Ruff is unavailable in this sandbox."
fiLength of output: 2512
@dhiraj113 Yes — it is still an issue. The current branch still has the Yoda condition at Line 1993, and Ruff reports SIM300 for it.
The minimal fix remains:
if cutlass.Int32(top_k) >= N:There is also a separate SIM300 violation at Line 735 (if N > i + ...), so both need addressing for the targeted Ruff check to pass.
dhiraj113
force-pushed
the
dhiraj113/gvr_topk
branch
2 times, most recently
from
July 22, 2026 03:42
78d299b to
5f9d3ae
Compare
Collaborator
Author
|
/bot run |
Collaborator
Collaborator
|
[FAILED] Pipeline #59041442: 1/20 passed |
Contributor
|
/bot run |
…topk_varlen - Rename flashinfer/topk_blackwell.py → flashinfer/topk_varlen.py and the public API top_k_decode → top_k_varlen to better reflect that the function operates on variable-length (ragged) sequences rather than being decode-step-specific. - Drop num_long_rows param and load_balance="auto" mode; load_balance is now bool only. - Add flashinfer/trace/templates/topk.py with top_k_varlen_trace (TraceTemplate) and wire it into @flashinfer_api(trace=top_k_varlen_trace). - Rename tests/topk/test_gvr_topk.py → tests/topk/test_topk_varlen.py; remove helper/auto tests that no longer apply; add trace correctness test. - Update benchmarks and all call sites. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…parity - Add single-pass multi-CTA radix CuTe DSL kernel with SM-aware chunking (fixes large-N shared-memory overflow, closes the uniform large-N gap) - Port the warp-redundant GVR decode kernel and feed per-CTA scan width to the load-balance launch heuristic (closes the mixed / N=131072 LB gaps) - Rename backends: radix -> radix_cutlass (masked CUDA fallback), radix_cutedsl -> radix (CuTe DSL); order radix, gvr, radix_cutlass - Fix top_k_varlen return-type annotations Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…code - Update benchmark registry, sampling routine, and bench_gvr_lb to the new backend names (radix = CuTe DSL, radix_cutlass = masked fallback); add an explicit radix_cutlass runner - Fix sampling refcheck to index top_k_varlen's (indices, values) tuple - Remove unused num_groups_eff in _run_radix (row->request mapping is in-kernel) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…CUDA graph) - Add explicit radix (CuTe DSL) backend tests, including the multi-CTA regime (N > max_chunk) that exercises the N=131072 shared-memory-overflow fix - Add CUDA-graph capture/replay tests (radix multi-CTA + GVR) - Add varlen-ragged, seq_len<=top_k degenerate, cross-backend consistency, backend-heuristic-priority, and input-validation tests - Strengthen _check_correct (require_all_checked); check correctness on all GPUs Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Line-wrapping only — ruff-format for the Python files, clang-format for csrc/topk.cu. No logic changes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
top_k_varlen is a sparse-attention KV-index selection primitive, not a vocabulary-sampling op, so it no longer lives under routines/sampling.py: - New routines/topk_varlen.py (parse_topk_varlen_args + run_topk_varlen_test) - Own benchmark_apis category + flashinfer_benchmark.py dispatch - Drop the "cuda" backend sentinel; --backends defaults to the cc-registry union (radix / gvr / radix_cutlass), narrowed per-GPU by the cc-filter - Rename --vocab_size -> --max_seq_len (arg, variable, and CSV column) - Remove testTopKVarlen and its top_k_varlen-specific --backends from sampling.py (reverted to the plain "cuda" default) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Lazy CuTe DSL imports: remove top-level `import cutlass` and
`from .cute_dsl.utils import ...` from topk_varlen.py; use
`importlib.util.find_spec` to compute `_CUTE_DSL_AVAILABLE` at
module load without importing cutlass, so `import flashinfer` does
not fail when nvidia-cutlass-dsl is absent.
- Fix order_row slice bug in _run_gvr (load_balance=False): remove
`seq_lens[::next_n]` — seq_lens is already request-level with shape
(num_rows // next_n,), so slicing made order_row next_n times too
short when next_n > 1, causing out-of-bounds kernel accesses.
- Fix compress_ratio division order in _run_radix: the next_n
per-row adjustment is in token units and must happen before dividing
by compress_ratio. Add compress_ratio as a const_expr parameter to
SinglePassMultiCTARadixTopKKernel and compute length as
(seq_len - next_n + nn + 1) // compress_ratio inside the kernel;
remove the Python-side pre-division.
- Use device-queried SMEM capacity in radix chunk config: replace
hardcoded _SM100_SMEM_CAPACITY=232448 with
get_shared_bytes_per_block_optin(device) so SM120/SM121 get the
correct SMEM bound rather than SM100's value.
- Strengthen _gvr_top_k_varlen_check: add guards for top_k not in
{512,1024,2048}, compress_ratio not in {1,4}, misaligned row width
(N % elem_align != 0), and load_balance batch > 1024; backend="auto"
now falls back to radix instead of crashing deep in the kernel.
Remove the now-unreachable alignment check from top_k_varlen body.
- Guard out_values in _run_radix: pass
`out_values if return_output_values else None` to the compiled
kernel; previously a caller-supplied out_values buffer was forwarded
unconditionally into a kernel compiled to expect None.
- Fix _top_k_varlen_heuristic signature: spell out the full parameter
list (logits, seq_lens, top_k, ...) matching top_k_varlen's
signature; the old **kwargs caused TypeError on the skip_check=True
path where the decorator forwards positional args directly.
- Remove pdl_utils.py: import griddepcontrol_wait /
griddepcontrol_launch_dependents from cutlass.cute.arch (already
used in gemm/kernels/dense_blockscaled_gemm_sm100.py) instead of
maintaining a third local copy.
- Fix stale docstrings in config.py: remove references to a
nonexistent config= argument on top_k_varlen; point GvrTopKLBConfig
at top_k_varlen(backend="gvr", load_balance=True) instead of
unexported internal functions.
- Eliminate masked_fill in _run_radix_cutlass: replace
masked_fill + radix_topk (which allocated its own output buffer)
with radix_topk_ragged_transform, which accepts per-row lengths
natively and writes directly into the pre-allocated out_indices;
gather values with torch.gather when return_values=True (~2x speedup
at 128x131072 bf16).
- Allow caller-provided workspace in _run_gvr_lb: add optional
workspace dict to top_k_varlen (keys: "gvr_order_row",
"gvr_counters") so decode-loop callers can reuse the LB workspace
across steps without per-call allocation; change counters from
torch.zeros to torch.empty since the prepare kernel overwrites both
entries. Document that the same workspace must not be shared across
concurrent streams.
- Restrict GVR to SM100/103: split _BLACKWELL_PLUS_CCS into
_GVR_CCS=[100,103] for GVR (B200-class where cluster_size=4 CTA
clusters are validated) and keep _BLACKWELL_PLUS_CCS=[100,103,110,
120,121] for the CuTe DSL radix backend; SM120/SM121 do not support
programmatic multicast (1x1x1 cluster only) so GVR LB would be
incorrect there.
- Tests: add test_gvr_no_lb_next_n, test_radix_next_n_compress_ratio,
test_out_values_ignored_when_return_values_false (all 3 backends +
both GVR LB paths), test_skip_check_auto_backend,
test_gvr_lb_workspace_reuse; expand test_load_balance_modes to
parametrize over next_n in {1,2}; expand test_radix_compress_ratio
to parametrize over (top_k, next_n); update
test_gvr_row_width_alignment to verify auto fallback on misaligned N.
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
The four CuTe-DSL compile helpers (_compile_gvr, _compile_gvr_lb, _compile_gvr_lb_prepare, _compile_radix) called cute.compile() directly. cute.compile() has no persistent cache, and the @functools.cache on each helper only dedupes within a single process, so every fresh process recompiled the kernels from scratch. Route all four through build_and_load_cute_dsl_kernel(), which exports each compiled kernel as an object file and reloads it on later runs. - kernel_name encodes every @functools.cache specializer, so distinct specializations get distinct on-disk artifacts (no cache collision). - The three GVR helpers share the "gvr_topk" module dir and pass an identical union source-key via _gvr_kernel_source_files(); radix uses _radix_kernel_source_files(). Both list all transitive kernel sources (decode modules + block_scan.py) so a source edit invalidates the cache. The first run per specialization still compiles; subsequent runs load the cached artifact from disk. Verified on SM90 and SM100; behavior unchanged (identical cute.compile arguments). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…dsl 4.6+) nvidia-cutlass-dsl 4.6 removed the deprecated free function cute.make_fragment; the topk CuTe DSL kernels called it in 9 places (GvrTopKKernel x7, radix x2), so both the GVR and radix backends raised "AttributeError: module 'cutlass.cute' has no attribute 'make_fragment'" at compile time on Blackwell under cutlass-dsl 4.6+. Replace all 9 cute.make_fragment((shape,), dtype) calls with cute.make_rmem_tensor((shape,), dtype) — the successor named in 4.5.x's own deprecation warning, with an identical (shape, dtype) signature and present in both 4.5.x and 4.6.x. Verified with nvidia-cutlass-dsl 4.6.1 on SM80, SM89, SM90, SM100, and SM120: tests/topk/test_topk_varlen.py::test_basic_decode passes 16/16 on every arch, exercising both the GVR and radix backends. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
_run_radix_cutlass fetches output values with torch.gather(logits, 1, out_indices). radix_topk_ragged_transform writes the -1 sentinel into surplus slots when a row's length <= top_k (topk.cuh Ragged branch), so the gather received a negative index and tripped a device-side bounds assert. Reachable via radix_cutlass + return_values=True + any row with seq_len < top_k; the prior masked_fill path never gathered, so this was introduced with the gather optimization. Clamp the gather index to >= 0 so it stays in-bounds, then zero the sentinel value slots (matches the kernel's DType(0) convention). No-op when every row has seq_len >= top_k, so the common path is unchanged. Extend test_seq_len_less_than_top_k to request return_values=True and assert valid-slot values match logits[row, idx] and that radix_cutlass zeroes the -1 sentinel slots. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gvr is gated to _GVR_CCS = [100, 103], so listing it for 12.0/12.1 in the top_k_varlen benchmark backend table raised BackendSupportedError when the default backend set was exercised on SM120/121. Drop it from both 12.x entries (radix and radix_cutlass remain — both build there). Also correct the _GVR_CCS comment: the non-LB (cluster_size=1) GVR CuTe-DSL kernel does not merely get restricted 'for consistency' — it fails to build on sm_120a (libNVVM rejects the generated device IR, verified on an RTX 5080), so consumer Blackwell cannot use GVR even without load balancing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… regression test_radix_next_n_compress_ratio relied on _make_inputs, which fills seq_lens with N*compress_ratio (16384). That value is divisible by compress_ratio, so adjust-before-divide and the buggy divide-before-adjust produce the same N_eff (4095) and the test could not catch the regression it describes. Override seq_lens to N*compress_ratio + 1 (16385), where the two orders diverge (4096 vs 4095), and boost the last column to a guaranteed top-1 value so the off-by-one is observable: correct N_eff == N selects column N-1 in every row, while the buggy bound drops it from the ofs=0 rows. Assert column N-1 is present in every row's top-k — deterministic and tie-immune (a plain k-th-value check misses it at seed=17, verified by simulation). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
test_repeated_calls asserted two identical calls return the same top-K index set per row. On non-Blackwell GPUs top_k_varlen runs radix_cutlass with deterministic=False, so BF16 values that tie at the K-th boundary let two correct calls select different (equally valid) tied indices — the sets can legitimately differ. This made the test flaky (failing intermittently in the H100 CI job). Assert each call returns a valid top-K via _check_correct(require_all_checked= True) instead of requiring bit-identical sets. Verified stable across 5 repeats on H100 (SM90) and once each on SM80/89/100/120. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…s extraction _build_axis_extractors matched the first Tensor input mentioning an axis and stopped, so the top_k axis was extracted from the optional pre_idx tensor instead of the always-present top_k scalar. Generated trace samples omit optional inputs, so the extractor returned None and the trace-consistency test reported: Const axes missing values: ['top_k']. Reorder the per-axis source preference to always-present sources first: required tensor shape -> scalar input named after the axis -> optional tensor -> scalar-kwarg fallback. Fixes the top_k_varlen template and generalizes to any axis whose only tensor source is optional but which is also a scalar. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Move the CuTe-DSL top-k kernels and the top_k_varlen API out of the shared flashinfer/cute_dsl/ tree into a dedicated flashinfer/top_k/ package, mirroring flashinfer/quantization/ (kernels under kernels/). Addresses PR review comments about not accumulating feature-specific code under flashinfer/cute_dsl/. flashinfer/cute_dsl/top_k/* -> flashinfer/top_k/kernels/* flashinfer/topk_varlen.py -> flashinfer/top_k/topk_varlen.py tests/topk/ -> tests/top_k/ (name consistency) Kernel sources are byte-identical (git mv only); only import paths and the ruff exclude list in pyproject.toml change. The public API flashinfer.top_k_varlen is unchanged (re-exported from __init__.py). Verified: pre-commit clean; tests pass on sm_80/89/90/100/120 (115 passed on sm_100, exercising GVR + radix from the new path); trace-registry inventory green. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ltering Drop the hard-coded top_k_varlen compute-capability table in the benchmark's routine_cc_to_supported_backends and instead resolve per-CC backend support at runtime via flashinfer.top_k_varlen.is_backend_supported(backend, cc), mirroring how the GEMM routines rely on their @backend_requirement support checkers (mm_fp4 / bmm_fp8). The decorator is the single source of truth, so the two can no longer drift. Verified end-to-end: sm_80 runs radix_cutlass only (radix/gvr skipped with warnings), sm_100 runs all three backends; refcheck passes on both. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…owing top_k()) The flashinfer/top_k/ package introduced by the relocation shadowed the existing flashinfer.top_k sampling function (from flashinfer/topk.py): importing the submodule rebinds flashinfer.top_k to the module, so every flashinfer.top_k(logits, k) call raised "TypeError: 'module' object is not callable" -- 429 failures in tests/utils/test_topk.py (the CI JIT Unittest sampling shard). Rename the package and its tests dir to the collision-free topk_varlen, mirroring the top_k_varlen API name. Note top_k_varlen (underscored) would collide identically with the top_k_varlen function, so topk_varlen is used. flashinfer/top_k/ -> flashinfer/topk_varlen/ tests/top_k/ -> tests/topk_varlen/ Verified: flashinfer.top_k is a callable function again; tests/utils/test_topk.py 962 passed (was 429 failed); tests/topk_varlen/test_topk_varlen.py 115 passed on sm_100; trace-registry inventory green; pre-commit clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…h failure) The optional-output_values change instantiated collect_indices twice in the Basic-mode deterministic path (a values lambda and a no-values lambda), so each carried its own __shared__ BlockScan temp storage. That doubled the kernel's static shared memory, overflowing the per-block opt-in cap on Ada/sm_89 (~99KB; Hopper's 227KB masked it) and failing the launch with cudaErrorInvalidValue -- breaking top_k_top_p_sampling_from_* at large vocab (e.g. 128256) on A10G CI. Collapse to a single collect_indices instantiation with a runtime null-check on the value pointer, restoring the original static-smem footprint while keeping the indices-only fast path. Verified: main passes / branch failed (branch regression); sm_80/89/90/100 test_sampling.py + test_topk.py pass (2416 on L40S, was 8 failing). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
dhiraj113
force-pushed
the
dhiraj113/gvr_topk
branch
from
August 6, 2026 22:47
6ff26fb to
55f1c5c
Compare
Collaborator
|
[FAILED] Pipeline #61432149 — 5/18 executed test jobs passed Compared with nightly #61182354. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
Collaborator
|
/bot run tests/topk_varlen |
Collaborator
Collaborator
|
[FAILED] Pipeline #61643724 — 16/18 executed test jobs passed Compared with nightly #61367193. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
No individual test or infrastructure failures could be extracted. |
5 tasks
5 tasks
bkryu
pushed a commit
that referenced
this pull request
Aug 9, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description @HumansAnd SGLang's DeepSeek V4 indexer uses a compact page table in which one entry represents 64 score positions. Its score rows may also have padding between rows, and CUDA graph capture owns the translated and raw-index output buffers. The current `top_k_page_table_transform` contract assumes one page-table entry per score and allocates only the translated output, so SGLang has to split this into `top_k`, score re-gathering, compact-page translation, and output copies. This PR extends the existing fused page-table transform for that layout: - Add `page_size`, defaulting to `1` for backward compatibility. - Add optional caller-owned `out` and `out_raw_indices` buffers. Raw indices remain positionally aligned with translated indices, including deterministic post-sort and `-1` padding; the two buffers must be disjoint. - Honor the input row stride instead of requiring tightly packed rows, while preserving last-dimension contiguity and alignment-safe vectorization. - Propagate the contract through the Python API, TVM FFI binding, Radix and Filtered implementations, graph-safe dispatch, deterministic post-sort, and the trivial `length <= k` path. - Extend only the operation-local trace template for `page_size`. Its core one-output definition excludes destination buffers; raw-output calls are deliberately not emitted and use a distinct routing identity so Trace Apply falls back to the API. For each selected local score index `idx`, the compact transform is: ```text physical_page = src_page_table[ batch_idx, page_table_row_start + idx // page_size ] output = physical_page * page_size + idx % page_size ``` `page_table_row_starts` is measured in page-table entries, while `row_starts` is measured in score elements. When `page_size > 1` and `row_starts` is supplied, callers must therefore provide `page_table_row_starts` explicitly rather than relying on the existing shared-start behavior. The C++ path uses policy types rather than a second boolean mode. The FFI boundary selects `DirectPageTableKernelPolicy` or `ConfigurablePageTableKernelPolicy` once, and the kernel ABI is derived structurally from the policy type: an empty policy contributes zero arguments, while a stateful trivially-copyable policy contributes one object. Host dispatch carries one typed policy value through the selection stack; only the terminal launch forms the zero-or-one argument pack. The direct policy therefore adds no kernel argument or device branch, while the configurable policy owns score-row layout, logical-to-physical translation, and the optional raw-index sink. Future transforms that share the flat row-selection contract can extend the configurable policy without multiplying kernel variants or changing the selection kernels. Direct translation sites retain their original expressions so supported-path SASS is preserved exactly. This is a clean extension of the API introduced in #4169: `page_size=1`, omitted output buffers, and tightly packed inputs retain the existing behavior and cluster fast path. [SGLang #33237](sgl-project/sglang#33237) uses this API to replace its DeepSeek V4 unfused workaround with one graph-safe FlashInfer call. ## ⚡ Performance Fresh performance validation compares the exact rebase base `29196cf437778906c72630dc5d9850de547501de` with head `cf0319a3497e252219544c0a8b4168c6ba598f88` on the same NVIDIA B200 (driver 580.126.09), CUDA 13.2.78, and PyTorch 2.13.0+cu132. Base and head used separate editable source trees and JIT workspaces. `benchmarks/bench_topk.py` is byte-identical on both sides (SHA-256 `57cd1ca61b38120380cb9ea7cf81ae3ee972484724bb2bb785649ae05cc199d9`). The script uses CUPTI (`cupti-python` 13.2.0 and `nvidia-cuda-cupti` 13.2.75), 10 dry runs, 100 measured iterations, cold L2, and the median. After #4295 it already sets `use_cuda_graph=False`, so no temporary benchmark-source edit was needed and the CUPTI plus CUDA graph instability is excluded. Exact PR-body commands: ```bash python3 benchmarks/bench_topk.py \ --op dsa_topk --dtype bf16 --dsa-input-pattern dsa_relu \ --dsa-case all --dsa-topk 2048 --tie-break python3 benchmarks/bench_topk.py \ --op varlen --dtype bf16 --length-dist causal \ --varlen-k 2048 --varlen-q-len 128 --tie-break ``` After #4295 these exact commands keep `deterministic=False`: they measure the default nondeterministic path plus SMALL/LARGE tie selection without the canonical output-order sort. For fair DSA pairing, the confirmation runs used the same command after `torch.manual_seed(1234)` and `torch.cuda.manual_seed_all(1234)`. Current-mode DSA ran base/head/head/base; varlen used base/head/head/base and its built-in per-case seeds. Canonical-output coverage repeated both commands with `--deterministic` on base and head. All comparisons use matched per-case medians; negative PR delta means the head is faster. | workload / mode | default PR delta (worst) | deterministic PR delta (worst) | tie-small PR delta (worst) | tie-large PR delta (worst) | |---|---:|---:|---:|---:| | DSA, current ABBA | `+0.05%` (`+0.13%`) | n/a | `-0.09%` (`+0.08%`) | `-0.01%` (`+0.12%`) | | page-table varlen, current ABBA | `+0.01%` (`+0.10%`) | n/a | `+0.06%` (`+0.19%`) | `+0.05%` (`+0.24%`) | | ragged varlen, current ABBA | `+0.02%` (`+0.08%`) | n/a | `+0.01%` (`+0.13%`) | `-0.02%` (`+0.14%`) | | DSA, explicit deterministic | `+0.03%` (`+0.15%`) | `-0.09%` (`+0.10%`) | `+0.01%` (`+0.41%`) | `-0.09%` (no regressed point) | | page-table varlen, explicit deterministic | `+0.01%` (`+0.10%`) | `-0.18%` (`+0.03%`) | `-0.09%` (no regressed point) | `-0.05%` (`+0.09%`) | | ragged varlen, explicit deterministic | `-0.01%` (`+0.06%`) | `-0.10%` (`+0.01%`) | `-0.01%` (`+0.03%`) | `-0.13%` (`+0.08%`) | All 12 varlen rows had identical `len_min`, `len_mean`, `len_max`, and `triv%` across paired runs. Current-mode suite geomeans are within `+0.06%`, the largest individual delta is `+0.24%`, and deterministic geomeans are flat or faster apart from a `+0.01%` DSA tie-small geomean. This supports no measurable kernel-latency regression after the final rebase. These sweeps exercise the existing direct compatibility paths (`page_size=1`, contiguous rows, no caller-owned outputs). The configured V4 path has no pre-PR API equivalent. Its policy cleanup was separately checked by an eager ABBA host audit using `page_size=64`, padded row stride, raw output, and all three production shapes: its configured 12-case batched end-to-end geomean was `+0.053%`, with a worst point of `+0.380%`. The policy design also has complementary binary evidence from the exhaustive audit performed after #4295: ```text direct: 354/354 affected kernels, 0 normalized SASS/resource/KPARAM mismatches radix=180, filtered=132, finalizer=42 configured: 116 kernels radix=60, filtered=44, finalizer=12 exactly one trailing 24-byte policy object; PageTable mode only ``` The direct variants retained their upstream parameter counts and constant-bank spans. The configured variants add one policy object without a second policy family or boolean template axis. ## 🔍 Related Issues - SGLang integration: sgl-project/sglang#33237 - Independent score/page-table starts: #4169 - SGLang packed-PAGED workaround and backend-selection fix: sgl-project/sglang#32490 - SGLang DeepSeek V4 Top-K backend integration: sgl-project/sglang#31087 ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). Validation used an editable source build on one NVIDIA B200 with CUDA 13.2.78 and PyTorch 2.13.0+cu132. The rebased head was synced to an isolated source tree and JIT workspace on `flashinfer-pr4366-cu132`: ```text pre-commit run --all-files Passed python3 -m pytest -q \ tests/utils/test_topk.py::test_top_k_page_table_transform_misaligned_scores_without_row_starts \ tests/utils/test_topk.py::test_top_k_page_table_transform_compact_pages_cuda_graph_replay 18 passed, 3 warnings in 93.98s python3 -m pytest -q \ tests/topk_varlen/test_topk_varlen.py::test_radix_preallocated_outputs \ tests/topk_varlen/test_topk_varlen.py::test_out_values_ignored_when_return_values_false 6 passed, 39 warnings in 4.41s python3 -m pytest -q \ tests/trace/test_fi_trace_template_consistency.py \ tests/trace/test_template_init.py \ -k top_k_page_table_transform 6 passed, 1 skipped, 972 deselected, 3 warnings in 0.27s python3 -m pytest -q tests/utils/test_topk.py 1495 passed, 3 warnings in 7.56s ``` Final CUDA 13.2 validation was rerun on rebased head `cf0319a3497e252219544c0a8b4168c6ba598f88`. It covers the page-table changes plus the optional-output overlap from #3901 after conflict resolution. `git range-diff`, `git diff --check`, and `pre-commit run --all-files` also passed. The Top-K matrix covers Radix multi-CTA and Filtered dispatch, graph-safe mode, deterministic mode, optional raw output for default and compact page sizes, independent score/page-table starts, compact and default page sizes, padded row strides, misaligned input bases, trivial and selected rows, and CUDA graph replay with mutated inputs. The trace checks cover the operation-local schema and default-argument initialization. A separate smoke check also verified that the Python `page_size <= 2**30` validation matches the native contract. The warnings are existing CUTLASS DSL deprecations from `tests/conftest.py` and `flashinfer/cute_dsl/utils.py`. ## Reviewer Notes Review focus is welcome on the positional pairing of raw and translated outputs across deterministic post-sort, the structural zero-or-one policy ABI, and the page-table transform in the Radix multi-CTA and graph-safe Filtered epilogues. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for compact page tables with configurable page sizes. * Added optional raw-index outputs alongside translated physical indices. * Added reusable output buffers for flexible result storage. * Improved support for empty rows, padding, non-contiguous inputs, and physical-page remapping. * **Bug Fixes** * Added validation for page metadata and output compatibility. * Improved CUDA graph and algorithm-path support. * **Documentation** * Updated API and tracing documentation for page sizes, physical indices, and raw-index outputs. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
Sparse attention has to pick the top-K KV positions for each request on every decode step, and requests in a batch have different sequence lengths. This adds flashinfer.top_k_varlen to do that, with three backends:
radix – CuTe DSL single-pass multi-CTA radix. Handles variable lengths natively (no logit masking) and is Blackwell-only. This is the default when no pre_idx hint is passed.
gvr – Guess-Verify-Refine. Uses the previous step's top-K indices (pre_idx) as a warm start, since a layer's attention pattern barely moves from one step to the next. Blackwell-only, with a load-balance mode for ragged batches.
radix_cutlass – masked fallback on the existing CUTLASS radix kernel. Runs on any GPU.
backend="auto" uses GVR when a pre_idx hint is available, otherwise the CuTe DSL radix on Blackwell, otherwise the CUTLASS fallback.
The GVR and radix kernels are ported from TRT-LLM and tuned to match it. Across the shapes we measured (batch × seq-len × K, bf16 and fp32 on B200) latency is on par with the equivalent TRT-LLM path. Two changes did most of that: SM-aware multi-CTA chunking on the radix path, which also fixes a shared-memory overflow at large N, and feeding the per-CTA scan width into the GVR launch heuristic.
Also in here:
tests/topk/test_topk_varlen.py – covers the three backends, the multi-CTA path, CUDA-graph capture/replay, ragged and seq_len <= top_k cases, and that the backends agree with torch.topk.
benchmarks/routines/topk_varlen.py – compares the backends on a given shape and reference-checks the results.
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
Unit tests — covers all three backends, the multi-CTA path, CUDA-graph capture/replay, ragged and seq_len <= top_k cases, and cross-backend agreement with torch.topk:
pytest tests/topk/test_topk_varlen.py
Trace-template reference correctness:
pytest tests/trace/test_top_k_varlen_reference_correctness.py
Benchmark with a reference check across backends (radix / gvr / radix_cutlass):
python benchmarks/flashinfer_benchmark.py --routine top_k_varlen
--batch_size 16 --max_seq_len 8192 --top_k 512 --refcheck
The GVR and CuTe-DSL radix paths need Blackwell (sm_100+); radix_cutlass runs anywhere.
Reviewer Notes
Summary by CodeRabbit
New Features
top_k_decodeAPI withradix,gvr, orautobackend selection, plusload_balancefor Blackwell-optimized hybrid GVR decoding.return_values,next_n,compress_ratio, and support for preallocated output buffers.top_k_decodeand new GVR tuning configurations (including load-balanced settings).Tests
next_n), compression, preallocation, determinism, fallback behavior, and configuration validation.Benchmarks
top_k_decodeto sampling benchmarks with backend-aware runs and result verification.