[TRTLLM-13767][chore] upgrade CUTLASS DSL stack to 4.6.1 - #17274
Conversation
|
/bot run |
|
PR_Github #63913 [ run ] triggered by Bot. Commit: |
|
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:
WalkthroughThe pull request updates CuTe DSL and FlashAttention dependencies, migrates kernel register allocations to ChangesCuTe DSL and FlashAttention compatibility
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
|
The The main question. Your own comments say the skew is " Can you state which it is? If FA4 needs a newer build to stay working on 4.6.1, that bump belongs in this PR; if Same question for Scope of validation. Smaller ones:
|
|
Following up on my own question — I pulled the wheel metadata, and it answers it. My FlashAttention-4 hypothesis was wrong; disregard that part.
But the same metadata shows a real side effect the PR doesn't mention.
So a one-line CUTLASS bump labelled Separately, My remaining asks are unchanged and both cheap: run the VisualGen attention-backend CI matrix your description still calls for (the PR isn't a draft), and confirm the in-tree CuTe DSL consumers outside VisualGen — MLA decode, the NVFP4 GEMM autotuner — are exercised somewhere on 4.6.1. |
|
Thanks for tracing the dependency metadata and the broader CuTe impact. I addressed this as one supported dependency/API set rather than extending the fallback: quack-kernels is now pinned to 0.6.2 alongside CUTLASS DSL 4.6.1; the in-tree VSA |
|
/bot run |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
requirements.txt (1)
77-79: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueDocument the
quack-kernelscoupling.
quack-kernels==0.6.2declaresnvidia-cutlass-dsl==4.6.1, so the exact pin is required. Update the inline comment to document this coupling;flash-attn-4alone requires onlyquack-kernels>=0.4.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 `@requirements.txt` around lines 77 - 79, Update the inline comment on the quack-kernels==0.6.2 requirement to document that it declares and therefore requires nvidia-cutlass-dsl==4.6.1; distinguish this coupling from flash-attn-4, which only requires quack-kernels>=0.4.0.
🤖 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.
Nitpick comments:
In `@requirements.txt`:
- Around line 77-79: Update the inline comment on the quack-kernels==0.6.2
requirement to document that it declares and therefore requires
nvidia-cutlass-dsl==4.6.1; distinguish this coupling from flash-attn-4, which
only requires quack-kernels>=0.4.0.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 41364dc7-9c6e-4f65-bc76-b65fcbf8416d
📒 Files selected for processing (4)
requirements.txttensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/__init__.pytensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/block_sparse_attn_dsl_fwd.pytensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py
|
PR_Github #63929 [ run ] triggered by Bot. Commit: |
|
PR_Github #63913 [ run ] completed with state
|
|
Updated the inline dependency note to make the exact quack-kernels 0.6.2 / CUTLASS DSL 4.6.1 coupling explicit. Commit: 323fc25. |
|
/bot run |
|
PR_Github #63934 [ run ] triggered by Bot. Commit: |
|
PR_Github #63929 [ run ] completed with state |
|
PR_Github #63934 [ run ] completed with state
|
|
Re-reviewed at The transitive coupling is declared: The import guard went the right direction rather than wider: Privatizing Still not approving as first reviewer — this moves two pinned dependencies and there are no approvals yet. The one ask left over from before is unchanged: the description still calls for the VisualGen attention-backend matrix, and it'd be good to confirm the non-VisualGen CuTe DSL consumers were exercised on 4.6.1, since |
|
/bot run |
|
PR_Github #64018 [ run ] triggered by Bot. Commit: |
…tlass-dsl version skew The visual_gen FlashAttn4 / parallel attention backends import flash_attn.cute (which pulls quack); both already guard that optional import but only catch (ImportError, OSError). A newer nvidia-cutlass-dsl (e.g. 4.6, required by the NVFP4 conv kernel) removes cute.core.ThrMma, which flash_attn.cute / quack reference at import time, raising AttributeError and hard-breaking `import tensorrt_llm` for every code path — including the VAE, which never uses these backends. Add AttributeError to the two guards so the optional cute-dsl backends degrade gracefully (re-imported lazily only when explicitly selected) instead of breaking the whole package import under a cutlass-dsl version mismatch. Signed-off-by: Taian Zhang <taianz@nvidia.com> (cherry picked from commit 87f0ef71a43549e69f8109d621e0fe9ad665da8b) Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
2cef3fc to
4419f29
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #67233 [ run ] triggered by Bot. Commit: |
|
PR_Github #67233 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #67333 [ run ] triggered by Bot. Commit: |
|
PR_Github #67333 [ run ] completed with state |
|
LGTM. |
…4 optional-import contract test_ltx2_example fails on DGX_B200 with the example subprocess exiting 1. Two independent product defects, both still present on main. 1) Kernel compile. NVIDIA#17274 moved the CuTe DSL stack to 4.6.1, and NVIDIA#17610 flipped ltx2-fp4-1gpu.yaml from VANILLA to FA4, so this test began exercising FlashAttention 4's sm100 kernel for the first time against a DSL release that rejects it: error[TYPE_UNSTABLE_JOIN]: `n_block_first` has type `None` on one path and `Int32` on another -- flash_attn/cute/flash_fwd_sm100.py:1476 n_block_first is assigned only inside the `not is_split_kv or n_block_min < n_block_max` branch, so the unstable join is emitted only when split-KV is live. flash_attn4.py requests num_splits=0, i.e. "ask the heuristic", which resolves per shape: 1 for long-query self-attention, 18-37 for LTX-2's low-occupancy cross-attention. Upstream fixed the type checker in 4.6.2, so the bump repairs the path rather than avoiding it. Pinning num_splits=1 was rejected: it only forces is_split_kv False, silently reverting the heuristic a02214a (NVIDIA#15399) added for exactly this workload. quack-kernels pins the DSL by equality in its wheel metadata (0.6.2 -> ==4.6.1, 0.6.4 -> ==4.6.2), so the two move together. flash-attn-4 stays at b19 and no FA4 config changes. 2) Optional-import contract. Two sites break when the optional CuTe DSL is absent or incomplete, and neither failure is confined to its own feature. cute_dsl_custom_ops defines the four Sm100BlockScaledContiguous*Runner classes inside its `if IS_CUTLASS_DSL_AVAILABLE:` block, which has no `else`. fused_moe_cute_dsl imported them at module scope, so importing it raised ImportError: cannot import name 'Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner'. That import sits on the create_moe -> modeling_utils -> models path, so every model import died, including the visual_gen pipeline loader this test drives. The four names are used at exactly one place, the isinstance in runner_tactic_comb_checker, and reaching that line means a CuteDSL runner is already being autotuned, so the DSL is necessarily installed. GroupedGemmInputsHelper is defined outside the guard and stays at module scope. flash_attn4.py and parallel.py each duplicated `except (ImportError, OSError)` around their flash_attn.cute.interface import. That covers a missing dependency but not a broken one: an FA4 built against a different CuTe DSL is found, starts executing, and dies dereferencing a symbol the DSL removed, which surfaces as AttributeError and escapes the tuple. Since attention_backend/__init__.py imports both modules unconditionally, the escape propagates out of `import tensorrt_llm` itself. Rather than adding one more exception type to two enumerations, both sites now share one resolver so the contract is defined once and cannot drift. Nothing is silenced: each site keeps its recorded error and first use still raises the pre-existing descriptive ImportError chained from the original, and the module-level _flash_attn_fwd / _flash_attn_combine and their *_import_error names are preserved, so the `is not None` availability contract that the visual_gen tests rely on is unchanged. Verified on DGX_B200. Holding the tree fixed and varying only DSL visibility confirms the guard honours optionality without disabling the feature: unpatched, DSL absent -> ImportError (the observed signature) patched, DSL absent -> models and VisualGen import, 0 runners bound, checker(non-CuteDSL comb) still True patched, DSL present -> models import, all 4 runners bound The third row is the anti-regression check: with the DSL present the real runner classes still bind, so the tile-size tactic filter keeps working. Removes the corresponding waiver (exactly one line). Note for reviewers: security_scanning/poetry.lock still records 4.6.1 / 0.6.2. It is generated by scripts/generate_lock_file.py (Poetry 2.4.1) and jenkins/TensorRT_LLM_PLC.groovy regenerates and auto-commits it nightly. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…4 optional-import contract test_ltx2_example fails on DGX_B200 with the example subprocess exiting 1. Two independent product defects, both still present on main. 1) Kernel compile. NVIDIA#17274 moved the CuTe DSL stack to 4.6.1, and NVIDIA#17610 flipped ltx2-fp4-1gpu.yaml from VANILLA to FA4, so this test began exercising FlashAttention 4's sm100 kernel for the first time against a DSL release that rejects it: error[TYPE_UNSTABLE_JOIN]: `n_block_first` has type `None` on one path and `Int32` on another -- flash_attn/cute/flash_fwd_sm100.py:1476 n_block_first is assigned only inside the `not is_split_kv or n_block_min < n_block_max` branch, so the unstable join is emitted only when split-KV is live. flash_attn4.py requests num_splits=0, i.e. "ask the heuristic", which resolves per shape: 1 for long-query self-attention, 18-37 for LTX-2's low-occupancy cross-attention. Upstream fixed the type checker in 4.6.2, so the bump repairs the path rather than avoiding it. Pinning num_splits=1 was rejected: it only forces is_split_kv False, silently reverting the heuristic a02214a (NVIDIA#15399) added for exactly this workload. quack-kernels pins the DSL by equality in its wheel metadata (0.6.2 -> ==4.6.1, 0.6.4 -> ==4.6.2), so the two move together. flash-attn-4 stays at b19 and no FA4 config changes. 2) Optional-import contract. Two sites break when the optional CuTe DSL is absent or incomplete, and neither failure is confined to its own feature. cute_dsl_custom_ops defines the four Sm100BlockScaledContiguous*Runner classes inside its `if IS_CUTLASS_DSL_AVAILABLE:` block, which has no `else`. fused_moe_cute_dsl imported them at module scope, so importing it raised ImportError: cannot import name 'Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner'. That import sits on the create_moe -> modeling_utils -> models path, so every model import died, including the visual_gen pipeline loader this test drives. The four names are used at exactly one place, the isinstance in runner_tactic_comb_checker, and reaching that line means a CuteDSL runner is already being autotuned, so the DSL is necessarily installed. GroupedGemmInputsHelper is defined outside the guard and stays at module scope. flash_attn4.py and parallel.py each duplicated `except (ImportError, OSError)` around their flash_attn.cute.interface import. That covers a missing dependency but not a broken one: an FA4 built against a different CuTe DSL is found, starts executing, and dies dereferencing a symbol the DSL removed, which surfaces as AttributeError and escapes the tuple. Since attention_backend/__init__.py imports both modules unconditionally, the escape propagates out of `import tensorrt_llm` itself. Rather than adding one more exception type to two enumerations, both sites now share one resolver so the contract is defined once and cannot drift. Nothing is silenced: each site keeps its recorded error and first use still raises the pre-existing descriptive ImportError chained from the original, and the module-level _flash_attn_fwd / _flash_attn_combine and their *_import_error names are preserved, so the `is not None` availability contract that the visual_gen tests rely on is unchanged. Verified on DGX_B200. Holding the tree fixed and varying only DSL visibility confirms the guard honours optionality without disabling the feature: unpatched, DSL absent -> ImportError (the observed signature) patched, DSL absent -> models and VisualGen import, 0 runners bound, checker(non-CuteDSL comb) still True patched, DSL present -> models import, all 4 runners bound The third row is the anti-regression check: with the DSL present the real runner classes still bind, so the tile-size tactic filter keeps working. Removes the corresponding waiver (exactly one line). Note for reviewers: security_scanning/poetry.lock still records 4.6.1 / 0.6.2. It is generated by scripts/generate_lock_file.py (Poetry 2.4.1) and jenkins/TensorRT_LLM_PLC.groovy regenerates and auto-commits it nightly. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…ersion ATTRIBUTIONS-Python.md still names 4.5.0, which the 4.6.1 bump in NVIDIA#17274 left behind. Point it at the version this PR pins so the attribution matches what ships. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
12 participants
Summary
Upgrade the supported Blackwell CuTe DSL stack to
nvidia-cutlass-dsl[cu13]==4.6.1. This version is required by the FP4 Conv3d integration in #17262, whose branch-merged shared-memory allocation uses the newerLaunchConfigAPI.The upgrade is a coordinated dependency and source migration rather than a one-line package bump:
quack-kernels==0.6.2is the release coupled exactly to CUTLASS DSL 4.6.1.apache-tvm-ffi==0.1.13.post2providesmap_dataclass_to_tuple, which CUTLASS DSL 4.6.1 passes while compiling NVFP4 GEMM; the previous 0.1.6 pin fails at compile time. This pin also matches the version already selected by currentmain's security lock and remains inside every consumer's declared<0.2compatibility range.flash-attn-4==4.0.0b19removes the oldcute.core.ThrMmausage while remaining compatible with the released CUTLASS DSL 4.6.1 package. Its private forward API can append diagnostic values, so the VisualGen wrapper now consumes the output/LSE prefix only.torch-c-dlpack-ext==0.1.3remains pinned as part of the same FA4 runtime stack.Scope
requirements.txt, constraints, the FA4 image installer, and the security-scanning Poetry project/lock file on the supported versions above.cute.make_fragmentalias tocute.make_rmem_tensor, and fromcute.core.ThrMmatocute.ThrMma.This PR changes the shared CuTeDSL runtime stack and compatibility code; it does not add the FP4 VAE/Conv3d feature itself.
Validation
On an NSC B200 with CUTLASS DSL 4.6.1 and TVM-FFI 0.1.11 (the feature-bearing version used for compatibility validation), using real JIT compilation and kernel execution:
The final TVM-FFI 0.1.13.post2 pin is newer than the NSC compatibility run above; it satisfies the declared CUTLASS DSL/Quack/FA4 ranges and is exercised by the final standard CI run.
Repository checks:
poetry check --lockpassed forsecurity_scanning/.main(33c6270c35) before the final validation run.No new test implementation is needed: the relevant runtime/correctness tests already exist. This PR adds the missing B200 L0 registration for FA4 key-padding-mask behavior and expands the dependency-upgrade validation across the existing suites.
Compatibility notes
(output, lse); those areNoneon the VisualGen inference path used here and are intentionally ignored.AuxData cannot be converted to JitArgumentwarning during first-shape JIT compilation. The tested kernels compile and execute successfully; this warning is non-fatal and does not occur in steady-state execution.Review focus
Please review this as one supported CUTLASS/TVM-FFI/quack/FA4 version set. In particular, focus on the mechanical
make_rmem_tensorAPI migration, VisualGen FA4 return compatibility, security-scanning lock consistency, and coverage of CuTeDSL consumers outside VisualGen.Dev Engineer Review
4.6.1.quack-kernelsto0.6.2.apache-tvm-ffito0.1.13.post2.4.0.0b19.cute.ThrMmaandcute.make_rmem_tensor.VideoSparseAttentionForwardexport.cute.make_rmem_tensor.Constexprhandling.cu13CUTLASS DSL extra to the devel image.QA Engineer Review
tests/integration/test_lists/test-db/l0_b200.yml.unittest/_torch/visual_gen/test_fa4_key_padding_mask.pyto the B200 pre-merge test list.