Skip to content

[TRTLLM-13767][chore] upgrade CUTLASS DSL to 4.6.1 - #17274

Open
taianz-nv wants to merge 6 commits into
NVIDIA:mainfrom
taianz-nv:user/taianz/TRTLLM-13767-cutlass-dsl-4-6-1
Open

[TRTLLM-13767][chore] upgrade CUTLASS DSL to 4.6.1#17274
taianz-nv wants to merge 6 commits into
NVIDIA:mainfrom
taianz-nv:user/taianz/TRTLLM-13767-cutlass-dsl-4-6-1

Conversation

@taianz-nv

@taianz-nv taianz-nv commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Summary

Upgrade the supported Blackwell CuTe DSL dependency set to nvidia-cutlass-dsl[cu13]==4.6.1 and quack-kernels==0.6.2. CUTLASS DSL 4.6.1 is required by the FP4 Conv3d integration in #17262, whose kernel uses the newer LaunchConfig shared-memory allocation API.

Pinning quack alongside CUTLASS is intentional. Older quack releases satisfy the previous open lower bound but reference cute.core.ThrMma, which CUTLASS DSL 4.6 removed. An existing environment could otherwise upgrade CUTLASS while retaining an incompatible quack wheel. The repository's VSA kernel used the same removed name, so this PR migrates it to cute.ThrMma as part of the dependency upgrade.

Scope

  • Upgrade CUTLASS DSL from 4.5.0 to 4.6.1.
  • Pin quack-kernels to the matching 0.6.2 release used by MiniMax M3 MSA and FlashAttention 4.
  • Migrate the Blackwell VSA kernel's ThrMma annotations to the CUTLASS DSL 4.6 API.
  • Keep the VSA package import lazy when optional packages or shared libraries are unavailable, while allowing genuine CuTe API regressions to fail loudly.
  • Stop exporting the internal VSA kernel class from the package; callers use the guarded functional interface.

Validation

  • Full changed-file pre-commit suite passed.
  • Python syntax compilation passed for all changed Python files.
  • CUTLASS DSL 4.6.1 has been exercised on NSC B200 by the dependent FP4 Conv3d test suite, including real kernel compilation and execution.
  • Standard CI is re-triggered on the complete CUTLASS 4.6.1 + quack 0.6.2 dependency set.

Review focus

This dependency change affects CuTe consumers beyond FP4 Conv3d. CI and reviewer focus should include VisualGen FA4/parallel-combine/VSA, MiniMax M3 MSA, and representative non-VisualGen CuTe paths rather than treating an AttributeError as an optional-backend fallback.

Dev Engineer Review

  • Upgraded nvidia-cutlass-dsl[cu13] to 4.6.1.
  • Pinned quack-kernels to 0.6.2 to match the supported CUTLASS DSL dependency set.
  • Updated CuTe DSL type annotations from cute.core.ThrMma to cute.ThrMma.
  • Made optional VSA imports tolerate unavailable Python modules and shared libraries.
  • Preserved the original import exception as the cause of backend failure.
  • Removed the public VideoSparseAttentionForward export.
  • CUTLASS DSL 4.6.1 passed the dependent FP4 Conv3d validation on NSC B200.
  • Broader VisualGen and other CuTe DSL validation was requested.

QA Engineer Review

No test changes.

…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>
@taianz-nv
taianz-nv marked this pull request as ready for review August 5, 2026 03:00
@taianz-nv
taianz-nv requested review from a team as code owners August 5, 2026 03:00
@taianz-nv

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63913 [ run ] triggered by Bot. Commit: 95a8c32 Link to invocation

@coderabbitai

coderabbitai Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5d59bf53-3d6b-481c-8c6e-30c9d410a03a

📥 Commits

Reviewing files that changed from the base of the PR and between b5049f8 and 323fc25.

📒 Files selected for processing (1)
  • requirements.txt
🚧 Files skipped from review as they are similar to previous changes (1)
  • requirements.txt

Walkthrough

Changes

CUDA backend compatibility

Layer / File(s) Summary
Optional backend imports and diagnostics
requirements.txt, tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py
The CuTe DSL requirement is upgraded to 4.6.1, and quack-kernels is pinned to 0.6.2. CuTe import handling records ImportError and OSError causes and includes them in backend errors.
Video Sparse Attention kernel API
tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/__init__.py, .../interface.py, .../block_sparse_attn_dsl_fwd.py
The internal kernel reference is renamed, its package export is removed, kernel construction uses the renamed reference, and ThrMma annotations use cute.ThrMma.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: karljang

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly identifies the dependency upgrade and follows the required ticket and type format.
Description check ✅ Passed The description explains the dependency changes, rationale, scope, validation, and review focus in sufficient detail.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@BowenFu

BowenFu commented Aug 5, 2026

Copy link
Copy Markdown

The AttributeError handling reads as defensive hygiene, but taken with the pins in the same file it may be describing a live break, and I couldn't settle that from the diff alone.

The main question. Your own comments say the skew is "flash_attn.cute / quack referencing a cute API removed in a newer cutlass-dsl". flash-attn-4 is a hard pin at 4.0.0b11 (requirements.txt:87) and this PR doesn't move it. If 4.0.0b11 really does touch an API that 4.6.1 removed, then after this merge the import fails in every environment, not just skewed ones — and FlashAttn4Attention.__init__ raises ImportError at construction (flash_attn4.py:107-110). That would take the FA4 backend and the parallel-combine path from working on 4.5.0 to hard-failing on 4.6.1, with the widened except making it look like a tolerated environment issue rather than a dependency that needs bumping alongside.

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 4.0.0b11 is fine and the catch is purely for out-of-tree installs, saying so in the comment would keep the next reader from having to re-derive it.

Same question for quack-kernels>=0.2.10 (requirements.txt:79, "required for MinimaxM3 MSA"). That's a floor rather than a pin, so a fresh resolve may pick up a 4.6.1-compatible build — but a pinned/cached 0.2.10 would not, and MSA has no equivalent graceful-degradation guard that I could find.

Scope of validation. nvidia-cutlass-dsl reaches well past VisualGen — the CuTe DSL MLA decode kernels, the NVFP4 GEMM autotuner tactic pruning (nvidia-matmul-heuristics right below it), and the DSv4 CuTe DSL paths all sit on it. "Exercised on NSC B200 by the dependent FP4 Conv3d test suite" covers the kernel that motivated the bump, not those. Your own PR body still says "Draft follow-up: Run the standard VisualGen import and attention-backend CI matrix before marking ready" — but the PR isn't a draft. Worth either running that matrix or moving the line out of the description.

Smaller ones:

  • The VSA __init__.py now re-exports VideoSparseAttentionForward from .interface, where it is None on import failure. A previously loud ImportError at import time becomes a None that surfaces later as a confusing AttributeError at the use site. block_sparse_attn_from_indices_cute guards on CUTE_AVAILABLE first, so the in-tree path is covered — it's direct importers of the symbol that lose the signal.
  • Heads-up on a merge conflict: [None][chore] Update flashinfer-python from 0.6.15 to 0.6.16 #17175 (approved) moves flashinfer-python to 0.6.16 in this same file, and this PR is based on 0.6.15. Whichever lands second will need a rebase.

@BowenFu

BowenFu commented Aug 5, 2026

Copy link
Copy Markdown

Following up on my own question — I pulled the wheel metadata, and it answers it. My FlashAttention-4 hypothesis was wrong; disregard that part.

flash-attn-4 4.0.0b11   Requires-Dist: nvidia-cutlass-dsl>=4.4.2
                        Requires-Dist: quack-kernels>=0.4.0
quack-kernels 0.2.10    Requires-Dist: nvidia-cutlass-dsl>=4.4.0
quack-kernels 0.4.0     Requires-Dist: nvidia-cutlass-dsl>=4.4.2
quack-kernels 0.6.2     Requires-Dist: nvidia-cutlass-dsl==4.6.1

4.0.0b11 declares an open lower bound, so 4.6.1 is inside its supported range. The widened except is defensive hygiene, not cover for a live break, and there's nothing for you to fix there.

But the same metadata shows a real side effect the PR doesn't mention. quack-kernels 0.6.2 pins nvidia-cutlass-dsl==4.6.1 exactly. So which quack gets installed is determined by your pin:

  • on main (==4.5.0), the 0.6.x line is excluded outright and the resolver settles on an older quack;
  • after this PR (==4.6.1), the 0.6.x line becomes the only one that satisfies both constraints.

So a one-line CUTLASS bump labelled [chore] also moves quack-kernels — the package requirements.txt:79 marks as "required for MinimaxM3 MSA" — across several minor versions, as a transitive consequence. That's worth saying out loud in the PR description, and it's a second reason the "FP4 Conv3d on B200" validation doesn't cover the blast radius: MSA is on a different kernel package version afterwards, and nothing in the diff exercises it.

Separately, requirements.txt:79 floors quack at >=0.2.10 while the flash-attn-4==4.0.0b11 you pin at :87 itself requires quack-kernels>=0.4.0. The repo's own floor is looser than what its pinned FA4 needs — worth tightening while you're in this file.

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.

Signed-off-by: Taian Zhang <taianz@nvidia.com>
@taianz-nv

Copy link
Copy Markdown
Collaborator Author

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 cute.core.ThrMma uses are migrated to cute.ThrMma; the broad AttributeError catches were removed; and the nullable internal VSA class export was removed in favor of the guarded functional interface. I also updated the PR description to call out the dependency transition and cross-consumer validation scope. Commit: b5049f8.

@taianz-nv

Copy link
Copy Markdown
Collaborator Author

/bot run

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
requirements.txt (1)

77-79: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low value

Document the quack-kernels coupling.

quack-kernels==0.6.2 declares nvidia-cutlass-dsl==4.6.1, so the exact pin is required. Update the inline comment to document this coupling; flash-attn-4 alone requires only quack-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

📥 Commits

Reviewing files that changed from the base of the PR and between 95a8c32 and b5049f8.

📒 Files selected for processing (4)
  • requirements.txt
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/__init__.py
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/block_sparse_attn_dsl_fwd.py
  • tensorrt_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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63929 [ run ] triggered by Bot. Commit: b5049f8 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63913 [ run ] completed with state ABORTED. Commit: 95a8c32
/LLM/main/L0_MergeRequest_PR pipeline #51852 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Taian Zhang <taianz@nvidia.com>
@taianz-nv

Copy link
Copy Markdown
Collaborator Author

Updated the inline dependency note to make the exact quack-kernels 0.6.2 / CUTLASS DSL 4.6.1 coupling explicit. Commit: 323fc25.

@taianz-nv

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63934 [ run ] triggered by Bot. Commit: 323fc25 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63929 [ run ] completed with state ABORTED. Commit: b5049f8

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63934 [ run ] completed with state FAILURE. Commit: 323fc25
/LLM/main/L0_MergeRequest_PR pipeline #51869 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@BowenFu

BowenFu commented Aug 5, 2026

Copy link
Copy Markdown

Re-reviewed at 323fc2589. Both of my points are addressed, and the guard is now better than what it replaced.

The transitive coupling is declared: quack-kernels>=0.2.10==0.6.2 with the comment naming the CUTLASS DSL 4.6.1 pin. That was the substantive concern — on main the >=0.2.10 floor plus nvidia-cutlass-dsl==4.5.0 excluded the 0.6.x line entirely, so bumping CUTLASS silently made 0.6.x the only satisfiable resolution. Pinning it exactly makes the move visible in the diff instead of in a resolver, and it also closes the gap where the repo floored quack below what the pinned FA4 requires.

The import guard went the right direction rather than wider: except (ImportError, OSError) as error instead of adding AttributeError. That matters — an AttributeError from cutlass.cute is exactly the signal that the API moved under us, which is the thing you'd want to fail loudly on, and swallowing it would have turned a version-skew break into a silent fallback. Capturing _cute_import_error and re-raising with from also means the eventual RuntimeError carries the real cause instead of a generic "not importable".

Privatizing VideoSparseAttentionForward_VideoSparseAttentionForward and dropping it from the package __all__ is fine: it's under _torch/visual_gen/ (internal, not the tensorrt_llm/visual_gen/ public surface), and nothing outside that package imports it.

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 nvidia-cutlass-dsl isn't only used here. Also still textually conflicts with #17175 in requirements.txt.

@zhenhuaw-me zhenhuaw-me left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thanks!

@taianz-nv

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64018 [ run ] triggered by Bot. Commit: 323fc25 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64018 [ run ] completed with state FAILURE. Commit: 323fc25
/LLM/main/L0_MergeRequest_PR pipeline #51950 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants