diff --git a/constraints.txt b/constraints.txt index a99d65019d38..06d203ce95f6 100644 --- a/constraints.txt +++ b/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. diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index 7164715ee7d3..5d5c94aee580 100644 --- a/docker/Dockerfile.multi +++ b/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 diff --git a/requirements.txt b/requirements.txt index abcaef09250a..c3a62d7ee170 100644 --- a/requirements.txt +++ b/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 diff --git a/security_scanning/pyproject.toml b/security_scanning/pyproject.toml index 24b5f57d6c9d..4db3fe49667d 100644 --- a/security_scanning/pyproject.toml +++ b/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)", diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index eb98bc2a7b27..6463363ab778 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/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 diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index b89a16f3a959..b1c9e901d457 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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) diff --git a/tests/unittest/_torch/visual_gen/test_fa4_split_kv.py b/tests/unittest/_torch/visual_gen/test_fa4_split_kv.py new file mode 100644 index 000000000000..faa9fdf909be --- /dev/null +++ b/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"])