[TRTLLM-13767][chore] upgrade CUTLASS DSL stack to 4.6.1 - #17274
[TRTLLM-13767][chore] upgrade CUTLASS DSL stack to 4.6.1#17274taianz-nv wants to merge 10 commits into
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: |
|
/bot run |
|
PR_Github #64296 [ run ] triggered by Bot. Commit: |
|
PR_Github #64296 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Clean, well-documented upgrade — the description matches the diff closely and the validation list is unusually thorough for a dependency bump. Two non-blocking asks:
-
The
make_fragment→make_rmem_tensorrename is only checked at JIT-compile time, and three migrated kernels aren't in the validation list:dense_gemm_persistent.py,filtered_top_k_varlen_util.py, andsingle_pass_multi_cta_radix_topk.py. Can you confirm existing B200 CI exercises their compile paths, or run them once? The rename is mechanical so risk is low, but an uncovered kernel would fail loudly on first customer use rather than in CI. -
See the inline note on the missing
[cu13]extra inconstraints.txt— with 4.6.1 this now determines which GPU runtime libs land in the devel image.
For the record I verified no bare cute.make_fragment( calls remain in the tree (the make_fragment_A/B/C/_like hits are different, still-supported APIs), and that fmha_helpers.py's exp2 emulation uses the string-based rnd="rm" form rather than the removed nvvm.RoundingModeKind enum, so the ptx.py deletion leaves no broken sibling.
| # Upgrade base image nvidia-cutlass-dsl 4.3.5 to 4.4.2 | ||
| nvidia-cutlass-dsl>=4.4.2 | ||
| # Upgrade the base-image nvidia-cutlass-dsl to the supported 4.6 API. | ||
| nvidia-cutlass-dsl==4.6.1 |
There was a problem hiding this comment.
requirements.txt pins nvidia-cutlass-dsl[cu13]==4.6.1 but this pin drops the [cu13] extra. That mattered less before, but per the 4.6.1 metadata (visible in the poetry.lock in this PR) the base package now unconditionally depends on nvidia-cutlass-dsl-libs-cu12, with the cu13 libs only behind the extra — and Dockerfile.multi consumes this file with pip3 install -r, i.e. it actually installs these packages into the devel image after uninstalling the base image's copy. So the devel image gets cu12 runtime libs on a CUDA 13 base and no cu13 libs until requirements.txt is installed later. Suggest nvidia-cutlass-dsl[cu13]==4.6.1 here to keep the two files interchangeable.
There was a problem hiding this comment.
Good catch—the devel stage was installing the constraints as requirements, so the cu13 extra was lost. I did not add the extra to constraints.txt because that file is also consumed via -c, where pip rejects extras. In 1336c2e I request nvidia-cutlass-dsl[cu13]==4.6.1 at this Docker install site instead.
|
@brnguyen2 Thanks for calling out the three kernels. I verified all three on an NSC B200 with CUTLASS DSL 4.6.1: the production dense GEMM native/TVM-FFI paths, filtered varlen top-k, and single-pass multi-CTA top-k all passed (job 1636364). The standalone dense runner did expose stale Tensor/Pointer and runtime- |
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 `@docker/Dockerfile.multi`:
- Around line 55-60: Add the repository-standard NVIDIA copyright header at the
beginning of docker/Dockerfile.multi, using 2026 as the copyright year; leave
the existing Docker build commands unchanged.
In `@tests/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py`:
- Around line 346-383: Add the return annotation -> cute.testing.JitArguments to
both generate_tensors functions, including the function shown and the second
generate_tensors definition nearby. Keep their existing implementations and
workspace_generator usage unchanged.
- Around line 221-233: Register run_dense_bf16_gemm_persistent.py in the
appropriate CI or QA test list. Configure the entry to exercise the run entry
point and both generate_tensors branches, reusing the test-list conventions for
similar CuTe DSL kernel scripts.
🪄 Autofix
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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b9e5ee51-3863-4b4b-9281-078cebc21d0b
📒 Files selected for processing (2)
docker/Dockerfile.multitests/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py
| _, a_torch = cutlass_torch.cute_tensor_like( | ||
| a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 | ||
| ) | ||
| b_tensor, b_torch = cutlass_torch.cute_tensor_like( | ||
| _, b_torch = cutlass_torch.cute_tensor_like( | ||
| b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 | ||
| ) | ||
| c_tensor, c_torch = cutlass_torch.cute_tensor_like( | ||
| _, c_torch = cutlass_torch.cute_tensor_like( | ||
| c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 | ||
| ) | ||
|
|
||
| # Mark tensor to be byte aligned | ||
| a_tensor.mark_compact_shape_dynamic( | ||
| mode=1 if a_major == "k" else 0, | ||
| stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), | ||
| divisibility=1, | ||
| ) | ||
| b_tensor.mark_compact_shape_dynamic( | ||
| mode=1 if b_major == "k" else 0, | ||
| stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), | ||
| divisibility=1, | ||
| ) | ||
| c_tensor.mark_compact_shape_dynamic( | ||
| mode=1 if c_major == "n" else 0, | ||
| stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), | ||
| divisibility=1, | ||
| a_ptr = make_ptr(ab_dtype, a_torch.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) | ||
| b_ptr = make_ptr(ab_dtype, b_torch.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) | ||
| c_tensor = from_dlpack(c_torch, assumed_align=16).mark_layout_dynamic( | ||
| leading_dim=1 if c_major == "n" else 0 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
fd -t f -e yml . tests/integration/test_lists 2>/dev/null | sort || true
rg -n -C 2 \
'run_dense_bf16_gemm_persistent|dense_bf16_gemm_persistent|PersistentDenseGemmKernel' \
tests/integration/test_lists/test-db \
tests/integration/test_lists/qa || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 4305
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- changed files ---'
git diff --name-status
printf '%s\n' '--- relevant source structure ---'
ast-grep outline tests/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py
printf '%s\n' '--- reviewed source ranges ---'
sed -n '180,255p' tests/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py
sed -n '320,435p' tests/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py
printf '%s\n' '--- repository-wide references ---'
rg -n -C 3 \
'run_dense_bf16_gemm_persistent|dense_bf16_gemm_persistent|PersistentDenseGemmKernel' \
. \
-g '!build/**' -g '!dist/**' -g '!node_modules/**' || true
printf '%s\n' '--- test-list entry shapes ---'
for f in $(fd -t f . tests/integration/test_lists | sort | head -n 8); do
printf '\n### %s\n' "$f"
sed -n '1,80p' "$f"
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- all test-list files ---'
git ls-files tests/integration/test_lists | sort
printf '%s\n' '--- direct script and directory references in test lists ---'
rg -n -i \
'cute[_ -]?dsl|tests/scripts|run_dense|dense_gemm_persistent|gemm_persistent' \
tests/integration/test_lists || true
printf '%s\n' '--- references outside test lists ---'
rg -n -i \
'run_dense_bf16_gemm_persistent\.py|tests/scripts/cute_dsl_kernels|cute_dsl_kernels' \
.github tests scripts tests/integration \
-g '!tests/integration/test_lists/**' \
-g '!**/__pycache__/**' || true
printf '%s\n' '--- function definitions in the reviewed script ---'
rg -n '^(async[[:space:]]+def|def)[[:space:]]+' \
tests/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 29471
Register the CuTe DSL test script.
No CI or QA test-list entry invokes tests/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py. Add an appropriate entry that exercises run and both generate_tensors branches.
Test coverage summary: modified entry point run; modified workspace generators: both generate_tensors definitions; no test functions added or removed; coverage verdict: insufficient.
🤖 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/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py` around
lines 221 - 233, Register run_dense_bf16_gemm_persistent.py in the appropriate
CI or QA test list. Configure the entry to exercise the run entry point and both
generate_tensors branches, reusing the test-list conventions for similar CuTe
DSL kernel scripts.
Source: Path instructions
There was a problem hiding this comment.
I don't plan to register this standalone benchmark runner directly in a CI/QA test list. The lists contain pytest test nodes, while this is one of 12 manual CuTe DSL run_*.py utilities and none is registered directly. The production BF16 GEMM path is covered by TestDeepSeekV3Lite::test_cute_dsl_bf16_gemm[...] in tests/integration/test_lists/qa/llm_function_core.txt; for this change, both the contiguous and strided runner branches were also executed with numerical checks on an NSC B200 (job 1636394). Adding a new recurring hardware test for the manual runner would be disproportionate to this dependency-upgrade follow-up.
|
/bot kill |
|
/bot run |
|
/bot kill |
…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>
f82cb2c to
8be2042
Compare
|
/bot run |
|
PR_Github #64521 [ run ] triggered by Bot. Commit: |
|
PR_Github #64521 [ run ] completed with state
|
8 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.