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:
📝 WalkthroughWalkthroughAdds Blackwell-specific CUDA softmax kernels, deterministic route selection, PDL-aware launches, capability-gated JIT and API dispatch, fallback handling, input validation, and CUDA route and numerical tests. ChangesBlackwell softmax
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant SoftmaxAPI
participant blackwell_softmax
participant BlackwellKernel
participant OnlineSoftmax
Caller->>SoftmaxAPI: request softmax
SoftmaxAPI->>blackwell_softmax: dispatch supported Blackwell input
blackwell_softmax->>BlackwellKernel: select and launch route
BlackwellKernel-->>blackwell_softmax: write output
blackwell_softmax-->>SoftmaxAPI: return status
blackwell_softmax->>OnlineSoftmax: fallback for unsupported route
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (5)
flashinfer/sampling.py (1)
805-807: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring is now out of sync with the dispatch behavior.
On SM100/SM103 this API routes to the Blackwell kernels, which ignore
enable_pdl(seecsrc/blackwell_softmax.cu:171-179), yet the docstring still promises automatic PDL for compute capability ≥ 9.0 and describes only the online-softmax kernel. Add a short note about the Blackwell route and theenable_pdlcaveat.As per coding guidelines, "Keep documentation synchronized with code changes, including infrastructure examples, skill files, conventions, deprecated approaches, and new error-handling patterns."
🤖 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/sampling.py` around lines 805 - 807, Update the enable_pdl parameter docstring in the sampling API to mention that SM100/SM103 dispatches to Blackwell kernels, where enable_pdl is ignored, and clarify that the automatic PDL behavior applies only to supported non-Blackwell online-softmax kernels.Source: Coding guidelines
tests/utils/test_sampling.py (2)
91-93: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the
flashinfer.utilscapability helper for the arch skip.
torch.cuda.get_device_capability() not in ((10, 0), (10, 3))duplicates the gate that already lives inflashinfer/sampling.py:113-115and skips the CUDA-version part of the check.As per coding guidelines, "Skip tests on unsupported CUDA architectures using the appropriate
flashinfer.utilschecks or backend capability APIs, rather than running unsupported kernels." Based on learnings,flashinfer.utils.is_sm100a_supported(device)already covers SM103 (compute capability 10.3) in addition to SM100.♻️ Proposed change
- if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): - pytest.skip("Loom Softmax routes require SM100 or SM103") + device = torch.device("cuda:0") + if not is_sm100a_supported(device): + pytest.skip("Blackwell Softmax routes require SM100 or SM103")Add the import alongside the existing test imports:
from flashinfer.utils import is_sm100a_supported🤖 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_sampling.py` around lines 91 - 93, Update test_softmax_blackwell_routes to use flashinfer.utils.is_sm100a_supported for the architecture and CUDA-version capability check, replacing the direct torch.cuda.get_device_capability comparison. Add the helper import with the existing test imports and preserve the current skip behavior for unsupported devices.Sources: Coding guidelines, Learnings
79-110: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffConsider asserting the route and trimming the largest case.
Two optional improvements: (1) the test name promises route coverage but a silent
cudaErrorNotSupportedfallback toOnlineSoftmaxwould still pass every assertion — spying onget_blackwell_softmax_op(or asserting on a route counter) would make the dispatcher itself the thing under test; (2)(989, 128256)allocates ~1.5 GB of fp32 acrosslogits,probs, andprobs_refplustorch.softmaxtemporaries, which is the one parametrization at real OOM risk on smaller SM103 parts.🤖 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_sampling.py` around lines 79 - 110, Strengthen test_softmax_blackwell_routes by asserting that get_blackwell_softmax_op is invoked, so unsupported-route fallback cannot satisfy the test. Reduce or remove the (989, 128256, "per_row") parametrization while retaining representative coverage of the intended Blackwell routes and avoiding excessive memory use.Source: Coding guidelines
csrc/blackwell_softmax.cu (1)
44-64: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRecord the provenance of these dispatch thresholds.
rows <= 128 && vocab_size <= 257,vocab_size >= 24576 && vocab_size <= 256000,measured_large_odd, etc. are sweep-derived constants with no note on how they were obtained or when they should be re-tuned. A one-line comment citing the benchmark/checkpoint (as the generated payload files do) makes future re-tuning safe.🤖 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/blackwell_softmax.cu` around lines 44 - 64, Add concise comments in use_warp_kernel and use_rowwise_kernel documenting that the dispatch thresholds, including the measured_large_odd range, were derived from the relevant benchmark/checkpoint sweep and should be re-tuned when that data changes. Follow the provenance-reference style used by generated payload files.csrc/blackwell_softmax_warp.cu (1)
72-88: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the implicit 288-column capacity of this kernel.
The register tile (
row_values[9]× 32 lanes) caps this kernel at 288 columns, but the enforcing predicate lives incsrc/blackwell_softmax.cu:44-48(vocab_size <= 257). If that gate is ever widened, columns ≥288 are silently never written and the output buffer keeps uninitialized data rather than failing. A short comment next to the tile (and mirroring it inuse_warp_kernel) makes the coupling explicit without touching the generated payload 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/blackwell_softmax_warp.cu` around lines 72 - 88, The warp kernel’s 9-element-per-lane register tile supports only 288 columns, but this capacity is implicit. Add concise comments beside row_values[9] and in use_warp_kernel documenting the 288-column limit and its coupling to the existing vocab_size <= 257 gate, without changing generated payload behavior.
🤖 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/blackwell_softmax.cu`:
- Around line 171-179: Add a brief comment near the Blackwell fast-path
invocation explaining that enable_pdl is intentionally not applied there and
only affects the OnlineSoftmax fallback. Update the softmax documentation in
sampling.py to note this Blackwell-specific behavior, preserving the existing
fallback handling.
- Line 1: Run the repository clang-format hook on csrc/blackwell_softmax.cu and
include/flashinfer/blackwell_softmax.cuh, then commit the formatter’s changes:
update boolean-expression parenthesization in
use_warp_kernel/use_rowwise_kernel, wrap cudaLaunchCooperativeKernel arguments,
and align the kernel signature declarations in the extern "C" block.
- Around line 151-168: Validate inputs in the softmax entry point before
deriving pointers or launching kernels: when maybe_temperature_arr is present,
require a one-dimensional tensor with exactly logits.size(0) elements, matching
logits dtype and device; also require output to match logits shape, dtype, and
device. Reuse the existing CHECK_INPUT/CHECK_DIM validation conventions and
reject invalid inputs before any per-row parameter indexing.
In `@flashinfer/jit/blackwell_softmax.py`:
- Around line 22-34: Update the JIT specification returned by the blackwell
softmax setup to enable CUDA relocatable device code/device linking for the
cross-translation-unit kernel launch. Extend the existing extra CUDA flags in
the relevant function around gen_jit_spec and preserve the current source list
and compiler-version flags.
In `@flashinfer/sampling.py`:
- Around line 842-853: Update the Blackwell dispatch in the sampling path to
require logits.is_cuda before resolving a CUDA device index or calling
_supports_blackwell_softmax. Keep CPU tensors on the existing fallback path,
avoiding torch.cuda.current_device() and the Blackwell custom op for non-CUDA
inputs.
---
Nitpick comments:
In `@csrc/blackwell_softmax_warp.cu`:
- Around line 72-88: The warp kernel’s 9-element-per-lane register tile supports
only 288 columns, but this capacity is implicit. Add concise comments beside
row_values[9] and in use_warp_kernel documenting the 288-column limit and its
coupling to the existing vocab_size <= 257 gate, without changing generated
payload behavior.
In `@csrc/blackwell_softmax.cu`:
- Around line 44-64: Add concise comments in use_warp_kernel and
use_rowwise_kernel documenting that the dispatch thresholds, including the
measured_large_odd range, were derived from the relevant benchmark/checkpoint
sweep and should be re-tuned when that data changes. Follow the
provenance-reference style used by generated payload files.
In `@flashinfer/sampling.py`:
- Around line 805-807: Update the enable_pdl parameter docstring in the sampling
API to mention that SM100/SM103 dispatches to Blackwell kernels, where
enable_pdl is ignored, and clarify that the automatic PDL behavior applies only
to supported non-Blackwell online-softmax kernels.
In `@tests/utils/test_sampling.py`:
- Around line 91-93: Update test_softmax_blackwell_routes to use
flashinfer.utils.is_sm100a_supported for the architecture and CUDA-version
capability check, replacing the direct torch.cuda.get_device_capability
comparison. Add the helper import with the existing test imports and preserve
the current skip behavior for unsupported devices.
- Around line 79-110: Strengthen test_softmax_blackwell_routes by asserting that
get_blackwell_softmax_op is invoked, so unsupported-route fallback cannot
satisfy the test. Reduce or remove the (989, 128256, "per_row") parametrization
while retaining representative coverage of the intended Blackwell routes and
avoiding excessive memory 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: 2f0d159a-b8f5-4a76-8134-f89b16627dd3
📒 Files selected for processing (9)
csrc/blackwell_softmax.cucsrc/blackwell_softmax_bootstrap.cucsrc/blackwell_softmax_rowwise.cucsrc/blackwell_softmax_warp.cuflashinfer/aot.pyflashinfer/jit/blackwell_softmax.pyflashinfer/sampling.pyinclude/flashinfer/blackwell_softmax.cuhtests/utils/test_sampling.py
| nvcc_flags = current_compilation_context.get_nvcc_flags_list( | ||
| supported_major_versions=[10] | ||
| ) | ||
| return gen_jit_spec( | ||
| "blackwell_softmax", | ||
| [ | ||
| jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax.cu", | ||
| jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax_bootstrap.cu", | ||
| jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax_rowwise.cu", | ||
| jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax_warp.cu", | ||
| ], | ||
| extra_cuda_cflags=nvcc_flags, | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Look for an existing precedent of launching a __global__ defined in another TU, and for rdc/device-link flag support.
rg -n -C3 'cudaLaunchCooperativeKernel|reinterpret_cast<const void\*>\(kernel_' csrc | head -40
rg -n -C3 'rdc|dlink|relocatable-device-code' flashinfer/jit/core.py flashinfer/jit/compilation_context.py
ast-grep outline flashinfer/jit/core.py --match 'gen_jit_spec' --view expandedRepository: flashinfer-ai/flashinfer
Length of output: 3729
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== repo files around blackwell_softmax =="
git ls-files | rg '(^csrc/blackwell_softmax|flashinfer/aot\.py|flashinfer/jit/core\.py|flashinfer/jit/compilation_context\.py|csrc/mixed_comm\.cu|csrc/fmha_v2/templates/kernel\.jinja)' || true
echo
echo "== flashinfer/jit/core.py gen_jit_spec and nvcc flags =="
if [ -f flashinfer/jit/core.py ]; then
nl -ba flashinfer/jit/core.py | sed -n '460,620p'
echo
rg -n -C4 'nvcc_flags|rdc|dlink|device-link|cuda|extra_cuda_cflags|_get_nvcc_flags' flashinfer/jit/core.py
fi
echo
echo "== compilation_context files =="
fd -i 'compilation_context|compiler|jit' flashinfer flashinfer/jit -t f
if [ -e flashinfer/jit/compilation_context.py ]; then
rg -n -C4 'get_nvcc_flags_list|nvcc|rdc|dlink|device-link|cuda|extra_cuda_cflags' flashinfer/jit/compilation_context.py
fi
echo
echo "== AOT blackwell_softmax spec usage =="
if [ -f flashinfer/aot.py ]; then
nl -ba flashinfer/aot.py | sed -n '620,655p'
rg -n -C5 'blackwell_softmax|has_sm100|has_sm103|nvcc|rdc|dlink|device-link' flashinfer/aot.py
fi
echo
echo "== blackwell_softmax source outline =="
ast-grep outline csrc/blackwell_softmax.cuh --view expanded 2>/dev/null || true
ast-grep outline csrc/blackwell_softmax.cu --view expanded 2>/dev/null || true
nl -ba csrc/blackwell_softmax.cuh | sed -n '1,180p'
nl -ba csrc/blackwell_softmax.cu | sed -n '1,220p'Repository: flashinfer-ai/flashinfer
Length of output: 526
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== flashinfer/jit/core.py gen_jit_spec and nvcc flags =="
sed -n '460,620p' flashinfer/jit/core.py
echo
grep -n -C4 'nvcc_flags\|rdc\|dlink\|device-link\|cuda\|extra_cuda_cflags\|_get_nvcc_flags' flashinfer/jit/core.py || true
echo
echo "== AOT blackwell_softmax spec usage =="
sed -n '620,655p' flashinfer/aot.py
grep -n -C5 'blackwell_softmax\|has_sm100\|has_sm103\|nvcc\|rdc\|dlink\|device-link' flashinfer/aot.py || true
echo
echo "== blackwell_softmax source snippets =="
sed -n '1,220p' csrc/blackwell_softmax.cuh 2>/dev/null || true
sed -n '1,220p' csrc/blackwell_softmax.cu 2>/dev/null || trueRepository: flashinfer-ai/flashinfer
Length of output: 31308
🌐 Web query:
CUDA device-linking -rdc=true separate compilation __global__ function address launch cross translation unit reinterpret_cast
💡 Result:
In CUDA, referencing the address of a global function or using it via reinterpret_cast across translation units requires the use of separate compilation and device linking (enabled via the -rdc=true compiler option) [1][2][3]. Without separate compilation (the default whole-program mode), each translation unit is compiled into a self-contained device program [1][3]. In this mode, attempting to resolve addresses or function pointers across file boundaries will result in incorrect behavior, invalid addresses, or linker errors, as the host-side stub functions and device-side kernels are not unified across object files [1][4]. Key technical constraints and mechanisms include: 1. Requirement of Relocatable Device Code (RDC): You must compile with -rdc=true to allow the NVCC device linker to resolve references to device symbols (kernels and device functions) across different object files [1][2]. 2. Limitations of Function Pointers: While you can take the address of a device or global function, doing so in host code returns a pointer that is only valid within the context of the CUDA runtime's internal mapping [4][5]. If you attempt to manipulate these addresses using reinterpret_cast or pass them across different compilation units without proper device-linking support, the CUDA runtime will fail to map the address to the correct device kernel, often leading to invalid device function errors or undefined behavior [4][6][5]. 3. Kernel Launch Stubs: When you launch a global function, you are actually calling a host-side stub function generated by the compiler [1][4]. This stub is responsible for interacting with the CUDA runtime to launch the correct device entry point [1]. If multiple translation units define kernels or use shared logic, the host linker may conflict if symbols are not properly managed or if the RDC mode is not consistently applied, leading to scenarios where the wrong kernel or no kernel is found at the expected address [1][4]. 4. Best Practices: To use function pointers reliably across files, define the function pointer in device code (e.g., as a device variable) and use cudaMemcpyFromSymbol to retrieve its address on the host for setup, or maintain both the caller and the target function within the same linkable module (e.g., the same static library) [7][8][9]. Using shared libraries (.so/.dll) for cross-translation-unit device code is significantly more restrictive and generally only works effectively with static linking [8]. Top results: [1], [2], [4], [8]
Citations:
- 1: https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html
- 2: https://developer.nvidia.com/blog/separate-compilation-linking-cuda-device-code/
- 3: https://docs.nvidia.com/cuda/archive/11.7.1/cuda-compiler-driver-nvcc/index.html
- 4: https://nvidia.github.io/cccl/unstable/cccl/development/visibility/device_kernel_visibility.html
- 5: https://stackoverflow.com/questions/26738079/cuda-kernel-with-function-pointer-and-variadic-templates
- 6: https://forums.developer.nvidia.com/t/using-a-device-function-pointer-problems-using-a-pointer-to-a-device-function/26896/3
- 7: https://stackoverflow.com/questions/22202726/passing-cuda-function-pointers-with-libraries
- 8: https://stackoverflow.com/questions/35354822/load-device-function-from-shared-library-with-dlopen
- 9: https://stackoverflow.com/questions/39329905/function-pointers-with-cuda-device-functions
Enable device linking for the cross-TU kernel launch.
blackwell_softmax.cu launches kernels defined in other translation units via reinterpret_cast<const void*>(...), which requires CUDA RDC/device linking (-rdc=true). Add the device-link flags through this JIT spec instead of compiling each .cu as an isolated object.
🤖 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/blackwell_softmax.py` around lines 22 - 34, Update the JIT
specification returned by the blackwell softmax setup to enable CUDA relocatable
device code/device linking for the cross-translation-unit kernel launch. Extend
the existing extra CUDA flags in the relevant function around gen_jit_spec and
preserve the current source list and compiler-version flags.
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
csrc/blackwell_softmax_mr515_exp2.cu (1)
86-91: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the overloaded
temperatureparameter, or pass an explicit PDL flag.This kernel derives
enable_pdlfromtemperature == nullptr. The three sibling kernels use bit 2 ofparameter_kind. The polarity is correct today: the launcher passesnullptrexactly whenlaunch_with_pdlis true, andTEMP_KINDis0so the kernel never dereferencestemperature.The coupling is fragile. If the mr515 route later reads
temperature, or if a caller passes a null temperature array for any other reason, PDL behavior changes silently. The declaration ininclude/flashinfer/blackwell_softmax.cuhat Lines 49-50 gives no indication that the second parameter carries a control flag.Add a comment at this site that states the encoding and its dependency on
TEMP_KIND == 0. Also add a matching note to the header declaration. The project guidelines require documented rationale for special choices on kernel hot paths.Suggested comment
`#if` defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 + // Integration-only PDL encoding. TEMP_KIND==0 never dereferences the + // `temperature` slot, so the launcher passes nullptr to request the + // wait/signal pair. Revisit if this specialization ever reads temperature. const bool enable_pdl = temperature == nullptr; if (enable_pdl) { asm volatile("griddepcontrol.wait;" ::: "memory"); } `#endif`🤖 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/blackwell_softmax_mr515_exp2.cu` around lines 86 - 91, Document the overloaded temperature parameter at the __CUDA_ARCH__ == 1030 enable_pdl logic, stating that temperature == nullptr encodes PDL enabled and is valid only because TEMP_KIND == 0 means the kernel does not dereference temperature. Add the same rationale to the corresponding declaration in the blackwell_softmax header, without changing the existing behavior.Source: Coding guidelines
csrc/blackwell_softmax_bootstrap.cu (1)
375-386: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRecord why splitting the 256-bit store is safe here.
The fallback replaces one
st.global.v8.b32with twost.global.v4.b32. For loads that substitution is unconditionally equivalent. For stores it is equivalent only when no concurrent reader can observe the 32-byte region in a partially written state.The current code satisfies that condition. The Line 375 stores run only in the
splits == 1branch, which has no cross-block reader ofoutput. The Line 541 stores run aftercooperative_groups::this_grid().sync()at Line 415, and no later code in the kernel readsoutput.Add a short comment at these two sites that states the assumption. A future change that reads
outputfrom another block while the grid runs would make the split store observable as torn, and the provenance note at Lines 19-21 does not capture that.Also applies to: 541-552
🤖 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/blackwell_softmax_bootstrap.cu` around lines 375 - 386, Add concise comments at both the fallback store sites around the split v4 stores (the blocks near the current 375 and 541 locations) documenting that splitting the 256-bit store is safe because no concurrent block can read output: the first is in the splits == 1 path, and the second follows grid synchronization with no later output reads. Note that introducing concurrent readers would make the split observable as torn.
🤖 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/blackwell_softmax_bootstrap.cu`:
- Around line 87-93: Strip the PDL ABI bit before any parameter_kind checks in
the adapted kernels: in csrc/blackwell_softmax_bootstrap.cu lines 87-93,
csrc/blackwell_softmax_warp.cu lines 67-73, and
csrc/blackwell_softmax_rowwise.cu lines 85-91, move enable_pdl extraction and
parameter_kind masking outside and before the __CUDA_ARCH__ == 1030 guard. Keep
only the griddepcontrol.wait assembly behind the architecture guard.
In `@csrc/blackwell_softmax_mr515_exp2.cu`:
- Line 197: Align the non-mr515 softmax routes selected by
select_softmax_route—kWarp, kRowwise, and kBootstrap—with the mr515 zero-sum
behavior by guarding reciprocal normalization when the row sum is zero and
returning zero probabilities instead of multiplying by infinity. Alternatively,
add coverage that exercises all-−inf rows through each route and verifies 0.0f
outputs.
In `@csrc/blackwell_softmax.cu`:
- Around line 303-321: Update query_softmax_route to return
SoftmaxRoute::kFallback when the current device capability is not supported by
_supports_blackwell_softmax, before invoking select_softmax_route. Keep the
existing capability queries and route selection for allowed SM100/SM103 devices,
and use the queried device_major/device_minor values for the gate.
---
Nitpick comments:
In `@csrc/blackwell_softmax_bootstrap.cu`:
- Around line 375-386: Add concise comments at both the fallback store sites
around the split v4 stores (the blocks near the current 375 and 541 locations)
documenting that splitting the 256-bit store is safe because no concurrent block
can read output: the first is in the splits == 1 path, and the second follows
grid synchronization with no later output reads. Note that introducing
concurrent readers would make the split observable as torn.
In `@csrc/blackwell_softmax_mr515_exp2.cu`:
- Around line 86-91: Document the overloaded temperature parameter at the
__CUDA_ARCH__ == 1030 enable_pdl logic, stating that temperature == nullptr
encodes PDL enabled and is valid only because TEMP_KIND == 0 means the kernel
does not dereference temperature. Add the same rationale to the corresponding
declaration in the blackwell_softmax header, without changing the existing
behavior.
🪄 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: 348ec8e8-f10c-4c47-ae66-c77954ee58cb
📒 Files selected for processing (9)
csrc/blackwell_softmax.cucsrc/blackwell_softmax_bootstrap.cucsrc/blackwell_softmax_mr515_exp2.cucsrc/blackwell_softmax_rowwise.cucsrc/blackwell_softmax_warp.cuflashinfer/jit/blackwell_softmax.pyflashinfer/sampling.pyinclude/flashinfer/blackwell_softmax.cuhtests/utils/test_sampling.py
🚧 Files skipped from review as they are similar to previous changes (1)
- flashinfer/jit/blackwell_softmax.py
| #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 | ||
| const bool enable_pdl = (parameter_kind & 4) != 0; | ||
| parameter_kind &= 3; | ||
| if (enable_pdl) { | ||
| asm volatile("griddepcontrol.wait;" ::: "memory"); | ||
| } | ||
| #endif |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check whether an sm_100a Blackwell softmax binary can run on an sm_103 device.
set -uo pipefail
echo "=== Blackwell softmax JIT module: arch and gencode flags ==="
fd -t f 'blackwell_softmax.py' --exec rg -n -C 8 'sm_10|103|100|arch|gencode|nvcc|flags|capability'
echo "=== AOT build registration ==="
fd -t f 'aot.py' --exec rg -n -C 8 'blackwell_softmax'
echo "=== is_sm103 producer and consumers ==="
rg -nP --type=cu --type=cpp --type=py -C 6 '\bis_sm103\b'
echo "=== Python capability gate for the softmax dispatch ==="
fd -t f 'sampling.py' --exec rg -n -C 10 'blackwell|get_device_capability|major|minor'
echo "=== Repository conventions for sm_103 arch selection ==="
rg -nP -C 4 "sm_103a?|'103'|\"103\"|\b103\b" --type=py -g '!**/tests/**' | head -80Repository: flashinfer-ai/flashinfer
Length of output: 12995
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== Files and top-level build context ==="
git ls-files | rg '(^|/)env\.py$|(^|/)blackwell_softmax\.py$|(^|/)sampling\.py$|(^|/)aot\.py$|(^|/)blackwell_softmax.*\.cu$|(^|/)build' | head -200
echo
echo "=== blackwell_softmax related snippets ==="
for f in \
python/flashinfer/jit/blackwell_softmax.py \
python/flashinfer/sampling.py \
csrc/blackwell_softmax.cu \
csrc/blackwell_softmax_bootstrap.cu \
csrc/blackwell_softmax_warp.cu \
csrc/blackwell_softmax_rowwise.cu \
csrc/blackwell_softmax_mr515_exp2.cu
do
[ -f "$f" ] || continue
echo "---- $f ----"
sed -n '1,230p' "$f"
done
echo
echo "=== Search exact PDL/flag patterns in cu/cpp ==="
rg -n --glob '*.cu' --glob '*.cpp' --glob '*.cc' --glob '*.cuh' -C 5 'parameter_kind|max_temperature|scalar_temperature|enable_pdl|griddepcontrol.wait|is_sm103|__CUDA_ARCH__|1030|1000' .Repository: flashinfer-ai/flashinfer
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== Python JIT env and compilation context ==="
fd -t f 'env.py|context.py|compilation.*\.py' -x sh -c 'echo "--- $1"; rg -n -C 6 "CUDA|CUDAARCH|__cuda|nvcc_flags|supported_major_versions|sm_10|capability|arch" "$1"' sh {}
echo
echo "=== Blackwell softmax C++ entrypoints ==="
for f in \
csrc/blackwell_softmax.cu \
csrc/blackwell_softmax_bootstrap.cu \
csrc/blackwell_softmax_warp.cu \
csrc/blackwell_softmax_rowwise.cu
do
echo "---- $f ----"
wc -l "$f"
rg -n -C 8 'softmax_dispatch|parameter_kind|scalar_temperature|max_temperature|template|__device__|softmax_route|parameter_kind' "$f" | head -260
done
echo
echo "=== Narrow cu search for relevant symbols ==="
rg -n --glob '*.cu' '(__CUDA_ARCH__|1030|1000|parameter_kind|scalar_temperature|max_temperature|is_sm103)' csrc include python -g '!**/nv_internal/**' -g '!**/test**/**' | sed -n '1,220p'Repository: flashinfer-ai/flashinfer
Length of output: 49311
Strip the PDL ABI bit unconditionally before reading parameter_kind.
parameter_kind_i is passed with bit 2 on when is_sm103 && enable_pdl, but each adapted kernel strips it only while compiling for __.cuda_arch__ == 1030. Keep the griddepcontrol.wait path behind the arch guard and move const bool enable_pdl plus parameter_kind &= 3 above the guard at blackwell_softmax_bootstrap.cu#87, blackwell_softmax_warp.cu#67, and blackwell_softmax_rowwise.cu#85. Otherwise parameter_kind == 2 checks will fail and per-row temperatures will use scalar_temperature.
📍 Affects 3 files
csrc/blackwell_softmax_bootstrap.cu#L87-L93(this comment)csrc/blackwell_softmax_warp.cu#L67-L73csrc/blackwell_softmax_rowwise.cu#L85-L91
🤖 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/blackwell_softmax_bootstrap.cu` around lines 87 - 93, Strip the PDL ABI
bit before any parameter_kind checks in the adapted kernels: in
csrc/blackwell_softmax_bootstrap.cu lines 87-93, csrc/blackwell_softmax_warp.cu
lines 67-73, and csrc/blackwell_softmax_rowwise.cu lines 85-91, move enable_pdl
extraction and parameter_kind masking outside and before the __CUDA_ARCH__ ==
1030 guard. Keep only the griddepcontrol.wait assembly behind the architecture
guard.
|
/bot run tests/utils |
|
/bot run tests/utils |
|
/bot run tests/utils |
|
Update on the refreshed The bot's earlier “0/0 executed test jobs” comment was generated by the report job before the long GPU jobs finished, so it is now stale. The completed jobs show:
I also ran an independent four-GB300 complete-denominator audit. It accounts for all 20,165 The strict LogitsPipe regression remains green (original node IDs 9/9, full temperature matrix 54/54, independent Torch oracle 3/3), and the full GB300 kernel contract remains 376/376 correctness plus 41/41 performance floors. |
|
/bot run tests/utils |
|
[FAILED] Pipeline #61373025 — 16/18 executed test jobs passed No usable JUnit artifact was available; individual tests and nightly comparison could not be recovered. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
4 participants
Description
This PR adds an optimized Blackwell FP32 softmax dispatcher behind the existing
flashinfer.sampling.softmaxAPI for SM100 and SM103. The public ABI and output semantics stay unchanged: 2-D CUDA logits, temperature absent/scalar/per-row, PDL support, and a fresh non-aliasing output. Unsupported devices and launches retain the existing OnlineSoftmax fallback.The current head is
f0edac69e1c4d299c14fa95b7a0073a172203cf9. It contains the warp, rowwise, cooperative-bootstrap, and MR515-derived exp2/512-thread/vec4 routes. The device kernels and dispatcher are unchanged from the source-bound082c7f54GB300 qualification;f0edac69fixes the strict LogitsPipe regression by makingFusedTemperatureSoftmaxOpcall the public architecture-dispatchedsampling.softmaxpath with the same temperature and PDL arguments, instead of bypassing it through the generic module. It also removes the now-unused workspace and adds an independent Torch-oracle regression for the three failing per-row-temperature shapes.Dispatch policy
The original PR4282 warp, rowwise, and bootstrap regions are preserved. Route 4 is restricted to the SM103 cases where same-session paired CUPTI measurements proved a local win:
vocab_size == 32000,temperature=None, PDL disabled, androws in {16, 32, 64, 128, 512, 1024};rows == 64,vocab_size == 32000, scalar temperature1.0, and PDL enabled.All other audited shapes remain on the prior dispatcher. The route observer exposes IDs
1..4; the 376-row audit has no external fallback.Current GB300 evidence
The latest source-bound requalification used Slurm job
376374on four GB300 GPUs (sm_103a). Timing is GPU-only CUPTI activity tracing with cold-L2 flushing and candidate/baseline blocks interleaved in the same process and GPU session.Ratified 41-row performance set
All 41/41 declared performance rows passed correctness and their contract floor. Geometric-mean speedup is 2.099151447x and the worst row is 1.328849819x. Every candidate sample contains exactly one kernel and no row is within 2% of the floor.
b1, v64000, temperature=Noneb16, v256000, temperature=Noneb1024, v128000, temperature=Noneb64, v32000, temperature=1.0, PDLThe formal run used 100 ms warmup and 1000 ms arms with three same-session interleaved CUPTI pairs. Its aggregate evidence SHA-256 is
3a19cfdcc5686393c41a4213ac0fd5b358649f933b856b992813efd2ad91084e.Supplemental full-376 timing map
Step
376374.331measured the exact 376-row regression manifest across four GB300s (94 rows/GPU) with same-session interleaved CUPTI timing.baseline_ms / candidate_ms >= 1.01:76, 2:42, 3:251, 4:7Every row retained its raw baseline/candidate timing arrays, produced a fresh output, used a custom route, and launched exactly one candidate kernel. This sweep used 25 ms warmup, 100 ms measurement arms, and three interleaved pairs, so it is a supplemental broad performance map; the 41-row 100/1000 ms run remains the formal contract-duration performance gate. Aggregate artifact SHA-256:
7b7cf9fa28250b5651161efcbaa703f7d43276a9bb3a31f1d69924effbd3a0d6.Correctness and implementation evidence
atol=rtol=1e-3.f0edac69:test_sampling.py1834 passed / 121 skipped;test_logits_processor.py591 passed / 12 skipped.376374.482: bootstrap vec8 retainsLDG.256/STG.256; route 4 retains 128-bit vec4 traffic; audited kernels contain noLDL/STLspill. SM103 retains the expected PDLACQBULK/PREEXITinstructions.376374.485(b64 route 4): 512 threads, 99 registers/thread, 23.68% active warps, 38.09% long-scoreboard stall, zero local loads/stores.376374.496(b1024 bootstrap): 256 threads, 48 registers/thread, 47.55% active warps, 72.47% DRAM throughput, 86.59% long-scoreboard stall, zero local loads/stores.ERROR SUMMARY: 0 errors. Racecheck is not applicable because the final routes use no DSM/cluster execution.tests/utils CI status
The refreshed internal pipeline was built from this exact head after a successful prepare retry.
sampling1809 passed / 146 skipped andlogits_processor591 passed / 12 skipped.quantized_allreduceTCPStore port race: the test chooses an unreserved random port while concurrent mpirun instances run. A retry reproduced the sameEADDRINUSEon a different port.1765105before checkout; CUDA 13.0 remains pending without a runner.No completed job in this refreshed pipeline has a softmax/operator failure. The infrastructure-blocked and pending lanes are intentionally not represented as green.
Independent GB300 complete-denominator audit
A fresh allocation (
377417) collected all 20,165 tests across the 25tests/utilsfiles. The two changed files are fully source-bound as reported above. The remaining 23-file sweep was followed by a source-bound replay of every failure using a non-editable 0.6.16 wheel built from exact headf0edac69; installed JIT-cache/cubin packages were blocked and all csrc/include/Jinja data came from the wheel. Two previously skipped IPC cases were rerun with four visible GB300s.The effective complete-denominator accounting is 19,920 passed / 243 skipped / 2 non-softmax failed / 0 errors. The only failures are the independent TopK reference-index bug (
torch.gatherreceives int32 instead of int64) and an FP4 logging replay that requires cuDNN backend >=91002 while the pinned image provides 91001. Softmax/kernel failures are 0. The receipt explicitly records that this is a combined complete-denominator audit, not a single-shot 20,165-case source-bound run. Receipt SHA-256:aad5289a20bf84045d531cc14800eb77054ed2edb443083021e747722eaa74f1.Checklist
f0edac69.Related to #4254
Summary by CodeRabbit
New Features
Bug Fixes
Tests