feat(cake_tinygemm2): add CAKE-generated SM100/SM103 tinygemm2 variants with bit-identical outputs - #4274
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 SM100/SM103 TinyGEMM2 CUDA kernels with stage4, stage8, and PDL variants, JIT compilation support, architecture- and shape-aware bias dispatch, an environment override, and CUDA correctness and routing tests. ChangesSM100 TinyGEMM2
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller as tinygemm_bf16
participant Router as routergemm dispatch
participant JIT as SM100 module loader
participant Kernel as tinygemm2_sm100 variant
Caller->>Router: submit biased BF16 GEMM
Router->>Router: check SM100a support, SM count, shape, and environment
Router->>JIT: load generated SM100 module
JIT-->>Router: return stage4/stage8 and PDL operations
Router->>Kernel: launch selected variant
Kernel-->>Caller: return BF16 output
Possibly related issues
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 |
![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 (4)
csrc/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuh (2)
181-205: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winOptional: hoist the per-invocation driver queries out of the launch path.
cudaFuncSetAttributehere plus the twocudaDeviceGetAttributecalls inCheckSm100Familyrun on every call. Both results are stable per (device, kernel), so caching them (e.g. a function-localstatictable keyed by device id) removes a few microseconds of host work from a kernel selected specifically for its low latency.🤖 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/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuh` around lines 181 - 205, Cache the stable per-device results used by CheckSm100Family and the cudaFuncSetAttribute setup in the launch path, using a function-local static table keyed by device ID. Update the relevant query and attribute-setting logic so repeated invocations reuse cached values while preserving existing behavior for each device and kernel.
110-118: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd an explicit
in_featuresalignment guard before TMA encoding.
EncodeWeightTmaandEncodeActivationTmapassstride(0) * sizeof(__nv_bfloat16)as the global TMA stride, which must be 16B-aligned. Since stride(0) equalsin_featuresfor these 2D contiguous bfloat16 tensors, rejectin_features % 8 != 0here instead of failing later insidecuTensorMapEncodeTiledwith an opaqueCUresult=error.♻️ Proposed guard
TVM_FFI_ICHECK(in_features >= kTileK) << "in_features (" << in_features << ") must be at least " << kTileK << " (one TMA box)"; + TVM_FFI_ICHECK(in_features % 8 == 0) + << "in_features (" << in_features + << ") must be a multiple of 8 so the TMA global stride is 16B-aligned";🤖 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/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuh` around lines 110 - 118, Add an explicit validation in the existing dimension-check block before TMA encoding, requiring in_features to be divisible by 8 and reporting the received value when it is not. Keep the current checks unchanged and anchor the guard alongside the in_features validation in the binding setup.csrc/tinygemm2_sm100/tinygemm2_sm100_stage4_binding.cu (1)
26-42: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRename the generated device TUs to
.cuhso they cannot be compiled standalone. The JIT source list only includes the*_binding.cufiles, but the generated device units currently carry a.cuextension and are only meant to be included with per-boundary macros defined. Renaming them to.cuh/.incremoves the stalecsrc/tinygemm2_sm100/*.cuglob risk and makes the include contracts explicit.🤖 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/tinygemm2_sm100/tinygemm2_sm100_stage4_binding.cu` around lines 26 - 42, Rename the generated device translation units included by the binding files to .cuh or .inc, and update the include targets in csrc/tinygemm2_sm100/tinygemm2_sm100_stage4_binding.cu:26-42, csrc/tinygemm2_sm100/tinygemm2_sm100_stage8_binding.cu:26-42, csrc/tinygemm2_sm100/tinygemm2_sm100_stage4_pdl_binding.cu:26-42, and csrc/tinygemm2_sm100/tinygemm2_sm100_stage8_pdl_binding.cu:26-42. Preserve the existing per-boundary macro definitions and undefinitions while ensuring the generated units are no longer standalone .cu files.flashinfer/jit/tinygemm2_sm100.py (1)
63-72: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: use iterable unpacking instead of list concatenation.
Static analysis flags
sm100a_nvcc_flags + [...]for RUF005;[*sm100a_nvcc_flags, "-gencode=arch=compute_103a,code=sm_103a"]is the idiomatic equivalent.🧹 Proposed nit
- extra_cuda_cflags=sm100a_nvcc_flags - + ["-gencode=arch=compute_103a,code=sm_103a"], + extra_cuda_cflags=[ + *sm100a_nvcc_flags, + "-gencode=arch=compute_103a,code=sm_103a", + ],🤖 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/tinygemm2_sm100.py` around lines 63 - 72, Update the extra_cuda_cflags argument in gen_jit_spec to use iterable unpacking for sm100a_nvcc_flags and append the compute_103a gencode flag without list concatenation, preserving the resulting flag order and values.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuh`:
- Around line 124-127: Correct the weight descriptor comment near tmap_wt: state
that the reduction/K axis may exceed K for non-multiple-of-1024 sizes and
correctness depends on TMA zero-filling via FLOAT_OOB_FILL_NONE, while only the
out_features axis is guaranteed in bounds. Remove the inaccurate claim that both
boxed axes remain in bounds.
- Around line 170-185: Add a CUDA device guard for input.device().device_id at
the RunStage* entry path, before stream creation/configuration and before
LaunchVariant performs cudaFuncSetAttribute or kernel launch. Match the existing
csrc binding guard pattern and keep the guard active through the full stream and
kernel execution flow.
In `@tests/model_optimizations/test_tinygemm2_sm100.py`:
- Around line 16-21: Update _skip_if_not_sm100_family to use
flashinfer.utils.is_sm100a_supported with the CUDA device, while preserving the
existing no-CUDA skip behavior. Skip the tests whenever that utility reports the
device or CUDA toolkit is unsupported, and remove the manual
get_compute_capability tuple check.
---
Nitpick comments:
In `@csrc/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuh`:
- Around line 181-205: Cache the stable per-device results used by
CheckSm100Family and the cudaFuncSetAttribute setup in the launch path, using a
function-local static table keyed by device ID. Update the relevant query and
attribute-setting logic so repeated invocations reuse cached values while
preserving existing behavior for each device and kernel.
- Around line 110-118: Add an explicit validation in the existing
dimension-check block before TMA encoding, requiring in_features to be divisible
by 8 and reporting the received value when it is not. Keep the current checks
unchanged and anchor the guard alongside the in_features validation in the
binding setup.
In `@csrc/tinygemm2_sm100/tinygemm2_sm100_stage4_binding.cu`:
- Around line 26-42: Rename the generated device translation units included by
the binding files to .cuh or .inc, and update the include targets in
csrc/tinygemm2_sm100/tinygemm2_sm100_stage4_binding.cu:26-42,
csrc/tinygemm2_sm100/tinygemm2_sm100_stage8_binding.cu:26-42,
csrc/tinygemm2_sm100/tinygemm2_sm100_stage4_pdl_binding.cu:26-42, and
csrc/tinygemm2_sm100/tinygemm2_sm100_stage8_pdl_binding.cu:26-42. Preserve the
existing per-boundary macro definitions and undefinitions while ensuring the
generated units are no longer standalone .cu files.
In `@flashinfer/jit/tinygemm2_sm100.py`:
- Around line 63-72: Update the extra_cuda_cflags argument in gen_jit_spec to
use iterable unpacking for sm100a_nvcc_flags and append the compute_103a gencode
flag without list concatenation, preserving the resulting flag order and values.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: f130ab3b-7310-4a56-a3f4-a28a3a3601f3
📥 Commits
Reviewing files that changed from the base of the PR and between a02d94d and 48f225a1dd1ccf1dc35dff9167a8c5b0038d3885.
📒 Files selected for processing (13)
csrc/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuhcsrc/tinygemm2_sm100/tinygemm2_sm100_stage4.cucsrc/tinygemm2_sm100/tinygemm2_sm100_stage4_binding.cucsrc/tinygemm2_sm100/tinygemm2_sm100_stage4_pdl.cucsrc/tinygemm2_sm100/tinygemm2_sm100_stage4_pdl_binding.cucsrc/tinygemm2_sm100/tinygemm2_sm100_stage8.cucsrc/tinygemm2_sm100/tinygemm2_sm100_stage8_binding.cucsrc/tinygemm2_sm100/tinygemm2_sm100_stage8_pdl.cucsrc/tinygemm2_sm100/tinygemm2_sm100_stage8_pdl_binding.cuflashinfer/gemm/routergemm.pyflashinfer/jit/__init__.pyflashinfer/jit/tinygemm2_sm100.pytests/model_optimizations/test_tinygemm2_sm100.py
| // 2D TMA descriptor for the weight matrix — field-for-field the descriptor | ||
| // the Loom host shim encodes for 'tmap_wt': box (kTileK, kTileM), 128B | ||
| // swizzle, no L2 promotion, no OOB fill. Both boxed axes stay in bounds | ||
| // (CheckInputs guarantees in_features >= kTileK and out_features >= kTileM). |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
The in-bounds claim for the weight descriptor is inaccurate.
The reduction axis goes out of bounds whenever K is not a multiple of 1024: the loaders address k_base = (ki * 4 + wslot) * 256 for ki < ceil(K/1024), so e.g. K = 1536 issues boxes ending at 1984. Correctness relies on TMA zero-filling those boxes (FLOAT_OOB_FILL_NONE ⇒ zero fill), exactly like the activation batch axis. Only the out_features axis is guaranteed in bounds (multiple of kTileM). Please correct the comment so the zero-fill dependency is explicit — a future reader could otherwise "harden" the descriptor and silently break non-multiple-of-1024 K.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@csrc/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuh` around lines 124 -
127, Correct the weight descriptor comment near tmap_wt: state that the
reduction/K axis may exceed K for non-multiple-of-1024 sizes and correctness
depends on TMA zero-filling via FLOAT_OOB_FILL_NONE, while only the out_features
axis is guaranteed in bounds. Remove the inaccurate claim that both boxed axes
remain in bounds.
| template <typename TMap> | ||
| inline void LaunchVariant(void (*kernel)(TMap, TMap, __nv_bfloat16*, __nv_bfloat16*, int, int, int), | ||
| int smem_bytes, bool pdl, const CUtensorMap& weight_map, | ||
| const CUtensorMap& activation_map, __nv_bfloat16* out, | ||
| __nv_bfloat16* bias, const ProblemDims& dims, cudaStream_t stream) { | ||
| static_assert(sizeof(TMap) == sizeof(CUtensorMap), | ||
| "generated tensor-map parameter must be layout-compatible with CUtensorMap"); | ||
| TMap wt_param, act_param; | ||
| std::memcpy(&wt_param, &weight_map, sizeof(TMap)); | ||
| std::memcpy(&act_param, &activation_map, sizeof(TMap)); | ||
|
|
||
| CheckCuda(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes), | ||
| "cudaFuncSetAttribute(tinygemm2_sm100 dynamic smem)"); | ||
|
|
||
| const dim3 grid((dims.out_features + kTileM - 1) / kTileM, (dims.batch + kTileN - 1) / kTileN); | ||
| const dim3 block(kThreads); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# How do other csrc operators bind the device before launching?
rg -nP -C3 '(CUDADeviceGuard|cudaSetDevice|DeviceGuard)\b' csrc | head -60
# Definition of get_stream to see whether it implies a device binding
rg -nP -C6 '\bget_stream\s*\(' --iglob '*tvm_ffi_utils*'Repository: flashinfer-ai/flashinfer
Length of output: 4550
🏁 Script executed:
#!/bin/bash
set -u
echo "== target file outline/size =="
wc -l csrc/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuh
ast-grep outline csrc/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuh --view expanded || true
echo
echo "== target section =="
cat -n csrc/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuh | sed -n '1,240p'
echo
echo "== tinygemm2 sm100 bindings list =="
fd -t f . csrc/tinygemm2_sm100 -E 'tinygemm2_sm100_binding_common.cuh' | sort | xargs -r -n1 sh -c 'echo "--- $0"; wc -l "$0"; rg -n "RunStage|CheckInputs|LaunchVariant|get_stream|CUDADeviceGuard|cudaSetDevice|CheckCuda" "$0" -C 3 || true'Repository: flashinfer-ai/flashinfer
Length of output: 16267
🏁 Script executed:
#!/bin/bash
set -u
echo "== tvm_ffi_utils declarations/usages =="
rg -n -C 5 '\bget_stream\s*\(' | head -120
fd -t f 'tvm_ffi_utils.*' csrc | sort | xargs -r -n1 sh -c 'echo "--- $0"; wc -l "$0"; rg -n -C 8 "get_stream|CUDADeviceGuard|class CUDADeviceGuard|struct CUDADeviceGuard|namespace ffi|inline.*device" "$0" || true'
echo
echo "== cudaFuncSetAttribute device-stream semantics references in repo comments =="
rg -n "cudaFuncSetAttribute|cudaLaunchKernelEx|programmatic stream|invalid resource|cudaErrorInvalidResourceHandle|SetDevice" csrc/include csrc | head -120Repository: flashinfer-ai/flashinfer
Length of output: 202
🌐 Web query:
CUDA cudaFuncSetAttribute target of attribute is current device stream invalid resource handle documentation
💡 Result:
The error message "target of attribute is current device stream" is not a standard, documented error string for the cudaFuncSetAttribute API in official NVIDIA CUDA documentation [1][2][3][4]. cudaFuncSetAttribute (and its Driver API counterpart cuFuncSetAttribute) is used to set attributes for a global function (kernel) [3][4]. If the call fails due to an invalid handle or invalid resource, it typically returns specific error codes defined by the CUDA API, most commonly cudaErrorInvalidValue (or CUDA_ERROR_INVALID_VALUE in the Driver API) if the attribute cannot be written, the value is incorrect, or the function pointer is invalid [2][3][5]. In the Driver API, CUDA_ERROR_INVALID_HANDLE is also a possibility if the CUfunction handle is invalid [2]. If you are encountering an error message referring to a "current device stream," it is highly likely that you are conflating cudaFuncSetAttribute with a different API that operates on streams, such as cudaStreamSetAttribute [6]. Key distinctions: 1. cudaFuncSetAttribute: Sets attributes for a kernel (function) [3][4]. It does not take a stream as an argument [3]. 2. cudaStreamSetAttribute: Sets attributes for a stream [6]. If this API is misused or called with an invalid stream handle, it may produce errors related to stream resources [6]. To troubleshoot: - Verify that you are calling the correct API for your intent. If you intend to modify kernel properties, use cudaFuncSetAttribute [3]. If you intend to modify stream properties, use cudaStreamSetAttribute [6]. - Ensure your handles (function pointers or stream handles) are valid and have been correctly initialized before use [3][6]. - Check for cudaErrorInvalidValue or cudaErrorInvalidResourceHandle return codes, which indicate the specific nature of the configuration or handle issue [3][4].
Citations:
- 1: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__EXEC.html
- 2: https://docs.nvidia.com/cuda/archive/11.0/cuda-driver-api/group__CUDA__EXEC.html
- 3: https://docs.nvidia.com/cuda/archive/13.2.1/cuda-runtime-api/group__CUDART__EXECUTION.html
- 4: https://docs.nvidia.com/cuda/developer-preview/13.4/cuda-runtime-api/cuda_runtime_api/group__CUDART__EXECUTION.html
- 5: https://docs.nvidia.com/cuda/archive/13.1.1/cuda-runtime-api/group__CUDART__EXECUTION.html
- 6: https://docs.nvidia.com/cuda/archive/13.2.1/cuda-runtime-api/group__CUDART__STREAM.html
🌐 Web query:
CUDA driver API cuFuncSetAttribute current context launch stream foreign stream invalid resource handle
💡 Result:
In CUDA, an "invalid resource handle" error (often associated with the CUDA_ERROR_INVALID_HANDLE return code) during API calls like cuFuncSetAttribute, cuLaunchKernel, or memory operations typically indicates that a resource—such as a stream, context, or memory handle—is being used in an invalid context [1][2][3]. While cuFuncSetAttribute itself manages kernel attributes, the error does not necessarily mean the function handle is invalid. It most frequently signifies a mismatch between the current CUDA device/context and the resource handle provided [2][4][3]. Key causes and troubleshooting steps include: Device/Context Mismatch: CUDA resources (streams, memory, kernels) are bound to specific devices and contexts at the time of creation [2][4][3]. If you attempt to use a stream or other handle that was created while a different device was active, the driver will return an invalid resource handle [4][3]. Ensure that cudaSetDevice() or cuCtxSetCurrent() is called correctly for the thread before any resources are created and before they are used [5][4]. Thread-Safety: If your application is multi-threaded, each CPU thread must properly manage its current context or device [5][4]. If a thread spawns or switches devices without re-associating its context, it may inadvertently attempt to access resources owned by a different device's context [5][4][3]. Stream Lifecycle: Ensure the stream handle being passed is valid and has not been destroyed [6]. A stream created on one device cannot be used to launch work on another device [2][4][3]. "Foreign" Stream Context: While the term "foreign stream" is not a standard CUDA technical term, it often refers to streams imported from other APIs (such as Graphics interop via Direct3D or OpenGL) or streams created in a different process/context. If you are performing interop, you must ensure the foreign handle is correctly mapped and synchronized according to the specific interop requirements (e.g., using cuExternalMemory or similar mechanisms) before use [2]. Verification: To debug, verify that the device ID is consistently set for all operations involving a specific resource [4][3]. If the code works on the default device but fails on others, it is almost certainly a device-association issue [4]. Always check the return value of every CUDA driver API call to isolate exactly which function is reporting the error [1][2].
Citations:
- 1: https://docs.nvidia.com/cuda/archive/11.0/cuda-driver-api/group__CUDA__EXEC.html
- 2: https://forums.developer.nvidia.com/t/invalid-resource-handle-on-culaunchkernel/277931/2
- 3: https://forums.developer.nvidia.com/t/cudamemcpyasync-returns-invalid-resource-handle/338750/1
- 4: https://stackoverflow.com/questions/56137558/invalid-resource-handle-returned-from-cudalaunchkernel-when-using-non-default-gp
- 5: Error Code 1: Cuda Runtime (invalid resource handle) NVIDIA/TensorRT#3028
- 6: https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html
🌐 Web query:
CUDA error cudaErrorInvalidResourceHandle launch kernel stream created on other device cudaStreamCreateWithFlags
💡 Result:
The error cudaErrorInvalidResourceHandle typically indicates that a CUDA stream is being accessed or used in a context that does not recognize it, often because it is bound to a different device or context than the one currently active [1][2]. CUDA streams are created specifically on the device context that is active at the time of creation [3][4]. If you attempt to launch a kernel, synchronize, or query a stream on a thread where a different device (or context) is currently active, the runtime may fail with cudaErrorInvalidResourceHandle because the stream handle is not valid for that specific device context [1][2][5]. Key causes and solutions include: 1. Device Context Mismatch: If you switch devices using cudaSetDevice or create a new context (e.g., when initializing libraries like TensorRT), a previously created stream may become invalid if it is not associated with the newly active context [1]. Always ensure that the correct device is set via cudaSetDevice before performing any operations on a specific stream [6][5][7]. 2. Cross-Thread Access: CUDA streams are associated with a specific host thread's context state [4][8]. While they can technically be used across threads, you must ensure that the thread attempting to use the stream has the correct device context active [2][9]. 3. Verification: To prevent these issues, it is recommended to keep track of which device a stream was created on [2][6]. You can verify the current device using cudaGetDevice and ensure it matches the device expected for that stream [6]. 4. Best Practices: If you are working in a multi-GPU environment, explicitly manage your stream-to-device association [6][5]. Some developers use a wrapper class to encapsulate both the stream and its associated device ID to ensure that cudaSetDevice is called before any operations are performed on that stream [2][6]. For further debugging, note that while some API calls (like cudaMemcpyPeerAsync or cudaStreamWaitEvent) are designed to work across devices, standard kernel launches will fail if issued to a stream not associated with the current device [5][10].
Citations:
- 1: after tensorrt init finished, my own stream invalid NVIDIA/TensorRT#2236
- 2: https://stackoverflow.com/questions/31474784/are-cuda-streams-device-associated-and-how-do-i-get-a-streams-device
- 3: https://docs.nvidia.com/cuda/archive/12.6.3/cuda-runtime-api/group__CUDART__STREAM.html
- 4: https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html
- 5: https://docs.nvidia.com/cuda/archive/13.1.1/cuda-programming-guide/03-advanced/multi-gpu-systems.html
- 6: https://forums.developer.nvidia.com/t/cudaerrorinvalidresourcehandle-after-dispatching-inclusive-scan-kernel-in-cuda12/326349
- 7: https://forums.developer.nvidia.com/t/on-implicit-synchronization-of-streams-on-separate-devices/115573
- 8: https://docs.nvidia.com/cuda/archive/13.0.2/cuda-runtime-api/group__CUDART__STREAM.html
- 9: https://forums.developer.nvidia.com/t/cudastreamquery-returning-cudaerrorinvalidresourcehandle-from-another-thread-works-fine-from-kernel/11903/4
- 10: https://docs.nvidia.com/cuda/cuda-programming-guide/03-advanced/multi-gpu-systems.html.md
Guard the device before configuring/launching the kernels.
RunStage* builds the stream and calls LaunchVariant from CheckInputs, but CheckInputs only compares device_id values. If this function is called while a different CUDA device is current, cudaFuncSetAttribute/cudaLaunchKernelEx can operate against the wrong device/stream context. Introduce a device guard for input.device().device_id before stream config, kernel attribute setup, and launch, matching the existing csrc binding pattern.
🤖 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/tinygemm2_sm100/tinygemm2_sm100_binding_common.cuh` around lines 170 -
185, Add a CUDA device guard for input.device().device_id at the RunStage* entry
path, before stream creation/configuration and before LaunchVariant performs
cudaFuncSetAttribute or kernel launch. Match the existing csrc binding guard
pattern and keep the guard active through the full stream and kernel execution
flow.
There was a problem hiding this comment.
hi @xslingcn I sugguested to encode. tma descriptor from a platform indepenty repr :
raw c pointer, layout of shape, stride
So the codes can be either neatly ported w or w/o TVM-FFI interfaces.
xslingcn
changed the title
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
flashinfer/jit/tinygemm2_sm100.py (1)
34-35: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse iterable unpacking for CUDA flag construction.
Ruff RUF005 flags concatenating
sm100a_nvcc_flagswith a single-item list. This preserves the same flag order while avoiding the warning:Proposed fix
- extra_cuda_cflags=sm100a_nvcc_flags - + ["-gencode=arch=compute_103a,code=sm_103a"], + extra_cuda_cflags=[ + *sm100a_nvcc_flags, + "-gencode=arch=compute_103a,code=sm_103a", + ],🤖 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/tinygemm2_sm100.py` around lines 34 - 35, Update the CUDA flags passed via extra_cuda_cflags in the relevant build configuration to use iterable unpacking of sm100a_nvcc_flags followed by the compute_103a flag, preserving the existing order and values while removing list concatenation.Source: Linters/SAST tools
csrc/tinygemm2_sm100/tinygemm2_sm100.cu (2)
1539-1540: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winPer-invocation CUDA setup queries on a microsecond-scale path. Both sites re-query immutable state (kernel smem opt-in, device compute capability) on every
RunStage*call; cache each once instead.
csrc/tinygemm2_sm100/tinygemm2_sm100.cu#L1539-L1540: issuecudaFuncSetAttributeonce per kernel function pointer.csrc/tinygemm2_sm100/tinygemm2_sm100.cu#L1409-L1419: cache the compute capability per device id and validate from the cache.🤖 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/tinygemm2_sm100/tinygemm2_sm100.cu` around lines 1539 - 1540, The RunStage* CUDA setup currently repeats immutable queries on every invocation. In csrc/tinygemm2_sm100/tinygemm2_sm100.cu lines 1539-1540, cache successful cudaFuncSetAttribute state once per kernel function pointer and skip subsequent calls; in lines 1409-1419, cache compute capability by device ID and validate using the cached value rather than querying CUDA each time.
1467-1470: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winValidate the TMA global-stride constraint explicitly.
The 128B-swizzled descriptors require the row stride in bytes to be a multiple of 16, i.e.
in_features % 8 == 0. Today an oddin_features(e.g. 100) passesCheckInputsand surfaces later as a barecuTensorMapEncodeTiled failed ... CUresult=N. A direct check gives a usable message.♻️ Suggested check
TVM_FFI_ICHECK(in_features >= kTileK) << "in_features (" << in_features << ") must be at least " << kTileK << " (one TMA box)"; + TVM_FFI_ICHECK(in_features % 8 == 0) + << "in_features (" << in_features + << ") must be a multiple of 8 so the row stride is 16B-aligned for the 128B-swizzled TMA " + "descriptors"; TVM_FFI_ICHECK(out_features >= kTileM && out_features % kTileM == 0)🤖 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/tinygemm2_sm100/tinygemm2_sm100.cu` around lines 1467 - 1470, Update the input validation near the existing in_features check in CheckInputs to require in_features % 8 == 0, with a clear message stating that in_features must be a multiple of 8 for the TMA global stride. Preserve the existing minimum-size validation and out_features 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 `@csrc/tinygemm2_sm100/tinygemm2_sm100.cu`:
- Around line 1480-1483: Update the TMA descriptor comments near the
weight-matrix descriptor to state that the kernel’s K-axis traversal may extend
to ceil(K/1024)*1024 beyond in_features and relies on TMA zero-fill for the
overrun, rather than claiming both boxed axes remain in bounds or that there is
no OOB fill. Preserve the existing description of the descriptor’s box
dimensions and swizzle behavior.
- Around line 1568-1612: Update RunStage4, RunStage4Pdl, RunStage8, and
RunStage8Pdl to create an ffi::CUDADeviceGuard for input.device().device_id
before TMA encoding and LaunchVariant execution. Keep the guard active through
stream retrieval, cuTensorMap encoding, and kernel launch so all SM100 setup and
launch calls run on input’s device.
---
Nitpick comments:
In `@csrc/tinygemm2_sm100/tinygemm2_sm100.cu`:
- Around line 1539-1540: The RunStage* CUDA setup currently repeats immutable
queries on every invocation. In csrc/tinygemm2_sm100/tinygemm2_sm100.cu lines
1539-1540, cache successful cudaFuncSetAttribute state once per kernel function
pointer and skip subsequent calls; in lines 1409-1419, cache compute capability
by device ID and validate using the cached value rather than querying CUDA each
time.
- Around line 1467-1470: Update the input validation near the existing
in_features check in CheckInputs to require in_features % 8 == 0, with a clear
message stating that in_features must be a multiple of 8 for the TMA global
stride. Preserve the existing minimum-size validation and out_features checks.
In `@flashinfer/jit/tinygemm2_sm100.py`:
- Around line 34-35: Update the CUDA flags passed via extra_cuda_cflags in the
relevant build configuration to use iterable unpacking of sm100a_nvcc_flags
followed by the compute_103a flag, preserving the existing order and values
while removing list concatenation.
🪄 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: beae5624-ed36-47bb-bdcf-cb15e8c26b41
📥 Commits
Reviewing files that changed from the base of the PR and between 48f225a1dd1ccf1dc35dff9167a8c5b0038d3885 and 57cf3f00aed451faa3d7c7d460d140f16754168d.
📒 Files selected for processing (3)
csrc/tinygemm2_sm100/tinygemm2_sm100.cuflashinfer/jit/tinygemm2_sm100.pytests/model_optimizations/test_tinygemm2_sm100.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/model_optimizations/test_tinygemm2_sm100.py
| // 2D TMA descriptor for the weight matrix — field-for-field the descriptor | ||
| // the Loom host shim encodes for 'tmap_wt': box (kTileK, kTileM), 128B | ||
| // swizzle, no L2 promotion, no OOB fill. Both boxed axes stay in bounds | ||
| // (CheckInputs guarantees in_features >= kTileK and out_features >= kTileM). |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Comment contradicts the kernels' K-axis access pattern.
Both loaders iterate k_loops = (K + 1023) / 1024 and issue boxes at k_base + i*64 (lines 406-421, 986-1001), so the K axis is read up to ceil(K/1024)*1024 — out of bounds whenever in_features % 1024 != 0 (the test suite's in_features=720 hits this). Correctness there relies on TMA zero-fill, exactly the property this comment denies; the activation comment at Line 1501-1503 documents zero-fill only for the batch axis. Please restate that the K axis may also overrun and is zero-filled, otherwise a future reader may "harden" the descriptor and silently break results.
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 `@csrc/tinygemm2_sm100/tinygemm2_sm100.cu` around lines 1480 - 1483, Update the
TMA descriptor comments near the weight-matrix descriptor to state that the
kernel’s K-axis traversal may extend to ceil(K/1024)*1024 beyond in_features and
relies on TMA zero-fill for the overrun, rather than claiming both boxed axes
remain in bounds or that there is no OOB fill. Preserve the existing description
of the descriptor’s box dimensions and swizzle behavior.
Source: Coding guidelines
| void RunStage4(TensorView input, TensorView weight, TensorView bias, TensorView out) { | ||
| const ProblemDims dims = CheckInputs(input, weight, bias, out); | ||
| const CUtensorMap weight_map = EncodeWeightTma(weight); | ||
| const CUtensorMap activation_map = EncodeActivationTma(input); | ||
| const cudaStream_t stream = get_stream(input.device()); | ||
| LaunchVariant(&kernel_tinygemm2_sm100_stage4, kSmemBytesStage4, /*pdl=*/false, weight_map, | ||
| activation_map, reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), | ||
| reinterpret_cast<__nv_bfloat16*>(bias.data_ptr()), dims, stream); | ||
| } | ||
|
|
||
| // out = input @ weight.T + bias (bf16, fp32 accumulation), column-major | ||
| // epilogue identical to csrc/tinygemm2.cu. | ||
| void RunStage4Pdl(TensorView input, TensorView weight, TensorView bias, TensorView out) { | ||
| const ProblemDims dims = CheckInputs(input, weight, bias, out); | ||
| const CUtensorMap weight_map = EncodeWeightTma(weight); | ||
| const CUtensorMap activation_map = EncodeActivationTma(input); | ||
| const cudaStream_t stream = get_stream(input.device()); | ||
| LaunchVariant(&kernel_tinygemm2_sm100_stage4_pdl, kSmemBytesStage4Pdl, /*pdl=*/true, weight_map, | ||
| activation_map, reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), | ||
| reinterpret_cast<__nv_bfloat16*>(bias.data_ptr()), dims, stream); | ||
| } | ||
|
|
||
| // out = input @ weight.T + bias (bf16, fp32 accumulation), column-major | ||
| // epilogue identical to csrc/tinygemm2.cu. | ||
| void RunStage8(TensorView input, TensorView weight, TensorView bias, TensorView out) { | ||
| const ProblemDims dims = CheckInputs(input, weight, bias, out); | ||
| const CUtensorMap weight_map = EncodeWeightTma(weight); | ||
| const CUtensorMap activation_map = EncodeActivationTma(input); | ||
| const cudaStream_t stream = get_stream(input.device()); | ||
| LaunchVariant(&kernel_tinygemm2_sm100_stage8, kSmemBytesStage8, /*pdl=*/false, weight_map, | ||
| activation_map, reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), | ||
| reinterpret_cast<__nv_bfloat16*>(bias.data_ptr()), dims, stream); | ||
| } | ||
|
|
||
| // out = input @ weight.T + bias (bf16, fp32 accumulation), column-major | ||
| // epilogue identical to csrc/tinygemm2.cu. | ||
| void RunStage8Pdl(TensorView input, TensorView weight, TensorView bias, TensorView out) { | ||
| const ProblemDims dims = CheckInputs(input, weight, bias, out); | ||
| const CUtensorMap weight_map = EncodeWeightTma(weight); | ||
| const CUtensorMap activation_map = EncodeActivationTma(input); | ||
| const cudaStream_t stream = get_stream(input.device()); | ||
| LaunchVariant(&kernel_tinygemm2_sm100_stage8_pdl, kSmemBytesStage8Pdl, /*pdl=*/true, weight_map, | ||
| activation_map, reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), | ||
| reinterpret_cast<__nv_bfloat16*>(bias.data_ptr()), dims, stream); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the guard type/ctor signature used elsewhere in csrc.
rg -nP --type=cu -C2 '\bCUDADeviceGuard\b' | head -50
fd -t f 'tvm_ffi_utils.h' --exec rg -nP -C4 'class\s+CUDADeviceGuard|struct\s+CUDADeviceGuard'Repository: flashinfer-ai/flashinfer
Length of output: 193
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== file exists =="
git ls-files | rg '(^|/)csrc/tinygemm2_sm100/tinygemm2_sm100\.cu$|tvm_ffi_utils\.h$' || true
echo "== CUDADeviceGuard occurrences in tracked source =="
rg -n -C3 '\bCUDADeviceGuard\b' --glob '*.cu' --glob '*.h' --glob '*.cuh' .
echo "== get_stream / LaunchVariant / EncodeWeightTma definitions =="
rg -n -C5 'CUtensorMap Encode(Weight|Activation)Tma|CUDADeviceGuard|get_stream\(|LaunchVariant\b|cudaLaunchKernelEx|cudaFuncSetAttribute|cudaSetDevice\b' --glob '*.cu' --glob '*.h' --glob '*.cuh' .Repository: flashinfer-ai/flashinfer
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== target file outline / relevant lines =="
wc -l csrc/tinygemm2_sm100/tinygemm2_sm100.cu
sed -n '1,80p' csrc/tinygemm2_sm100/tinygemm2_sm100.cu
sed -n '1520,1635p' csrc/tinygemm2_sm100/tinygemm2_sm100.cu
echo "== locate tinygemm2_sm100 helpers =="
rg -n -C4 'std::variant|struct DeviceGuard|class DeviceGuard|struct CUDADeviceGuard|struct CUDAStreamGuard|inline cudaStream_t get_stream|TensorView device\(\)|cuTensorMapEncodeTiled|LaunchVariant\b' csrc/tinygemm2_sm100 csrc/include csrc -g '*tinygemm2_sm100*' --glob '*.h' --glob '*.cu' --glob '*.cuh'Repository: flashinfer-ai/flashinfer
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== tvm_ffi_utils.h relevant definitions =="
sed -n '300,355p' csrc/tvm_ffi_utils.h
rg -n -C4 'struct CUDADeviceGuard|struct CUDAStreamGuard|inline cudaStream_t get_stream\(const TensorView|class TensorView|struct TensorView' csrc/tvm_ffi_utils.h
echo "== compact CUDADeviceGuard usage examples =="
rg -n --glob '*.cu' --glob '*.h' --glob '*.cuh' 'ffi::CUDADeviceGuard device_guard\(' | head -30
echo "== target launcher/calls compact =="
sed -n '1480,1615p' csrc/tinygemm2_sm100/tinygemm2_sm100.cuRepository: flashinfer-ai/flashinfer
Length of output: 2731
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== CUDADeviceGuard definition and surrounding helpers =="
rg -n -C6 'CUDADeviceGuard|CUDAStreamGuard|TVMFFIEnvGetStream|kDLCUDA' csrc/tvm_ffi_utils.h csrc | head -160
echo "== compact target launcher/code =="
sed -n '1480,1615p' csrc/tinygemm2_sm100/tinygemm2_sm100.cuRepository: flashinfer-ai/flashinfer
Length of output: 19489
Add the device guard around the SM100 TMA/lane setup.
RunStage4/4Pdl/8/8Pdl retrieve input.device()’s stream but do not make that device current before cuTensorMapEncodeTiled, cudaFuncSetAttribute, or the kernel launch. Add ffi::CUDADeviceGuard device_guard(input.device().device_id); before the encoding/launcher path and keep LaunchVariant device-scoped for the smem/launch calls, matching the existing csrc pattern.
🤖 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/tinygemm2_sm100/tinygemm2_sm100.cu` around lines 1568 - 1612, Update
RunStage4, RunStage4Pdl, RunStage8, and RunStage8Pdl to create an
ffi::CUDADeviceGuard for input.device().device_id before TMA encoding and
LaunchVariant execution. Keep the guard active through stream retrieval,
cuTensorMap encoding, and kernel launch so all SM100 setup and launch calls run
on input’s device.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
flashinfer/jit/tinygemm2.py (1)
30-31: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse iterable unpacking for the CUDA flags.
Ruff RUF005 flags concatenating a one-element list here. This preserves the same flags while matching the repository’s lint configuration.
Proposed fix
- extra_cuda_cflags=sm100a_nvcc_flags - + ["-gencode=arch=compute_103a,code=sm_103a"], + extra_cuda_cflags=[ + *sm100a_nvcc_flags, + "-gencode=arch=compute_103a,code=sm_103a", + ],🤖 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/tinygemm2.py` around lines 30 - 31, Update the extra_cuda_cflags construction in the tinygemm2 configuration to use iterable unpacking for sm100a_nvcc_flags and the additional CUDA flag instead of list concatenation, preserving the exact flag order and values.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/tinygemm2_sm100.cu`:
- Around line 1466-1474: Add an ICHECK alongside the existing dimension
validation in the kernel entry-point checks, requiring in_features to be
divisible by 8 and reporting the received value and alignment requirement. Keep
the existing minimum-size and scalar-range validations unchanged so misaligned
inputs fail with a clear shape diagnostic before EncodeWeightTma or
EncodeActivationTma runs.
In `@flashinfer/gemm/routergemm.py`:
- Around line 518-522: The documentation for _use_tinygemm2_sm100 inaccurately
claims all SM100/SM103 devices use the generated kernel. Update its docstring to
state that dispatch requires compute capability 10.x and CUDA 12.8 or newer via
is_sm100a_supported, while preserving the FLASHINFER_DISABLE_TINYGEMM2_SM100
override description.
---
Nitpick comments:
In `@flashinfer/jit/tinygemm2.py`:
- Around line 30-31: Update the extra_cuda_cflags construction in the tinygemm2
configuration to use iterable unpacking for sm100a_nvcc_flags and the additional
CUDA flag instead of list concatenation, preserving the exact flag order and
values.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 21b4cd41-c96d-4fe2-b0c2-354287d2ec42
📥 Commits
Reviewing files that changed from the base of the PR and between 57cf3f00aed451faa3d7c7d460d140f16754168d and 80d4ef57adb9e023dd35a47d6087166af96ebaed.
📒 Files selected for processing (4)
csrc/tinygemm2_sm100.cuflashinfer/gemm/routergemm.pyflashinfer/jit/__init__.pyflashinfer/jit/tinygemm2.py
| TVM_FFI_ICHECK(batch > 0) << "batch must be positive, got " << batch; | ||
| TVM_FFI_ICHECK(in_features >= kTileK) | ||
| << "in_features (" << in_features << ") must be at least " << kTileK << " (one TMA box)"; | ||
| TVM_FFI_ICHECK(out_features >= kTileM && out_features % kTileM == 0) | ||
| << "out_features (" << out_features << ") must be a positive multiple of " << kTileM; | ||
| TVM_FFI_ICHECK(batch <= std::numeric_limits<int>::max() && | ||
| in_features <= std::numeric_limits<int>::max() && | ||
| out_features <= std::numeric_limits<int>::max()) | ||
| << "problem dimensions exceed the kernel's i32 scalar range"; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Add the in_features alignment precondition.
cuTensorMapEncodeTiled requires each globalStrides entry to be a multiple of 16 bytes, and both encoders pass stride(0) * sizeof(__nv_bfloat16) = in_features * 2. So in_features must be a multiple of 8; otherwise callers get an opaque CUresult= failure from EncodeWeightTma/EncodeActivationTma instead of a shape diagnostic.
🛡️ Proposed guard
TVM_FFI_ICHECK(in_features >= kTileK)
<< "in_features (" << in_features << ") must be at least " << kTileK << " (one TMA box)";
+ TVM_FFI_ICHECK(in_features % 8 == 0)
+ << "in_features (" << in_features
+ << ") must be a multiple of 8 so the TMA global stride is 16B-aligned";📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| TVM_FFI_ICHECK(batch > 0) << "batch must be positive, got " << batch; | |
| TVM_FFI_ICHECK(in_features >= kTileK) | |
| << "in_features (" << in_features << ") must be at least " << kTileK << " (one TMA box)"; | |
| TVM_FFI_ICHECK(out_features >= kTileM && out_features % kTileM == 0) | |
| << "out_features (" << out_features << ") must be a positive multiple of " << kTileM; | |
| TVM_FFI_ICHECK(batch <= std::numeric_limits<int>::max() && | |
| in_features <= std::numeric_limits<int>::max() && | |
| out_features <= std::numeric_limits<int>::max()) | |
| << "problem dimensions exceed the kernel's i32 scalar range"; | |
| TVM_FFI_ICHECK(batch > 0) << "batch must be positive, got " << batch; | |
| TVM_FFI_ICHECK(in_features >= kTileK) | |
| << "in_features (" << in_features << ") must be at least " << kTileK << " (one TMA box)"; | |
| TVM_FFI_ICHECK(in_features % 8 == 0) | |
| << "in_features (" << in_features | |
| << ") must be a multiple of 8 so the TMA global stride is 16B-aligned"; | |
| TVM_FFI_ICHECK(out_features >= kTileM && out_features % kTileM == 0) | |
| << "out_features (" << out_features << ") must be a positive multiple of " << kTileM; | |
| TVM_FFI_ICHECK(batch <= std::numeric_limits<int>::max() && | |
| in_features <= std::numeric_limits<int>::max() && | |
| out_features <= std::numeric_limits<int>::max()) | |
| << "problem dimensions exceed the kernel's i32 scalar range"; |
🤖 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/tinygemm2_sm100.cu` around lines 1466 - 1474, Add an ICHECK alongside
the existing dimension validation in the kernel entry-point checks, requiring
in_features to be divisible by 8 and reporting the received value and alignment
requirement. Keep the existing minimum-size and scalar-range validations
unchanged so misaligned inputs fail with a clear shape diagnostic before
EncodeWeightTma or EncodeActivationTma runs.
| On SM100/SM103 (B200/B300 class) devices the bias path dispatches to | ||
| ``tinygemm2_sm100`` — generated variants of the same kernel with | ||
| bit-identical outputs and lower latency (see | ||
| ``csrc/tinygemm2_sm100.cu``). Set ``FLASHINFER_DISABLE_TINYGEMM2_SM100=1`` | ||
| to force the reference implementation everywhere. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document the CUDA 12.8 eligibility gate.
_use_tinygemm2_sm100 delegates to is_sm100a_supported, so the optimized path is used only on compute capability 10.x with CUDA 12.8 or newer. The current docstring says all SM100/SM103 devices dispatch to the generated kernel.
- On SM100/SM103 (B200/B300 class) devices the bias path dispatches to
+ When ``is_sm100a_supported(input.device)`` is true (SM100/SM103 with
+ CUDA 12.8 or newer), the bias path dispatches toAs per coding guidelines, documentation must remain synchronized with code changes.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| On SM100/SM103 (B200/B300 class) devices the bias path dispatches to | |
| ``tinygemm2_sm100`` — generated variants of the same kernel with | |
| bit-identical outputs and lower latency (see | |
| ``csrc/tinygemm2_sm100.cu``). Set ``FLASHINFER_DISABLE_TINYGEMM2_SM100=1`` | |
| to force the reference implementation everywhere. | |
| When ``is_sm100a_supported(input.device)`` is true (SM100/SM103 with | |
| CUDA 12.8 or newer), the bias path dispatches to | |
| ``tinygemm2_sm100`` — generated variants of the same kernel with | |
| bit-identical outputs and lower latency (see | |
| ``csrc/tinygemm2_sm100.cu``). Set ``FLASHINFER_DISABLE_TINYGEMM2_SM100=1`` | |
| to force the reference implementation everywhere. |
🤖 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/gemm/routergemm.py` around lines 518 - 522, The documentation for
_use_tinygemm2_sm100 inaccurately claims all SM100/SM103 devices use the
generated kernel. Update its docstring to state that dispatch requires compute
capability 10.x and CUDA 12.8 or newer via is_sm100a_supported, while preserving
the FLASHINFER_DISABLE_TINYGEMM2_SM100 override description.
Source: Coding guidelines
|
/bot run |
Add four frozen device TUs generated from Loom schedules that exactly port csrc/tinygemm2.cu (deep/shallow pipeline ring x PDL on/off), with bit-identical outputs and lower measured latency on B200. The bias path of tinygemm_bf16 dispatches to them on SM100/SM103; set FLASHINFER_DISABLE_TINYGEMM2_SM100=1 to force the reference implementation. - csrc/tinygemm2_sm100/: generated device TUs (verbatim, clang-format off) + one binding TU per variant (typedef isolation + per-variant kernel symbol rename), linked into a single JIT module - stage selection mirrors the measured B200 crossover axes: shallow ring for K <= 1024 or grids past 2x SM count - activation TMA descriptor allows an out-of-bounds box on the batch axis (TMA zero-fills), so batch 1..7 decode shapes are covered - tests: bitwise parity vs csrc/tinygemm2.cu across the batch/K/stage axes, per-variant direct launches, PDL back-to-back, dispatch escape hatch
…le JIT module
Follow the repo convention set by the incumbent csrc/tinygemm2.cu (one
translation unit holding the whole kernel family plus its TVM-FFI
binding, one JitSpec listing that single file): the four generated
SM100/SM103 variants (stage4/stage8 pipeline ring x PDL on/off) now
live in a single csrc/tinygemm2_sm100/tinygemm2_sm100.cu instead of
four device TUs + four binding TUs + a shared binding header.
Mechanical transform (no kernel code touched):
- the extern "C" symbol kernel_flashinfer_tinygemm2 is renamed per
section to kernel_tinygemm2_sm100_{stage4,stage4_pdl,stage8,
stage8_pdl}
- each section keeps its generated #define block verbatim and #undef's
all 17 metadata macros plus the 3 in-kernel mbarrier address macros
at section end; SMEM_TOTAL/USE_PDL are captured into per-variant
constexprs carrying the same static_asserts the per-variant bindings
had
- the generated prelude's fixed-width typedefs and opaque CUtensorMap
typedef are dropped in favor of the host headers; the LoomTensorMap
struct and __device__ helpers (byte-identical across the four
generated TUs) are deduplicated to one copy
- binding_common.cuh and the four binding TUs fold into the same file;
LaunchVariant is de-templatized (one concrete LoomTensorMap now);
the exported op names (stage4_op, stage4_pdl_op, stage8_op,
stage8_pdl_op) are unchanged, so routergemm.py dispatch and the test
semantics needed no changes.
Kernel bodies are byte-identical to the generated output modulo the
symbol rename (asserted by the merge script). SASS identity was
verified with cuobjdump -sass on the JIT-built modules before/after
the merge: all 8 kernel/arch pairs
(4 variants x {sm_100a, sm_103a}) are byte-identical after normalizing only the kernel symbol name (1265/1281/1457/1473
SASS lines per stage4/stage4_pdl/stage8/stage8_pdl kernel).
Validation on GB200 (sm_100):
- tests/model_optimizations/test_tinygemm2_sm100.py: 23 passed
- tests/model_optimizations/test_tinygemm2.py: 59 passed (incumbent path
unaffected)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…mm2.py, reuse is_sm100a_supported - csrc/tinygemm2_sm100/tinygemm2_sm100.cu -> csrc/tinygemm2_sm100.cu: after the single-TU merge the subdirectory held one file; the incumbent convention is a flat csrc/tinygemm2.cu. Pure rename, content unchanged. - gen_tinygemm2_sm100_module moved into flashinfer/jit/tinygemm2.py next to gen_tinygemm2_module (family jit files hold all their variants — see jit/fused_moe.py with seven gens); flashinfer/jit/tinygemm2_sm100.py removed, package-level import path unchanged. - _use_tinygemm2_sm100 now gates on flashinfer.utils.is_sm100a_supported instead of a hand-rolled compute-capability tuple check — same devices, plus the util's CUDA>=12.8 toolchain floor the ad-hoc check was missing. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Same rationale as the dispatch change: reuse flashinfer's own capability util instead of a hand-rolled compute-capability tuple check. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
ceeb31a to
b38a0d6
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tests/model_optimizations/test_tinygemm2_sm100.py (1)
167-188: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the selected backend.
Both backends have the same output contract. The current assertions can pass if
tinygemm_bf16always selects the reference module. Wrap or monkeypatchget_tinygemm2_sm100_moduleandget_tinygemm2_module, then assert that the enabled case calls the SM100 op and the escape-hatch case calls the reference op.🤖 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/model_optimizations/test_tinygemm2_sm100.py` around lines 167 - 188, Strengthen test_tinygemm2_sm100_dispatch_and_escape_hatch by monkeypatching get_tinygemm2_sm100_module and get_tinygemm2_module with wrappers that record invocations while preserving behavior. Assert the enabled call selects the SM100 backend and, after setting FLASHINFER_DISABLE_TINYGEMM2_SM100, the escape-hatch call selects the reference backend, while retaining the existing output 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/gemm/routergemm.py`:
- Line 441: In flashinfer/gemm/routergemm.py at line 441, rename the private
wrapper parameter `input` to a domain-specific identifier like `activations`, or
add a narrow A002 suppression if the registered custom-op schema requires
retaining the `input` name. In
tests/model_optimizations/test_tinygemm2_sm100.py, rename all variables
shadowing the builtin `input` to `activations` or an appropriate domain-specific
name at lines 25, 75, 118, 146, 159, and 174. These renames will resolve the
Ruff A001/A002 linting errors caused by shadowing Python builtins.
In `@tests/model_optimizations/test_tinygemm2_sm100.py`:
- Around line 3-6: Correct the CUDA source path in the test documentation to
reference csrc/tinygemm2_sm100.cu, matching the JIT generator specification; do
not alter the parity tests or their torch.equal behavior.
---
Nitpick comments:
In `@tests/model_optimizations/test_tinygemm2_sm100.py`:
- Around line 167-188: Strengthen test_tinygemm2_sm100_dispatch_and_escape_hatch
by monkeypatching get_tinygemm2_sm100_module and get_tinygemm2_module with
wrappers that record invocations while preserving behavior. Assert the enabled
call selects the SM100 backend and, after setting
FLASHINFER_DISABLE_TINYGEMM2_SM100, the escape-hatch call selects the reference
backend, while retaining the existing output checks.
🪄 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: 0eb1ba4f-0ef5-4621-bd13-e05ab6918772
📥 Commits
Reviewing files that changed from the base of the PR and between ceeb31aee0f75e4d1e9076fa6088daa8c9d95545 and b38a0d6.
📒 Files selected for processing (5)
csrc/tinygemm2_sm100.cuflashinfer/gemm/routergemm.pyflashinfer/jit/__init__.pyflashinfer/jit/tinygemm2.pytests/model_optimizations/test_tinygemm2_sm100.py
🚧 Files skipped from review as they are similar to previous changes (2)
- flashinfer/jit/init.py
- csrc/tinygemm2_sm100.cu
| mutates_args=["out"], | ||
| ) | ||
| def tinygemm2_sm100_op_impl( | ||
| input: torch.Tensor, |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Resolve the new Ruff A001/A002 errors.
New code shadows the Python builtin input. Rename these identifiers to activations or another domain-specific name. If the registered custom-op schema must retain input, use a narrow A002 suppression with that rationale.
flashinfer/gemm/routergemm.py#L441-L441: rename the private wrapper parameter, or add a narrow documented A002 suppression if the schema requires this name.tests/model_optimizations/test_tinygemm2_sm100.py#L25-L25: rename the generated activation tensor.tests/model_optimizations/test_tinygemm2_sm100.py#L75-L75: rename the unpacked activation tensor.tests/model_optimizations/test_tinygemm2_sm100.py#L118-L118: rename the unpacked activation tensor.tests/model_optimizations/test_tinygemm2_sm100.py#L146-L146: rename the loop activation tensor.tests/model_optimizations/test_tinygemm2_sm100.py#L159-L159: rename the unpacked activation tensor.tests/model_optimizations/test_tinygemm2_sm100.py#L174-L174: rename the generated activation tensor.
🧰 Tools
🪛 Ruff (0.16.0)
[error] 441-441: Function argument input is shadowing a Python builtin
(A002)
📍 Affects 2 files
flashinfer/gemm/routergemm.py#L441-L441(this comment)tests/model_optimizations/test_tinygemm2_sm100.py#L25-L25tests/model_optimizations/test_tinygemm2_sm100.py#L75-L75tests/model_optimizations/test_tinygemm2_sm100.py#L118-L118tests/model_optimizations/test_tinygemm2_sm100.py#L146-L146tests/model_optimizations/test_tinygemm2_sm100.py#L159-L159tests/model_optimizations/test_tinygemm2_sm100.py#L174-L174
🤖 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/gemm/routergemm.py` at line 441, In flashinfer/gemm/routergemm.py
at line 441, rename the private wrapper parameter `input` to a domain-specific
identifier like `activations`, or add a narrow A002 suppression if the
registered custom-op schema requires retaining the `input` name. In
tests/model_optimizations/test_tinygemm2_sm100.py, rename all variables
shadowing the builtin `input` to `activations` or an appropriate domain-specific
name at lines 25, 75, 118, 146, 159, and 174. These renames will resolve the
Ruff A001/A002 linting errors caused by shadowing Python builtins.
Source: Linters/SAST tools
| The four frozen variants in ``csrc/tinygemm2_sm100/tinygemm2_sm100.cu`` are | ||
| generated Loom schedules exactly porting ``csrc/tinygemm2.cu``; the contract | ||
| is bit-identical outputs. Every parity test below therefore uses | ||
| ``torch.equal``, not a tolerance. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the documented CUDA source path.
The JIT generator compiles csrc/tinygemm2_sm100.cu. This docstring instead names csrc/tinygemm2_sm100/tinygemm2_sm100.cu. Update the path so the test documentation matches the JIT specification.
Proposed fix
-The four frozen variants in ``csrc/tinygemm2_sm100/tinygemm2_sm100.cu`` are
+The four frozen variants in ``csrc/tinygemm2_sm100.cu`` are📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| The four frozen variants in ``csrc/tinygemm2_sm100/tinygemm2_sm100.cu`` are | |
| generated Loom schedules exactly porting ``csrc/tinygemm2.cu``; the contract | |
| is bit-identical outputs. Every parity test below therefore uses | |
| ``torch.equal``, not a tolerance. | |
| The four frozen variants in ``csrc/tinygemm2_sm100.cu`` are | |
| generated Loom schedules exactly porting ``csrc/tinygemm2.cu``; the contract | |
| is bit-identical outputs. Every parity test below therefore uses | |
| ``torch.equal``, not a tolerance. |
🤖 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/model_optimizations/test_tinygemm2_sm100.py` around lines 3 - 6,
Correct the CUDA source path in the test documentation to reference
csrc/tinygemm2_sm100.cu, matching the JIT generator specification; do not alter
the parity tests or their torch.equal behavior.
Source: Coding guidelines
|
/bot run tests/model_optimizations/test_tinygemm2_sm100.py |
|
/bot run tests/model_optimizations |
…ge-K shapes (#4423) ## 📌 Description Adds a third ring-depth tier to the `tinygemm2_sm100` family (#4274) and moves ring selection into the binding. All six kernels are regenerated from the current CAKE generator state, which also updates the tensor-map parameter passing to a single by-value `__grid_constant__` pack. - **New STAGES=16 ring for single-wave large-K shapes.** The 4/8-stage rings were tuned for K up to ~4K; at larger K the 8-deep ring no longer covers the weight-stream latency when the working set is not L2-resident, and trails the reference kernel's own 16-deep configuration. The new tier closes that: at N=8/M=128/K=7168 (bias path) it measures **4.93 µs vs the reference's 5.22 µs on GB300, and 5.31 µs vs 5.56 µs on B200** (cold-L2 CUPTI). - **Ring selection in the binding**, following the reference `csrc/tinygemm2.cu` launcher convention: stage 4 for K <= 1024 or grids past 2x the SM count (unchanged), stage 16 for single-wave grids with K >= 4608 (measured crossover on both GB300 and B200), stage 8 otherwise. The Python dispatcher becomes a single call into the combined op. - **Dispatch gates on exact compute capabilities (10, 0)/(10, 3).** SM107 passes the previous `major == 10` predicate but must keep the reference path instead of erroring in the binding. - The dynamic-SMEM attribute is now set once per (kernel, device) instead of on every launch. **Correctness contract unchanged: bitwise equality with `csrc/tinygemm2.cu`** — verified per-variant and through the dispatcher (`torch.equal`, batch 1-64, K to 7168, M to 4096, on B200 and GB300), and end-to-end in SGLang serving with zero token flips (gpt-oss-120b, and Mistral-Large-3 whose router GEMM sits in the new tier). `tests/model_optimizations/test_tinygemm2_sm100.py` extends to the stage-16 variants and long-K parity shapes. compute-sanitizer synccheck is clean on all six variants. Limitations: unchanged from #4274 — the nobias path stays on the reference kernel. ## 📊 End-to-end serving validation — Mistral-Large-3 (675B FP8), SGLang, TP4, GB300 (CUDA graphs on, production defaults) **Correctness**: 32 fixed greedy prompts — token streams bitwise identical between arms; GSM8K-200: CAKE 0.950 / ref 0.945. **Performance** (CAKE = this PR's kernels via default dispatch, ref = reference `tinygemm2`; throughput in mean output tok/s, ITL is median in ms): | concurrency | throughput (CAKE) | throughput (ref) | ITL (CAKE) | ITL (ref) | |---|---|---|---|---| | 1 | 100.7 | 101.0 | 9.47 | 9.48 | | 8 | 488.6 | 487.9 | 15.47 | 15.48 | | 32 | 1068 | 1073 | 28.90 | 28.92 | | 128 | 3000 | 2997 | 41.31 | 41.44 | ## 🔍 Related Issues Follow-up to #4274. Related to #4254 (CAKE-generated kernel progress tracker). ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes The generated device sections are frozen artifacts (regeneration happens outside this repo, as with #4262/#4274); review focus is best spent on the binding section of `csrc/tinygemm2_sm100.cu` (pack construction, stage selection, launch attributes), the dispatcher in `flashinfer/gemm/routergemm.py`, and the test additions. The single-TU merge follows the mechanical transform documented in the file header. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for larger matrix workloads through new stage16 execution variants. * Automatically selects the appropriate execution stage based on workload size and hardware. * Added direct launch support for all available execution variants. * Expanded compatibility for supported compute capabilities and newer CUDA versions. * **Bug Fixes** * Improved handling of large reduction sizes and deep-ring workloads. * **Tests** * Added coverage for larger K dimensions, including 7168 and 14336, and stage16 variants. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
6 participants
📌 Description
This PR adds
tinygemm2_sm100: four CAKE-generated SM100/SM103 devicekernels specializing
csrc/tinygemm2.cu(the TensorRT-LLM tinygemm2 portbehind
tinygemm_bf16) for Blackwell, with bit-identical outputs andlower latency. The bias path
of
tinygemm_bf16automatically dispatches to them on compute capability10.0/10.3; every other path and architecture is unchanged, and
FLASHINFER_DISABLE_TINYGEMM2_SM100=1forces the reference implementation.The kernels are CAKE-generated Loom schedules, following the
frozen-generated-source pattern of the CAKE-generated FlashKDA prefill
export (#4262): generated kernel sources checked in verbatim
(
clang-format off, provenance headers, no host-library dependency),concatenated into a single translation unit
(
csrc/tinygemm2_sm100/tinygemm2_sm100.cu) with per-variant kernel symbolrenames and a small hand-written binding section doing validation, TMA
descriptor encode, and launch — one TU, one JIT module, mirroring the
incumbent
csrc/tinygemm2.culayout. The single-TU restructure isSASS-verified: all four kernels are byte-identical to per-variant builds on
both sm_100a and sm_103a (
cuobjdump -sass, symbol names normalized).The variants are {deep, shallow} pipeline ring x {PDL on, off}:
(shallow ring for
K <= 1024or grids past 2x the SM count);griddepcontrolpair in-kernel and launch withprogrammatic stream serialization, matching the reference kernel's
USE_PDL=trueinstantiation;(TMA zero-fills), covering batch 1-7 decode shapes.
Correctness contract: bitwise equality with
csrc/tinygemm2.cu, not atolerance. On B200:
torch.equalparity across batch {1,2,4,7,8,13,16,64}x five (M, K) pairs x PDL on/off, per-variant direct launches, PDL
back-to-back replay, dispatch/escape-hatch — 23/23; existing
test_tinygemm2.pyunchanged and green (59/59) with the bias path routed tothe new backend. Internally, the same generated kernels passed a 239-row
bitwise shape sweep and end-to-end gpt-oss serving with zero token flips on
both B200 and GB300 (sm_103a): greedy outputs bit-identical across arms on
gpt-oss-20b and -120b, GSM8K-200 scores identical.
Performance (torch.profiler kernel-time medians, B200, bias path, 200
launches): (1,128,720) 2.24us vs 2.34us reference; (16,1024,1024) 2.50 vs
2.58; (64,4096,3072) 18.22 vs 32.58 (1.79x). Small decode shapes are
launch-bound; a wider internal CUPTI sweep (cold L2, 35 canonical + 239
regression shapes) measured an 18-23% geometric-mean kernel-time reduction. In SGLang serving
(gpt-oss-120b) the swap measures up to +7.6% output throughput at high
concurrency (c128, TP1) over the reference path, and parity within noise at
TP4.
Limitations: the nobias path stays on the reference kernel (no generated
nobias variant yet).
🔍 Related Issues
Related to #4254 (long-term CAKE-generated kernel progress tracker).
Follows the frozen generated-kernel export pattern introduced in #4262.
🚀 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.).Reviewer Notes
The generated kernel sections are frozen artifacts (regeneration happens
outside this repo, as with FlashKDA); review focus is best spent on the
binding section at the tail of
csrc/tinygemm2_sm100/tinygemm2_sm100.cu, thedispatch in
flashinfer/gemm/routergemm.py, and the parity tests. The TMA descriptorencode in the binding is field-for-field identical to the reference kernel's
own
cuTensorMapEncodeTiledcalls.Summary by CodeRabbit
New Features
Bug Fixes
Tests