[https://nvbugs/6641122][fix] Bump CUTLASS DSL to 4.6.2 to unblock FA4 split-KV

constraints.txt

@@ -1,12 +1,12 @@
 # setup.py appends these constraints to wheel install_requires, so every entry
 # must also be a valid TensorRT-LLM runtime dependency outside the base image.
 # Keep the base-image and wheel CuTe DSL stack aligned with requirements.txt.
-nvidia-cutlass-dsl==4.6.1
+nvidia-cutlass-dsl==4.6.2
 # Keep the CuTe DSL and FlashAttention 4 runtime stack aligned with requirements.txt.
 apache-tvm-ffi==0.1.13.post2
-quack-kernels==0.6.2
+quack-kernels==0.6.4
 torch-c-dlpack-ext==0.1.3
-# CUTLASS DSL 4.6.1 libraries require protobuf 6.x.
+# The CUTLASS DSL libraries require protobuf 6.x.
 protobuf>=6.30.2,<7
 # The `nvidia-cutlass-dsl` package does not pin numpy at all, which can be problematic in certain CI
 # stages.

docker/Dockerfile.multi

@@ -72,7 +72,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
         nvidia-cutlass-dsl-libs-core nvidia-cutlass-dsl-libs-cu12 nvidia-cutlass-dsl-libs-cu13 numpy wandb || true && \
     # Remove any leftover namespace dirs or dist-info that pip missed
     rm -rf $(python3 -c "import site; print(site.getsitepackages()[0])")/nvidia_cutlass_dsl* && \
-    pip3 install -r /tmp/constraints.txt "nvidia-cutlass-dsl[cu13]==4.6.1" && \
+    pip3 install -r /tmp/constraints.txt "nvidia-cutlass-dsl[cu13]==4.6.2" && \
     rm /tmp/constraints.txt
 
 # Install UCX, NIXL, etcd

requirements.txt

@@ -78,15 +78,17 @@ xdsl>=0.59.0 # Optional: required for MLIR-based elementwise fusion in AutoDeplo
 tiktoken
 blobfile
 openai-harmony==0.0.4
-nvidia-cutlass-dsl[cu13]==4.6.1; python_version >= "3.10" # required by Blackwell CuTe DSL kernels
+# Do not go below 4.6.2: 4.6.1 fails to compile FlashAttention 4's sm100 split-KV
+# kernel (TYPE_UNSTABLE_JOIN), which every FA4 attention config reaches via num_splits=0.
+nvidia-cutlass-dsl[cu13]==4.6.2; python_version >= "3.10" # required by Blackwell CuTe DSL kernels
 nvidia-matmul-heuristics==0.1.0.27; python_version >= "3.10" # analytic GEMM heuristics for CuTe DSL autotuner tactic pruning
-quack-kernels==0.6.2; python_version >= "3.10" # required for MinimaxM3 MSA and FlashAttention 4; pins CUTLASS DSL 4.6.1
+quack-kernels==0.6.4; python_version >= "3.10" # required for MinimaxM3 MSA and FlashAttention 4; pins the CUTLASS DSL by equality, so bump both together
 jinja2 # required for MinimaxM3 MSA
 plotly
 numexpr
 partial_json_parser
 mcp<2.0.0 # 2.0.0 removed mcp.server.fastmcp; scaffolding still uses the 1.x APIs
-apache-tvm-ffi==0.1.13.post2 # required by CUTLASS DSL 4.6.1's TVM-FFI provider
+apache-tvm-ffi==0.1.13.post2 # required by the CUTLASS DSL TVM-FFI provider
 torch-c-dlpack-ext==0.1.3 # used for reduce nvidia-cutlass-dsl host overhead, optional package for improved torch tensor calling perf
 flash-attn-4==4.0.0b19
 mistral-common>=1.10.0

security_scanning/pyproject.toml

@@ -70,9 +70,9 @@ dependencies = [
     "tiktoken (>=0.14.0,<0.15.0)",
     "blobfile (>=3.3.0,<4.0.0)",
     "openai-harmony (==0.0.4)",
-    "nvidia-cutlass-dsl[cu13] (==4.6.1)",
+    "nvidia-cutlass-dsl[cu13] (==4.6.2)",
     "nvidia-matmul-heuristics (==0.1.0.27)",
-    "quack-kernels (==0.6.2)",
+    "quack-kernels (==0.6.4)",
     "jinja2 (>=3.1.6,<4.0.0)",
     "plotly (>=7.0.0,<8.0.0)",
     "partial-json-parser (>=0.2.1.1.post7,<0.3.0.0)",

tests/integration/test_lists/test-db/l0_b200.yml

@@ -241,6 +241,7 @@ l0_b200:
   - unittest/_torch/visual_gen/test_attention_trtllm_sage.py
   - unittest/_torch/visual_gen/test_attention_integration.py
   - unittest/_torch/visual_gen/test_fa4_key_padding_mask.py
+  - unittest/_torch/visual_gen/test_fa4_split_kv.py
   - unittest/_torch/visual_gen/test_attention_perf.py
   - unittest/_torch/visual_gen/test_qwen_image_layered_registry.py
   - unittest/_torch/visual_gen/test_trtllm_serve_e2e.py

tests/integration/test_lists/waives.txt

@@ -111,7 +111,6 @@ examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[f
 examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux2-nvfp4] SKIP (https://nvbugs/6572800)
 examples/visual_gen/test_visual_gen_glm.py::test_glm_image_feature_accuracy_against_golden[nvfp4] SKIP (https://nvbugs/6644450)
 examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_cuda_graph_trtllm_backend SKIP (https://nvbugs/6668775)
-examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_example SKIP (https://nvbugs/6641122)
 examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[cuda-graph] SKIP (https://nvbugs/6655986)
 examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[nvfp4] SKIP (https://nvbugs/6572800)
 examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_lpips_against_golden SKIP (https://nvbugs/6655986)

tests/unittest/_torch/visual_gen/test_fa4_split_kv.py

@@ -0,0 +1,91 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Numerical parity for FlashAttn4's sm100 split-KV kernel.
+
+FlashAttn4Attention passes ``num_splits=0``, so FA4's heuristic picks the split
+count and any count above 1 selects a separate split-KV kernel. That kernel is
+compiled by the CuTe DSL at first use, so a DSL that rejects it takes down every
+FA4 config at pipeline load rather than at import. The other FA4 tests use K/V
+short enough that the heuristic returns 1 (it short-circuits at
+``num_n_blocks <= 4``), leaving the split-KV kernel uncovered; the K/V length
+here is what reaches it.
+
+Requires CUDA (FA4 is GPU-only).
+"""
+
+import pytest
+import torch
+
+try:
+    from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd
+
+    FA4_AVAILABLE = _flash_attn_fwd is not None
+except ImportError:
+    FA4_AVAILABLE = False
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="FA4 requires CUDA"),
+    pytest.mark.skipif(not FA4_AVAILABLE, reason="FA4 kernel not available"),
+]
+
+# S_kv must stay well above FA4's `num_n_blocks <= 4` short-circuit for the
+# `auto` case to reach the split-KV kernel; the short S_q keeps occupancy low
+# enough that the heuristic prefers splitting, which is the LTX-2 cross-attn shape.
+B, S_Q, S_KV, H, D = 1, 64, 4096, 8, 128
+
+
+@pytest.mark.parametrize("num_splits", [0, 8], ids=["auto", "forced"])
+def test_split_kv_matches_sdpa(num_splits):
+    """FA4 split-KV output matches SDPA. Fails to compile on CUTLASS DSL < 4.6.2."""
+    torch.manual_seed(0)
+    device = "cuda"
+    q, k, v = (
+        torch.randn(B, s, H, D, dtype=torch.bfloat16, device=device) for s in (S_Q, S_KV, S_KV)
+    )
+
+    out, _, *_ = _flash_attn_fwd(
+        q,
+        k,
+        v,
+        seqused_k=None,
+        softmax_scale=D**-0.5,
+        causal=False,
+        window_size_left=None,
+        window_size_right=None,
+        learnable_sink=None,
+        softcap=0.0,
+        pack_gqa=None,
+        mask_mod=None,
+        block_sparse_tensors=None,
+        return_lse=True,
+        num_splits=num_splits,
+    )
+
+    # SDPA wants [B, H, S, D]; FA4 uses [B, S, H, D].
+    ref = torch.nn.functional.scaled_dot_product_attention(
+        q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+    ).transpose(1, 2)
+
+    torch.testing.assert_close(
+        out,
+        ref,
+        rtol=2e-2,
+        atol=2e-2,
+        msg=f"FA4 num_splits={num_splits} diverges from SDPA",
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])