feat(topk): support compact page table transforms - #4315
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
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:
📝 WalkthroughWalkthroughThe page-table top-k transform now supports compact pages, physical-slot translation, optional raw-index outputs, reusable buffers, expanded validation, policy-based CUDA dispatch, trace support, and broader regression coverage. ChangesPaged top-k transformation
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant PythonWrapper
participant TopKPageTableTransformDispatch
participant CUDAKernel
participant OutputBuffers
Caller->>PythonWrapper: page_size and optional buffers
PythonWrapper->>TopKPageTableTransformDispatch: validated transform request
TopKPageTableTransformDispatch->>CUDAKernel: policy, stride, and page bits
CUDAKernel->>OutputBuffers: physical indices and optional raw indices
OutputBuffers->>Caller: transformed results
Possibly related PRs
Suggested labels: 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 |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests/utils/test_topk.py (2)
424-424: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the no-op statement.
scores.size(1)has no effect. It is a leftover of a removed local variable.♻️ Proposed cleanup
num_rows = scores.size(0) - scores.size(1) device = scores.device🤖 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/utils/test_topk.py` at line 424, Remove the standalone no-op scores.size(1) statement from the affected test, leaving the surrounding test logic unchanged.
791-803: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a raw-output case for
page_size=1.The matrix exercises
out_raw_indicesonly withpage_size=64. Thepage_size=1layout is the default and legacy path, and it stays untested for raw indices. Add one case to cover it.💚 Proposed parameter addition
pytest.param(1, False, False, id="page_size_1_shared_starts"), pytest.param(1, True, False, id="page_size_1_separate_starts"), + pytest.param(1, True, True, id="page_size_1_raw_output"), pytest.param(64, True, False, id="compact_pages"), pytest.param(64, True, True, id="compact_pages_raw_output"),Note: the raw-index assertions at lines 954-977 read
page_table_row_starts, which is set forseparate_page_table_row_starts=True. The proposed case keeps that assertion valid.🤖 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/utils/test_topk.py` around lines 791 - 803, Add a parametrized test case in the matrix for page_size=1 with separate_page_table_row_starts=True and with_raw_output=True, ensuring the existing raw-index assertions over page_table_row_starts remain valid.csrc/topk.cu (1)
125-129: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAdd an overlap check between
output_page_tableandmaybe_output_raw_indices.The Python wrapper rejects overlapping buffers, but this binding is also reachable from other FFI callers. The kernel writes both buffers, so an overlap corrupts results silently. The buffers are contiguous and have equal shape, so a pointer-range check is cheap.
♻️ Proposed validation
if (maybe_output_raw_indices.has_value()) { CHECK_INPUT_AND_TYPE(maybe_output_raw_indices.value(), dl_int32); CHECK_DIM(2, maybe_output_raw_indices.value()); CHECK_SHAPE(maybe_output_raw_indices.value(), output_page_table); + const char* out_ptr = static_cast<const char*>(output_page_table.data_ptr()); + const char* raw_ptr = + static_cast<const char*>(maybe_output_raw_indices.value().data_ptr()); + const size_t out_bytes = static_cast<size_t>(output_page_table.numel()) * sizeof(int32_t); + TVM_FFI_ICHECK(out_ptr + out_bytes <= raw_ptr || raw_ptr + out_bytes <= out_ptr) + << "output_page_table and output_raw_indices must not overlap"; }🤖 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 `@csrc/topk.cu` around lines 125 - 129, Add a pointer-range overlap validation in the maybe_output_raw_indices validation block, comparing its contiguous storage range with output_page_table before the kernel runs. Reject any overlapping buffers while preserving the existing type, dimension, and shape checks.
🤖 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/trace/templates/sampling.py`:
- Around line 1108-1114: Update the trace setup for the axes mapping near
`top_k_page_table_transform` so the trace namespace explicitly provides the
default `page_size` value, allowing the `Const(abbrev="ps")` axis to produce
`ps1` in definition names. Regenerate the required test JSON using
`top_k_page_table_transform.init(...)`.
---
Nitpick comments:
In `@csrc/topk.cu`:
- Around line 125-129: Add a pointer-range overlap validation in the
maybe_output_raw_indices validation block, comparing its contiguous storage
range with output_page_table before the kernel runs. Reject any overlapping
buffers while preserving the existing type, dimension, and shape checks.
In `@tests/utils/test_topk.py`:
- Line 424: Remove the standalone no-op scores.size(1) statement from the
affected test, leaving the surrounding test logic unchanged.
- Around line 791-803: Add a parametrized test case in the matrix for
page_size=1 with separate_page_table_row_starts=True and with_raw_output=True,
ensuring the existing raw-index assertions over page_table_row_starts remain
valid.
🪄 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 Plus
Run ID: 3c7c13b4-9a10-44a2-bfb8-ed6fdc040771
📥 Commits
Reviewing files that changed from the base of the PR and between d020372 and 3c9627782a2345fddcae4aa51bb7de00ece837c0.
📒 Files selected for processing (8)
csrc/flashinfer_topk_binding.cucsrc/topk.cuflashinfer/topk.pyflashinfer/trace/template.pyflashinfer/trace/templates/sampling.pyinclude/flashinfer/topk.cuhtests/trace/test_fi_trace_template_consistency.pytests/utils/test_topk.py
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
csrc/topk.cu (1)
92-137: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRequire explicit layout contracts before passing raw pointers.
CHECK_INPUT_AND_TYPEandCHECK_LAST_DIM_CONTIGUOUS_INPUTalready require CUDA tensors with contiguous-layout semantics, but this code still passesoutput_page_table,lengths,row_to_batch,row_starts,page_table_row_starts, andoutput_raw_indicesas raw pointers. Assert the expected strided layouts as part of the binding contract, or pass strides to the dispatch path, to prevent a direct FFI caller from using tensors with the accepted shapes but incorrect underlying storage.🤖 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 `@csrc/topk.cu` around lines 92 - 137, Require explicit contiguous or expected-stride layout validation for every tensor later passed as a raw pointer in this binding, including output_page_table, lengths, optional row_to_batch, row_starts, page_table_row_starts, and output_raw_indices. Add the checks alongside the existing CHECK_INPUT_AND_TYPE/CHECK_DIM validation, or propagate each tensor’s strides through the dispatch path; ensure shape-valid tensors with incompatible storage layouts are rejected before pointer use.
🤖 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/trace/templates/__init__.py`:
- Around line 57-62: Update the Const behavior documentation near
build_fi_trace_fn to state that Const() writes an axis value only when tensor or
scalar extraction succeeds; if both return None and default is None, omit
axes[axis]["value"] from the generated JSON, while Const(default=N) still
supplies N.
---
Outside diff comments:
In `@csrc/topk.cu`:
- Around line 92-137: Require explicit contiguous or expected-stride layout
validation for every tensor later passed as a raw pointer in this binding,
including output_page_table, lengths, optional row_to_batch, row_starts,
page_table_row_starts, and output_raw_indices. Add the checks alongside the
existing CHECK_INPUT_AND_TYPE/CHECK_DIM validation, or propagate each tensor’s
strides through the dispatch path; ensure shape-valid tensors with incompatible
storage layouts are rejected before pointer use.
🪄 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 Plus
Run ID: cdf1c194-9586-4c82-81a8-7e66660fdb9f
📥 Commits
Reviewing files that changed from the base of the PR and between 3c9627782a2345fddcae4aa51bb7de00ece837c0 and b8f701b8da5f1c863e356a70b14c074d3691f3ae.
📒 Files selected for processing (7)
csrc/topk.cuflashinfer/trace/template.pyflashinfer/trace/templates/__init__.pyflashinfer/trace/templates/sampling.pytests/trace/test_fi_trace.pytests/trace/test_fi_trace_template_consistency.pytests/utils/test_topk.py
🚧 Files skipped from review as they are similar to previous changes (2)
- flashinfer/trace/templates/sampling.py
- tests/utils/test_topk.py
|
Follow-up on the outside-diff native-layout note: the listed contracts are already enforced at the native FFI boundary. CHECK_INPUT_AND_TYPE expands to CHECK_CUDA, CHECK_CONTIGUOUS (TensorView::IsContiguous), and the dtype check, and radix_topk_page_table_transform applies it to output_page_table, lengths, row_to_batch, row_starts, page_table_row_starts, output_raw_indices, and src_page_table. The score input is intentionally the only padded-row exception: it requires last-dimension stride 1 plus nonoverlapping rows, and its row stride is passed into CUDA. Duplicating explicit stride checks would reject safe singleton layouts without strengthening the pointer contract. Current-head validation after the overload cleanup: editable CUDA 13.2 source build, all 1,463 Top-K tests, and pre-commit passed. |
|
@zianglih, merging your other PR has caused merge conflicts in two files. Can you check and update the PR? |
|
Hi @bkryu , github ci passed. |
|
Hi @zianglih, the main branch and CI opened back to business, so I started merging the long backlog of PR and #3901 resulted in a merge conflict. Would you be kind to resolve the conflict one last time? This PR is now at the front of my tracking list of "PRs to merge". Will keep an eye on it once resolved to ensure it gets in. I appreciate your patience with the PR |
0bd5a75 to
cf0319a
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
csrc/topk.cu (1)
189-193: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the dispatch-policy rationale.
Document why the direct policy remains on the default path and why the configurable policy is required for compact pages, raw indices, padded rows, and unaligned input. State that the direct-policy alternative preserves the existing direct-kernel ABI and code generation.
As per coding guidelines, document the rationale for special algorithmic choices and relevant alternatives in performance-critical hot paths.
Proposed comment
+ // Keep the direct policy for the packed page_size=1 ABI-compatible fast path. + // Use the configurable policy for compact pages, raw outputs, padded rows, + // or inputs that cannot safely use the direct policy's vectorization. const bool use_configurable_page_table_policy =🤖 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 `@csrc/topk.cu` around lines 189 - 193, Document the dispatch-policy rationale adjacent to use_configurable_page_table_policy: explain that the direct policy remains the default because its alternative preserves the existing direct-kernel ABI and code generation, while the configurable policy is required for compact pages, raw indices, padded rows, and unaligned input. Keep the condition and behavior unchanged.Source: Coding guidelines
🤖 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 `@csrc/topk.cu`:
- Around line 135-140: Update the validation around maybe_output_raw_indices in
the top-k path to reject storage overlap with every kernel input and output,
including src_page_table and output_page_table, rather than relying only on
CHECK_SHAPE and CHECK_DEVICE. Reuse the project’s existing storage-overlap
validation utility if available, and preserve the current dtype, device,
dimension, and shape checks.
- Around line 189-193: Update the use_configurable_page_table_policy condition
to check input_stride multiplied by sizeof(dtype), not input_stride alone, when
detecting row-spanning alignment issues. Ensure non-aligned byte strides select
ConfigurablePageTableKernelPolicy so refine_vec_size() accounts for every row
base before LoadToSharedOrdered() vector loads.
---
Nitpick comments:
In `@csrc/topk.cu`:
- Around line 189-193: Document the dispatch-policy rationale adjacent to
use_configurable_page_table_policy: explain that the direct policy remains the
default because its alternative preserves the existing direct-kernel ABI and
code generation, while the configurable policy is required for compact pages,
raw indices, padded rows, and unaligned input. Keep the condition and behavior
unchanged.
🪄 Autofix
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 Plus
Run ID: 28b377be-132c-4f22-9ca9-a5aac5135d5a
📒 Files selected for processing (6)
csrc/flashinfer_topk_binding.cucsrc/topk.cuflashinfer/topk.pyflashinfer/trace/templates/sampling.pyinclude/flashinfer/topk.cuhtests/utils/test_topk.py
🚧 Files skipped from review as they are similar to previous changes (4)
- csrc/flashinfer_topk_binding.cu
- flashinfer/trace/templates/sampling.py
- flashinfer/topk.py
- include/flashinfer/topk.cuh
| if (maybe_output_raw_indices.has_value()) { | ||
| CHECK_INPUT_AND_TYPE(maybe_output_raw_indices.value(), dl_int32); | ||
| CHECK_DEVICE(maybe_output_raw_indices.value(), input); | ||
| CHECK_DIM(2, maybe_output_raw_indices.value()); | ||
| CHECK_SHAPE(maybe_output_raw_indices.value(), output_page_table); | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect public-buffer validation and existing aliasing tests.
rg -n -C6 'output_raw_indices|radix_topk_page_table_transform|data_ptr|overlap|alias' \
flashinfer/topk.py csrc/flashinfer_topk_binding.cu tests/utils/test_topk.pyRepository: flashinfer-ai/flashinfer
Length of output: 19152
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== topk.cu outlines =="
ast-grep outline csrc/topk.cu --view expanded || true
echo "== relevant topk.cu sections =="
wc -l csrc/topk.cu
sed -n '1,220p' csrc/topk.cu
echo "== raw index references in topk.cu =="
rg -n -C4 'output_raw_indices|maybe_output_raw_indices|output_raw|page_table|raw_indices' csrc/topk.cuRepository: flashinfer-ai/flashinfer
Length of output: 17521
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate TopKPageTableTransformDispatch and policies =="
rg -n -C8 'TopKPageTableTransformDispatch|ConfigurablePageTableKernelPolicy|DirectPageTableKernelPolicy|output_raw_indices_ptr|raw_indices_ptr' csrc flashinfer tests -g '*.cu' -g '*.cuh' -g '*.cpp' -g '*.h' -g '*.cuh' -g '*.py'
echo "== source snippets from candidate files =="
for f in $(rg -l 'ConfigurablePageTableKernelPolicy|TopKPageTableTransformDispatch' csrc flashinfer tests -g '*.cu' -g '*.cuh' -g '*.cpp'); do
echo "--- $f ---"
wc -l "$f"
rg -n 'ConfigurablePageTableKernelPolicy|TopKPageTableTransformDispatch|TopKPageTable' "$f" | while IFS=: read -r lineno rest; do
start=$((lineno-80)); [ "$start" -lt 1 ] && start=1; endl=$((lineno+120))
sed -n "${start},${endl}p" "$f"
done
done
echo "== validate_helper/overlap helpers =="
rg -n -C4 'CHECK_INPUT|data_ptr|range|overlap|alias|is_same|storage_offset|stride|size' csrc/tvm_ffi_utils.h csrc/topk.cu flashinfer/topk.pyRepository: flashinfer-ai/flashinfer
Length of output: 50380
Reject overlapping raw-index buffers.
CHECK_SHAPE and CHECK_DEVICE do not prevent output_raw_indices from overlapping other buffers such as src_page_table or output_page_table. If it overlaps src_page_table, kernel stores can modify page-table entries while other threads read them, producing nondeterministic translated indices. Reject storage-range overlap with all kernel inputs and outputs, or explicitly define safe aliasing semantics.
🤖 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 `@csrc/topk.cu` around lines 135 - 140, Update the validation around
maybe_output_raw_indices in the top-k path to reject storage overlap with every
kernel input and output, including src_page_table and output_page_table, rather
than relying only on CHECK_SHAPE and CHECK_DEVICE. Reuse the project’s existing
storage-overlap validation utility if available, and preserve the current dtype,
device, dimension, and shape checks.
| const bool use_configurable_page_table_policy = | ||
| page_bits != 0 || output_raw_indices_ptr != nullptr || | ||
| input_stride != static_cast<int64_t>(max_len) || | ||
| (row_starts_ptr == nullptr && !dsa_graph_safe && | ||
| reinterpret_cast<std::uintptr_t>(input.data_ptr()) % sampling::TOPK_MAX_VECTOR_BYTES != 0); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map policy definitions before inspecting vector-load and alignment behavior.
ast-grep outline include/flashinfer/topk.cuh --items all \
--match 'DirectPageTableKernelPolicy|ConfigurablePageTableKernelPolicy|TopKPageTableTransformDispatch'
rg -n -C8 \
'DirectPageTableKernelPolicy|ConfigurablePageTableKernelPolicy|TOPK_MAX_VECTOR_BYTES|input_stride|VecSize|vector' \
include/flashinfer/topk.cuh
# Locate coverage for non-vector-aligned row strides.
rg -n -C6 'misalign|stride|page_table_transform|page_size' tests/utils/test_topk.pyRepository: flashinfer-ai/flashinfer
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Files under csrc with topk:"
fd -a 'topk|TopK' csrc include tests 2>/dev/null | sed 's#^\./##' | head -80
echo
echo "csrc/topk.cu dispatch lines:"
sed -n '150,230p' csrc/topk.cu
echo
echo "Relevant csrc/topk.cu policy dispatch if present:"
rg -n -C5 'use_configurable_page_table_policy|input\.data_ptr\(\)|input_stride|DirectPageTableKernelPolicy|ConfigurablePageTableKernelPolicy|TopKPageTableTransformDispatch|top_k_page_table_transform' csrc/topk.cu || true
echo
echo "Top-level dispatch in include around TopKPageTableTransformDispatch:"
sed -n '3620,3740p' include/flashinfer/topk.cuh
echo
echo "Multi-CTA vec-size refinement:"
sed -n '2088,2110p' include/flashinfer/topk.cuh
echo
echo "Filtered vec-size dispatch:"
sed -n '3399,3425p' include/flashinfer/topk.cuhRepository: flashinfer-ai/flashinfer
Length of output: 20274
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Vector load definitions/uses in include/flashinfer/topk_common.cuh:"
rg -n -C4 'struct vec_t|cast_load|load|vectorized|__ldg' include/flashinfer/topk_common.cuh include/flashinfer/topk.cuh | head -240
echo
echo "Locate vec_t template and cast_load implementations:"
rg -n 'template .*vec_t|struct vec_t|cast_load|load_data|vectorize' include/flashinfer/topk_common.cuh include/flashinfer/topk.cuhRepository: flashinfer-ai/flashinfer
Length of output: 12352
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Find topk_common include path and vec_t definition in repository files:"
rg -n -C8 'struct vec_t|template <typename DType|alignas|__ldg|__align__' include/flashinfer/topk_common.cuh | head -220
echo
echo "Candidate topk common outline:"
ast-grep outline include/flashinfer/topk_common.cuh --match vec_t --view expanded || true
rg -n 'vec_t|struct vec_t|cast_load' include/flashinfer/topk_common.cuh include/flashinfer/topk.cuhRepository: flashinfer-ai/flashinfer
Length of output: 2200
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Repository-wide vec_t definitions:"
rg -n -C8 'struct vec_t|template .*vec_t|alignas.*vec|__align__.*vec|cast_load' . --glob '*.{cuh,cu,cpp,h,py,hpp}' | head -260
echo
echo "topk_common.cuh beginning:"
sed -n '1,80p' include/flashinfer/topk_common.cuh
echo
echo "topk.cuh vec_t occurrences with surrounding context:"
sed -n '290,340p' include/flashinfer/topk.cuh
sed -n '535,625p' include/flashinfer/topk.cuh
sed -n '2620,2670p' include/flashinfer/topk.cuh
echo
echo "TopK entry wrapper / input_shape validation if present:"
rg -n -C6 'top_k_page_table_transform|input_shape|stride|page_size|input_stride' python flashinfer python include csrc -g '*.py' -g '*.cuh' -g '*.cu' | head -260Repository: flashinfer-ai/flashinfer
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Check whether vec_t has include-only definition in generated build files / cache, excluding source data:"
git ls-files | rg '(^|/)(build|dist|third_party|venv|\.venv|flashinfer/_ffi|tmp_|_ffi|aot_output|torch|cuda)' | head -50 || true
find . -type f \( -name '*topk_common*.cuh' -o -name '*topk*.cuh' -o -name '*vec*.cuh' \) | head -100
echo
echo "topk.cuh includes at top:"
sed -n '1,40p' include/flashinfer/topk.cuh
echo
echo "Search for any cast_load/global load macro definitions in include/topk-common style files without excluding common topk.cuh definitions already scanned:"
rg -n -C12 'template <typename DType|template <uint32_t|struct vec_t|class vec_t|__forceinline__ vec_t|cast_load|__ldg|asm volatile|global_memory|read_via' include/flashinfer/topk_common.cuh include/flashinfer/topk.cuh | head -220
echo
echo "Python entry signature around top_k_page_table_transform:"
fd -a 'topk.py|flashinfer_topk' python flashinfer csrc include 2>/dev/null | head -100
rg -n -C12 'def top_k_page_table_transform|top_k_page_table_transform\(|top_k_page_table_transform_cython|__from_dlpack|input\.stride\(|input_stride' flashinfer csrc include --glob '*.{py,cu,cuh}' | head -200Repository: flashinfer-ai/flashinfer
Length of output: 38360
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "vec_dtypes.cuh outline and vector definitions:"
wc -l include/flashinfer/vec_dtypes.cuh
sed -n '1,220p' include/flashinfer/vec_dtypes.cuh
rg -n -C16 'template <typename T|template <typename T,|struct vec_t|alignas|(using|typedef).*load_value|cast_load|load_value|read_via|global_memory|__ldg|asm volatile' include/flashinfer/vec_dtypes.cuh include/flashinfer/topk.cuh include/flashinfer/topk_common.cuh | sed -n '1,260p'
echo
echo "Python topk dispatch and stride check:"
sed -n '300,530p' flashinfer/topk.py
sg -n 'input\.stride\(|input_stride|top_k_page_table_transform|page_size|DSA|row_starts' flashinfer/topk.py csrc/flashinfer_topk_binding.cu csrc/topk.cu include/flashinfer/*.cuh | head -160 || trueRepository: flashinfer-ai/flashinfer
Length of output: 33422
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Search vec_t definition in vec_dtypes.cuh sections excluding first 220 lines:"
rg -n -C24 'template <typename T, typename|template <typename T|template <typename T,\s*uint32_t|class vec_t|struct vec_t|using vec_t|__ldg|__align__|alignas\(.*vec|cast_load\(|load_value\(|global_memory|read_via' include/flashinfer/vec_dtypes.cuh | sed -n '1,260p'
echo
echo "Show around cast_load definitions in vec_dtypes.cuh:"
rg -n 'cast_load|load_value|global_memory|read_via|ld\.global|ld\.cg|__ldg' include/flashinfer/vec_dtypes.cuh
if [ -n "$(rg -n 'cast_load' include/flashinfer/vec_dtypes.cuh)" ]; then
rg -n -C16 'cast_load' include/flashinfer/vec_dtypes.cuh | sed -n '1,260p'
fi
echo
echo "TopK entry registration call site:"
rg -n -C12 'radix_topk_page_table_transform\(|top_k_page_table_transform|input_stride|input\.stride\(' csrc/flashinfer_topk_binding.cu flashinfer/topk.py python flashinfer csrc --glob '*.{cu,cuh,py}' | sed -n '1,220p'Repository: flashinfer-ai/flashinfer
Length of output: 36727
🌐 Web query:
CUDA unaligned global memory load vector float2 float4 undefined behavior
💡 Result:
In CUDA, accessing global memory with an unaligned address—specifically for vector types like float2 or float4—is classified as undefined behavior [1][2][3]. Hardware requirements dictate that an N-byte data object must be accessed at an address that is a multiple of N bytes (natural alignment) [4][3]. For vector types, this means: - float2 (8 bytes) requires 8-byte alignment [4]. - float4 (16 bytes) requires 16-byte alignment [5][4]. If these alignment requirements are violated, the behavior is explicitly undefined [1][6][3]. In practice, this may result in: - Silent data corruption (e.g., the hardware may mask off lower-order address bits to force alignment) [5][6]. - Incorrect results [5]. - Hardware exceptions or faults, which can crash the kernel or the application [5][6]. Because the GPU hardware does not support unaligned loads and stores, there is no "safe" way to perform an unaligned access using vector types [4][1]. Simply casting a float* pointer to a float4* pointer and dereferencing it will lead to undefined behavior if the base address is not a multiple of 16 [5][4]. To ensure correct and efficient operation, developers must manually manage data alignment or use multiple smaller, naturally aligned loads (e.g., loading individual floats) [5][4]. Tools such as cuda-memcheck can be used to detect these invalid memory access violations during development [2][3].
Citations:
- 1: https://stackoverflow.com/questions/47692070/can-one-make-unaligned-loads-and-stores-in-cuda
- 2: https://forums.developer.nvidia.com/t/question-about-cuda-memcheck-manual/23444
- 3: https://forums.developer.nvidia.com/t/atomicxor-misaligned-memory-access/117236/2
- 4: https://forums.developer.nvidia.com/t/float4-must-read-adjacent-element-can-we-modify-it-for-coalesced-reading/212691/3
- 5: https://stackoverflow.com/questions/16013758/float4-missaligned-access-debug-works-release-dont
- 6: https://forums.developer.nvidia.com/t/alignment-requirements/11062/5
Check input_stride * sizeof(dtype) before selecting the direct policy.
With input_stride = max_len + 1 and max_len odd, a contiguous fp16 row can be aligned at input.data_ptr() but unaligned at each subsequent row. The direct policy still gets a vector size of gcd(8, max_len+1) and then refines only on the first row pointer; LoadToSharedOrdered() casts the row address to a vector load. Route row-spanning non-aligned strides to ConfigurablePageTableKernelPolicy so refine_vec_size() can reduce the vector width for every row base.
🤖 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 `@csrc/topk.cu` around lines 189 - 193, Update the
use_configurable_page_table_policy condition to check input_stride multiplied by
sizeof(dtype), not input_stride alone, when detecting row-spanning alignment
issues. Ensure non-aligned byte strides select ConfigurablePageTableKernelPolicy
so refine_vec_size() accounts for every row base before LoadToSharedOrdered()
vector loads.
|
Hi @bkryu , I have finished rebasing and refreshed the perf numbers. Overhead is still nigligible. Thanks! |
|
/bot run tests/utils |
|
@zianglih thank you the CI result looks good. Merged! |
4 participants
📌 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_transformcontract assumes one page-table entry per score and allocates only the translated output, so SGLang has to split this intotop_k, score re-gathering, compact-page translation, and output copies.This PR extends the existing fused page-table transform for that layout:
page_size, defaulting to1for backward compatibility.outandout_raw_indicesbuffers. Raw indices remain positionally aligned with translated indices, including deterministic post-sort and-1padding; the two buffers must be disjoint.length <= kpath.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:page_table_row_startsis measured in page-table entries, whilerow_startsis measured in score elements. Whenpage_size > 1androw_startsis supplied, callers must therefore providepage_table_row_startsexplicitly 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
DirectPageTableKernelPolicyorConfigurablePageTableKernelPolicyonce, 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 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
29196cf437778906c72630dc5d9850de547501dewith headcf0319a3497e252219544c0a8b4168c6ba598f88on 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.pyis byte-identical on both sides (SHA-25657cd1ca61b38120380cb9ea7cf81ae3ee972484724bb2bb785649ae05cc199d9).The script uses CUPTI (
cupti-python13.2.0 andnvidia-cuda-cupti13.2.75), 10 dry runs, 100 measured iterations, cold L2, and the median. After #4295 it already setsuse_cuda_graph=False, so no temporary benchmark-source edit was needed and the CUPTI plus CUDA graph instability is excluded.Exact PR-body commands:
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 aftertorch.manual_seed(1234)andtorch.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--deterministicon base and head. All comparisons use matched per-case medians; negative PR delta means the head is faster.+0.05%(+0.13%)-0.09%(+0.08%)-0.01%(+0.12%)+0.01%(+0.10%)+0.06%(+0.19%)+0.05%(+0.24%)+0.02%(+0.08%)+0.01%(+0.13%)-0.02%(+0.14%)+0.03%(+0.15%)-0.09%(+0.10%)+0.01%(+0.41%)-0.09%(no regressed point)+0.01%(+0.10%)-0.18%(+0.03%)-0.09%(no regressed point)-0.05%(+0.09%)-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, andtriv%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 usingpage_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:
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
🚀 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
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: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, andpre-commit run --all-filesalso 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**30validation matches the native contract.The warnings are existing CUTLASS DSL deprecations from
tests/conftest.pyandflashinfer/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.
Summary by CodeRabbit