Revert PR 4122 - #4171
Merged
Merged
Revert PR 4122#4171
Conversation
jimmyzho
requested review from
Anerudhan,
Aneureka,
IwakuraRein,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jiahanc,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 27, 2026 08:00
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
📝 WalkthroughWalkthroughThis change removes SM107 architecture support across runtime dispatch, JIT/AOT compilation, quantization, GEMM/MoE/TRTLLM kernels, and CuTe DSL checks. It also updates tests, artifact references, sampling declarations, and CUDA architecture documentation. ChangesArchitecture support removal
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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 |
Collaborator
|
/bot run |
Collaborator
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/utils/test_fp4_quantize.py (1)
386-393: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winDo not replace GPU capability checks with importability-only checks.
Both helpers now treat an importable CuTe DSL stack as sufficient, so unsupported GPUs can execute architecture-specific tests instead of being skipped.
tests/utils/test_fp4_quantize.py#L386-L393: retain the CuTe DSL architecture predicate alongside the importability check.tests/utils/test_fp4_quantize_padding.py#L43-L50: restore the same architecture-aware gating behavior.🤖 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_fp4_quantize.py` around lines 386 - 393, Update _is_cute_dsl_available in tests/utils/test_fp4_quantize.py (386-393) and its corresponding helper in tests/utils/test_fp4_quantize_padding.py (43-50) to require both CuTe-DSL importability and the existing supported-GPU architecture predicate, preserving skips on unsupported GPUs.Source: Coding guidelines
🧹 Nitpick comments (1)
include/flashinfer/sampling.cuh (1)
293-294: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftDocument and benchmark the
__launch_bounds__removal.These sampling and softmax kernels are hot paths. Removing
__launch_bounds__(BLOCK_THREADS)changes the compiler’s register-allocation and occupancy constraints, but the diff gives no rationale or performance evidence. Add a short explanation of the intended trade-off and benchmark the affected paths, or retain the attribute where it remains beneficial.Also applies to: 410-412, 495-498, 724-726, 773-775, 837-840, 970-973, 1097-1100, 1192-1196, 1663-1664, 1860-1866
🤖 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 `@include/flashinfer/sampling.cuh` around lines 293 - 294, Review every affected sampling and softmax kernel, including OnlineSoftmaxFusedKernel, for the removed __launch_bounds__ attribute. Benchmark the relevant hot paths and document the intended register-allocation/occupancy trade-off, or restore __launch_bounds__(BLOCK_THREADS) on kernels where measurements show it remains beneficial.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/compilation_context.py`:
- Around line 98-101: Update the generic NVCC flag generation in the relevant
compilation-context method to remove SM107 from supported_cuda_archs before
sorting and constructing -gencode entries. Apply this centrally so both
device-detected and FLASHINFER_CUDA_ARCH_LIST architectures are filtered, while
preserving all other architecture flags and COMMON_NVCC_FLAGS.
- Line 84: Update the supported_major_versions parameter annotation in the
relevant constructor or function from list[int] to list[int] | None while
preserving its default value of None.
In `@flashinfer/cute_dsl/gemm_allreduce_two_shot.py`:
- Around line 1191-1195: Update the CooperativeGroup construction in the
epilogue producer setup to pass the agent and only one group size/alignment
value. Remove the duplicated 32 * len(self.epilog_warp_id) positional argument
while preserving the existing group-size calculation.
In `@flashinfer/mamba/ssd_combined.py`:
- Around line 292-296: Update the architecture validation in the SSDCombined
capability check to allow only SM100, SM103, and SM110, rejecting SM107 and all
other major/minor combinations before kernel compilation. Keep the existing
ValueError and its supported-architecture message aligned with this explicit
allowlist.
In `@tests/gemm/test_cute_dsl_gemm_allreduce_two_shot.py`:
- Around line 490-491: Update the pytest.skip message in the architecture check
around get_compute_capability to mention both supported architectures, SM100 and
SM103, matching the allowed (10, 0) and (10, 3) condition.
In `@tests/utils/test_fp8_quantize.py`:
- Around line 16-21: Restore the shared exact architecture policy using the
capability utilities, keeping _is_available in
tests/utils/test_fp8_quantize.py:16-21 limited to dependency availability. In
tests/gemm/test_cute_dsl_blockscaled_gemm.py:263-267 and :314-318, restore
architecture gates; narrow the predicates at
tests/gemm/test_groupwise_scaled_gemm_fp8.py:52, :98-106, :227, :288, and :336
and exclude SM107 at tests/gemm/test_groupwise_scaled_gemm_mxfp4.py:257-261.
Replace the major-version-only marker at
tests/mamba/test_chunk_scan_combined.py:146-151, restore the shared marker at
tests/moe/test_cute_dsl_fused_moe.py:36-37, and retain its gating at :862,
:1128, :1650, :1729, :1933, and :2129 for the respective architecture-specific
tests.
---
Outside diff comments:
In `@tests/utils/test_fp4_quantize.py`:
- Around line 386-393: Update _is_cute_dsl_available in
tests/utils/test_fp4_quantize.py (386-393) and its corresponding helper in
tests/utils/test_fp4_quantize_padding.py (43-50) to require both CuTe-DSL
importability and the existing supported-GPU architecture predicate, preserving
skips on unsupported GPUs.
---
Nitpick comments:
In `@include/flashinfer/sampling.cuh`:
- Around line 293-294: Review every affected sampling and softmax kernel,
including OnlineSoftmaxFusedKernel, for the removed __launch_bounds__ attribute.
Benchmark the relevant hot paths and document the intended
register-allocation/occupancy trade-off, or restore
__launch_bounds__(BLOCK_THREADS) on kernels where measurements show it remains
beneficial.
🪄 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: 4d561f28-968b-42ac-be82-c497ec3e28d0
📒 Files selected for processing (80)
README.mdbenchmarks/routines/flashinfer_benchmark_utils.pycsrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inlcsrc/trtllm_batched_gemm_runner.cucsrc/trtllm_gemm_runner.cudocs/installation.rstflashinfer/aot.pyflashinfer/artifacts.pyflashinfer/compilation_context.pyflashinfer/cute_dsl/add_rmsnorm_fp4quant.pyflashinfer/cute_dsl/attention/compat.pyflashinfer/cute_dsl/attention/mla_decode.pyflashinfer/cute_dsl/attention/mla_decode_fp8.pyflashinfer/cute_dsl/attention/prefill.pyflashinfer/cute_dsl/attention/wrappers/batch_decode.pyflashinfer/cute_dsl/attention/wrappers/batch_mla.pyflashinfer/cute_dsl/attention/wrappers/batch_prefill.pyflashinfer/cute_dsl/gemm_allreduce_two_shot.pyflashinfer/cute_dsl/rmsnorm_fp4quant.pyflashinfer/cute_dsl/utils.pyflashinfer/cutile/cutile_common.pyflashinfer/decode.pyflashinfer/deep_gemm.pyflashinfer/fp4_quantization.pyflashinfer/fused_moe/core.pyflashinfer/fused_moe/cute_dsl/fused_moe.pyflashinfer/fused_moe/fused_routing_dsv3.pyflashinfer/fused_moe/hash_topk.pyflashinfer/gemm/gemm_base.pyflashinfer/gemm/routergemm.pyflashinfer/jit/__init__.pyflashinfer/jit/attention/modules.pyflashinfer/jit/core.pyflashinfer/jit/cpp_ext.pyflashinfer/jit/fp4_quantization.pyflashinfer/jit/fp8_quantization.pyflashinfer/jit/fused_moe.pyflashinfer/jit/gemm/core.pyflashinfer/jit/gemm/cutlass/generate_kernels.pyflashinfer/jit/mla.pyflashinfer/jit/xqa.pyflashinfer/kda_decode.pyflashinfer/kda_kernels/__init__.pyflashinfer/mamba/ssd_combined.pyflashinfer/mla/_core.pyflashinfer/norm/__init__.pyflashinfer/prefill.pyflashinfer/quantization/fp4_quantization.pyinclude/flashinfer/sampling.cuhinclude/flashinfer/trtllm/common.hinclude/flashinfer/trtllm/fmha/fmhaKernels.cuhinclude/flashinfer/trtllm/fmha/fmhaRunner.cuhtests/attention/test_cute_dsl_decode.pytests/attention/test_cute_dsl_mla_decode.pytests/attention/test_deepseek_mla.pytests/attention/test_modular_fmha_prefill.pytests/attention/test_trtllm_gen_attention_decode.pytests/attention/test_trtllm_gen_attention_prefill.pytests/attention/test_trtllm_gen_mla.pytests/attention/test_vsa_block_sparse.pytests/comm/test_mnnvl_custom_comm.pytests/gdn/test_decode_delta_rule.pytests/gdn/test_decode_pretranspose_bf16_padded_pool.pytests/gdn/test_decode_pretranspose_noncontiguous_pool.pytests/gemm/test_cute_dsl_blockscaled_gemm.pytests/gemm/test_cute_dsl_gemm_allreduce_two_shot.pytests/gemm/test_groupwise_scaled_gemm_fp8.pytests/gemm/test_groupwise_scaled_gemm_mxfp4.pytests/mamba/conftest.pytests/mamba/test_chunk_scan_combined.pytests/model_optimizations/test_router_gemms.pytests/moe/test_cute_dsl_fused_moe.pytests/moe/test_trtllm_cutlass_fused_moe.pytests/moe/test_unified_moe.pytests/norm/test_add_rmsnorm_fp4_quant_cute_dsl.pytests/norm/test_rmsnorm_fp4_quant_cute_dsl.pytests/utils/test_fp4_quantize.pytests/utils/test_fp4_quantize_padding.pytests/utils/test_fp4_tensor_torch_cute.pytests/utils/test_fp8_quantize.py
💤 Files with no reviewable changes (31)
- flashinfer/jit/init.py
- flashinfer/jit/core.py
- tests/utils/test_fp4_tensor_torch_cute.py
- tests/norm/test_add_rmsnorm_fp4_quant_cute_dsl.py
- include/flashinfer/trtllm/common.h
- tests/attention/test_cute_dsl_decode.py
- tests/gdn/test_decode_pretranspose_noncontiguous_pool.py
- tests/attention/test_modular_fmha_prefill.py
- flashinfer/cute_dsl/add_rmsnorm_fp4quant.py
- flashinfer/cute_dsl/attention/wrappers/batch_mla.py
- flashinfer/cutile/cutile_common.py
- tests/attention/test_trtllm_gen_attention_prefill.py
- flashinfer/fused_moe/cute_dsl/fused_moe.py
- tests/attention/test_trtllm_gen_attention_decode.py
- flashinfer/kda_kernels/init.py
- tests/comm/test_mnnvl_custom_comm.py
- flashinfer/mla/_core.py
- tests/norm/test_rmsnorm_fp4_quant_cute_dsl.py
- tests/gdn/test_decode_pretranspose_bf16_padded_pool.py
- flashinfer/jit/fp4_quantization.py
- tests/attention/test_vsa_block_sparse.py
- flashinfer/aot.py
- flashinfer/cute_dsl/attention/wrappers/batch_decode.py
- flashinfer/cute_dsl/attention/wrappers/batch_prefill.py
- flashinfer/fp4_quantization.py
- tests/attention/test_cute_dsl_mla_decode.py
- tests/attention/test_trtllm_gen_mla.py
- flashinfer/cute_dsl/rmsnorm_fp4quant.py
- benchmarks/routines/flashinfer_benchmark_utils.py
- tests/mamba/conftest.py
- flashinfer/cute_dsl/attention/compat.py
| self, | ||
| supported_major_versions: list[int] = None, | ||
| map_sm107_to_100f: bool = False, | ||
| self, supported_major_versions: list[int] = None |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
ruff check flashinfer/compilation_context.py --select RUF013Repository: flashinfer-ai/flashinfer
Length of output: 767
Use an explicit nullable annotation.
Ruff checks RUF013 on flashinfer/compilation_context.py:84; replace list[int] = None with list[int] | None = None.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 84-84: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
🤖 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/compilation_context.py` at line 84, Update the
supported_major_versions parameter annotation in the relevant constructor or
function from list[int] to list[int] | None while preserving its default value
of None.
Source: Linters/SAST tools
Comment on lines
+98
to
+101
| return [ | ||
| f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}" | ||
| for major, minor in sorted(supported_cuda_archs) | ||
| ] + self.COMMON_NVCC_FLAGS |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Exclude SM107 before forming generic NVCC flags.
TARGET_CUDA_ARCHS can still contain SM107 from device detection or FLASHINFER_CUDA_ARCH_LIST; filtering only by major version emits compute_107a/sm_107a for every caller requesting major 10. That leaves SM107 JIT targeting active despite this revert. Reject or filter SM107 centrally before generating flags.
🤖 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/compilation_context.py` around lines 98 - 101, Update the generic
NVCC flag generation in the relevant compilation-context method to remove SM107
from supported_cuda_archs before sorting and constructing -gencode entries.
Apply this centrally so both device-detected and FLASHINFER_CUDA_ARCH_LIST
architectures are filtered, while preserving all other architecture flags and
COMMON_NVCC_FLAGS.
Comment on lines
1191
to
1195
| c_producer_group = pipeline.CooperativeGroup( | ||
| pipeline.Agent.Thread, | ||
| 32 * len(self.epilog_warp_id), | ||
| 32 * len(self.epilog_warp_id), | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate file =="
git ls-files | rg '^flashinfer/cute_dsl/gemm_allreduce_two_shot\.py$|cutlass|CooperativeGroup|flashinfer/cute_dsl'
echo
echo "== target lines =="
if [ -f flashinfer/cute_dsl/gemm_allreduce_two_shot.py ]; then
sed -n '1160,1215p' flashinfer/cute_dsl/gemm_allreduce_two_shot.py | nl -ba -v1160
fi
echo
echo "== CooperativeGroup usages in repo =="
rg -n "CooperativeGroup\s*\(" . -g '*.py' -g '*.cu' -g '*.cpp' -g '*.h' -g '*.hpp' -g '*.cuh' -g '*.jinja'
echo
echo "== imports around cooperative group/pipeline =="
rg -n "from cutlass|import cutlass|class CooperativeGroup|def CooperativeGroup|CooperativeGroup" flashinfer/cute_dsl -g '*.py' || trueRepository: flashinfer-ai/flashinfer
Length of output: 22081
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== target lines =="
sed -n '1160,1215p' flashinfer/cute_dsl/gemm_allreduce_two_shot.py
echo
echo "== CooperativeGroup constructor/class definitions =="
rg -n "class CooperativeGroup|def CooperativeGroup|CooperativeGroup" 3rdparty flashinfer/cute_dsl -g '*.py' -g '*.cu' -g '*.cpp' -g '*.h' -g '*.hpp' -g '*.cuh' -g '*.inl' | head -200
echo
echo "== cooperative group constructor snippets =="
python3 - <<'PY'
import re, os
for root, dirs, files in os.walk('flashinfer/cute_dsl'):
for fname in files:
if not fname.endswith('.py'):
continue
path = os.path.join(root, fname)
text = open(path, encoding='utf-8', errors='ignore').read()
for m in re.finditer(r'class\s+CooperativeGroup\b|def\s+CooperativeGroup\b', text):
line = text.count('\n',0,m.start())+1
print(f'-- {path}:{line} --')
PY
echo
echo "== all CooperativeGroup call expressions with line numbers =="
rg -n "CooperativeGroup\s*\(" flashinfer/cute_dsl -g '*.py' -C 2Repository: flashinfer-ai/flashinfer
Length of output: 27230
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== imports in target file =="
sed -n '1,90p' flashinfer/cute_dsl/gemm_allreduce_two_shot.py
echo
echo "== CooperativeGroup definitions in 3rdparty/cutlass =="
rg -n "class CooperativeGroup|CooperativeGroup\s*=" 3rdparty/cutlass -g '*.h' -g '*.hpp' -g '*.inl' -g '*.h' -g '*.cuh' | head -80
echo
echo "== surrounding target file section around earlier cooperative groups =="
sed -n '710,746p' flashinfer/cute_dsl/gemm_allreduce_two_shot.pyRepository: flashinfer-ai/flashinfer
Length of output: 2531
🌐 Web query:
CUTLASS pipeline CooperativeGroup constructor Agent Thread count
💡 Result:
In the CUTLASS library's Python DSL (CuTeDSL), the CooperativeGroup constructor is defined within the cutlass.pipeline module [1][2][3]. It is used to specify the participation of threads in a pipeline's producer or consumer agents [1][4]. The constructor signature is: class CooperativeGroup(agent: Agent, size: int = 1, alignment: int | None = None) [2][3][5] Parameters: agent (Agent): Defines the type of thread grouping, typically selected from Agent.Thread, Agent.Warp, Agent.ThreadBlock, or Agent.ThreadBlockCluster [1][3]. size (int): Specifies the number of participating agents (e.g., the number of threads if Agent.Thread is used) [3][4]. It must be greater than 0 [3]. alignment (int | None): An optional parameter that is deprecated and should be removed from code [3]. The size parameter serves as the thread count (or agent count) for the group [3][4]. When Agent.Thread is used, the size represents the specific number of threads participating in the producer or consumer group [4]. For example, developers often set this to 1 for a single-threaded producer or scale it based on the number of warps or threads involved in specific operations like TMA (Tensor Memory Accelerator) store or multicast [4]. Top results: [1], [3], [4]
Citations:
- 1: https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_api/pipeline.html
- 2: https://docs.nvidia.com/cutlass/4.5.3/media/docs/pythonDSL/cute_dsl_api/pipeline.html
- 3: https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/pipeline/helpers.py
- 4: https://veitner.bearblog.dev/blackwell-pipelining-with-cutedsl/
- 5: https://docs.nvidia.com/cutlass/4.4.1/media/docs/pythonDSL/cute_dsl_api/pipeline.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== CooperativeGroup implementation in installed/available cutlass package =="
python3 - <<'PY'
from pathlib import Path
try:
import cutlass
print("cutlass", getattr(cutlass, "__version__", "unknown"), getattr(cutlass, "__file__", "no __file__"))
except Exception as e:
print("cutlass import unavailable:", type(e).__name__, e)
PY
find /opt /home /usr -path '*cutlass*helpers.py' -o -path '*cute_dsl_api*/pipeline.html' 2>/dev/null | head -20 || trueRepository: flashinfer-ai/flashinfer
Length of output: 313
Remove the duplicated producer-group count.
pipeline.CooperativeGroup takes the agent and a single group size/alignment, so the extra positional argument makes this producer group construction invalid.
Proposed fix
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread,
32 * len(self.epilog_warp_id),
- 32 * len(self.epilog_warp_id),
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| c_producer_group = pipeline.CooperativeGroup( | |
| pipeline.Agent.Thread, | |
| 32 * len(self.epilog_warp_id), | |
| 32 * len(self.epilog_warp_id), | |
| ) | |
| c_producer_group = pipeline.CooperativeGroup( | |
| pipeline.Agent.Thread, | |
| 32 * len(self.epilog_warp_id), | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/cute_dsl/gemm_allreduce_two_shot.py` around lines 1191 - 1195,
Update the CooperativeGroup construction in the epilogue producer setup to pass
the agent and only one group size/alignment value. Remove the duplicated 32 *
len(self.epilog_warp_id) positional argument while preserving the existing
group-size calculation.
Comment on lines
+292
to
296
| if major not in (10, 11): | ||
| raise ValueError( | ||
| f"SSDCombined requires datacenter Blackwell (SM100/SM103/SM110) " | ||
| f"for tcgen05 MMA. Got SM{major}{minor}." | ||
| f"SSDCombined requires SM100-SM110 (tcgen05): Blackwell datacenter " | ||
| f"GPUs (SM100/SM103/SM110). Got SM{major}{minor}." | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Check no longer excludes SM107, undermining its own error message.
major not in (10, 11) lets SM107 (major=10, minor=7) through, even though the comment and error message explicitly state only SM100/SM103/SM110 are supported. On actual SM107 hardware this will now fall through to kernel compilation and hit the cryptic cute-dsl OpError that this check exists to avoid.
🐛 Proposed fix
- if major not in (10, 11):
+ if major not in (10, 11) or (major, minor) == (10, 7):
raise ValueError(
f"SSDCombined requires SM100-SM110 (tcgen05): Blackwell datacenter "
f"GPUs (SM100/SM103/SM110). Got SM{major}{minor}."
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if major not in (10, 11): | |
| raise ValueError( | |
| f"SSDCombined requires datacenter Blackwell (SM100/SM103/SM110) " | |
| f"for tcgen05 MMA. Got SM{major}{minor}." | |
| f"SSDCombined requires SM100-SM110 (tcgen05): Blackwell datacenter " | |
| f"GPUs (SM100/SM103/SM110). Got SM{major}{minor}." | |
| ) | |
| if major not in (10, 11) or (major, minor) == (10, 7): | |
| raise ValueError( | |
| f"SSDCombined requires SM100-SM110 (tcgen05): Blackwell datacenter " | |
| f"GPUs (SM100/SM103/SM110). Got SM{major}{minor}." | |
| ) |
🤖 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/mamba/ssd_combined.py` around lines 292 - 296, Update the
architecture validation in the SSDCombined capability check to allow only SM100,
SM103, and SM110, rejecting SM107 and all other major/minor combinations before
kernel compilation. Keep the existing ValueError and its supported-architecture
message aligned with this explicit allowlist.
Comment on lines
+490
to
+491
| if get_compute_capability(torch.device("cuda")) not in [(10, 0), (10, 3)]: | ||
| pytest.skip("cute_dsl_gemm_allreduce_two_shot requires SM100") |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the architecture skip message.
The condition now allows SM100 and SM103, but the message says only “requires SM100.” Change it to mention both supported architectures so skipped-test diagnostics match the actual condition.
🤖 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/gemm/test_cute_dsl_gemm_allreduce_two_shot.py` around lines 490 - 491,
Update the pytest.skip message in the architecture check around
get_compute_capability to mention both supported architectures, SM100 and SM103,
matching the allowed (10, 0) and (10, 3) condition.
Comment on lines
16
to
21
| try: | ||
| from flashinfer.cute_dsl import is_cute_dsl_available as _is_available | ||
|
|
||
| if not _is_available(): | ||
| return False | ||
| import torch as _torch | ||
|
|
||
| if _torch.cuda.is_available(): | ||
| from flashinfer.cute_dsl.utils import is_cute_dsl_arch_supported | ||
|
|
||
| return is_cute_dsl_arch_supported(*_torch.cuda.get_device_capability(0)) | ||
| return True | ||
| return _is_available() | ||
| except ImportError: | ||
| return False |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Restore one exact architecture policy across the test suite. The current changes either remove architecture checks or retain predicates that accept all major-version-10 devices, including SM107.
tests/utils/test_fp8_quantize.py#L16-L21: keep dependency availability separate from exact GPU support.tests/gemm/test_cute_dsl_blockscaled_gemm.py#L263-L267: restore the grouped-GEMM architecture gate.tests/gemm/test_cute_dsl_blockscaled_gemm.py#L314-L318: restore the output-layout test architecture gate.tests/gemm/test_groupwise_scaled_gemm_fp8.py#L52-L52: narrow the blockscale predicate.tests/gemm/test_groupwise_scaled_gemm_fp8.py#L98-L106: narrow the TRTLLM and CUTLASS predicates.tests/gemm/test_groupwise_scaled_gemm_fp8.py#L227-L227: narrow the grouped-GEMM predicate.tests/gemm/test_groupwise_scaled_gemm_fp8.py#L288-L288: narrow the DeepGEMM predicate.tests/gemm/test_groupwise_scaled_gemm_fp8.py#L336-L336: narrow the batch-DeepGEMM predicate.tests/gemm/test_groupwise_scaled_gemm_mxfp4.py#L257-L261: exclude SM107 in the predicate.tests/mamba/test_chunk_scan_combined.py#L146-L151: replace the major-version-only marker.tests/moe/test_cute_dsl_fused_moe.py#L36-L37: restore the shared architecture marker.tests/moe/test_cute_dsl_fused_moe.py#L862-L862: retain gating for functional tests.tests/moe/test_cute_dsl_fused_moe.py#L1128-L1128: retain gating for wrapper tests.tests/moe/test_cute_dsl_fused_moe.py#L1650-L1650: retain gating for API consistency tests.tests/moe/test_cute_dsl_fused_moe.py#L1729-L1729: retain gating for expert-parallelism tests.tests/moe/test_cute_dsl_fused_moe.py#L1933-L1933: retain gating for poisoned-buffer tests.tests/moe/test_cute_dsl_fused_moe.py#L2129-L2129: retain gating for all-tactics tests.
As per coding guidelines, architecture-specific tests must skip unsupported GPUs using capability utilities.
📍 Affects 6 files
tests/utils/test_fp8_quantize.py#L16-L21(this comment)tests/gemm/test_cute_dsl_blockscaled_gemm.py#L263-L267tests/gemm/test_cute_dsl_blockscaled_gemm.py#L314-L318tests/gemm/test_groupwise_scaled_gemm_fp8.py#L52-L52tests/gemm/test_groupwise_scaled_gemm_fp8.py#L98-L106tests/gemm/test_groupwise_scaled_gemm_fp8.py#L227-L227tests/gemm/test_groupwise_scaled_gemm_fp8.py#L288-L288tests/gemm/test_groupwise_scaled_gemm_fp8.py#L336-L336tests/gemm/test_groupwise_scaled_gemm_mxfp4.py#L257-L261tests/mamba/test_chunk_scan_combined.py#L146-L151tests/moe/test_cute_dsl_fused_moe.py#L36-L37tests/moe/test_cute_dsl_fused_moe.py#L862-L862tests/moe/test_cute_dsl_fused_moe.py#L1128-L1128tests/moe/test_cute_dsl_fused_moe.py#L1650-L1650tests/moe/test_cute_dsl_fused_moe.py#L1729-L1729tests/moe/test_cute_dsl_fused_moe.py#L1933-L1933tests/moe/test_cute_dsl_fused_moe.py#L2129-L2129
🤖 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_fp8_quantize.py` around lines 16 - 21, Restore the shared
exact architecture policy using the capability utilities, keeping _is_available
in tests/utils/test_fp8_quantize.py:16-21 limited to dependency availability. In
tests/gemm/test_cute_dsl_blockscaled_gemm.py:263-267 and :314-318, restore
architecture gates; narrow the predicates at
tests/gemm/test_groupwise_scaled_gemm_fp8.py:52, :98-106, :227, :288, and :336
and exclude SM107 at tests/gemm/test_groupwise_scaled_gemm_mxfp4.py:257-261.
Replace the major-version-only marker at
tests/mamba/test_chunk_scan_combined.py:146-151, restore the shared marker at
tests/moe/test_cute_dsl_fused_moe.py:36-37, and retain its gating at :862,
:1128, :1650, :1729, :1933, and :2129 for the respective architecture-specific
tests.
Source: Coding guidelines
jiahanc
approved these changes
Jul 27, 2026
jiahanc
enabled auto-merge (squash)
July 27, 2026 15:14
bkryu
approved these changes
Jul 27, 2026
kahyunnam
approved these changes
Jul 27, 2026
aleozlx
pushed a commit
that referenced
this pull request
Jul 28, 2026
…4189) ## 📌 Description Fixes the B300 sampling performance regression tracked in **NVBug 6517769**, targeting `release-v0.6.16` directly. [#4122](#4122) annotated 11 sampling kernels with `__launch_bounds__(BLOCK_THREADS)`. That annotation is effectively `__launch_bounds__(1024)` on every cc ≥ 8, which forces ptxas under a 64-reg/thread ceiling and causes measurable spills / slowdowns on B300 and H100. ### Why not a blanket removal? The annotations were not unexplained noise. Internal GitLab [MR !611](https://gitlab-master.nvidia.com/dl/flashinfer/flashinfer/-/merge_requests/611) (`Fix Rubin (SM 10.7) launch failures in sampling kernels via __launch_bounds__`, merged into `feat_sm107` 2026-04-28) documented the Rubin failure mode: > The Top-K sampling kernel was using >64 regs per thread with 1024 threads per block, exceeding the 65,536 register budget. The [Rubin Kernel Requirements Tracker](https://docs.google.com/spreadsheets/d/1Cn0_F6n_HMrzXPYreYbVIWbKQmfp-zzKne3Lz7akysU) recorded the pre-fix failures as `too many resources requested for launch` on `tests/utils/test_sampling.py` / logits sampling. That rationale never made the public #4122 description, but the release branch still carries SM107 support, so a full drop would re-break Rubin native launches. ### Fix Gate the annotation to native SM107 compiles only: ```cuda #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1070) #define FLASHINFER_SAMPLING_LAUNCH_BOUNDS(block_threads) __launch_bounds__(block_threads) #else #define FLASHINFER_SAMPLING_LAUNCH_BOUNDS(block_threads) #endif ``` All 11 kernels use `FLASHINFER_SAMPLING_LAUNCH_BOUNDS(BLOCK_THREADS)` instead of unconditional `__launch_bounds__(BLOCK_THREADS)`. When today's JIT maps SM107 → `sm_100f` (`map_sm107_to_100f=True` while CUTLASS lacks `Sm107`), `__CUDA_ARCH__` is 1000 and the gate is inactive; that path inherits sm_100 register counts, which already fit a 1024-thread launch on the arches we measured. The gate is the safety net for **native `sm_107a`** compiles — the path MR !611 actually fixed. ## 🔍 Related Issues - NVBug 6517769 — B300 sampling perf drop after 290c091 - Internal MR !611 / commit `a8c10c24` on `feat_sm107` (Rubin launch-failure root cause) - Caused by #4122; #4171 reverted #4122 wholesale on `main` (no cherry-pick possible) ## 🧪 Tests ### ptxas (CUDA 13.0) with the gated macro | Target | Kernel | Result | |---|---|---| | sm_90a | MinP / TopK / TopP / TopKTopP | ≤48 regs, **no spills** (matches unconstrained) | | sm_103a | same | ≤48 regs, **no spills** | | sm_107a | — | **cannot compile** on this toolkit (`Unsupported gpu architecture`) | ### Wall-clock (NVBug repro on H100 NVL) ``` python3 flashinfer_benchmark.py --routine min_p_sampling_from_probs \ --batch_size 32 --vocab_size 129280 --min_p 0.1 -vv ``` | Config | median | |---|---| | release (unconditional `__launch_bounds__`) | 0.051 ms | | this PR (SM107-gated) | **0.045 ms** (same as full removal) | Correctness: earlier full-removal run of `tests/utils/test_sampling.py` on this header family was 1320 passed / 0 failed; the gated form is a no-op on H100 device code, so that result still applies. ## Reviewer Notes - Please confirm on **B300** that the NVBug case is restored, and on **Rubin** that `tests/utils/test_sampling.py` still launches when compiled for native `sm_107a` (and/or under the current `sm_100f` JIT mapping). - `#4171` already removed these annotations from `main` along with all SM107 enablement; this PR is release-only and preserves SM107 launchability.
3 tasks
5 tasks
8 tasks
feih-nv
added a commit
that referenced
this pull request
Jul 29, 2026
Expose the existing TRTLLM-gen `MxFP4xMxFP8` (W4A8) and `MxFP4xBf16` (W4A16) kernels through the unified MoE API. ## 📌 Description This is PR 3 in the unified MoE quantization series for FP8 support: 1. #4026 — unified block-scale FP8 (merged) 2. #4091 — unified per-tensor FP8 (merged) 3. This PR — unified TRTLLM MXFP4×MXFP8/W4A8 and MXFP4×BF16/W4A16 ### Changes - Generalize `TrtllmFp4RoutedRunner` beyond NVFP4: - `QuantVariant.MXFP4`: `MxE2m1` weights × `MxE4m3` activations - `QuantVariant.W4A16`: `MxE2m1` weights × BF16 activations - Add variant-aware TRTLLM FP4 preparation: - MXFP4 weights with 32-element UE8M0 scales - MXFP8 activation preparation for W4A8 - BF16 activation preparation for W4A16 - Add shape, dtype, and scale-layout validation. - Add unified conformance and fuzzer coverage for packed and `FromLogits` routing. ### Support matrix - NVFP4 and MXFP4/W4A8: SM100, SM103 - W4A16: SM100 only (remains disabled on SM103, matching upstream xfail #1754) - SM120/121: separate b12x backends where available - SM107 is unsupported after #4171 ### Scope - No CUDA/C++ kernel changes; both modes already exist in the TRTLLM flat API. - CUTLASS W4A8 is out of scope. ## 🚀 Pull Request Checklist ### ✅ 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 SM100, CUDA 13 CI container: - `tests/moe/test_unified_moe_mxfp4.py` - 21 passed - Unified fuzzer, packed routing: - seeds `900017,900018` - 2 passed - Unified fuzzer, `FromLogits`: - seeds `900019,900020` - 2 passed
5 tasks
aleozlx
added a commit
that referenced
this pull request
Jul 30, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description This PR relands SM 107 support to main branch (reverted in #4171) as well as some other release fixes. #### Cherry Picks - #4191 - #4189 - #4200 - #4215 - #4225 - #4230 - #4235 - #4226 - #4257 - #4258 - #4261 #### Other Changes - Rubin guards from #4252's conflict resolution (`TLLM_RUBIN_FEATURES`: SiTuGlu static_asserts + tile-192 advertisement, compiled out for the Rubin BMM pin) - Test-contract update: `test_unified_moe.py` arch assertions written post-revert (#4159) flipped to the restored contract (FP4/BF16 claim 107; FP8 stays 100/103) <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> #4107, #4164, reverts #4171 ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for Rubin/SM107 GPUs across GEMM, MoE, attention, quantization, sampling, and DeepGEMM workflows. * Added architecture-aware kernel selection, memory sizing, compilation, and artifact handling. * **Bug Fixes** * Improved validation and error messages for incompatible GPU architectures and invalid kernel configurations. * Clearly rejects unsupported NVFP4 KV-cache operations on SM107. * **Documentation** * Updated installation guidance with the SM107 architecture target. * **Tests** * Expanded architecture coverage and compatibility checks across GPU test suites. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Vinnie6167 <Vinnie6167@users.noreply.github.com> Co-authored-by: Ka-Hyun Nam <knam@nvidia.com> Co-authored-by: Alex Yang <aleyang@nvidia.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Jimmy Zhou <79552142+jimmyzho@users.noreply.github.com>
Vinnie6167
added a commit
to Vinnie6167/flashinfer
that referenced
this pull request
Aug 21, 2026
…t it is_cute_dsl_arch_supported's docstring said the probe "pins the DSL's default target via the CUTE_DSL_ARCH environment variable". That was true when flashinfer-ai#4122 landed -- it called os.environ.setdefault("CUTE_DSL_ARCH", family) -- but the line went away with the flashinfer-ai#4171 revert and was not restored by the flashinfer-ai#4280 re-land, leaving the wording stale. The passive behaviour is the correct one, and FlashInfer should not set the variable: as the comment a few lines below notes, an env var set after cutlass is imported does not retarget the DSL, and the setdefault sat after `from cutlass.base_dsl.arch import Arch` so it was a no-op anyway. Retargeting is the user's call, since it affects every DSL kernel in the process. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Compatibility
Bug Fixes
Documentation