Reviewing files that changed from the base of the PR and between da01b1b and 6c1079eacf0ff8b8c1a221c05ea4bb6eb497af8f.

• flashinfer/gemm.py (2 hunks)
• flashinfer/utils.py (2 hunks)
• tests/utils/test_decorators.py (1 hunks)

• GitHub Check: Deploy Docs

tests/utils/test_decorators.py (1)

194-229: LGTM! Well-structured test for auto-backend selection.

The test correctly validates the new suitable_auto_backends mechanism by:

• Using shape-based constraints to filter backends
• Verifying the suitable_auto_backends attribute is populated correctly
• Testing the error path when no backends are suitable

Note: The static analysis warnings about unused backend parameters are false positives—these requirement check functions must match the decorated function's signature.

flashinfer/gemm.py (4)

2012-2023: LGTM! Proper backend requirement gate for cuDNN.

The function correctly validates cuDNN availability and is decorated with appropriate compute capabilities (SM89, 90, 100, 103, 120).

Note: The unused parameter warnings from static analysis are false positives—requirement check functions must match the decorated function signature.

---

2026-2036: LGTM! Compute capability gate for cuBLAS backend.

The function serves as a pure compute capability gate for the cuBLAS backend (SM89, 90, 100, 103, 120). No additional validation is needed since cuBLAS is the most permissive backend.

---

2039-2051: LGTM! Proper validation for CUTLASS backend constraints.

The function correctly:

• Gates on newer compute capabilities (SM100, 103, 110, 120, 121)
• Validates that e5m2 is not used (CUTLASS doesn't support it for FP8 BMM)

Note: Ruff flags the inline error message as a style issue (TRY003), but this is acceptable for this use case.

---

2054-2064: LGTM! Common validation shared across all backends.

The function serves as the common_check parameter in the @backend_requirement decorator, validating that the output dtype is either bf16 or fp16—a constraint common to all FP8 BMM backends.

flashinfer/utils.py (1)

986-995: Race condition with wrapper.suitable_auto_backends mutation is real; attribute access pattern requires careful review.

The mutation at line 994 sets wrapper.suitable_auto_backends as a side effect, creating actual concurrency hazards:

1. Concurrent calls: If the decorated function is called from multiple threads, the attribute is overwritten between validation and usage. In flashinfer/gemm.py:2152, backends = bmm_fp8.suitable_auto_backends accesses this attribute after the wrapper returns, where it could reflect a different thread's result.

2. Skip-check stale data: When skip_check=True is passed, suitable_auto_backends() is never called, so wrapper.suitable_auto_backends is not set. The decorated function accessing this attribute would receive stale data from a prior call or hit an AttributeError.

Consider storing the suitable backends in a local variable within the wrapper and passing it to the decorated function (e.g., as a kwarg or via a thread-local context) rather than relying on attribute mutation.

Reviewing files that changed from the base of the PR and between f85c63d744d4adfedf62d6cfda4ee31215242a18 and 392cb79c2af6f7e76240d0ba855df6aba1ed1899.

• flashinfer/gemm.py (3 hunks)
• flashinfer/utils.py (3 hunks)
• tests/utils/test_decorators.py (1 hunks)

• GitHub Check: Deploy Docs

tests/utils/test_decorators.py (2)

194-229: LGTM! Test correctly validates auto-backend selection.

The test properly validates:

• Backend filtering based on problem size constraints
• Correct population of suitable_auto_backends attribute
• Error handling when no suitable backends are found

Note: The static analysis warnings about unused backend parameters (lines 204, 208, 212) are false positives - these parameters are required by the backend_requirement decorator API to match the decorated function's signature.

---

231-263: LGTM! Test correctly validates heuristic function behavior.

The test properly verifies:

• Heuristic function is called during auto-backend selection
• Backend ordering is determined by the heuristic function
• Output shape validation after backend selection

Note: The static analysis warnings about unused parameters in _heuristic_func (lines 248) are false positives for this test fixture - the function signature must match the decorator's expected interface.

flashinfer/utils.py (1)

994-1010: LGTM! Common check validation is correctly implemented.

The function properly validates both common_check (lines 995-996) and individual backend checks, addressing the concern from previous reviews. The heuristic function integration (lines 1007-1008) is also correct.

However, note that this function sets wrapper.suitable_auto_backends (line 1009), which is later accessed when backend="auto". This creates a potential AttributeError if skip_check=True - see separate comment on lines 1046-1060.

flashinfer/gemm.py (6)

2000-2039: LGTM! Backend requirement functions correctly specify compute capability support.

The three requirement functions properly gate backend support:

• cudnn: SM89, 90, 100, 103, 120
• cublas: SM89, 90, 100, 103, 120
• cutlass: SM100, 103, 110, 120, 121 (includes SM110 and SM121 that were missing in earlier iterations)

The e5m2 rejection for cutlass (lines 2037-2038) is appropriate.

Note: Static analysis warnings about unused arguments are false positives - these functions must match the decorated function signature for the backend_requirement decorator to work correctly.

---

2042-2052: LGTM! Common check validates output dtype.

The function performs the essential validation of output dtype that applies to all backends. The unused arguments are required for signature matching with the decorator API.

---

2055-2084: LGTM! Heuristic function correctly orders backends and handles SM-specific constraints.

The function properly:

• Filters out cutlass for e5m2 inputs (line 2072)
• Selects appropriate cutlass variant based on SM architecture (lines 2073-2078)
• Applies k_dim >= 128 constraint for SM120/121 (lines 2076-2078), which matches the blockwise scaling requirement at line 244
• Preserves the expected backend ordering: cutlass variants → cublas → cudnn

---

2087-2095: LGTM! Decorator correctly applied with all backend requirements.

The @backend_requirement decorator properly integrates:

• Backend-specific requirement functions for cudnn, cublas, and cutlass
• Common problem size validation
• Heuristic function for backend ordering during auto-selection

---

367-375: LGTM! SM-specific CUTLASS runner support correctly added.

The changes properly handle:

• cutlass_sm10x runner for SM100/103/110 architectures
• cutlass_sm12x runner for SM120/121 architectures

This integrates correctly with the heuristic function that selects these runners based on device capability and problem size.

---

2172-2175: Verify auto-backend selection with skip_check handling.

The logic correctly uses bmm_fp8.suitable_auto_backends for automatic backend selection. However, this attribute is only set during validation. If users call bmm_fp8(... backend="auto", skip_check=True), this will raise AttributeError.

This issue should be addressed by the fix proposed in flashinfer/utils.py (lines 1046-1061), which will prevent the invalid combination.

Once the fix in utils.py is applied, verify that:

1. backend="auto" with skip_check=True raises a clear error
2. backend="auto" with skip_check=False (default) works correctly
3. Explicit backends work with both skip_check values

Reviewing files that changed from the base of the PR and between 392cb79c2af6f7e76240d0ba855df6aba1ed1899 and 2cc760441a8e320f010c346183d52d3e55f28736.

• flashinfer/gemm.py (3 hunks)
• flashinfer/utils.py (3 hunks)
• tests/utils/test_decorators.py (1 hunks)

• GitHub Check: Deploy Docs

Reviewing files that changed from the base of the PR and between 2cc760441a8e320f010c346183d52d3e55f28736 and f76b98106a6050de30229b1e36522b99d128d20a.

• flashinfer/gemm.py (3 hunks)
• flashinfer/utils.py (3 hunks)

flashinfer/gemm.py (5)

2000-2039: LGTM: Backend requirement functions properly structured.

The three requirement functions (_cudnn_bmm_fp8_requirement, _cublas_bmm_fp8_requirement, _cutlass_bmm_fp8_requirement) are correctly decorated with @supported_compute_capability to specify their SM support. The unused parameters flagged by static analysis are intentional—the backend_requirement decorator requires all checker functions to have matching signatures.

The CUTLASS requirement correctly validates that e5m2 format is not supported (line 2037-2038).

Note: The static analysis warnings about unused parameters can be safely ignored in this context.

---

2055-2083: Good implementation of dynamic backend heuristic.

The heuristic function correctly addresses the requirement from past review comments to dynamically order backends based on device architecture and problem characteristics. Key features:

• Filters out CUTLASS for e5m2 format (not supported)
• Maps to concrete runner names (cutlass_sm10x, cutlass_sm12x) for different SM architectures
• Applies the k_dim >= 128 constraint for SM12x CUTLASS variant
• Preserves preferred ordering: CUTLASS (when suitable), cuBLAS, cuDNN

Based on learnings

---

2171-2183: LGTM: Backend selection correctly handles auto, explicit cutlass, and other backends.

The three-way branching properly addresses the different backend selection scenarios:

1. Auto mode (line 2171-2172): Uses the decorator-computed suitable_auto_backends list
2. Explicit CUTLASS (line 2173-2176): Applies the heuristic function to derive concrete runner names (cutlass_sm10x or cutlass_sm12x), fixing the issue from past comments where explicit CUTLASS would produce zero runnable tactics
3. Other explicit backends (line 2177-2180): Passes through as singleton lists

Based on learnings

---

367-375: LGTM: Runner selection updated for concrete backend names.

The function correctly handles the concrete runner names (cutlass_sm10x, cutlass_sm12x) produced by the heuristic function, replacing the previous generic backend names. The assertion at line 375 ensures at least one runner is available, providing a clear failure point if backend selection goes wrong.

---

2086-2094: SM110 exclusion in cuDNN/cuBLAS is intentional and justified.

The compute capability annotations exclude SM110 from cuDNN and cuBLAS FP8 BMM support deliberately. Based on available evidence:

1. cuDNN has a known FP8 bug: cuDNN has a known FP8 scaled‑dot‑product attention deadlock/bug on sm10.0 (Blackwell), which explains the exclusion.

2. cuBLAS theoretically supports SM110, but is likely excluded for consistency and to avoid related compatibility issues with the broader FP8 stack.

3. Pattern in codebase confirms intentionality: Other backends explicitly block SM110 (e.g., TRTLLM backend raises ValueError("TRTLLM FP8 GEMM is not supported on SM110.")), indicating deliberate SM110 FP8 exclusions are standard practice in this codebase due to known issues.

No changes are needed.

Reviewing files that changed from the base of the PR and between f76b98106a6050de30229b1e36522b99d128d20a and 259a3d3a3b89c31ac66f66891da398b577435992.

• flashinfer/gemm.py (3 hunks)
• flashinfer/utils.py (3 hunks)
• tests/utils/test_decorators.py (1 hunks)

tests/utils/test_decorators.py (2)

330-365: LGTM! Test structure validates auto backend selection.

The test properly exercises the new suitable_auto_backends feature with both success and error paths.

Note: Static analysis warnings about unused backend arguments in the requirement checker functions (lines 340, 344, 348) are false positives. These functions must match the signature of the decorated function even when not all parameters are used.

---

367-399: LGTM! Test validates heuristic function integration.

The test correctly verifies that the heuristic function can reorder the suitable backends list.

Note: Static analysis warnings about unused arguments (lines 377, 381, 384, 391) are false positives. The requirement checker and heuristic functions must match expected signatures.

flashinfer/gemm.py (5)

2000-2011: Verify SM110 support for cuDNN FP8.

The compute capability list includes [89, 90, 100, 103, 120, 121] but excludes 110. Based on past review comments about ensuring complete SM coverage, please confirm whether SM110 should be included in the supported list for cuDNN FP8 BMM.

Note: Static analysis warnings about unused parameters are false positives—requirement checker functions must match the decorated function signature.

Based on learnings from past reviews.

---

2014-2024: LGTM! cuBLAS requirement checker is correct.

The function correctly defines cuBLAS support for the specified compute capabilities. Same note about SM110 applies here as mentioned in the cuDNN requirement checker.

Note: Unused parameter warnings are false positives.

---

2027-2039: LGTM! CUTLASS requirement checker correctly validates e5m2 restriction.

The function appropriately restricts CUTLASS to SM100+ architectures and rejects float8_e5m2 inputs, which aligns with CUTLASS FP8 kernel limitations.

Note: Static analysis warning about long error message (TRY003) is acceptable here for clarity. Unused parameter warnings are false positives.

---

2042-2083: LGTM! Heuristic function correctly maps backends to SM-specific runners.

The heuristic properly transforms the logical "cutlass" backend into architecture-specific runner names ("cutlass_sm10x", "cutlass_sm12x") based on device capability and problem size (k dimension check for SM120/121). This aligns with the updated fp8_gemm_sm100 function that expects these specific runner names.

Note: Unused parameter warnings are false positives—these functions must match the decorated function signature.

---

2086-2094: LGTM! Decorator and implementation correctly handle auto and explicit backends.

The decorator integrates the new per-backend requirements and heuristic function. The implementation correctly:

• Uses bmm_fp8.suitable_auto_backends for backend="auto" (line 2172)
• Applies the heuristic to explicit "cutlass" to resolve runner names (lines 2174-2176)
• Handles "cudnn" with availability check (lines 2177-2178)
• Falls through to other explicit backends (line 2180)

This addresses past review comments about explicit cutlass needing heuristic transformation.

Also applies to: 2171-2183

Reviewing files that changed from the base of the PR and between 259a3d3a3b89c31ac66f66891da398b577435992 and 330f89b.

• flashinfer/gemm.py (3 hunks)
• flashinfer/utils.py (3 hunks)
• tests/utils/test_decorators.py (1 hunks)

• GitHub Check: Deploy Docs

tests/utils/test_decorators.py (2)

330-365: LGTM - comprehensive test coverage for auto-backend selection.

The test correctly validates that:

1. suitable_auto_backends is populated with backends passing their requirement checks
2. The error path raises BackendSupportedError when no suitable backends exist
3. The suitable_auto_backends attribute is accessible on the decorated function

Note: Static analysis warnings about unused backend parameters (lines 340, 344, 348) are false positives—these parameters must match the decorated function's signature even if unused in the requirement check.

---

367-419: LGTM - thorough validation of heuristic function behavior.

This test correctly validates that:

1. The heuristic function receives the filtered list of suitable backends (after capability/requirement checks)
2. The heuristic can reorder backends dynamically based on runtime parameters
3. The suitable_auto_backends attribute reflects the heuristic's ordering
4. Different input shapes produce different backend orderings as expected

The assertion on line 390 that "cutlass" is not in suitable_backends is particularly good—it confirms that unsuitable backends are filtered before the heuristic runs.

Note: Static analysis warnings about unused parameters (lines 377, 381, 385, 388, 403) are false positives—these are required by the decorator pattern.

flashinfer/gemm.py (6)

367-375: LGTM - cleaner backend selection logic.

The refactoring from SM-specific conditionals to a generic runner_names list is cleaner and more maintainable. The new structure supports distinct "cutlass_sm10x" and "cutlass_sm12x" runners, aligning well with the heuristic function's output.

The assertion on line 375 provides a clear error if no runners are configured, which is better than silently failing later.

---

2000-2039: Verify SM110 exclusion from cuDNN and cuBLAS backends is intentional.

Comparing the @supported_compute_capability decorators across the three backend requirement functions:

• cuDNN (line 2000): [89, 90, 100, 103, 120, 121] — missing SM110
• cuBLAS (line 2014): [89, 90, 100, 103, 120, 121] — missing SM110
• CUTLASS (line 2027): [100, 103, 110, 120, 121] — has SM110, but missing SM89/90

This asymmetry may be intentional (reflecting actual backend/architecture support), but given the past review comment asking to verify SM110 and SM121 aren't missing, please confirm:

1. Is SM110 truly unsupported by cuDNN and cuBLAS for FP8 BMM?
2. Is the absence of SM89/90 from CUTLASS intentional?

If these exclusions are correct, consider adding a brief comment in the code or commit message explaining the architecture support matrix to prevent future confusion.

---

2042-2052: LGTM - appropriate common validation.

The common check correctly validates the output dtype constraint that applies to all backends. Backend-specific constraints (like e5m2 support) are appropriately handled in the individual requirement functions.

Note: Static analysis warnings about unused parameters are false positives—the signature must match the decorated function even if not all parameters are used in the validation.

---

2055-2083: Good heuristic logic with SM-specific CUTLASS variant selection.

The heuristic correctly handles:

1. e5m2 exclusion for CUTLASS (line 2066, 2072) — aligns with the requirement check on line 2037-2038
2. SM-specific CUTLASS variants (lines 2073-2078) — maps to the concrete runner names expected by fp8_gemm_sm100
3. K-dimension constraint for SM12x (line 2077) — the k >= 128 check aligns with the SM120 groupwise scaling kernel requirement (see line 242-245 earlier in this file)

The backend ordering (CUTLASS → cuBLAS → cuDNN) implies a performance preference, which is reasonable for auto-selection.

Note: Static analysis warnings about unused parameters (lines 2059-2063) are false positives—the heuristic signature must match the decorated function.

---

2086-2094: LGTM - decorator correctly configured.

The decorator integrates all the new components:

• Backend-specific requirement functions with compute capability decorators
• Common validation for output dtype
• Heuristic function for dynamic backend ordering

This provides a complete auto-backend selection framework for bmm_fp8.

---

2171-2180: Backend selection logic correctly handles all cases.

The implementation properly handles:

1. backend="auto" (line 2171-2172): Uses the decorator-populated suitable_auto_backends list
2. backend="cutlass" (lines 2173-2176): Explicitly calls the heuristic to map to SM-specific variants ("cutlass_sm10x" or "cutlass_sm12x"), resolving the past issue where explicit cutlass produced zero runnable tactics
3. backend="cudnn" (lines 2177-2178): Guards with CUDNN_AVAILABLE check
4. Other backends (lines 2179-2180): Passes through as-is (e.g., "cublas")

The explicit cutlass path (lines 2173-2176) is particularly important—it ensures that users who explicitly request "cutlass" get the correct SM-specific variant, not an empty list.

Reviewing files that changed from the base of the PR and between 330f89b and 9def7ba.

• flashinfer/utils.py (3 hunks)

• GitHub Check: Deploy Docs

flashinfer/utils.py (6)

857-859: LGTM: Function signature enhanced with heuristic support.

The addition of heuristic_func parameter is properly typed and aligns with the PR's goal of enabling dynamic backend ordering for auto-selection.

---

1045-1062: LGTM: Capability extraction is correctly implemented.

The helper function appropriately extracts compute capability from the first tensor argument. The assumption that all tensors share the same device (line 1049) is reasonable for performance.

---

1081-1081: LGTM: Capability extraction refactored to helper.

Good refactoring that extracts capability once for reuse across validation paths.

---

1088-1106: LGTM: Auto and explicit backend validation paths are correctly implemented.

The logic appropriately handles both auto-selection (lines 1088-1094) and explicit backend validation (lines 1095-1106), with proper error messages for unsupported backends.

---

1108-1116: LGTM: Backend-agnostic validation correctly implemented.

The comment on line 1108 improves clarity, and the validation logic properly handles functions with implicit single backends.

Based on past review feedback.

---

1117-1121: Consider applying defaults before calling suitable_auto_backends in skip_check path.

On line 1121, suitable_auto_backends is called with *args, **kwargs without applying defaults, unlike the main validation path (line 1090) which uses kwargs_with_defaults. If the heuristic function or backend requirement checkers rely on default parameter values, they could fail or behave incorrectly.

Trade-off: Per past review feedback, keeping the skip_check path lightweight is important for performance. However, this creates an inconsistency that could cause subtle bugs.

Recommendation: Document that when using skip_check=True, callers must explicitly provide all parameters that the heuristic or backend checks depend on, or apply defaults here:

 elif skip_check and heuristic_func is not None:
     if kwargs.get("backend") == "auto":
+        # Apply defaults for heuristic function
+        bound_args = sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+        kwargs_with_defaults = dict(bound_args.arguments)
         # This needs to be called for heuristic function
         capability = _get_capability(*args, **kwargs)
-        suitable_auto_backends(capability, *args, **kwargs)
+        suitable_auto_backends(capability, **kwargs_with_defaults)

Based on past review feedback.