feat(cake_kda): share recurrent prefill kernels across SM100 family - #4313
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
📝 WalkthroughWalkthroughRecurrent KDA prefill now supports SM100a and SM103a. Target selection depends on compute capability and CUDA version. JIT/AOT loading, CUDA bindings, beta packing, tests, documentation, and benchmark metadata were updated. ChangesRecurrent KDA prefill
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant recurrent_kda
participant RecurrentKDAPrefillWorkspace
participant flash_kda
participant FlashKDABinding
Caller->>recurrent_kda: submit prefill tensors
recurrent_kda->>RecurrentKDAPrefillWorkspace: validate and select target
RecurrentKDAPrefillWorkspace->>flash_kda: load variant and target module
flash_kda-->>RecurrentKDAPrefillWorkspace: return compiled module
RecurrentKDAPrefillWorkspace->>FlashKDABinding: launch with beta_tma
FlashKDABinding-->>Caller: return output and final state
Possibly related PRs
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 |
|
What's the difference between this sm103 kernel compared to the sm100 kernel, any B300 specific features being used here? Also, where does the host-side overhead coming from, we should remove that overhead. |
|
/bot run |
|
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: 3
🧹 Nitpick comments (8)
benchmarks/bench_recurrent_kda_prefill.py (1)
168-177: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider detecting untracked files in the provenance check.
--untracked-files=nohides new untracked sources in the checkout. An editable FlashKDA build can compile such files. The reportedsource_committhen does not fully describe the measured peer. If you want a strict provenance claim, drop the flag or add a separate untracked-file check.🤖 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/bench_recurrent_kda_prefill.py` around lines 168 - 177, Update the provenance check around _git_output to detect untracked files as well as tracked modifications by removing "--untracked-files=no" or adding a separate untracked-file validation. Ensure any untracked checkout contents cause the existing RuntimeError, so source_commit fully identifies the measured FlashKDA checkout.csrc/kda/flashkda_bf16_fused_m64_binding.cu (1)
76-80: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCall
PackBetaForTmaIfNeededfor symmetry with the M128 binding.The M128 binding calls
PackBetaForTmaIfNeeded(beta, beta_tma, num_heads, stream)at line 72 ofcsrc/kda/flashkda_bf16_fused_m128_binding.cu. This binding omits the call.The omission is correct today. Line 59 enforces
num_heads == 64, andPackBetaForTmaIfNeededreturns immediately fornum_heads >= kBetaTmaMinHeads. The call is therefore a no-op here.Add the call anyway. It costs nothing at runtime and it keeps the two launch paths identical, so a future change to the head constraint or to the packing rule cannot silently skip packing in the M64 path.
♻️ Proposed change
const TmaPointers tma = EncodeTmaPointers<64>(q, k, v, g, beta_tma, out, descriptor_storage, prepare_descriptors, stream); + // No-op while the M64 variant pins H=64, kept for parity with the M128 path. + PackBetaForTmaIfNeeded(beta, beta_tma, num_heads, stream);🤖 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/kda/flashkda_bf16_fused_m64_binding.cu` around lines 76 - 80, Add a call to PackBetaForTmaIfNeeded(beta, beta_tma, num_heads, stream) in the M64 binding before EncodeTmaPointers, matching the M128 binding’s launch setup. Keep the existing num_heads validation and stream usage unchanged.csrc/kda/flashkda_binding_common.cuh (3)
127-137: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winEvery prefill launch repeats five CUDA driver calls for setup that never changes. The shared root cause is that device properties and the kernel shared-memory attribute are queried and set per launch instead of once per device and once per kernel. The compute capability and the opt-in shared-memory limit are fixed for the process lifetime.
cudaFuncSetAttributeis idempotent for a given kernel and value. The PR reviewer asked to remove host-side overhead, and these calls sit directly on the prefill launch path.
csrc/kda/flashkda_binding_common.cuh#L127-L137: cache the compute capability perdevice_idinCheckExactFlashKDATargetso later launches skip bothcudaDeviceGetAttributecalls.csrc/kda/flashkda_binding_common.cuh#L504-L511: cachecudaDevAttrMaxSharedMemoryPerBlockOptinperdevice_idinCheckDynamicSmemCapacityso later launches skip the query.csrc/kda/flashkda_bf16_fused_m64_binding.cu#L67-L69: runcudaFuncSetAttributeforkernel_flashkda_bf16_fused_m64once per device, for example behind a static per-device flag, instead of on everyRunM64call.csrc/kda/flashkda_bf16_fused_m128_binding.cu#L60-L62: runcudaFuncSetAttributeforkernel_flashkda_bf16_fused_m128once per device, instead of on everyRunM128call.Keep the validation semantics unchanged. A cached value must still produce the same
TVM_FFI_ICHECKfailure message when the device does not match the compiled target.🤖 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/kda/flashkda_binding_common.cuh` around lines 127 - 137, Cache compute capability per device in CheckExactFlashKDATarget while preserving the existing validation and failure message; cache opt-in shared-memory capacity per device in CheckDynamicSmemCapacity. In csrc/kda/flashkda_bf16_fused_m64_binding.cu lines 67-69 and csrc/kda/flashkda_bf16_fused_m128_binding.cu lines 60-62, guard each kernel’s cudaFuncSetAttribute call with a per-device one-time flag so it runs once per device rather than per launch.
352-399: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd the trailing-dimension guards that
EncodeQkTmaalready has.
EncodeValueTmaat line 357 andEncodeGateTmaat line 381 computetensor.numel() / (d1 * d2)without first checkingd1 > 0 && d2 > 0, and without checkingtensor.ndim() >= 2before callingtensor.size(tensor.ndim() - 2).EncodeQkTmaincludes both guards at lines 327 and 330.
CheckCommonInputsvalidates thevandgshapes before these functions run, so no division by zero occurs today. Add the guards so the helpers stay safe if a future caller skipsCheckCommonInputs.🛡️ Proposed guards
template <int ValueRows> inline CUtensorMap EncodeValueTma(const TensorView& tensor) { static_assert(ValueRows == 64 || ValueRows == 128); + TVM_FFI_ICHECK(tensor.ndim() >= 2) << "v must have at least two dimensions"; const int64_t d1 = tensor.size(tensor.ndim() - 1); const int64_t d2 = tensor.size(tensor.ndim() - 2); + TVM_FFI_ICHECK(d1 > 0 && d2 > 0) << "v has invalid trailing dimensions"; const int64_t outer2 = tensor.numel() / (d1 * d2);inline CUtensorMap EncodeGateTma(const TensorView& tensor) { + TVM_FFI_ICHECK(tensor.ndim() >= 2) << "g must have at least two dimensions"; const int64_t d1 = tensor.size(tensor.ndim() - 1); const int64_t d2 = tensor.size(tensor.ndim() - 2); + TVM_FFI_ICHECK(d1 > 0 && d2 > 0) << "g has invalid trailing dimensions"; const int64_t outer2 = tensor.numel() / (d1 * d2);🤖 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/kda/flashkda_binding_common.cuh` around lines 352 - 399, Add the same ndim and trailing-dimension validation used by EncodeQkTma to both EncodeValueTma and EncodeGateTma before accessing d2 or computing outer2: require tensor.ndim() >= 2, then require d1 > 0 and d2 > 0 before dividing. Preserve the existing TMA encoding and shape checks after these guards.
220-225: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReplace the literal
8and32with named constants.Line 220 uses the literal
8whilekBetaTmaMinHeadsholds the same value and is used at line 227. Line 223 uses the literal32for the minimum padded token count.EncodeBetaTmarepeats both literals at line 405, and_beta_tma_sourceinflashinfer/kda_prefill.pyrepeats32at line 384.Use
kBetaTmaMinHeadsat line 220. Add akBetaTmaMinTokens = 32constant and use it at line 223 and at line 405.♻️ Proposed change
- const int64_t beta_tma_heads = std::max<int64_t>(num_heads, 8); + const int64_t beta_tma_heads = std::max<int64_t>(num_heads, kBetaTmaMinHeads); TVM_FFI_ICHECK(beta_tma.ndim() >= 2 && beta_tma.size(beta_tma.ndim() - 1) == beta_tma_heads && beta_tma.numel() % beta_tma_heads == 0 && - beta_tma.numel() / beta_tma_heads >= std::max<int64_t>(token_count, 32)) + beta_tma.numel() / beta_tma_heads >= std::max<int64_t>(token_count, kBetaTmaMinTokens))Add the constant next to
kBetaTmaMinHeads:constexpr int64_t kBetaTmaMinTokens = 32;🤖 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/kda/flashkda_binding_common.cuh` around lines 220 - 225, Replace the literal 8 in the beta_tma_heads calculation and validation with the existing kBetaTmaMinHeads constant. Define kBetaTmaMinTokens as an int64_t constant alongside kBetaTmaMinHeads, then use it for the minimum token validation and the corresponding padding logic in EncodeBetaTma; preserve the existing thresholds and behavior.flashinfer/kda.py (1)
17-25: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe module description is not the module docstring.
The license block at lines 1-15 is the module docstring. The second string literal at lines 17-25 is a bare expression statement.
flashinfer.kda.__doc__andhelp(flashinfer.kda)show only the license text. Merge the description into the first string literal, or convert lines 17-25 into a comment block.The same pattern exists in
flashinfer/kda_prefill.pyat lines 17-24.🤖 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/kda.py` around lines 17 - 25, Merge the KDA facade description into the leading module docstring in kda.py so flashinfer.kda.__doc__ and help() include it, rather than leaving it as a separate string expression. Apply the same correction to the module description in kda_prefill.py, preserving the existing license text and descriptive content.flashinfer/kda_prefill.py (2)
606-627: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
final_state_argis assigned in two places with a duplicated condition.On the path where
initial_state is None,output_final_stateisTrue, andprefill_workspaceis notNone, this block leavesfinal_state_argunbound. Line 668 binds it, guarded by the same condition restated. The code is correct today. If either condition changes, line 710 raisesUnboundLocalError.Bind
final_state_argonce. Set it todummy_stateat line 621 and overwrite it at line 668. Also remove the redundantif initial_state is None:at line 673, because line 667 already establishes it.♻️ Proposed restructure
elif output_final_state: initial_state_arg = dummy_state if prefill_workspace is None: final_state_arg = torch.empty( state_shape, dtype=torch.bfloat16, device=q.device ) returned_state = final_state_arg else: # Assigned to caller-owned stable state scratch under its lock. + final_state_arg = dummy_state returned_state = None store_final_state = Trueif initial_state is None and output_final_state and explicit_workspace: final_state_arg = _state_scratch( workspace=workspace, device=q.device, shape=state_shape, ) - if initial_state is None: - returned_state = final_state_arg + returned_state = final_state_arg🤖 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/kda_prefill.py` around lines 606 - 627, Initialize final_state_arg to dummy_state in the output_final_state branch before the prefill_workspace check, then overwrite it where the caller-owned stable state is assigned. Remove the redundant initial_state is None guard around that later assignment, relying on the existing branch structure to establish the condition and ensure final_state_arg is always bound.
254-267: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winThe global tensor cache and the stream workspace map grow without bound.
_flash_kda_tensor_cacheadds one entry per distinct(kind, device_index, stream_ptr, shape...)tuple and never removes an entry._flash_kda_stream_workspacesadds one_FlashKDAStreamWorkspaceper stream and never removes an entry. A long-running server that uses many batch sizes, sequence lengths, or short-lived CUDA streams accumulates device allocations for the process lifetime. Each workspace also retains beta padding and two 768-byte descriptor buffers.The stream key uses the raw
cuda_streamhandle. CUDA can reuse a handle value after a stream is destroyed, so a later stream can hit an entry created for a destroyed stream.Add a bound or an eviction path. One option is an LRU with a size cap. Another option is to key the stream workspace map by a
weakrefto thetorch.cuda.Streamobject so the entry is dropped when the stream is released.Based on learnings: "ensure that every tensor whose
data_ptr()is included in the key also participates in cache eviction/invalidation... verify the cache invalidation triggers when those key tensors go out of scope."🤖 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/kda_prefill.py` around lines 254 - 267, Bound and add eviction/invalidation for both _flash_kda_tensor_cache and _flash_kda_stream_workspaces so allocations do not grow for the process lifetime. Ensure stream workspaces cannot be reused solely from stale raw cuda_stream handles, and invalidate every cached tensor entry when any tensor whose data_ptr() contributes to its key is released or otherwise becomes invalid. Preserve cache reuse for live resources while removing entries and associated workspace buffers when their key tensors or streams go out of scope.Source: 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 `@benchmarks/bench_recurrent_kda_prefill.py`:
- Around line 238-241: Update the state rotation sizing around _make_state_pool
and state_cursors so state_rotations covers every iteration executed by the full
bench_gpu_time block, including warmup_ms and bench_ms across both timing
backends. Derive the pool size from a measured post-warmup per-call duration, or
increase and document DEFAULT_STATE_ROTATIONS with its memory cost, while
preserving the existing cursor behavior.
In `@csrc/kda/flashkda_binding_common.cuh`:
- Around line 464-494: Add a tensormap release fence at the end of
PublishTensorMaps after all destination writes, using the
fence.proxy.tensormap::generic.release.gpu instruction so descriptor updates
through the generic proxy are visible to TMA consumers. Keep the existing launch
and descriptor preparation flow unchanged.
In `@flashinfer/jit/flash_kda.py`:
- Around line 40-70: Update gen_flash_kda_module to follow the standard JIT
generator flow: compute a unique URI, create its directory under
FLASHINFER_GEN_SRC_DIR, copy the sources from _get_flash_kda_csrc_dir() into
that directory, and pass the copied paths to gen_jit_spec. Preserve the
per-schedule translation-unit behavior and ensure no generated or compiled
artifacts are written to package source directories.
---
Nitpick comments:
In `@benchmarks/bench_recurrent_kda_prefill.py`:
- Around line 168-177: Update the provenance check around _git_output to detect
untracked files as well as tracked modifications by removing
"--untracked-files=no" or adding a separate untracked-file validation. Ensure
any untracked checkout contents cause the existing RuntimeError, so
source_commit fully identifies the measured FlashKDA checkout.
In `@csrc/kda/flashkda_bf16_fused_m64_binding.cu`:
- Around line 76-80: Add a call to PackBetaForTmaIfNeeded(beta, beta_tma,
num_heads, stream) in the M64 binding before EncodeTmaPointers, matching the
M128 binding’s launch setup. Keep the existing num_heads validation and stream
usage unchanged.
In `@csrc/kda/flashkda_binding_common.cuh`:
- Around line 127-137: Cache compute capability per device in
CheckExactFlashKDATarget while preserving the existing validation and failure
message; cache opt-in shared-memory capacity per device in
CheckDynamicSmemCapacity. In csrc/kda/flashkda_bf16_fused_m64_binding.cu lines
67-69 and csrc/kda/flashkda_bf16_fused_m128_binding.cu lines 60-62, guard each
kernel’s cudaFuncSetAttribute call with a per-device one-time flag so it runs
once per device rather than per launch.
- Around line 352-399: Add the same ndim and trailing-dimension validation used
by EncodeQkTma to both EncodeValueTma and EncodeGateTma before accessing d2 or
computing outer2: require tensor.ndim() >= 2, then require d1 > 0 and d2 > 0
before dividing. Preserve the existing TMA encoding and shape checks after these
guards.
- Around line 220-225: Replace the literal 8 in the beta_tma_heads calculation
and validation with the existing kBetaTmaMinHeads constant. Define
kBetaTmaMinTokens as an int64_t constant alongside kBetaTmaMinHeads, then use it
for the minimum token validation and the corresponding padding logic in
EncodeBetaTma; preserve the existing thresholds and behavior.
In `@flashinfer/kda_prefill.py`:
- Around line 606-627: Initialize final_state_arg to dummy_state in the
output_final_state branch before the prefill_workspace check, then overwrite it
where the caller-owned stable state is assigned. Remove the redundant
initial_state is None guard around that later assignment, relying on the
existing branch structure to establish the condition and ensure final_state_arg
is always bound.
- Around line 254-267: Bound and add eviction/invalidation for both
_flash_kda_tensor_cache and _flash_kda_stream_workspaces so allocations do not
grow for the process lifetime. Ensure stream workspaces cannot be reused solely
from stale raw cuda_stream handles, and invalidate every cached tensor entry
when any tensor whose data_ptr() contributes to its key is released or otherwise
becomes invalid. Preserve cache reuse for live resources while removing entries
and associated workspace buffers when their key tensors or streams go out of
scope.
In `@flashinfer/kda.py`:
- Around line 17-25: Merge the KDA facade description into the leading module
docstring in kda.py so flashinfer.kda.__doc__ and help() include it, rather than
leaving it as a separate string expression. Apply the same correction to the
module description in kda_prefill.py, preserving the existing license text and
descriptive content.
🪄 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: 9fe5a4b1-fccd-4c66-88ce-57659c012d10
📒 Files selected for processing (18)
benchmarks/bench_recurrent_kda_prefill.pycsrc/kda/flashkda_bf16_fused_m128.cucsrc/kda/flashkda_bf16_fused_m128_binding.cucsrc/kda/flashkda_bf16_fused_m64.cucsrc/kda/flashkda_bf16_fused_m64_binding.cucsrc/kda/flashkda_binding_common.cuhdocs/api/kda.rstdocs/api/kda_prefill.rstdocs/index.rstflashinfer/__init__.pyflashinfer/aot.pyflashinfer/jit/__init__.pyflashinfer/jit/flash_kda.pyflashinfer/kda.pyflashinfer/kda_kernels/__init__.pyflashinfer/kda_prefill.pytests/jit/test_flash_kda_jit.pytests/kda/test_recurrent_kda_prefill.py
| candidate_state_pool = _make_state_pool(initial_state, state_rotations) | ||
| candidate_output = torch.empty_like(q) | ||
| candidate_workspace = RecurrentKDAPrefillWorkspace(q.device) | ||
| state_cursors = {"pr": [0], "adapted": [0]} |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect bench_gpu_time to confirm how dry_run_time_ms/repeat_time_ms map to iteration counts.
fd -t f 'profiler.py|utils.py|__init__.py' flashinfer/testing | while IFS= read -r f; do
rg -n -C 20 'def bench_gpu_time' "$f"
doneRepository: flashinfer-ai/flashinfer
Length of output: 7041
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== target context =="
sed -n '220,280p' benchmarks/bench_recurrent_kda_prefill.py
echo "== bench implementation =="
sed -n '1546,1660p' flashinfer/testing/profiler.py
echo "== _make_state_pool and cursor use =="
sed -n '198,270p' benchmarks/bench_recurrent_kda_prefill.py
echo "== searches for candidate_run / bench_gpu_time in benchmark =="
rg -n 'candidate_run|state_rotations|bench_gpu_time|dry_run|repeat' benchmarks/bench_recurrent_kda_prefill.py flashinfer/testing/profiler.pyRepository: flashinfer-ai/flashinfer
Length of output: 2064
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== files =="
git ls-files | rg '(^|/)bench_recurrent_kda_prefill\.py$|profiler\.py$|testing'
echo
echo "== target context =="
file_target="$(git ls-files | rg '(^|/)bench_recurrent_kda_prefill\.py$' | head -n1)"
sed -n '180,310p' "$file_target"
echo
echo "== search profiler symbols =="
rg -n -C 12 'def bench_gpu_time|def bench_gpu_time_with_cupti|def bench_gpu_time_with_cugraph|repeat_time_ms|dry_run_time_ms|for iter_idx in range\(repeat_iters\)' .
echo
echo "== _make_state_pool and cursor use complete =="
rg -n -C 15 '_make_state_pool|state_rotations|state_cursors|candidate_run|bench_gpu_time' "$file_target"Repository: flashinfer-ai/flashinfer
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target=$(git ls-files | rg '(^|/)bench_recurrent_kda_prefill\.py$' | head -n1)
profiler_file=$(git ls-files | rg '(^|/)profiler\.py$' | head -n1)
utils_file=$(git ls-files | rg '(^|/)utils\.py$' | head -n1)
echo "== benchmark function invocation =="
sed -n '440,530p' "$target"
echo
echo "== _make_state_pool implementation references =="
rg -n -C 8 '_make_state_pool|state_rotations\s*=\s*|default.*state|--state|state-rotations' benchmarks docs .
echo
echo "== profiler wrapper sequence =="
line=$(rg -n 'def bench_gpu_time\(' "$profiler_file" | head -n1 | cut -d: -f1)
sed -n "${line},$((line+180))p" "$profiler_file"
echo
echo "== profiling backend wrappers relevant sections =="
for name in bench_gpu_time_with_cuda_event bench_gpu_time_with_cupti bench_gpu_time_with_cudagraph; do
line=$(rg -n "def $name" "$profiler_file" | head -n1 | cut -d: -f1)
echo "--- $name: $line"
sed -n "${line},$((line+160))p" "$profiler_file"
done
echo
echo "== argument parser for state rotations =="
rg -n -C 8 'StateRotations|state_rotations|add_argument\(--.*state|add_argument\('--state' "$target"Repository: flashinfer-ai/flashinfer
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target=$(git ls-files | rg '(^|/)bench_recurrent_kda_prefill\.py$' | head -n1)
echo "target=$target"
echo "== benchmark function invocation =="
sed -n '440,530p' "$target"
echo
echo "== _make_state_pool implementation references =="
rg -n -C 8 '_make_state_pool|state_rotations\s*=|default.*state|--state-rotations|add_argument.*state' "$target" docs .
echo
echo "== profiler wrapper sequence =="
profiler_file=$(git ls-files | rg '(^|/)profiler\.py$' | head -n1)
echo "profiler_file=$profiler_file"
line=$(rg -n 'def bench_gpu_time\(' "$profiler_file" | head -n1 | cut -d: -f1)
sed -n "${line},$((line+180))p" "$profiler_file"
echo
echo "== profiling backend wrappers relevant sections =="
for name in bench_gpu_time_with_cuda_event bench_gpu_time_with_cupti bench_gpu_time_with_cudagraph; do
line=$(rg -n "def $name" "$profiler_file" | head -n1 | cut -d: -f1)
echo "--- $name: $line"
sed -n "${line},$((line+160))p" "$profiler_file"
doneRepository: flashinfer-ai/flashinfer
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target=$(git ls-files | rg '(^|/)bench_recurrent_kda_prefill\.py$' | head -n1)
echo "== DEFAULT_STATE_ROTATIONS and constants =="
rg -n -C 4 'DEFAULT_STATE_ROTATIONS|DEFAULT.*ROTATION|class Case|seq_lens|NUM.*HEADS' "$target"
echo
echo "== profiler imports =="
rg -n 'from .*profiler import|import .*profiler|bench_gpu_time_with_cuda_event|bench_gpu_time_with_cupti|bench_gpu_time_with_cudagraph' flashinfer/testing/__init__.py flashinfer/testing/utils.py "$target"Repository: flashinfer-ai/flashinfer
Length of output: 4967
🌐 Web query:
CUDA profiler CUPTI activity synchronization per iteration timing loop bench_gpu_time_with_cupti
💡 Result:
To accurately measure GPU kernel execution time while utilizing CUPTI activity synchronization, it is essential to understand that the CUPTI Activity API operates asynchronously [1][2][3]. CUPTI does not provide direct "per-iteration" timing hooks in the way a simple start/stop timer would; instead, it records activity events—such as kernel executions and synchronization calls—with high-precision hardware timestamps [4][5]. To benchmark or profile a loop involving kernel execution and synchronization, follow these practices: 1. Enable Synchronization Tracking: Synchronization operations (like cudaDeviceSynchronize or stream synchronization) can be tracked by enabling CUPTI_ACTIVITY_KIND_SYNCHRONIZATION [4][6]. You can use cuptiActivityEnableAllSyncRecords(1) to ensure these records are captured [7][8]. 2. Correlation via IDs: CUPTI activity records, such as CUpti_ActivityKernel* and CUpti_ActivitySynchronization*, contain fields (like contextId, streamId, and often unique synchronization or event IDs) that allow you to correlate synchronization events with the specific kernel or memory operation they are waiting for [4][5][9]. 3. Asynchronous Buffer Handling: Because CUPTI is asynchronous, you must register callbacks (using cuptiActivityRegisterCallbacks) to process activity records as they are written to buffers [10][1][2]. Attempting to "synchronize" timing inside the loop by directly reading these buffers will introduce significant overhead and break the asynchronous nature of the profiler [1][2]. Instead, the typical workflow is: - Launch your kernels and perform synchronization in your application loop normally [10]. - Allow CUPTI to asynchronously populate buffers via the registered callback [10][1]. - Parse the buffers after the loop or at a later point to extract the start and end timestamps (in nanoseconds) from the kernel activity records and compare them against the synchronization activity records to derive the effective duration [10][5][9]. 4. Flushing Buffers: To ensure all records are captured for a specific benchmark iteration, use cuptiActivityFlushAll after your timed loop completes [10][1][2]. This forces CUPTI to deliver the remaining records in the buffers so they can be processed and correlated with your iteration data [1][2]. For precise timing, rely on the nanosecond-precision start and end timestamps provided within the activity records themselves rather than relying on host-side timing, which is prone to skew due to the asynchronous nature of kernel launches [5][9].
Citations:
- 1: https://docs.nvidia.com/cupti/main/main.html
- 2: https://docs.nvidia.com/cupti/13.1.1/main/main.html
- 3: https://developer.download.nvidia.com/compute/DevZone/docs/html/C/doc/CUPTI_Users_Guide.pdf
- 4: https://docs.nvidia.com/cupti/api/structCUpti__ActivitySynchronization2.html
- 5: https://docs.nvidia.com/cupti/api/structCUpti__ActivityKernel12.html
- 6: https://docs.nvidia.com/cupti/api/structCUpti__ActivitySynchronization.html
- 7: https://docs.nvidia.com/cupti/api/group__CUPTI__ACTIVITY__API.html
- 8: https://docs.nvidia.com/cupti/13.3.1/api/group__CUPTI__ACTIVITY__API.html
- 9: https://docs.nvidia.com/cupti/13.1.1/api/structCUpti__ActivityKernel10.html
- 10: https://docs.nvidia.com/cupti/tutorial/tutorial.html
Make state_rotations account for the full timed block.
state_cursors["pr"][0] += 1 once per bench_gpu_time() call, and the default call covers warmup_ms + bench_ms via both CUPTI and CUDA graph backends. With DEFAULT_STATE_ROTATIONS = 512, a 120 ms block fails with PR state rotations exhausted if iterations average under ~235 us. Size the pool from an measured per-call time after warmup, or document a sufficient default and its memory cost.
| static __global__ void PublishTensorMaps(uint64_t* destination, TensorMapWords source) { | ||
| const uint32_t index = threadIdx.x; | ||
| if (index < TensorMapWords::kWordCount) { | ||
| destination[index] = source.words[index]; | ||
| } | ||
| } | ||
|
|
||
| template <int ValueRows> | ||
| inline TmaPointers EncodeTmaPointers(const TensorView& q, const TensorView& k, const TensorView& v, | ||
| const TensorView& g, const TensorView& beta_tma, | ||
| const TensorView& out, const TensorView& descriptor_storage, | ||
| int64_t prepare_descriptors, cudaStream_t stream) { | ||
| if (prepare_descriptors != 0) { | ||
| cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; | ||
| CheckCuda(cudaStreamIsCapturing(stream, &capture_status), | ||
| "cudaStreamIsCapturing(TMA descriptor preparation)"); | ||
| TVM_FFI_ICHECK(capture_status == cudaStreamCaptureStatusNone) | ||
| << "prepare_descriptors must be 0 during CUDA graph capture; warm " | ||
| "this exact workspace and tensor signature before capture"; | ||
|
|
||
| const std::array<CUtensorMap, kTensorMapCount> host_maps = { | ||
| EncodeQkTma(q, "q"), EncodeQkTma(k, "k"), EncodeValueTma<ValueRows>(v), | ||
| EncodeGateTma(g), EncodeBetaTma(beta_tma), EncodeOutputTma<ValueRows>(out), | ||
| }; | ||
| static_assert(sizeof(host_maps) == kDescriptorStorageBytes); | ||
| TensorMapWords words{}; | ||
| std::memcpy(words.words, host_maps.data(), sizeof(host_maps)); | ||
| PublishTensorMaps<<<1, 128, 0, stream>>>( | ||
| reinterpret_cast<uint64_t*>(descriptor_storage.data_ptr()), words); | ||
| CheckCuda(cudaGetLastError(), "PublishTensorMaps launch"); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
PTX fence.proxy.tensormap generic release required device-side tensormap update kernel boundary
💡 Result:
The PTX instruction fence.proxy.tensormap::generic.release is a uni-directional fence used to establish ordering between memory accesses performed via different proxies [1][2]. Specifically, it acts as a release fence that synchronizes with an acquire sequence containing the fence.proxy.tensormap::generic.acquire instruction [1][2]. These instructions are primarily used on NVIDIA GPU architectures with compute capability 9.0 or higher (Hopper and later) to manage the coherence of Tensor Map (TMA) descriptors [3][4]. When a tensor map object is updated—such as when modifying its base address, dimensions, or strides in global or shared memory—these fence operations are required to ensure that the modifications (performed via the generic proxy) are visible to the hardware TMA engine (which reads through the tensormap proxy) [5][6]. Key technical aspects include: - Functionality: The release fence establishes a release sequence, and the acquire fence consumes it, ensuring that memory updates become visible to the subsequent TMA operations [1][2]. - Scope: Both release and acquire fences support various scopes (e.g.,.cta,.cluster,.gpu,.sys) to define the extent of the visibility [7][8]. - Implementation: In low-level development, these are often utilized alongside tensormap.cp_fenceproxy, which performs a fused copy and fence operation when updating TMA descriptors from shared memory to global memory [9][3][8]. - Usage Context: Failure to correctly use these fence sequences—or choosing an insufficient scope—can lead to stale data being consumed by the TMA engine, as the hardware may otherwise rely on cached or inconsistent descriptor states [6]. These PTX instructions were introduced starting in PTX ISA 8.3 [7][8]. Developers typically access them via specialized intrinsics provided in the CUDA C++ Core Compute Libraries (CCCL) or through inline assembly wrappers in libraries like CUTLASS [7][3][8].
Citations:
- 1: [LLVM][NVPTX] Add NVPTX codegen support for fence.proxy.tensormap llvm/llvm-project#100748
- 2: https://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20240805/1421919.html
- 3: https://github.com/NVIDIA/cutlass/blob/0ce648f5/include/cute/arch/copy_sm90_desc.hpp
- 4: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__TENSOR__MEMORY.html
- 5: https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/async-copies.html
- 6: https://gh.evko.io/nvopen-tools/tileiras/lowering/nvgpu-and-gpu-to-nvvm.html
- 7: https://nvidia.github.io/cccl/unstable/libcudacxx/ptx/instructions/fence.html
- 8: https://nvidia.github.io/cccl/unstable/libcudacxx/ptx/instructions/tensormap_cp_fenceproxy.html
- 9: [LLVM][NVPTX] Add support for tensormap.cp_fenceproxy llvm/llvm-project#107555
Add the tensormap release fence after updating the descriptors.
PublishTensorMaps writes CUtensorMap data through the generic proxy, while the TMA kernels acquire through the tensormap proxy. The kernel boundary does not establish this cross-proxy ordering, so the consumer can acquire a stale descriptor state. Add fence.proxy.tensormap::generic.release.gpu after the global-store sequence.
🤖 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/kda/flashkda_binding_common.cuh` around lines 464 - 494, Add a tensormap
release fence at the end of PublishTensorMaps after all destination writes,
using the fence.proxy.tensormap::generic.release.gpu instruction so descriptor
updates through the generic proxy are visible to TMA consumers. Keep the
existing launch and descriptor preparation flow unchanged.
| def _get_flash_kda_csrc_dir() -> Path: | ||
| """Locate frozen FlashKDA sources in installed and source checkouts.""" | ||
|
|
||
| installed = jit_env.FLASHINFER_CSRC_DIR / "kda" | ||
| if installed.exists(): | ||
| return installed | ||
|
|
||
| checkout = Path(__file__).resolve().parents[2] / "csrc" / "kda" | ||
| if checkout.exists(): | ||
| return checkout | ||
|
|
||
| raise FileNotFoundError( | ||
| "FlashKDA CUDA sources were not found. Checked:\n" | ||
| f" - {installed}\n" | ||
| f" - {checkout}" | ||
| ) | ||
|
|
||
|
|
||
| def _get_flash_kda_include_dir() -> Path: | ||
| """Locate FlashInfer headers in installed and source checkouts.""" | ||
|
|
||
| if jit_env.FLASHINFER_INCLUDE_DIR.exists(): | ||
| return jit_env.FLASHINFER_INCLUDE_DIR | ||
| checkout = Path(__file__).resolve().parents[2] / "include" | ||
| if checkout.exists(): | ||
| return checkout | ||
| raise FileNotFoundError( | ||
| "FlashInfer headers were not found. Checked:\n" | ||
| f" - {jit_env.FLASHINFER_INCLUDE_DIR}\n" | ||
| f" - {checkout}" | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Sources are compiled directly from csrc_dir instead of being copied to FLASHINFER_GEN_SRC_DIR.
gen_flash_kda_module passes binding (a path inside _get_flash_kda_csrc_dir(), which can resolve to the installed, potentially read-only FLASHINFER_CSRC_DIR) straight into gen_jit_spec(sources=[binding], ...). Other JIT generators in this repository compute a unique URI, create a directory under FLASHINFER_GEN_SRC_DIR, copy the csrc sources there, and pass the copies as sources, even when no Jinja rendering is needed.
The docstring explains why each schedule compiles in its own translation unit, but it does not explain this departure from the standard copy-to-FLASHINFER_GEN_SRC_DIR step. Align with the established pattern, or add an explicit comment documenting why this frozen-source case intentionally reads sources in place.
♻️ Proposed alignment with the established JIT generator pattern
+import shutil
+
`@functools.cache`
def gen_flash_kda_module(
variant: FlashKDAVariant, arch: FlashKDAArch = "sm100a"
) -> JitSpec:
...
csrc_dir = _get_flash_kda_csrc_dir()
include_dir = _get_flash_kda_include_dir()
uri = get_flash_kda_uri(variant, arch)
- binding = csrc_dir / f"flashkda_bf16_fused_{variant}_binding.cu"
- if not binding.exists():
- raise FileNotFoundError(f"FlashKDA binding source not found: {binding}")
+ binding_name = f"flashkda_bf16_fused_{variant}_binding.cu"
+ binding_src = csrc_dir / binding_name
+ if not binding_src.exists():
+ raise FileNotFoundError(f"FlashKDA binding source not found: {binding_src}")
+ gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri
+ gen_directory.mkdir(parents=True, exist_ok=True)
+ binding = gen_directory / binding_name
+ shutil.copy(binding_src, binding)As per path instructions: "JIT generators should compute a unique URI, generate files under FLASHINFER_GEN_SRC_DIR, copy source files there, and return a JitSpec; do not write generated or compiled artifacts into package source directories."
Also applies to: 83-118
🤖 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/jit/flash_kda.py` around lines 40 - 70, Update
gen_flash_kda_module to follow the standard JIT generator flow: compute a unique
URI, create its directory under FLASHINFER_GEN_SRC_DIR, copy the sources from
_get_flash_kda_csrc_dir() into that directory, and pass the copied paths to
gen_jit_spec. Preserve the per-schedule translation-unit behavior and ensure no
generated or compiled artifacts are written to package source directories.
Source: Coding guidelines
|
/bot run |
|
/bot run tests/kda |
Compile the frozen M64 and M128 prefill kernels once as sm_100f for CUDA 12.9+ and route both CC 10.0 and CC 10.3 through the shared modules. Preserve the exact sm_100a CUDA 12.8 compatibility path.
aafaad1 to
8d3bbb2
Compare
yyihuang
changed the title
|
/bot run tests/kda |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/kda_prefill.py`:
- Around line 167-172: Update _flash_kda_prefill_is_eligible to reject SM103a
devices when _is_cuda_version_at_least("12.9") is false, matching the
CUDA-version requirement in _select_flash_kda_prefill_target. Preserve
eligibility for other supported compute capabilities and ensure this combination
falls back instead of reaching _run_flash_kda_prefill and raising.
🪄 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: 5c8dd0cf-bd3b-41d7-8940-f74a82ae368b
📒 Files selected for processing (11)
benchmarks/bench_recurrent_kda_prefill.pycsrc/kda/flashkda_bf16_fused_m128_binding.cucsrc/kda/flashkda_bf16_fused_m64_binding.cucsrc/kda/flashkda_binding_common.cuhdocs/api/kda_prefill.rstflashinfer/aot.pyflashinfer/jit/flash_kda.pyflashinfer/kda.pyflashinfer/kda_prefill.pytests/jit/test_flash_kda_jit.pytests/kda/test_recurrent_kda_prefill.py
🚧 Files skipped from review as they are similar to previous changes (6)
- flashinfer/kda.py
- docs/api/kda_prefill.rst
- csrc/kda/flashkda_bf16_fused_m64_binding.cu
- csrc/kda/flashkda_binding_common.cuh
- csrc/kda/flashkda_bf16_fused_m128_binding.cu
- tests/kda/test_recurrent_kda_prefill.py
| if ( | ||
| not q.is_cuda | ||
| or get_compute_capability(q.device) != _FLASH_KDA_B200_COMPUTE_CAPABILITY | ||
| or get_compute_capability(q.device) | ||
| not in _FLASH_KDA_SUPPORTED_COMPUTE_CAPABILITIES | ||
| ): | ||
| return False |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Eligibility check does not require CUDA 12.9 for SM103a, so unsupported calls crash instead of falling back.
_flash_kda_prefill_is_eligible only checks that compute capability is in _FLASH_KDA_SUPPORTED_COMPUTE_CAPABILITIES (line 170). It does not check the CUDA toolkit version. For compute capability (10, 3), _select_flash_kda_prefill_target requires CUDA 12.9 or newer (lines 542-546); otherwise it raises RuntimeError.
If a call is on an SM103a device with an older CUDA toolkit available for JIT (_is_cuda_version_at_least("12.9") is False), _flash_kda_prefill_is_eligible still returns True, so recurrent_kda routes into _run_flash_kda_prefill. That function then raises RuntimeError instead of falling back to the CuTe DSL backend. This breaks the PR's stated goal of retaining existing fallback behavior for this specific combination.
Add the same CUDA-version condition to the eligibility check so unsupported combinations fall back instead of raising.
🐛 Proposed fix to align eligibility with target-selection requirements
if (
not q.is_cuda
- or get_compute_capability(q.device)
- not in _FLASH_KDA_SUPPORTED_COMPUTE_CAPABILITIES
):
return False
+ compute_capability = get_compute_capability(q.device)
+ if compute_capability not in _FLASH_KDA_SUPPORTED_COMPUTE_CAPABILITIES:
+ return False
+ if compute_capability == (10, 3) and not _is_cuda_version_at_least("12.9"):
+ # SM103a only has the sm_100f family target, which requires CUDA
+ # 12.9+. Without this check, an eligible-but-unbuildable call falls
+ # through to `_run_flash_kda_prefill`, which raises instead of
+ # falling back to the CuTe DSL backend.
+ return FalseAlso applies to: 530-550
🤖 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/kda_prefill.py` around lines 167 - 172, Update
_flash_kda_prefill_is_eligible to reject SM103a devices when
_is_cuda_version_at_least("12.9") is false, matching the CUDA-version
requirement in _select_flash_kda_prefill_target. Preserve eligibility for other
supported compute capabilities and ensure this combination falls back instead of
reaching _run_flash_kda_prefill and raising.
Follow-up to merged #4262 and #4313. ## What changed - pad the beta-only TMA source to `round_up(H, 8)` instead of only padding `H < 8`; - teach the beta pack kernel and binding validation to use the dynamic padded-head stride; - retain the caller-visible `H`, state shape, launch grid, and frozen M64/M128 CUDA bodies unchanged; - add H=12 eager, packed, full-chunk-plus-tail, final-state, and CUDA graph replay coverage. The frozen kernels load beta in 8-head TMA boxes. For H=12, the original BF16 row stride is 24 bytes and `cuTensorMapEncodeTiled` rejects it. Padding the descriptor source to 16 heads gives a 32-byte row stride; heads 12–15 are padding only and are never assigned CTAs. This generalizes the fix to every positive head count that is not divisible by eight, while aligned head counts retain the existing zero-copy beta path. ## Correctness Both runs used the public `flashinfer.recurrent_kda` facade and compared BF16 output plus the complete final state against the PyTorch reference with `atol=rtol=1e-2`. | GPU | Compute capability | CUDA | Result | | --- | --- | --- | --- | | NVIDIA B200 | 10.0 | 12.9 | `62 passed, 0 skipped, 0 failed` | | NVIDIA GB300 | 10.3 | 12.9 | `62 passed, 0 skipped, 0 failed` | The test gate runs: ```text tests/jit/test_flash_kda_jit.py tests/kda/test_recurrent_kda_prefill.py ``` H=12 coverage includes fixed T=32, fixed T=33 (one full TMA chunk plus the direct-load tail), packed sequence lengths `[32, 3]`, in-place initial/final state, and CUDA graph replay after beta is changed. The JIT contract test also verifies that the frozen generated M64/M128 bodies remain unchanged. `pre-commit run --files <changed files>` passes. ## Performance The H=12 path was benchmarked through the public `flashinfer.recurrent_kda` facade with fallback forbidden. Speedup is the official FlashKDA raw GPU span divided by this PR's public-API GPU span. The baseline is the same official FlashKDA source used for #4262: [`MoonshotAI/FlashKDA@d2ff19a`](MoonshotAI/FlashKDA@d2ff19a), with CUTLASS `5c149f5`. Measurements use strict CUPTI first-to-last correlated compute-kernel span, cold L2, no CUDA Graph, and two independent 128-sample blocks in symmetric ABCCBA order. The PR span includes both the beta pack and frozen M128 recurrence kernels. All six benchmark shapes passed output and complete-final-state correctness against the official peer with BF16 `atol=rtol=1e-2`. | H=12 shape | SM100 / B200: PR / baseline | Speedup | SM103 / GB300: PR / baseline | Speedup | | --- | ---: | ---: | ---: | ---: | | packed `[512] x 32` | 136.159 / 240.239 us | **1.7644x** | 128.264 / 233.496 us | **1.8204x** | | packed `[128] x 8` | 23.760 / 46.448 us | **1.9549x** | 25.712 / 52.904 us | **2.0576x** | | fixed `[512]` | 46.184 / 76.383 us | **1.6539x** | 47.712 / 82.240 us | **1.7237x** | | fixed `[8192]` | 514.197 / 814.435 us | **1.5839x** | 487.161 / 779.426 us | **1.5999x** | | mixed `[1300, 547, 2048, 963, 271, 3063]` | 208.647 / 351.550 us | **1.6849x** | 198.216 / 340.913 us | **1.7199x** | | uniform `[1024] x 8` | 82.080 / 162.703 us | **1.9823x** | 78.440 / 162.505 us | **2.0717x** | | **Six-shape geometric mean** | | **1.7645x** | | **1.8238x** | | Comparison | Result | | --- | --- | | FlashInfer upstream main at H=12 | `unsupported / N/A` | Upstream main fails before kernel launch with `cuTensorMapEncodeTiled failed for beta_tma with CUresult=1`, so there is no valid upstream H=12 timing or speedup claim. This PR leaves the frozen compute kernel unchanged; it adds the required beta packing only for non-8-aligned head counts. Related to #4254. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved beta padding for head counts that are not divisible by eight. - Ensured packed inputs correctly handle larger, non-aligned head counts. - Added validation for padded storage requirements. - **Documentation** - Clarified beta padding behavior and public tensor shapes. - **Tests** - Expanded coverage for 12-head inputs, varied sequence lengths, chunk boundaries, packed inputs, and CUDA graph updates. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Yingyi Huang <averyh@nvidia.com>
5 participants
Important
Follow-up to merged #4262. This PR contains only SM100-family packaging,
routing, validation, and GB300 enablement; it does not add another frozen
prefill kernel body.
Description
Compile the existing frozen recurrent-KDA prefill M64 and M128 kernels once as
sm_100fon CUDA 12.9 and newer, and use those modules on both CC 10.0(B200/GB200) and CC 10.3 (B300/GB300). CUDA 12.8 B200 keeps the exact
sm_100acompatibility modules because that toolkit predatessm_100f.The M64 and M128 generated CUDA bodies remain byte-identical to #4262. JIT and
AOT use target-bearing module identities, and the binding accepts the family
module only on CC 10.0 or CC 10.3.
Validation
Publication commit:
8d3bbb27e26795609638f8002d345917c6d802c5(treec859b6d5241b5994cc9ccf4ceb5faacce15f4057), based directly on upstreammain@4433996eafter #4279 merged. The incremental prefill implementation isidentical to the candidate validated on both GPUs: the ten non-AOT files are
blob-identical, and their stable patch-id is unchanged. The only replay
resolution was mechanical in
flashinfer/aot.py: it preserves the merged#4279 decode registration while applying the same prefill family registration.
Strict CUPTI, cold-L2, balanced-process A/B below compares the family target
with the same frozen bodies compiled for the exact device target. Speedup is
exact / sm100f; values near 1 mean target consolidation is performanceneutral.
The only stable tiny-shape difference is
h6_packed_n8_t128:sm100f25.664 us versus exact
sm103a25.360 us (0.9882x, 0.304 us absolute).Keeping an extra exact M128 target for this one coordinate would reintroduce a
third target and shape-specific routing, so this PR retains the simpler shared
family module.
The existing six-shape comparison with pinned official FlashKDA remains the
performance baseline documented by this PR; the new A/B isolates only the
physical target change.
Related to #4254.
Summary by CodeRabbit
New Features
Bug Fixes
Documentation