diff --git a/.claude/skills/kernel-cute-writing/references/api-core.md b/.claude/skills/kernel-cute-writing/references/api-core.md index 4945971a6a88..9b640b236f1b 100644 --- a/.claude/skills/kernel-cute-writing/references/api-core.md +++ b/.claude/skills/kernel-cute-writing/references/api-core.md @@ -126,7 +126,8 @@ cute.is_major(x, d) # True if dimension d is stride-1 ```python cute.make_tensor(ptr, layout) # From pointer + layout cute.make_identity_tensor(shape) # Coordinate mapping tensor -cute.make_rmem_tensor(layout_or_shape, dtype) # Register memory tensor +cute.make_rmem_tensor(layout) # Register memory tensor +cute.make_fragment(shape, dtype) # Register fragment ``` ### Data Initialization diff --git a/.claude/skills/kernel-cute-writing/references/concepts-tensors.md b/.claude/skills/kernel-cute-writing/references/concepts-tensors.md index 09e3e775bdef..5b7c9c4d67d7 100644 --- a/.claude/skills/kernel-cute-writing/references/concepts-tensors.md +++ b/.claude/skills/kernel-cute-writing/references/concepts-tensors.md @@ -81,7 +81,8 @@ mA = cute.make_tensor(ptr, layout=layout) ### Register/Fragment Tensors (Owning) ```python -rmem = cute.make_rmem_tensor(layout_or_shape, dtype) +rmem = cute.make_rmem_tensor(layout) +frag = cute.make_fragment(shape, dtype) ``` ## Accessing Tensors diff --git a/constraints.txt b/constraints.txt index a99d65019d38..7de1c603d547 100644 --- a/constraints.txt +++ b/constraints.txt @@ -1,18 +1,10 @@ -# 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 -# 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 -torch-c-dlpack-ext==0.1.3 -# CUTLASS DSL 4.6.1 libraries require protobuf 6.x. -protobuf>=6.30.2,<7 +# These vulnerabilities were inherited from the base image (pytorch:25.12-py3) and should be removed when the base image +# is updated. +# Upgrade base image nvidia-cutlass-dsl 4.3.5 to 4.4.2 +nvidia-cutlass-dsl>=4.4.2 # The `nvidia-cutlass-dsl` package does not pin numpy at all, which can be problematic in certain CI # stages. numpy>=2.0.0,<2.4 # numba 0.63.1 requires numpy<2.4 -# The remaining vulnerabilities were inherited from the base image -# (pytorch:25.12-py3) and can be removed after the base image is updated. # WAR against https://github.com/advisories/GHSA-qcq2-496w-v96p mistune>=3.3.0 # WAR against https://github.com/advisories/GHSA-rch3-82jr-f9w9 diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index 7164715ee7d3..14798a5a354a 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -66,13 +66,11 @@ RUN --mount=type=bind,source=docker/common,target=/opt/docker/common \ # Install constraints after install.sh so cleanup() doesn't delete the file mid-RUN COPY constraints.txt /tmp/constraints.txt -# constraints.txt is also consumed as a pip constraint file, which forbids extras. RUN --mount=type=cache,target=/root/.cache/pip \ - pip3 uninstall -y tornado black nbconvert pillow nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base \ - nvidia-cutlass-dsl-libs-core nvidia-cutlass-dsl-libs-cu12 nvidia-cutlass-dsl-libs-cu13 numpy wandb || true && \ + pip3 uninstall -y tornado black nbconvert pillow nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base 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 && \ rm /tmp/constraints.txt # Install UCX, NIXL, etcd diff --git a/docker/common/install_fa4.sh b/docker/common/install_fa4.sh index e70083220dc7..5138b19b2925 100644 --- a/docker/common/install_fa4.sh +++ b/docker/common/install_fa4.sh @@ -2,7 +2,7 @@ set -ex -FLASH_ATTN_4_VERSION="4.0.0b19" +FLASH_ATTN_4_VERSION="4.0.0b11" if [ -n "${GITHUB_MIRROR}" ]; then export PIP_INDEX_URL="https://urm.nvidia.com/artifactory/api/pypi/pypi-remote/simple" diff --git a/docs/source/installation/installation-guide.md b/docs/source/installation/installation-guide.md index 6ff9c295211e..dca8d522fb44 100644 --- a/docs/source/installation/installation-guide.md +++ b/docs/source/installation/installation-guide.md @@ -68,8 +68,7 @@ Before installing the latest version, uninstall any previous CUTLASS DSL install [CUTLASS DSL installation guide](https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/quick_start.html#installation): ```bash -pip3 uninstall nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base \ - nvidia-cutlass-dsl-libs-core nvidia-cutlass-dsl-libs-cu12 nvidia-cutlass-dsl-libs-cu13 +pip3 uninstall nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base nvidia-cutlass-dsl-libs-cu13 ``` ```bash diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index c942fff785cb..c4a12ceb9089 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -6464,11 +6464,7 @@ def launchTestJobs(pipeline, testFilter, globalVars) trtllm_utils.llmExecStepWithRetry(pipeline, script: "apt-get remove -y python3-pygments") // Remove stale nvidia-cutlass-dsl from the base image to prevent namespace // directory corruption when pip upgrades to the version required by tensorrt_llm. - trtllm_utils.llmExecStepWithRetry( - pipeline, - script: "pip3 uninstall -y nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base " + - "nvidia-cutlass-dsl-libs-core nvidia-cutlass-dsl-libs-cu12 " + - "nvidia-cutlass-dsl-libs-cu13 || true") + trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 uninstall -y nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base || true") trtllm_utils.llmExecStepWithRetry(pipeline, script: 'rm -rf $(python3 -c "import site; print(site.getsitepackages()[0])")/nvidia_cutlass_dsl*') } trtllm_utils.llmExecStepWithRetry(pipeline, script: "apt-get update && apt-get install -y python3-pip git rsync curl wget") diff --git a/requirements.txt b/requirements.txt index 8fff36a3a30e..24de148b3138 100644 --- a/requirements.txt +++ b/requirements.txt @@ -78,17 +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 +nvidia-cutlass-dsl[cu13]==4.5.0; python_version >= "3.10" 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.2.10; python_version >= "3.10" # required for MinimaxM3 MSA 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.6 # used for reduce nvidia-cutlass-dsl host overhead 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 +flash-attn-4==4.0.0b11 mistral-common>=1.10.0 torchao>=0.14.1,<0.16.0 cuda-core diff --git a/security_scanning/poetry.lock b/security_scanning/poetry.lock index 2ca48a23b4c3..060d7cdf6ee6 100644 --- a/security_scanning/poetry.lock +++ b/security_scanning/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.4.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. [[package]] name = "accelerate" @@ -1314,7 +1314,7 @@ files = [ [package.dependencies] datasets = ">=2.0.0" dill = "*" -fsspec = {version = ">=2021.5.0", extras = ["http"]} +fsspec = {version = ">=2021.05.0", extras = ["http"]} huggingface-hub = ">=0.7.0" multiprocess = "*" numpy = ">=1.17" @@ -1368,27 +1368,27 @@ files = [ [[package]] name = "flash-attn-4" -version = "4.0.0b19" +version = "4.0.0b11" description = "Flash Attention CUTE (CUDA Template Engine) implementation" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "flash_attn_4-4.0.0b19-py3-none-any.whl", hash = "sha256:bc3856e018fa32e2b833726566641e51d02d45f2d444849937276a46c01653a2"}, - {file = "flash_attn_4-4.0.0b19.tar.gz", hash = "sha256:c904c9da0387c2ac0420cffdf61d6201d94cfeb354b6768fde794b04e9fc4e89"}, + {file = "flash_attn_4-4.0.0b11-py3-none-any.whl", hash = "sha256:f028871f46a63d466d05762876506a12957147374b17f66a52376beed4238dc5"}, + {file = "flash_attn_4-4.0.0b11.tar.gz", hash = "sha256:b2ceede17eea2dfe9c62e8ddef454566702b119421bd098f61faa9aab4cd5885"}, ] [package.dependencies] apache-tvm-ffi = ">=0.1.5,<0.2" einops = "*" -nvidia-cutlass-dsl = ">=4.5.2" -quack-kernels = ">=0.5.0" +nvidia-cutlass-dsl = ">=4.4.2" +quack-kernels = ">=0.4.0" torch = "*" torch-c-dlpack-ext = "*" typing_extensions = "*" [package.extras] -cu13 = ["nvidia-cutlass-dsl[cu13] (>=4.5.2)"] +cu13 = ["nvidia-cutlass-dsl[cu13] (>=4.4.2)"] dev = ["pytest", "pytest-xdist", "ruff"] [[package]] @@ -2146,7 +2146,7 @@ files = [ [package.dependencies] attrs = ">=22.2.0" -jsonschema-specifications = ">=2023.3.6" +jsonschema-specifications = ">=2023.03.6" referencing = ">=0.28.4" rpds-py = ">=0.25.0" @@ -3242,19 +3242,6 @@ files = [ {file = "nvidia_cuda_cupti-13.0.85-py3-none-win_amd64.whl", hash = "sha256:683f58d301548deeefcb8f6fac1b8d907691b9d8b18eccab417f51e362102f00"}, ] -[[package]] -name = "nvidia-cuda-nvdisasm" -version = "13.3.73" -description = "CUDA nvdisasm" -optional = false -python-versions = ">=3" -groups = ["main"] -files = [ - {file = "nvidia_cuda_nvdisasm-13.3.73-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:dd4751884f9016b9b6dbf007abdeb5681d0a2edc731dd3d2fda9d6d878e88f73"}, - {file = "nvidia_cuda_nvdisasm-13.3.73-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fa17084b07c0dca68a42892f771b4b1b40fbe9b91660209623e61cea611cae8c"}, - {file = "nvidia_cuda_nvdisasm-13.3.73-py3-none-win_amd64.whl", hash = "sha256:da2fab133c3d095d83f13587eb87149beabd199b34a3cc270d0aa99449a628c3"}, -] - [[package]] name = "nvidia-cuda-nvrtc" version = "13.0.88" @@ -3446,130 +3433,72 @@ files = [ [[package]] name = "nvidia-cutlass-dsl" -version = "4.6.1" +version = "4.5.0" description = "NVIDIA CUTLASS Python DSL" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "nvidia_cutlass_dsl-4.6.1-py3-none-any.whl", hash = "sha256:93135a9d48e1bedf584828e0a021f174ac31591ef21f6ea53c206169ccbfab26"}, + {file = "nvidia_cutlass_dsl-4.5.0-py3-none-any.whl", hash = "sha256:3b051fe02ca69422ab840e64d9865667aba288a3984a7ca4ccd038a82aef1344"}, ] [package.dependencies] -nvidia-cutlass-dsl-libs-base = "4.6.1" -nvidia-cutlass-dsl-libs-cu12 = "4.6.1" -nvidia-cutlass-dsl-libs-cu13 = {version = "4.6.1", optional = true, markers = "extra == \"cu13\""} +nvidia-cutlass-dsl-libs-base = "4.5.0" +nvidia-cutlass-dsl-libs-cu13 = {version = "4.5.0", optional = true, markers = "extra == \"cu13\""} [package.extras] -cu13 = ["nvidia-cutlass-dsl-libs-cu13 (==4.6.1)"] +cu13 = ["nvidia-cutlass-dsl-libs-cu13 (==4.5.0)"] [[package]] name = "nvidia-cutlass-dsl-libs-base" -version = "4.6.1" +version = "4.5.0" description = "NVIDIA CUTLASS Python DSL" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:df002a6d01624bacfef3bb0b0096ef190a324e3951870ad9685ba9a3daf7ec78"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:e68ac160f75f79b690c267be3785cd35631d5bc147fc5c5a4e1902454bdd0d4c"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:617ff12b9a1fcbc3ec7397dea99380a72855a477c69c2421ce8a773966c127bb"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:dcbbf471839801501030f1097ce13dbf5082ad34f26d6b503459c6fed078e4e9"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:e244956cb5097c0f9595a23fb204c7b6b1ef9ecb23db6389e5103605f7d324da"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b1ffc932923fc7871be9d9890622c4a827fe64ab291f3660792b53f808f08801"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:7ca4e5ca9bd0caa2ec797ebb1e404bb3eb0f76d235823611795cb14d8d341e4c"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:09a6aace841594ea39cf75f7cee68dc31b8cf3cdbadffcbd494f08fc73c6466d"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:96815af98b19b8045802fd8b0a53d5e197e0d3bce42c7038e2861af7c8483f5d"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:27461ff52338280b6c02475f87851273b5a11c387d9daa383376225a3f767cc2"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:5fd2c4051a9251009a978ecdb61381271f904dd44704fc7f2a8995777d87dbfa"}, - {file = "nvidia_cutlass_dsl_libs_base-4.6.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:0a1bd67e0a387834389b81c1742a3af8c583fe2499fd01d8da9fc4354d4c322d"}, + {file = "nvidia_cutlass_dsl_libs_base-4.5.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:c78b18f2b44ca10a91bc76380ebd65bb7b86aa97a9330bae9b73eb0a1bc51d55"}, + {file = "nvidia_cutlass_dsl_libs_base-4.5.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:5cfdf52bea8feede5e512a094484956693cb87adaafa310991d2876653b1a88e"}, + {file = "nvidia_cutlass_dsl_libs_base-4.5.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:f8635ad1e0a670323cc729f167067fa880cb56577ec2e79afb80a35ab371912e"}, + {file = "nvidia_cutlass_dsl_libs_base-4.5.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:7bb6de91b00a2b392cd834fec174a1461bf0f10a9b6d28086c8f4885aed27218"}, + {file = "nvidia_cutlass_dsl_libs_base-4.5.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:3f7c133d31fa82ae7db697fd6943a5f9a2c97c8a40ee1056c67ef29fe00974d8"}, + {file = "nvidia_cutlass_dsl_libs_base-4.5.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:bd18322d9247f8c033a10ed4e519c4985ca6b4fb578ade382e5a264422ebd915"}, + {file = "nvidia_cutlass_dsl_libs_base-4.5.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:90a4d802a03963fa36eb287fbc9b40a1374590fc7e8cc1b9673dee8872f75713"}, + {file = "nvidia_cutlass_dsl_libs_base-4.5.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:8e58b016da5bb09bd1d809d0c025433edb36b279adfbcd107e96361b214bd8bc"}, + {file = "nvidia_cutlass_dsl_libs_base-4.5.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:2e121b20f0a48122c9b48227d00a7d681189e1de2fd4d211f9661a4e1658f066"}, + {file = "nvidia_cutlass_dsl_libs_base-4.5.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:0a60dfce3349984315306ef719ed1edf0e225527158f26019a5cf266e06cc45d"}, ] [package.dependencies] cuda-python = ">=12.8" numpy = "*" -nvidia-cuda-nvdisasm = ">=13.3,<14" -nvidia-cutlass-dsl-libs-core = "4.6.1" -protobuf = ">=6.30.2,<7" -typing-extensions = ">=4.10.0" - -[[package]] -name = "nvidia-cutlass-dsl-libs-core" -version = "4.6.1" -description = "NVIDIA CUTLASS Python DSL" -optional = false -python-versions = ">=3.10" -groups = ["main"] -files = [ - {file = "nvidia_cutlass_dsl_libs_core-4.6.1-py3-none-any.whl", hash = "sha256:f1d895ee24b1ba711b2b9d4a43c62fa1f6fe3a50634e25eeffe88fa17f6c5e47"}, -] - -[package.dependencies] -cuda-python = ">=12.8" -numpy = "*" -nvidia-cuda-nvdisasm = ">=13.3,<14" -protobuf = ">=6.30.2,<7" -typing-extensions = ">=4.10.0" - -[[package]] -name = "nvidia-cutlass-dsl-libs-cu12" -version = "4.6.1" -description = "NVIDIA CUTLASS Python DSL" -optional = false -python-versions = ">=3.10" -groups = ["main"] -files = [ - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:4aa1c937d941218cc77c08145b65a6f9be876cc3cc2fe5e760f9a1921f8a9c7c"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:68a09e6e42d64588b2026e5a36eedc789233c200cbe1871a9e4bdf5371e980cc"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:271d2a586be0782a68a1c9cb5cd59794cedca002cc61f9e040e93086f0080aef"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:adef539dbee08d79825aac8603737cc5354c5021d6959fff0dfdace3c0d014c1"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:2999be5954805219611969c7e3f204d8a8d03e73cf7dd44c8dfba196e83142d7"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:7a0745abb6b65a2c43002d924949df979fcb2d2bb6aad6efa5af44ea611e0eb4"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:17d9d5bdb8f530266dc8424236a80523ee2a9807e852a24c22096438a12df3bd"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:2504b5bdb2f0210a6a1bc923acefe979006366694b73f1458a3c8571aef3b765"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:c84ec5d4f45b634abe16b2f111c5352fd21ec665ebb929e4031883dcf135d311"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:02f89f79bffca1f670dc3e19e5238f1208051630eb4ebd83571ccaaa34e7ea45"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:e212b17fd8e35c8ddbbf27288dc56dcce1ca0051fb3d466ced4afb47d3e97bf8"}, - {file = "nvidia_cutlass_dsl_libs_cu12-4.6.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:f4df623107d4f1e92fd534b8c49c8a3fe8bdd5a9ce47eb8ce4c186666e254ccf"}, -] - -[package.dependencies] -cuda-python = ">=12.8" -numpy = "*" -nvidia-cuda-nvdisasm = ">=13.3,<14" -nvidia-cutlass-dsl-libs-base = "4.6.1" -protobuf = ">=6.30.2,<7" -typing-extensions = ">=4.10.0" +typing-extensions = "*" [[package]] name = "nvidia-cutlass-dsl-libs-cu13" -version = "4.6.1" +version = "4.5.0" description = "NVIDIA CUTLASS Python DSL" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:da30f3b6f7673ee6fa1f5d62534c9419fb4927f9f6a83efd3f2e125b110ba194"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:062021189a027f2ff49f3eae2f33fef70aa5ac1b9cddddd05e77914d2401a684"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:7d327ad46acb442c09b9666ca51d28e834f4f49f22a2d1cd3e722722201d03a4"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a8483778ad75ae50efd5b8981adfe8d7bc0c9cbcf5fd16a02b3bd9a062b98e2d"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:68ef426fdc5d8081d620bfa5dcddad4eadd8a98700e5575c707b21f161670444"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:a6dd028303713201e51827ca8e27048537a7ed1056aaf748a0be7e4353d36c3f"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:81370511dfd9e1c08da2302cf140f4adec89df504ec348a23b08aa0174720ec3"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:7da99584f4f0d498ca70a647022ab295ca54a20949441181cd0c3289584da8a3"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:abb6cee63d20238bf8ff7e3dac12534b5e8bc45d4256273d979398aeb17a46bd"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:88f58027a254ac3e6b9592eaf36125cb6ecf77ad4e4026847804125546ec4b10"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:9aacf551cf570e4d7bd0fde04f3aff5e8cd9032fe2b386298e2c0a40d585a681"}, - {file = "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:ee772072c220363486e0ae0981d68ecf8b995e508c5df48fcc998efd738728f7"}, + {file = "nvidia_cutlass_dsl_libs_cu13-4.5.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:b95cd78ca5f3440cd488ea03c30d394ca07e491e5bd6bb274f8f26d4f4713439"}, + {file = "nvidia_cutlass_dsl_libs_cu13-4.5.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:145cf5563b2d0d347c722d84a3b74541da5a1ec5b46821b54acc202693ce56ae"}, + {file = "nvidia_cutlass_dsl_libs_cu13-4.5.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:dadddab5bfe11973f8a0dfa66645a42ccb587287abb0766992e0e69da621dc82"}, + {file = "nvidia_cutlass_dsl_libs_cu13-4.5.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:35726dddb6f6e63c767caf7533d583b4462e2c6b90df7fa67738d58b2eede742"}, + {file = "nvidia_cutlass_dsl_libs_cu13-4.5.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:a799eb9d65ba03095444907b8bf617b5a8bd7d03d1b95cec9637c558af6d0b60"}, + {file = "nvidia_cutlass_dsl_libs_cu13-4.5.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:fc0b5a81ff591db72489134ca206ae886f0cce43f20863010a7f30fcfe484a7b"}, + {file = "nvidia_cutlass_dsl_libs_cu13-4.5.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:f2b821671add2e69a1377e7fd87e6261995db281b2f8e516ddae2fd6b7a6c1d0"}, + {file = "nvidia_cutlass_dsl_libs_cu13-4.5.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:554b775069d093f308949a65880bf9c9bfd48b1f4b2fd0e0d97aa2f608e6eea1"}, + {file = "nvidia_cutlass_dsl_libs_cu13-4.5.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:652e34d8a78accceab321da0232157e493b113e417306269d700a297a7d62447"}, + {file = "nvidia_cutlass_dsl_libs_cu13-4.5.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:3d460d03d6ea0a463e262ee07964ccafffde8eeefe058c9615d77da1ddd6d003"}, ] [package.dependencies] cuda-python = ">=12.8" numpy = "*" -nvidia-cuda-nvdisasm = ">=13.3,<14" -nvidia-cutlass-dsl-libs-base = "4.6.1" -protobuf = ">=6.30.2,<7" -typing-extensions = ">=4.10.0" +nvidia-cutlass-dsl-libs-base = "4.5.0" +typing-extensions = "*" [[package]] name = "nvidia-matmul-heuristics" @@ -5138,29 +5067,27 @@ cffi = {version = "*", markers = "implementation_name == \"pypy\""} [[package]] name = "quack-kernels" -version = "0.6.2" +version = "0.4.1" description = "" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "quack_kernels-0.6.2-py3-none-any.whl", hash = "sha256:c193d5df51a56f5fdb681bee03f4042ed812f8190d022695cba9c14eb0e63cd6"}, - {file = "quack_kernels-0.6.2.tar.gz", hash = "sha256:83684ceb3ee4f5096035eb6f6e005e63f3f838bc9a44cf3c3e580f24c2b7233d"}, + {file = "quack_kernels-0.4.1-py3-none-any.whl", hash = "sha256:c1c8df2935bf5156ec47d2c5384ac08b411fd0ee702d80ae916dbf6d6f5ae813"}, + {file = "quack_kernels-0.4.1.tar.gz", hash = "sha256:9d7d6ba412bc0c8a9b1331c52a73db76280adb9dc2f2750df4851ddabef1466b"}, ] [package.dependencies] apache-tvm-ffi = ">=0.1.6,<0.2" einops = "*" -nvidia-cutlass-dsl = "4.6.1" +nvidia-cutlass-dsl = ">=4.4.2" torch = "*" torch-c-dlpack-ext = "*" [package.extras] -bench = ["pandas", "tyro"] -cu13 = ["nvidia-cutlass-dsl[cu13] (==4.6.1)"] +cu13 = ["nvidia-cutlass-dsl[cu13] (>=4.4.2)"] dev = ["pre-commit", "pytest", "pytest-xdist", "ruff"] heuristics = ["nvidia-matmul-heuristics"] -jax = ["jax", "jax-tvm-ffi"] [[package]] name = "referencing" @@ -6914,4 +6841,4 @@ type = ["pytest-mypy (>=1.0.1) ; platform_python_implementation != \"PyPy\""] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.13" -content-hash = "2754deb353c9e32d33a28599c69a1456d3dd71071906300ff0e276c46f20ea93" +content-hash = "637d5d2dcb9c329e72fed08068cf97d937adb3bd00cb716f9d2040823f9a527d" diff --git a/security_scanning/pyproject.toml b/security_scanning/pyproject.toml index 566045eb86a8..e73c1a6074fe 100644 --- a/security_scanning/pyproject.toml +++ b/security_scanning/pyproject.toml @@ -70,15 +70,15 @@ 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.5.0)", "nvidia-matmul-heuristics (==0.1.0.27)", - "quack-kernels (==0.6.2)", + "quack-kernels (>=0.2.10)", "jinja2 (>=3.1.6,<4.0.0)", "plotly (>=6.9.0,<7.0.0)", "partial-json-parser (>=0.2.1.1.post7,<0.3.0.0)", "mcp (<2.0.0)", "torch-c-dlpack-ext (==0.1.3)", - "flash-attn-4 (==4.0.0b19)", + "flash-attn-4 (==4.0.0b11)", "mistral-common (>=1.10.0)", "torchao (>=0.14.1,<0.16.0)", "cuda-core (>=1.1.1,<2.0.0)", diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py index ba4cca087517..a5c3b8dc00d4 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py @@ -20,34 +20,15 @@ MSA_REQUIRED_HEAD_DIM = 128 -def _install_msa_cutlass_compatibility() -> None: - """Provide the CUTLASS 4.5 names still referenced by the packaged MSA sources.""" - try: - import cutlass.cute as cute - except ImportError: - return - - # MSA has not yet migrated these two aliases to their CUTLASS DSL 4.6 - # names. Keep this shim local to the MSA import path and remove it once the - # packaged sources use cute.ThrMma and cute.make_rmem_tensor directly. - if not hasattr(cute.core, "ThrMma"): - setattr(cute.core, "ThrMma", cute.ThrMma) - if not hasattr(cute, "make_fragment"): - setattr(cute, "make_fragment", cute.make_rmem_tensor) - - @functools.lru_cache(maxsize=1) def msa_package_available() -> bool: - """Prepare and report whether the packaged fmha_sm100 module can be imported. + """Return whether the packaged fmha_sm100 module can be imported. Cached: each call scans sys.path until fmha_sm100 is first imported, and create_fmha_libs asks once per attention layer. Whether the package is installed cannot change within a process. """ - if importlib.util.find_spec("fmha_sm100") is None: - return False - _install_msa_cutlass_compatibility() - return True + return importlib.util.find_spec("fmha_sm100") is not None def require_msa_module() -> types.ModuleType: @@ -57,7 +38,6 @@ def require_msa_module() -> types.ModuleType: advertised in the config schema on systems where the kernels cannot load. A missing package is a hard error, never a silent fallback to another backend. """ - _install_msa_cutlass_compatibility() try: import fmha_sm100 except ImportError as exc: diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/argmax.py b/tensorrt_llm/_torch/cute_dsl_kernels/argmax.py index 0a019dec74b2..d76e52ec6297 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/argmax.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/argmax.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025, Tri Dao. # SPDX-FileCopyrightText: Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -117,7 +117,7 @@ def store_shared_remote( @cute.jit def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor: - tApA = cute.make_rmem_tensor( + tApA = cute.make_fragment( cute.make_layout( ( cute.size(tAcA, mode=[0, 1]), diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py index 06ccffdad867..0194cbc8359e 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py @@ -1394,7 +1394,7 @@ def kernel( self.epilog_tmem_copy_and_partition(epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs)) - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tTR_rC = cute.make_fragment(tTR_rAcc.shape, self.c_dtype) tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( tiled_copy_t2r, tTR_rC, epi_tidx, sC) tma_atom_c, bSG_sC, bSG_gC_partitioned = ( @@ -1862,7 +1862,7 @@ def epilog_tmem_copy_and_partition( # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) # (T2R, T2R_M, T2R_N) - tTR_rAcc = cute.make_rmem_tensor( + tTR_rAcc = cute.make_fragment( tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) return tiled_copy_t2r, tTR_tAcc, tTR_rAcc diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_gemm_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_gemm_persistent.py index 867bf02214c5..fcd137a38ffd 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_gemm_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_gemm_persistent.py @@ -659,7 +659,7 @@ class SharedStorage: gC_mnl_epi = cute.flat_divide(tCgC[((None, None), 0, 0, None, None, None)], epi_tile) tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) - tTR_rAcc = cute.make_rmem_tensor( + tTR_rAcc = cute.make_fragment( tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype ) @@ -670,7 +670,7 @@ class SharedStorage: tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) tRS_sC = thr_copy_r2s.partition_D(sC) - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tTR_rC = cute.make_fragment(tTR_rAcc.shape, self.c_dtype) tRS_rC = tiled_copy_r2s.retile(tTR_rC) # SMEM -> GMEM TMA store setup diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py index dee10d230030..e769566218fe 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py @@ -192,7 +192,6 @@ def kda_decode_mtp_kernel( cute.make_layout((KERNEL_WIDTH * vec_size,), stride=(1,)), cutlass.Float32 ) r_exp_A = cutlass.Float32(0.0) - run_precompute = False if cutlass.const_expr(USE_REGULAR_METADATA) or eos > bos: if cutlass.const_expr(FUSE_PRECOMPUTE or RUNTIME_PRECOMPUTE_FLAG): if cutlass.const_expr(RUNTIME_PRECOMPUTE_FLAG): diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py index 1516eafc22ec..94db5d143085 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py @@ -326,7 +326,7 @@ def utccp_required_smem_warp_transpose(smem_ptr) -> None: smem_ptr: cute.Pointer to int32, must be 128-int (= 512-byte) aligned. """ lane_idx = cute.arch.lane_idx() - values = cute.make_rmem_tensor(4, cutlass.Int32) + values = cute.make_fragment(4, cutlass.Int32) for i in cutlass.range_constexpr(4): offset = (i ^ (lane_idx >> 3)) * 32 + lane_idx values[i] = ld_shared_b32(smem_ptr + offset) @@ -1839,12 +1839,12 @@ def kernel( else: # fp32, 4-byte weights MAX_NUM_W_IN_REG = 56 if next_n == 3 else 64 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads) - w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype) # Batched STG: hold reduced result per t in register; the # actual STG happens once after the for-t loop to land all # STGs in one contiguous LSU phase. if cutlass.const_expr(self.use_batched_store): - result_arr = cute.make_rmem_tensor(next_n, self.output_dtype) + result_arr = cute.make_fragment(next_n, self.output_dtype) else: result_arr = None q_stage_local = cutlass.Int32(0) @@ -2102,12 +2102,12 @@ def kernel( else: MAX_NUM_W_IN_REG = 56 if next_n == 3 else 64 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads) - w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype) # Batched STG: hold reduced result per t in register; the # actual STG happens once after the for-t loop to land all # STGs in one contiguous LSU phase. if cutlass.const_expr(self.use_batched_store): - result_arr = cute.make_rmem_tensor(next_n, self.output_dtype) + result_arr = cute.make_fragment(next_n, self.output_dtype) else: result_arr = None q_stage_local = cutlass.Int32(0) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index 722ac73cebd0..4367f06f442a 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -1283,7 +1283,7 @@ def kernel( else: MAX_NUM_W_IN_REG = 64 if next_n == 1 else 40 if next_n >= 4 else 52 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads) - w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype) q_stage_local = cutlass.Int32(0) while has_work: @@ -1518,7 +1518,7 @@ def kernel( else: MAX_NUM_W_IN_REG = 64 if next_n == 1 else 40 if next_n >= 4 else 52 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads) - w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) + w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype) q_stage_local = cutlass.Int32(0) while has_work: diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py index 5ce3ab29cdcc..159d44cb3445 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py @@ -1551,8 +1551,8 @@ def _phase3_writeback(self, tidx, row_start, s_indices, score, indices, dst, dst nvec_per_thread = cutlass.const_expr( cute.ceil_div(self.top_k, vecsize_out * self.num_threads_per_cta) ) - topk_vals = cute.make_rmem_tensor((vecsize_out, nvec_per_thread), self.dtype) - topk_indices = cute.make_rmem_tensor((vecsize_out, nvec_per_thread), cutlass.Int32) + topk_vals = cute.make_fragment((vecsize_out, nvec_per_thread), self.dtype) + topk_indices = cute.make_fragment((vecsize_out, nvec_per_thread), cutlass.Int32) stride = self.num_threads_per_cta * vecsize_out for i in cutlass.range(nvec_per_thread, unroll_full=True): @@ -1706,7 +1706,7 @@ def filtered_topk_kernel_per_row( num_bits_per_copy=self.num_copy_bits, ) - scan_frag = cute.make_rmem_tensor((vec_size,), self.dtype) + scan_frag = cute.make_fragment((vec_size,), self.dtype) # Trivial case: length <= top_k. In SP multi-CTA cluster mode this # per-chunk shortcut is unsafe (a CTA taking it would skip the cluster diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py index b3725ef4e6c9..44275ab832c3 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py @@ -626,7 +626,7 @@ def load_slice_to_smem( # block_count_ge's fast path. ic_local indexes both GMEM # (input_row[slice_start + ic_local]) and SMEM (smem_input[ic_local]). if self.enable_unroll_4: - rng_frag = cute.make_rmem_tensor((vec_w,), self.dtype) + rng_frag = cute.make_fragment((vec_w,), self.dtype) big_iters = cutlass.Int32(0) if slice_len > i_local + cutlass.Int32(vec_w - 1): big_iters = (slice_len - i_local - cutlass.Int32(vec_w)) // cutlass.Int32( @@ -653,7 +653,7 @@ def load_slice_to_smem( i_local = i_local + big_iters * cutlass.Int32(step_elem) # 1-way tail vec loop (slice_len mod step_elem residual). - tail_frag = cute.make_rmem_tensor((vec_w,), self.dtype) + tail_frag = cute.make_fragment((vec_w,), self.dtype) while i_local + cutlass.Int32(vec_w - 1) < slice_len: src_ptr = cute.make_ptr( self.dtype, @@ -1022,7 +1022,7 @@ def phase1b_hspace_rungs( # cum_at). qfracs descending in h => thresholds ascending in m. if warp_id == cutlass.Int32(0): top = cutlass.Int32(NB - 1) - lane * cutlass.Int32(SEG) - seg_frag = cute.make_rmem_tensor((SEG,), cutlass.Int32) + seg_frag = cute.make_fragment((SEG,), cutlass.Int32) part = cutlass.Int32(0) for j in cutlass.range_constexpr(SEG): v8 = smem_hist[top - cutlass.Int32(j)] @@ -1095,7 +1095,7 @@ def phase1b_hspace_rungs_cached( if warp_id == cutlass.Int32(0): top = cutlass.Int32(NB - 1) - lane * cutlass.Int32(SEG) - seg_frag = cute.make_rmem_tensor((SEG,), cutlass.Int32) + seg_frag = cute.make_fragment((SEG,), cutlass.Int32) part = cutlass.Int32(0) for j in cutlass.range_constexpr(SEG): v8 = smem_hist[top - cutlass.Int32(j)] @@ -1206,7 +1206,7 @@ def block_count_ge( # Each iter loads 1 vec_w chunk; LLVM unrolls 4 iters at IR level # so 4 LDG.E.* stay in flight. if self.enable_unroll_4: - rng_frag = cute.make_rmem_tensor((vec_w,), self.dtype) + rng_frag = cute.make_fragment((vec_w,), self.dtype) # Number of complete vec_w-aligned loads this thread can do: # need: i + k*step_elem + (vec_w - 1) < N # max k: floor((N - i - vec_w) / step_elem) @@ -1250,7 +1250,7 @@ def block_count_ge( # full vec_w-stride or less). i is always vec_w-aligned here (it # advanced by multiples of step_elem = num_threads*vec_w), so the # same vec_align bytes hold. - tail_frag = cute.make_rmem_tensor((vec_w,), self.dtype) + tail_frag = cute.make_fragment((vec_w,), self.dtype) while i + cutlass.Int32(vec_w - 1) < N: if cutlass.const_expr(smem_input is not None): src_ptr = cute.make_ptr( @@ -1423,8 +1423,8 @@ def block_count_ge_multi( copy_atom = self._make_load_copy_atom() step_elem = cutlass.const_expr(num_threads * vec_w) - thr_frag = cute.make_rmem_tensor((M,), cutlass.Float32) - cnt_frag = cute.make_rmem_tensor((M,), cutlass.Int32) + thr_frag = cute.make_fragment((M,), cutlass.Float32) + cnt_frag = cute.make_fragment((M,), cutlass.Int32) for m in cutlass.range_constexpr(M): thr_frag[m] = s_mt_thr[m] cnt_frag[m] = cutlass.Int32(0) @@ -1436,7 +1436,7 @@ def block_count_ge_multi( step = cutlass.Int32(step_elem) if self.enable_unroll_4: - rng_frag = cute.make_rmem_tensor((vec_w,), self.dtype) + rng_frag = cute.make_fragment((vec_w,), self.dtype) big_iters = cutlass.Int32(0) if slice_end > i + cutlass.Int32(vec_w - 1): big_iters = (slice_end - i - cutlass.Int32(vec_w)) // cutlass.Int32( @@ -1461,7 +1461,7 @@ def block_count_ge_multi( cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vj >= thr_frag[m]) i = i + big_iters * cutlass.Int32(step_elem) - tail_frag = cute.make_rmem_tensor((vec_w,), self.dtype) + tail_frag = cute.make_fragment((vec_w,), self.dtype) while i + cutlass.Int32(vec_w - 1) < slice_end: src_ptr = cute.make_ptr( self.dtype, @@ -2171,7 +2171,7 @@ def phase3_collect_candidates( if self.enable_phase3_unroll: # Fast path: 4-way unrolled vec loop (4 loading instructions in flight). if self.enable_unroll_4: - rng_frag = cute.make_rmem_tensor((vec_w,), self.dtype) + rng_frag = cute.make_fragment((vec_w,), self.dtype) big_iters = cutlass.Int32(0) if N_local > ic + cutlass.Int32(vec_w - 1): big_iters = (N_local - ic - cutlass.Int32(vec_w)) // cutlass.Int32( @@ -2213,7 +2213,7 @@ def phase3_collect_candidates( ic = ic + big_iters * cutlass.Int32(step_elem) # Tail vec loop: 1-way, handles remainder < 2*step. - tail_frag = cute.make_rmem_tensor((vec_w,), self.dtype) + tail_frag = cute.make_fragment((vec_w,), self.dtype) while ic + cutlass.Int32(vec_w - 1) < N_local: if cutlass.const_expr(smem_input is not None): src_ptr = cute.make_ptr( @@ -2578,7 +2578,7 @@ def _kth_bin_search_rw(self, smem_hist, smem_wcnt, lo, binw, tidx, warp_id, lane # Step 3 (every warp, lane-parallel): lane l owns the contiguous # descending positions [l*ppl, (l+1)*ppl) of the target slice. ppl = cutlass.const_expr((bins_per_warp + self.WARP_SIZE - 1) // self.WARP_SIZE) - cnt_frag = cute.make_rmem_tensor((ppl,), cutlass.Int32) + cnt_frag = cute.make_fragment((ppl,), cutlass.Int32) my_sum = cutlass.Int32(0) for j3 in cutlass.range_constexpr(ppl): pos = lane * cutlass.Int32(ppl) + cutlass.Int32(j3) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_direct.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_direct.py index 66165a286729..5dbe90ced305 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_direct.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_direct.py @@ -147,7 +147,7 @@ def _bin_select( ): per = cutlass.const_expr(nbins // self.num_threads) num_warps = cutlass.const_expr(self.num_warps) - h = cute.make_rmem_tensor((per,), cutlass.Int32) + h = cute.make_fragment((per,), cutlass.Int32) s_sum = cutlass.Int32(0) for q in cutlass.range_constexpr(per): h[q] = smem_hist[tidx * cutlass.Int32(per) + cutlass.Int32(q)] @@ -349,7 +349,7 @@ def direct_topk_kernel( # ~1.03 cold-kernel gap at npad 8256. Keep the flat batch. n_batch = cutlass.const_expr((DKCMAX // 4 + num_threads - 1) // num_threads) frags = [ - cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(n_batch) + cute.make_fragment((4,), cutlass.Float32) for _ in range(n_batch) ] # Python-unrolled register batch for u in cutlass.range_constexpr(n_batch): i_vec = tidx + cutlass.Int32(u * num_threads) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_reg.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_reg.py index 4ee9beb8e745..472edfb0ef3d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_reg.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_reg.py @@ -162,8 +162,8 @@ def __init__( def count_reg(self, R: cutlass.Constexpr, a, tidx, s_rungs, s_ptcnt): TB = cutlass.const_expr(self.num_threads) MAXV = cutlass.const_expr(self.maxv) - tr = cute.make_rmem_tensor((R,), cutlass.Float32) - cnt = cute.make_rmem_tensor((R,), cutlass.Int32) + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) for r in cutlass.range_constexpr(R): tr[r] = s_rungs[r] cnt[r] = cutlass.Int32(0) @@ -391,8 +391,8 @@ def gvr_reg_kernel( v1 = V4 copy_atom = self._copy_atom() - a = cute.make_rmem_tensor((MAXV * 4,), cutlass.Float32) - frag4 = cute.make_rmem_tensor((4,), cutlass.Float32) + a = cute.make_fragment((MAXV * 4,), cutlass.Float32) + frag4 = cute.make_fragment((4,), cutlass.Float32) for u in cutlass.range_constexpr(MAXV): i = v0 + tidx + cutlass.Int32(u * TB) if i < v1: @@ -728,7 +728,7 @@ def gvr_reg_kernel( cute.arch.barrier() if tidx < cutlass.Int32(32): b8 = tidx * cutlass.Int32(8) - h = cute.make_rmem_tensor((8,), cutlass.Int32) + h = cute.make_fragment((8,), cutlass.Int32) Ssum = cutlass.Int32(0) for q in cutlass.range_constexpr(8): h[q] = s_hist[b8 + cutlass.Int32(q)] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_tp.py index 9bdb5e5a82cb..9c194a556546 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_tp.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_tp.py @@ -411,13 +411,13 @@ def count_pass( # Python-unrolled register batch. TB = cutlass.const_expr(self.num_threads) copy_atom = self._copy_atom() - tr = cute.make_rmem_tensor((R,), cutlass.Float32) - cnt = cute.make_rmem_tensor((R,), cutlass.Int32) + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) for r in cutlass.range_constexpr(R): tr[r] = s_rungs[r] cnt[r] = cutlass.Int32(0) frags = [ - cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(U) + cute.make_fragment((4,), cutlass.Float32) for _ in range(U) ] # Python-unrolled register batch # Mask hoist: float4s below vmain = min(v1, n_eff >> 2) are fully # valid (4i+3 < n_eff), so the main loop runs mask-free (the @@ -476,12 +476,12 @@ def count_pass( def sample_count(self, R: cutlass.Constexpr, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt): TB = cutlass.const_expr(self.num_threads) copy_atom = self._copy_atom() - tr = cute.make_rmem_tensor((R,), cutlass.Float32) - cnt = cute.make_rmem_tensor((R,), cutlass.Int32) + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) for r in cutlass.range_constexpr(R): tr[r] = s_rungs[r] cnt[r] = cutlass.Int32(0) - frag = cute.make_rmem_tensor((4,), cutlass.Float32) + frag = cute.make_fragment((4,), cutlass.Float32) j = cutlass.Int32(tidx) while v0 + j * cutlass.Int32(SS) < v1: self._ld_float4(copy_atom, row_addr, v0 + j * cutlass.Int32(SS), frag) @@ -559,7 +559,7 @@ def max_below_pass(self, row_addr, v0, v1, n_eff, t_hi_bound, par, tidx, s_fwred copy_atom = self._copy_atom() m = cutlass.Float32(-FLT_MAX) # Explicit U=4 batched loads (CUDA `float4 a[4]` idiom). - frags = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(4)] + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(4)] vmain = v1 vfull = n_eff >> cutlass.Int32(2) if vmain > vfull: @@ -635,8 +635,8 @@ def phase1( K = cutlass.const_expr(self.top_k) if tidx < cutlass.Int32(64): s_hist[tidx] = cutlass.Int32(0) - hv = cute.make_rmem_tensor((4,), cutlass.Float32) - hok = cute.make_rmem_tensor((4,), cutlass.Int32) + hv = cute.make_fragment((4,), cutlass.Float32) + hok = cute.make_fragment((4,), cutlass.Int32) mn = cutlass.Float32(FLT_MAX) mx = cutlass.Float32(-FLT_MAX) for jj in cutlass.range_constexpr(4): @@ -858,8 +858,8 @@ def fused_count_collect( kcap = cutlass.const_expr(self.kC) R = cutlass.const_expr(2) copy_atom = self._copy_atom() - tr = cute.make_rmem_tensor((R,), cutlass.Float32) - cnt = cute.make_rmem_tensor((R,), cutlass.Int32) + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) for r in cutlass.range_constexpr(R): tr[r] = s_rungs[r] cnt[r] = cutlass.Int32(0) @@ -872,7 +872,7 @@ def fused_count_collect( # vec-tail while-loops (same fix as count_pass). Mask # hoist as in count_pass: [v0, vmain) mask-free (gi computed only # inside the rare push branch), [vmain, v1) masked epilogue. - frags = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(U)] + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(U)] vmain = v1 vfull = n_eff >> cutlass.Int32(2) if vmain > vfull: @@ -950,7 +950,7 @@ def collect_at(self, row_addr, v0, v1, n_eff, thr, tidx, s_cand, s_isc): # CUDA collect_at loop 4x with 4 LDG.E.128 issued back-to-back; a # 1-deep loop was the dominant stall site (36% of all warp-stall # samples) on cells whose reuse check fails at big npad x big BS. - frags = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(4)] + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(4)] vmain = v1 vfull = n_eff >> cutlass.Int32(2) if vmain > vfull: @@ -1657,7 +1657,7 @@ def gvr_tp_kernel( cute.arch.barrier() nt = cutlass.Int32(K) - m_gt copy_atom = self._copy_atom() - frag = cute.make_rmem_tensor((4,), cutlass.Float32) + frag = cute.make_fragment((4,), cutlass.Float32) ii = v0 + tidx while ii < v1: self._ld_float4(copy_atom, row_addr, ii, frag) @@ -1732,7 +1732,7 @@ def gvr_tp_kernel( cute.arch.barrier() if tidx < cutlass.Int32(32): b8 = tidx * cutlass.Int32(8) - h = cute.make_rmem_tensor((8,), cutlass.Int32) + h = cute.make_fragment((8,), cutlass.Int32) Ssum = cutlass.Int32(0) for q in cutlass.range_constexpr(8): h[q] = s_hist[b8 + cutlass.Int32(q)] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk.py index e7b5262c7732..a277bf6439e5 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/single_pass_multi_cta_radix_topk.py @@ -306,7 +306,7 @@ def load_chunk_to_smem(self, input_row, shared_ordered, chunk_start, actual_chun self.dtype, num_bits_per_copy=cutlass.const_expr(self.num_copy_bits), ) - frag = cute.make_rmem_tensor((vec_size,), self.dtype) + frag = cute.make_fragment((vec_size,), self.dtype) stride = cutlass.const_expr(num_threads * vec_size) # Thread t covers [t*vs, t*vs+vs), [t*vs+stride, ...) in the aligned region. @@ -326,7 +326,7 @@ def load_chunk_to_smem(self, input_row, shared_ordered, chunk_start, actual_chun # Apply to_ordered into a correctly-typed uint32 fragment so # autovec_copy sees matching types (uint32 → uint32) and emits # STS.128 instead of 8 scalar STS instructions. - ordered_frag = cute.make_rmem_tensor((vec_size,), self.ordered_type) + ordered_frag = cute.make_fragment((vec_size,), self.ordered_type) for j in cutlass.range(vec_size, unroll_full=True): ordered_frag[j] = self.to_ordered(frag[j]) diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py index 352b44788d2c..b397ae552bd0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -138,8 +138,8 @@ def is_megamoe_cute_dsl_runtime_available() -> Tuple[bool, Optional[str]]: by ``kernel_fc12.py``, and the async-copy helpers used by ``dispatch_kernel.py``. PR https://github.com/NVIDIA/TensorRT-LLM/pull/14354 pins - ``nvidia-cutlass-dsl[cu13]==4.6.1``; version 4.5.0 was the first release - that shipped all of them, and older wheels return ``(False, reason)``. + ``nvidia-cutlass-dsl[cu13]==4.5.0`` which is the first release that + ships all of them; older wheels return ``(False, reason)``. Returns ``(True, None)`` on success or ``(False, reason)`` with an actionable message. The result is cached for the process lifetime. diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py index b65eb1e9a645..532e86e91e08 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py @@ -74,8 +74,7 @@ def _fwd( seqused_k: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Calls _flash_attn_fwd with torch.compile disabled. Returns (output, lse).""" - # FA4's private forward API may append diagnostics that this backend does not consume. - output, lse, *_ = _flash_attn_fwd( + output, lse = _flash_attn_fwd( q, k, v, diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/__init__.py index 1e882733b529..2ced031ae28c 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/__init__.py @@ -13,10 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .block_sparse_attn_dsl_fwd import ( + VideoSparseAttentionForwardGroup2QInterleaveKV as VideoSparseAttentionForward, +) from .interface import CUTE_AVAILABLE, block_sparse_attn_from_indices_cute, is_cute_supported __all__ = [ "CUTE_AVAILABLE", + "VideoSparseAttentionForward", "block_sparse_attn_from_indices_cute", "is_cute_supported", ] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/block_sparse_attn_dsl_fwd.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/block_sparse_attn_dsl_fwd.py index 10473119ed69..a2d814b2dbc8 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/block_sparse_attn_dsl_fwd.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/block_sparse_attn_dsl_fwd.py @@ -1818,7 +1818,7 @@ def softmax( num_q_blocks: cutlass.Int32, sP: cute.Tensor, tCtS: cute.Tensor, - thr_mma_qk: cute.ThrMma, + thr_mma_qk: cute.core.ThrMma, sm_scale_log2: cutlass.Float32, sScale: cute.Tensor, corr_pipeline: pipeline.PipelineAsync, @@ -1848,8 +1848,8 @@ def softmax( tCcS[(None, None), 0, None], (self.mma_tiler_qk[0], self.block_n) ) tCcS_ld = thr_copy_t2r.partition_D(cS_load) - tCrS_ld = cute.make_rmem_tensor(cute.select(tCcS_ld.shape, mode=[0, 1, 2]), self.acc_dtype) - tCrS_ld_half = cute.make_rmem_tensor(tCrS_ld.layout, sP.element_type) + tCrS_ld = cute.make_fragment(cute.select(tCcS_ld.shape, mode=[0, 1, 2]), self.acc_dtype) + tCrS_ld_half = cute.make_fragment(tCrS_ld.layout, sP.element_type) sP_cpy_slice = None tiled_copy_r2t = None @@ -2174,7 +2174,7 @@ def correction( st_O_pipeline: pipeline.PipelineAsync, st_O_producer_state: pipeline.PipelineState, tCtO: cute.Tensor, - thr_mma_pv: cute.ThrMma, + thr_mma_pv: cute.core.ThrMma, sO: cute.Tensor, corr_pipeline: pipeline.PipelineAsync, correction_consumer_state: pipeline.PipelineState, @@ -2225,7 +2225,7 @@ def correction( tCcO[(None, None), 0, None], (self.mma_tiler_pv[0], corr_ld_inst) ) corr_tCcO_ld = corr_thr_copy_t2r.partition_D(corr_cO_load) - corr_tCrO_ld = cute.make_rmem_tensor( + corr_tCrO_ld = cute.make_fragment( cute.select(corr_tCcO_ld.shape, mode=[0, 1, 2]), self.acc_dtype ) @@ -2252,10 +2252,10 @@ def correction( tCcO[(None, None), 0, None], (self.mma_tiler_pv[0], wb_ld_inst) ) wb_tCcO_ld = wb_thr_copy_t2r.partition_D(wb_cO_load) - wb_tCrO_ld = cute.make_rmem_tensor( + wb_tCrO_ld = cute.make_fragment( cute.select(wb_tCcO_ld.shape, mode=[0, 1, 2]), self.acc_dtype ) - wb_tCrO_ld_half = cute.make_rmem_tensor(wb_tCrO_ld.layout, sO.element_type) + wb_tCrO_ld_half = cute.make_fragment(wb_tCrO_ld.layout, sO.element_type) tv_layout = cute.make_ordered_layout( (self.threads_per_wg, self.mma_tiler_pv[1], self.s_stage), (0, 1, 2) ) @@ -2264,7 +2264,7 @@ def correction( wb_tCrO_reduction = None if cutlass.const_expr(self.o_stage != 1): - wb_tCrO_reduction = cute.make_rmem_tensor(wb_tCrO_ld.layout, self.acc_dtype) + wb_tCrO_reduction = cute.make_fragment(wb_tCrO_ld.layout, self.acc_dtype) _correction_loop = partial( self.correction_loop, diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py index 4bee4dad352a..093803572e59 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py @@ -28,24 +28,22 @@ import torch -_cute_import_error: Exception | None = None try: import cuda.bindings.driver as _cuda import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack from .block_sparse_attn_dsl_fwd import ( - VideoSparseAttentionForwardGroup2QInterleaveKV as _VideoSparseAttentionForward, + VideoSparseAttentionForwardGroup2QInterleaveKV as VideoSparseAttentionForward, ) CUTE_AVAILABLE = True -except (ImportError, OSError) as error: +except ImportError: # cuda-bindings / cutlass-dsl not installed _cuda = None cute = None from_dlpack = None - _VideoSparseAttentionForward = None + VideoSparseAttentionForward = None CUTE_AVAILABLE = False - _cute_import_error = error __all__ = [ @@ -108,14 +106,14 @@ def block_sparse_attn_from_indices_cute( if not CUTE_AVAILABLE: raise RuntimeError( "block_sparse_attn_from_indices_cute called but cuda.bindings or " - f"cutlass-dsl is not importable: {_cute_import_error}" - ) from _cute_import_error + "cutlass-dsl is not importable." + ) num_q_blk = variable_block_sizes.shape[0] - if num_q_blk > _VideoSparseAttentionForward.MAX_INDICES: + if num_q_blk > VideoSparseAttentionForward.MAX_INDICES: raise ValueError( f"variable_block_sizes has {num_q_blk} entries but the CuTe kernel " - f"supports at most {_VideoSparseAttentionForward.MAX_INDICES} " + f"supports at most {VideoSparseAttentionForward.MAX_INDICES} " "(SMEM-allocated sVariable_block_sizes buffer). Lower video " "resolution/length or fall back to dense SDPA." ) @@ -145,7 +143,7 @@ def block_sparse_attn_from_indices_cute( compile_key = (D, q.dtype, float(sm_scale)) + tuple(q2k_idx.shape) compiled = _COMPILE_CACHE.get(compile_key) if compiled is None: - fwd_kernel = _VideoSparseAttentionForward(block_m=64, block_n=64, headdim=D) + fwd_kernel = VideoSparseAttentionForward(block_m=64, block_n=64, headdim=D) compiled = cute.compile( fwd_kernel, q_packed, diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/ptx.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/ptx.py index 9a1cd7b7dc55..12fa15bc5930 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/ptx.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/ptx.py @@ -12,8 +12,12 @@ # 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. +from functools import partial +from typing import Tuple + import cutlass -from cutlass._mlir.dialects import llvm +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm, nvvm from cutlass.cute import typing as cutlass_typing from cutlass.cutlass_dsl import dsl_user_op @@ -251,3 +255,123 @@ def max3f( ip=ip, ) ) + + +sub_packed_f32x2 = partial( + cute.arch.calc_packed_f32x2_op, + src_c=None, + calc_func=nvvm.sub_packed_f32x2, + rnd=nvvm.RoundingModeKind.RN, +) + + +@dsl_user_op +@cute.jit +def evaluate_polynomial_2( + x: cutlass.Float32, + y: cutlass.Float32, + poly: Tuple[cutlass.Float32, ...], + *, + loc=None, + ip=None, +) -> Tuple[cutlass.Float32, cutlass.Float32]: + deg = len(poly) - 1 + out = (poly[deg], poly[deg]) + for i in cutlass.range_constexpr(deg - 1, -1, -1): + out = cute.arch.fma_packed_f32x2(out, (x, y), (poly[i], poly[i])) + return out + + +@dsl_user_op +def combine_int_frac_ex2( + x_rounded: cutlass.Float32, + frac_ex2: cutlass.Float32, + *, + loc=None, + ip=None, +) -> cutlass.Float32: + return cutlass.Float32( + llvm.inline_asm( + cutlass_typing.Float32.mlir_type, + [ + cutlass_typing.Float32(x_rounded).ir_value(loc=loc, ip=ip), + cutlass_typing.Float32(frac_ex2).ir_value(loc=loc, ip=ip), + ], + """{\n\t + .reg .s32 x_rounded_i, frac_ex_i, x_rounded_e, out_i; \n\t + mov.b32 x_rounded_i, $1; \n\t + mov.b32 frac_ex_i, $2; \n\t + shl.b32 x_rounded_e, x_rounded_i, 23; \n\t + add.s32 out_i, x_rounded_e, frac_ex_i; \n\t + mov.b32 $0, out_i; \n\t + \n\t}""", + "=f, f, f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def exp2_emulation_2( + x: cutlass.Float32, + y: cutlass.Float32, + *, + loc=None, + ip=None, +) -> Tuple[cutlass.Float32, cutlass.Float32]: + # assume x <= 127.0 and y <= 127.0 + poly_ex2_deg3 = ( + 1.0, + 0.695146143436431884765625, + 0.227564394474029541015625, + 0.077119089663028717041015625, + ) + fp32_round_int = float(2**23 + 2**22) + xy_clamped = (cute.arch.fmax(x, -127.0), cute.arch.fmax(y, -127.0)) + xy_rounded = cute.arch.add_packed_f32x2( + xy_clamped, (fp32_round_int, fp32_round_int), rnd=nvvm.RoundingModeKind.RM + ) + xy_rounded_back = sub_packed_f32x2(xy_rounded, (fp32_round_int, fp32_round_int)) + xy_frac = sub_packed_f32x2(xy_clamped, xy_rounded_back) + xy_frac_ex2 = evaluate_polynomial_2(*xy_frac, poly_ex2_deg3, loc=loc, ip=ip) + x_out = combine_int_frac_ex2(xy_rounded[0], xy_frac_ex2[0], loc=loc, ip=ip) + y_out = combine_int_frac_ex2(xy_rounded[1], xy_frac_ex2[1], loc=loc, ip=ip) + return x_out, y_out + + +@dsl_user_op +def exp2f_packed_f32x2( + x: cutlass.Float32, + y: cutlass.Float32, + *, + loc=None, + ip=None, +) -> Tuple[cutlass.Float32, cutlass.Float32]: + result = llvm.inline_asm( + llvm.StructType.get_literal( + [ + cutlass_typing.Float32.mlir_type, + cutlass_typing.Float32.mlir_type, + ] + ), + [ + cutlass_typing.Float32(x).ir_value(loc=loc, ip=ip), + cutlass_typing.Float32(y).ir_value(loc=loc, ip=ip), + ], + """{\n\t + ex2.approx.f32 $0, $2;\n\t + ex2.approx.f32 $1, $3;\n\t + \n\t}""", + # Keep constraints compact (no spaces) and in the correct order + "=f,=f,f,f", + loc=loc, + ip=ip, + ) + + # Extract struct fields + out0_val = llvm.extractvalue(cutlass_typing.Float32.mlir_type, result, [0], loc=loc, ip=ip) + out1_val = llvm.extractvalue(cutlass_typing.Float32.mlir_type, result, [1], loc=loc, ip=ip) + + # Wrap back into cutlass.Float32 + return cutlass.Float32(out0_val), cutlass.Float32(out1_val) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index e1acb231f60a..fbefac545b92 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -231,7 +231,6 @@ l0_b200: - unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py - 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_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/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py b/tests/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py index 978441f849ed..6ed163bfcc9c 100644 --- a/tests/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py +++ b/tests/scripts/cute_dsl_kernels/run_dense_bf16_gemm_persistent.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause # Redistribution and use in source and binary forms, with or without @@ -56,11 +56,9 @@ from tensorrt_llm._torch.cute_dsl_kernels.blackwell import ( dense_gemm_persistent as kernel_module, ) - from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr except (ModuleNotFoundError, ImportError): sys.path.insert(0, str(Path(__file__).parents[3] / "tensorrt_llm/_torch/cute_dsl_kernels")) from blackwell import dense_gemm_persistent as kernel_module - from blackwell.utils import make_ptr PersistentDenseGemmKernel = kernel_module.PersistentDenseGemmKernel @@ -218,19 +216,31 @@ def run( b_ref = cutlass_torch.matrix(batch, n, k, b_major == "n", cutlass.Float32) c_ref = cutlass_torch.matrix(batch, m, n, c_major == "m", cutlass.Float32) - _, a_torch = cutlass_torch.cute_tensor_like( + a_tensor, a_torch = cutlass_torch.cute_tensor_like( a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - _, b_torch = cutlass_torch.cute_tensor_like( + b_tensor, b_torch = cutlass_torch.cute_tensor_like( b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - _, c_torch = cutlass_torch.cute_tensor_like( + c_tensor, c_torch = cutlass_torch.cute_tensor_like( c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 ) - 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 + + # 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, ) # Configure gemm kernel @@ -267,8 +277,8 @@ def run( n, k, batch, - a_ptr, - b_ptr, + a_tensor, + b_tensor, c_tensor, a_stride_m, a_stride_batch, @@ -288,6 +298,7 @@ def run_kernel(a_t, b_t, c_t, stream): c_t, a_stride_m, a_stride_batch, + max_active_clusters, stream, ) else: @@ -298,8 +309,8 @@ def run_kernel(a_t, b_t, c_t, stream): n, k, batch, - a_ptr, - b_ptr, + a_tensor, + b_tensor, c_tensor, max_active_clusters, current_stream, @@ -315,13 +326,14 @@ def run_kernel(a_t, b_t, c_t, stream): a_t, b_t, c_t, + max_active_clusters, stream, ) # Compute reference result if not skip_ref_check: # Execute kernel once for reference checking - run_kernel(a_ptr, b_ptr, c_tensor, current_stream) + run_kernel(a_tensor, b_tensor, c_tensor, current_stream) print("Verifying results...") # Reference: C = einsum("mkl,nkl->mnl", A, B) @@ -339,86 +351,82 @@ def run_kernel(a_t, b_t, c_t, stream): torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) - # make_ptr stores only raw addresses, so retain each generated tensor until - # the benchmark has finished using its JIT arguments. - backing_tensors = [] - if use_strided: - def generate_tensors() -> cute.testing.JitArguments: - _, a_torch_new = cutlass_torch.cute_tensor_like( + def generate_tensors(): + a_tensor_new, _ = cutlass_torch.cute_tensor_like( a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - _, b_torch_new = cutlass_torch.cute_tensor_like( + b_tensor_new, _ = cutlass_torch.cute_tensor_like( b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - _, c_torch_new = cutlass_torch.cute_tensor_like( + c_tensor_new, _ = cutlass_torch.cute_tensor_like( c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 ) - a_ptr_new = make_ptr( - ab_dtype, - a_torch_new.data_ptr(), - cute.AddressSpace.gmem, - assumed_align=16, + a_tensor_new.mark_compact_shape_dynamic( + mode=1, + stride_order=(2, 0, 1), + divisibility=1, ) - b_ptr_new = make_ptr( - ab_dtype, - b_torch_new.data_ptr(), - cute.AddressSpace.gmem, - assumed_align=16, + b_tensor_new.mark_compact_shape_dynamic( + mode=1, + stride_order=(2, 0, 1), + divisibility=1, ) - c_tensor_new = from_dlpack(c_torch_new, assumed_align=16).mark_layout_dynamic( - leading_dim=1 if c_major == "n" else 0 + c_tensor_new.mark_compact_shape_dynamic( + mode=1, + stride_order=(2, 0, 1), + divisibility=1, ) - backing_tensors.append((a_torch_new, b_torch_new, c_torch_new)) return cute.testing.JitArguments( m, n, k, batch, - a_ptr_new, - b_ptr_new, + a_tensor_new, + b_tensor_new, c_tensor_new, a_stride_m, a_stride_batch, + max_active_clusters, current_stream, ) else: - def generate_tensors() -> cute.testing.JitArguments: - _, a_torch_new = cutlass_torch.cute_tensor_like( + def generate_tensors(): + a_tensor_new, _ = cutlass_torch.cute_tensor_like( a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - _, b_torch_new = cutlass_torch.cute_tensor_like( + b_tensor_new, _ = cutlass_torch.cute_tensor_like( b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - _, c_torch_new = cutlass_torch.cute_tensor_like( + c_tensor_new, _ = cutlass_torch.cute_tensor_like( c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 ) - a_ptr_new = make_ptr( - ab_dtype, - a_torch_new.data_ptr(), - cute.AddressSpace.gmem, - assumed_align=16, + a_tensor_new.mark_compact_shape_dynamic( + mode=1, + stride_order=(2, 0, 1), + divisibility=1, ) - b_ptr_new = make_ptr( - ab_dtype, - b_torch_new.data_ptr(), - cute.AddressSpace.gmem, - assumed_align=16, + b_tensor_new.mark_compact_shape_dynamic( + mode=1, + stride_order=(2, 0, 1), + divisibility=1, ) - c_tensor_new = from_dlpack(c_torch_new, assumed_align=16).mark_layout_dynamic( - leading_dim=1 if c_major == "n" else 0 + c_tensor_new.mark_compact_shape_dynamic( + mode=1, + stride_order=(2, 0, 1), + divisibility=1, ) - backing_tensors.append((a_torch_new, b_torch_new, c_torch_new)) return cute.testing.JitArguments( m, n, k, batch, - a_ptr_new, - b_ptr_new, + a_tensor_new, + b_tensor_new, c_tensor_new, + max_active_clusters, current_stream, ) diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index 9fab7b82ef79..9f1a9366c870 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -7,8 +7,6 @@ the Triton reference is covered by the SM100 integration accuracy test. """ -import sys -from types import ModuleType, SimpleNamespace from unittest.mock import Mock import pytest @@ -20,30 +18,6 @@ from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig -def test_msa_package_availability_installs_cutlass_46_compatibility_aliases(monkeypatch): - from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( - msa_package_available, - ) - - cute = ModuleType("cutlass.cute") - cute.core = SimpleNamespace() - cute.ThrMma = object() - cute.make_rmem_tensor = object() - cutlass = ModuleType("cutlass") - cutlass.cute = cute - monkeypatch.setitem(sys.modules, "cutlass", cutlass) - monkeypatch.setitem(sys.modules, "cutlass.cute", cute) - monkeypatch.setattr("importlib.util.find_spec", lambda unused_name: object()) - - msa_package_available.cache_clear() - try: - assert msa_package_available() - assert cute.core.ThrMma is cute.ThrMma - assert cute.make_fragment is cute.make_rmem_tensor - finally: - msa_package_available.cache_clear() - - def test_resolver_selects_msa_backend_when_available(monkeypatch): import tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_availability as avail