diff --git a/.claude/skills/kernel-cute-writing/references/api-core.md b/.claude/skills/kernel-cute-writing/references/api-core.md index 9b640b236f1b..4945971a6a88 100644 --- a/.claude/skills/kernel-cute-writing/references/api-core.md +++ b/.claude/skills/kernel-cute-writing/references/api-core.md @@ -126,8 +126,7 @@ 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) # Register memory tensor -cute.make_fragment(shape, dtype) # Register fragment +cute.make_rmem_tensor(layout_or_shape, dtype) # Register memory tensor ``` ### 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 5b7c9c4d67d7..09e3e775bdef 100644 --- a/.claude/skills/kernel-cute-writing/references/concepts-tensors.md +++ b/.claude/skills/kernel-cute-writing/references/concepts-tensors.md @@ -81,8 +81,7 @@ mA = cute.make_tensor(ptr, layout=layout) ### Register/Fragment Tensors (Owning) ```python -rmem = cute.make_rmem_tensor(layout) -frag = cute.make_fragment(shape, dtype) +rmem = cute.make_rmem_tensor(layout_or_shape, dtype) ``` ## Accessing Tensors diff --git a/constraints.txt b/constraints.txt index 7de1c603d547..a99d65019d38 100644 --- a/constraints.txt +++ b/constraints.txt @@ -1,10 +1,18 @@ -# 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 +# 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 # 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 14798a5a354a..7164715ee7d3 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -66,11 +66,13 @@ 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 numpy wandb || true && \ + 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 && \ # 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 && \ + pip3 install -r /tmp/constraints.txt "nvidia-cutlass-dsl[cu13]==4.6.1" && \ rm /tmp/constraints.txt # Install UCX, NIXL, etcd diff --git a/docker/common/install_fa4.sh b/docker/common/install_fa4.sh index 5138b19b2925..e70083220dc7 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.0b11" +FLASH_ATTN_4_VERSION="4.0.0b19" 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 dca8d522fb44..6ff9c295211e 100644 --- a/docs/source/installation/installation-guide.md +++ b/docs/source/installation/installation-guide.md @@ -68,7 +68,8 @@ 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-cu13 +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 ``` ```bash diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 3e8e4bf8391d..c4ad88465263 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -6325,7 +6325,11 @@ 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 || true") + 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: '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 c5ee0bc21191..a6246b3160d7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -76,17 +76,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.5.0; python_version >= "3.10" +nvidia-cutlass-dsl[cu13]==4.6.1; 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.2.10; python_version >= "3.10" # required for MinimaxM3 MSA +quack-kernels==0.6.2; python_version >= "3.10" # required for MinimaxM3 MSA and FlashAttention 4; pins CUTLASS DSL 4.6.1 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.6 # used for reduce nvidia-cutlass-dsl host overhead +apache-tvm-ffi==0.1.13.post2 # required by CUTLASS DSL 4.6.1's 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.0b11 +flash-attn-4==4.0.0b19 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 ce1cb998cfa3..dae7229dee26 100644 --- a/security_scanning/poetry.lock +++ b/security_scanning/poetry.lock @@ -274,41 +274,37 @@ trio = ["trio (>=0.32.0)"] [[package]] name = "apache-tvm-ffi" -version = "0.1.13.post3" +version = "0.1.13.post2" description = "tvm ffi" optional = false python-versions = ">=3.9" groups = ["main"] files = [ - {file = "apache_tvm_ffi-0.1.13.post3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:06431d9139c20e9e5654b8929a34d92ea63a6fb4ada1bee6612f2bafb77e7054"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5f63a19af2728b974a0b808fc71f4376896301969b2510258c991fa1f0c30d19"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fbbab8aa62ba0afe996197267c3d7ff999c0747ba479b636589e52baf754f79e"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp310-cp310-win_amd64.whl", hash = "sha256:42c4d438f19f7d946ff04e1626066c9d04e7625dfd2761a2df0145bc6ecf3c38"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f571b636af7ce906200dbd31402762eceff099013de17963a4a2d2ae9b05e400"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a624852616364a67b16ff90509a1025128aee29e84a15279956b8f7c524ac837"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b2c9be6c2422423d11e0e9ba2883f8fb35ae5d568d329b3cf5df5fee2b64d19f"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp311-cp311-win_amd64.whl", hash = "sha256:5d48a04a3ebc8463b0c34d8c0eba7eef23b188171fd00fef12f09d12f89d59e1"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e4853ab8283ea15a094887ad561115454a47ae2d9ac776c0570f3771b9a51e65"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c108c26935a609b636bb1a53df6d7c53ab776550ad04f17d6dcb35e883112f78"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:edaca70656f87831c3b4989169b1d85fe563ccd6d2e2cadc1f4f5880e1295f7a"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp312-cp312-win_amd64.whl", hash = "sha256:295900a3bf111f133be5dff9f4d92d943f01e5e4ce3da6234f02d0e51a8cf14d"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bffbcc229f25668eceab522706acf7085d665c246e92ef8f5e4508d4d07feffd"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7382278abb5622f271d50e97ad7ffbb2970fd0cb836be4f7095fbcb116941582"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:db3b8c3d18dae01c6962d7052b335ff77e97b309263c49ec92aea25836fa4b97"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp313-cp313-win_amd64.whl", hash = "sha256:701b85a1c14abed1cc92da8e7b76d0e4ff13802e1363b4f5111813ec1da04c12"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:8047715bce35b2a8a75d278d9a48aeec58f461f99f5a8fa57bd32b9589f23cc6"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d07bd9d7c1b78f3ffd321afef400e9959cf3254db3be5fcb9a2f736180903ba8"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:32c7b40994ecd7b5670a42ff723f31e018ffe7e0cea2e67f2eaed6a86615b9dd"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp314-cp314-win_amd64.whl", hash = "sha256:33a6b9f308870dd9b6405111319bbe22b9ba774037c5867e110b03409cc317df"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a6fee5dcfa6fd7a4d62311b41ccb5946e3ef9d9939ead875d2e192da8763f819"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:341799f61c9e37f63ca283e39ddf5a952d09b1c7e22a60f32510666cd20677f5"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c2d629b592e352a46a5ebffbe43c9e1ce509cf10940f1c750d70d29aa7452888"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp314-cp314t-win_amd64.whl", hash = "sha256:7727fd0370b634b4d69c255e11aa0361bc120020940012a084d0c825fa705784"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:c55700790af0db289ad4eff32a9cb22d30787074a629eb8942f45aa832076a89"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp39-cp39-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d56b55ef204351238374b00a5538e0178082928a7696a6fd82afd6ccd034a3b4"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp39-cp39-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ad3dd40cb94d225040ae7534717927b91a1dfe3d12a99e8ee983f22f896e17d9"}, - {file = "apache_tvm_ffi-0.1.13.post3-cp39-cp39-win_amd64.whl", hash = "sha256:753ddb3d11390126c0a049d328e239664979268a41843542e21ae4efd24090c8"}, - {file = "apache_tvm_ffi-0.1.13.post3.tar.gz", hash = "sha256:b8153f020418df39cd93c3a6beda9d8c9b6e49b4e1b96de26ea6889c32347a00"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:f5ceb6fb6122078281ac68c8d59fdcab9575b1cbc0773fd03af31957a46ddc25"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d6f77e8c601319a0d43410939212c46e338628b74f133ac298f0e807a13d5021"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:36e27c8c8b38ac061638896cc593e2328a2c04441c28ec51aa5cdf451279d5fd"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp310-cp310-win_amd64.whl", hash = "sha256:1b604ca681a1381f96639cff53d94861e9e4949b52e6f3b3f3252aa28c5b0a60"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:ba333bf376883f8feb6ae1adca24ab65b2b8cd68c11be23085510f4707c1471e"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6146e0adeef8e10679c3cde3efb2062c4fd00601bbd8b4e1c20c63d73e505b6e"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:38dfd42b3af5e23baf95aa5054eca3e23f183b09fa3394b56a5597317fa1fec4"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp311-cp311-win_amd64.whl", hash = "sha256:e581c54e29e96b79facbacd6b3bd3a992f6056b43c95a0bfd3551326db09f9ab"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:db4486a73d8131b12c3ac7160d8f5aa7be5408a84f64425ab603078e1bd7b8d0"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d7f350f14a946ee572b1087f8916a57a20d0c3105f94aafdd600b5f61cdd2905"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4a18d1cb67585952bc9502da6ee428f4c75934c8d3d7c0d03cdd59b82202f2e9"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp312-cp312-win_amd64.whl", hash = "sha256:0e77c79fe84f6221e58e5b820fffd7dc711c5636ff84b828f8977d821a0ef9d6"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:552ff0d059ae20b969d2c96a28e6a0e3e810123cb41f6d81673fc3cdb12c9a40"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c00d083b55d7fcb5b50683ff5c464fe224d1611d93cd5a6218251caa8c605f23"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0284905d97aa5f92a42de2ad61faa7e3c19679d1a0e47caaa3abbe368b0ba303"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp313-cp313-win_amd64.whl", hash = "sha256:6e0df6627ca2652abd175b6da1e4b60877684e67295419287a04ff0ddfb95372"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:0873badc25ec37d0d5b9ad00d95d63476cf2ef95069ec3406ef8c39552aaf3b2"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:057c431d56bfd172ed6795247cb25c144d0e49d8e4b03840055d85e7b662b0e9"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2587cf82d54d9e376ae9368b1ef3403f8e23c8270fa8819d46358ff5923754f1"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp314-cp314t-win_amd64.whl", hash = "sha256:93bdb9ca4e85c88b4a5c9a944d353bd7aaec49d7d2779b110137f2a4543ae033"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:ee176022d4fdbd2acc7a42c3ed604783f79999178967f9061887783bd2057a01"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp39-cp39-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:41db1e7b7db0f0d74728e58a867f29fe747cc6eb7b4d3a28ea2cc70b6afd8582"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp39-cp39-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f6cf7875dab8d77f9e4d46cbf446c0d0e2f35fd8a634ab35fbdc50254d279237"}, + {file = "apache_tvm_ffi-0.1.13.post2-cp39-cp39-win_amd64.whl", hash = "sha256:619a20d7cbe31b0c14f0d474ff5f54d5c1349c3ccfb4133fadc96d6f545d89bb"}, + {file = "apache_tvm_ffi-0.1.13.post2.tar.gz", hash = "sha256:7da4ede2392527c122afd36ac717d9a267627901881a38f5d66d214ff68f482e"}, ] [package.dependencies] @@ -1381,27 +1377,27 @@ files = [ [[package]] name = "flash-attn-4" -version = "4.0.0b11" +version = "4.0.0b19" description = "Flash Attention CUTE (CUDA Template Engine) implementation" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {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"}, + {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"}, ] [package.dependencies] apache-tvm-ffi = ">=0.1.5,<0.2" einops = "*" -nvidia-cutlass-dsl = ">=4.4.2" -quack-kernels = ">=0.4.0" +nvidia-cutlass-dsl = ">=4.5.2" +quack-kernels = ">=0.5.0" torch = "*" torch-c-dlpack-ext = "*" typing_extensions = "*" [package.extras] -cu13 = ["nvidia-cutlass-dsl[cu13] (>=4.4.2)"] +cu13 = ["nvidia-cutlass-dsl[cu13] (>=4.5.2)"] dev = ["pytest", "pytest-xdist", "ruff"] [[package]] @@ -3210,6 +3206,19 @@ 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" @@ -3401,72 +3410,130 @@ files = [ [[package]] name = "nvidia-cutlass-dsl" -version = "4.5.0" +version = "4.6.1" description = "NVIDIA CUTLASS Python DSL" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "nvidia_cutlass_dsl-4.5.0-py3-none-any.whl", hash = "sha256:3b051fe02ca69422ab840e64d9865667aba288a3984a7ca4ccd038a82aef1344"}, + {file = "nvidia_cutlass_dsl-4.6.1-py3-none-any.whl", hash = "sha256:93135a9d48e1bedf584828e0a021f174ac31591ef21f6ea53c206169ccbfab26"}, ] [package.dependencies] -nvidia-cutlass-dsl-libs-base = "4.5.0" -nvidia-cutlass-dsl-libs-cu13 = {version = "4.5.0", optional = true, markers = "extra == \"cu13\""} +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\""} [package.extras] -cu13 = ["nvidia-cutlass-dsl-libs-cu13 (==4.5.0)"] +cu13 = ["nvidia-cutlass-dsl-libs-cu13 (==4.6.1)"] [[package]] name = "nvidia-cutlass-dsl-libs-base" -version = "4.5.0" +version = "4.6.1" description = "NVIDIA CUTLASS Python DSL" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {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"}, + {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"}, ] [package.dependencies] cuda-python = ">=12.8" numpy = "*" -typing-extensions = "*" +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" [[package]] name = "nvidia-cutlass-dsl-libs-cu13" -version = "4.5.0" +version = "4.6.1" description = "NVIDIA CUTLASS Python DSL" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {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"}, + {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"}, ] [package.dependencies] cuda-python = ">=12.8" numpy = "*" -nvidia-cutlass-dsl-libs-base = "4.5.0" -typing-extensions = "*" +nvidia-cuda-nvdisasm = ">=13.3,<14" +nvidia-cutlass-dsl-libs-base = "4.6.1" +protobuf = ">=6.30.2,<7" +typing-extensions = ">=4.10.0" [[package]] name = "nvidia-matmul-heuristics" @@ -4286,20 +4353,22 @@ files = [ [[package]] name = "protobuf" -version = "7.35.1" +version = "6.33.6" description = "" optional = false -python-versions = ">=3.10" +python-versions = ">=3.9" groups = ["main"] files = [ - {file = "protobuf-7.35.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:24f857477359a85c0c235261b8ba905fd51b2562f4a64ca1df5473f29850cbf6"}, - {file = "protobuf-7.35.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:11d6b0ec246892d85215b0a13ca6e0233cf5284b68f0ac02646427f4ff88a799"}, - {file = "protobuf-7.35.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:b73f9489a4b8b1c9cb1f8ed951c736392592edb24b9d6819f36d2e10b171d5b4"}, - {file = "protobuf-7.35.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:74758715c53d7158fb76caf4f0cfdacc5329a4b1bb994f865d6cf302d413a1c4"}, - {file = "protobuf-7.35.1-cp310-abi3-win32.whl", hash = "sha256:353652e4efd0bca5b5fc2656abf8307ef351f0cf938c9eba09f0e09c20a25c30"}, - {file = "protobuf-7.35.1-cp310-abi3-win_amd64.whl", hash = "sha256:230a75ddfc2de4806e56696ce9640c1cdfdb6543b7cfce98d42a4c0a0e7bdb87"}, - {file = "protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9"}, - {file = "protobuf-7.35.1.tar.gz", hash = "sha256:ce115a26fe0c39a2c29973d914d327e516a6455464489fe3cd1e51a1b354f81a"}, + {file = "protobuf-6.33.6-cp310-abi3-win32.whl", hash = "sha256:7d29d9b65f8afef196f8334e80d6bc1d5d4adedb449971fefd3723824e6e77d3"}, + {file = "protobuf-6.33.6-cp310-abi3-win_amd64.whl", hash = "sha256:0cd27b587afca21b7cfa59a74dcbd48a50f0a6400cfb59391340ad729d91d326"}, + {file = "protobuf-6.33.6-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:9720e6961b251bde64edfdab7d500725a2af5280f3f4c87e57c0208376aa8c3a"}, + {file = "protobuf-6.33.6-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:e2afbae9b8e1825e3529f88d514754e094278bb95eadc0e199751cdd9a2e82a2"}, + {file = "protobuf-6.33.6-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:c96c37eec15086b79762ed265d59ab204dabc53056e3443e702d2681f4b39ce3"}, + {file = "protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:e9db7e292e0ab79dd108d7f1a94fe31601ce1ee3f7b79e0692043423020b0593"}, + {file = "protobuf-6.33.6-cp39-cp39-win32.whl", hash = "sha256:bd56799fb262994b2c2faa1799693c95cc2e22c62f56fb43af311cae45d26f0e"}, + {file = "protobuf-6.33.6-cp39-cp39-win_amd64.whl", hash = "sha256:f443a394af5ed23672bc6c486be138628fbe5c651ccbc536873d7da23d1868cf"}, + {file = "protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901"}, + {file = "protobuf-6.33.6.tar.gz", hash = "sha256:a6768d25248312c297558af96a9f9c929e8c4cee0659cb07e780731095f38135"}, ] [[package]] @@ -5041,27 +5110,29 @@ cffi = {version = "*", markers = "implementation_name == \"pypy\""} [[package]] name = "quack-kernels" -version = "0.4.1" +version = "0.6.2" description = "" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "quack_kernels-0.4.1-py3-none-any.whl", hash = "sha256:c1c8df2935bf5156ec47d2c5384ac08b411fd0ee702d80ae916dbf6d6f5ae813"}, - {file = "quack_kernels-0.4.1.tar.gz", hash = "sha256:9d7d6ba412bc0c8a9b1331c52a73db76280adb9dc2f2750df4851ddabef1466b"}, + {file = "quack_kernels-0.6.2-py3-none-any.whl", hash = "sha256:c193d5df51a56f5fdb681bee03f4042ed812f8190d022695cba9c14eb0e63cd6"}, + {file = "quack_kernels-0.6.2.tar.gz", hash = "sha256:83684ceb3ee4f5096035eb6f6e005e63f3f838bc9a44cf3c3e580f24c2b7233d"}, ] [package.dependencies] apache-tvm-ffi = ">=0.1.6,<0.2" einops = "*" -nvidia-cutlass-dsl = ">=4.4.2" +nvidia-cutlass-dsl = "4.6.1" torch = "*" torch-c-dlpack-ext = "*" [package.extras] -cu13 = ["nvidia-cutlass-dsl[cu13] (>=4.4.2)"] +bench = ["pandas", "tyro"] +cu13 = ["nvidia-cutlass-dsl[cu13] (==4.6.1)"] dev = ["pre-commit", "pytest", "pytest-xdist", "ruff"] heuristics = ["nvidia-matmul-heuristics"] +jax = ["jax", "jax-tvm-ffi"] [[package]] name = "referencing" @@ -6807,4 +6878,4 @@ type = ["pytest-mypy (>=1.0.1) ; platform_python_implementation != \"PyPy\""] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.13" -content-hash = "6332abd73e7f6be24069c8fe229e33098f950c05571fc1160fc87444bd642058" +content-hash = "beca1eccc889c464e4e997ccf558683ff259d61742d3cd635541e43c30810f81" diff --git a/security_scanning/pyproject.toml b/security_scanning/pyproject.toml index 7cb023d88f09..2cb0b15e8581 100644 --- a/security_scanning/pyproject.toml +++ b/security_scanning/pyproject.toml @@ -70,15 +70,16 @@ dependencies = [ "tiktoken (>=0.13.0,<0.14.0)", "blobfile (>=3.2.0,<4.0.0)", "openai-harmony (==0.0.4)", - "nvidia-cutlass-dsl[cu13] (==4.5.0)", + "nvidia-cutlass-dsl[cu13] (==4.6.1)", "nvidia-matmul-heuristics (==0.1.0.27)", - "quack-kernels (>=0.2.10)", + "quack-kernels (==0.6.2)", "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.0b11)", + "apache-tvm-ffi (==0.1.13.post2)", + "flash-attn-4 (==4.0.0b19)", "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 a5c3b8dc00d4..ba4cca087517 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,15 +20,34 @@ 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: - """Return whether the packaged fmha_sm100 module can be imported. + """Prepare and report 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. """ - return importlib.util.find_spec("fmha_sm100") is not None + if importlib.util.find_spec("fmha_sm100") is None: + return False + _install_msa_cutlass_compatibility() + return True def require_msa_module() -> types.ModuleType: @@ -38,6 +57,7 @@ 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 d76e52ec6297..0a019dec74b2 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 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 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_fragment( + tApA = cute.make_rmem_tensor( 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 0194cbc8359e..06ccffdad867 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_fragment(tTR_rAcc.shape, self.c_dtype) + tTR_rC = cute.make_rmem_tensor(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_fragment( + tTR_rAcc = cute.make_rmem_tensor( 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 fcd137a38ffd..867bf02214c5 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_fragment( + tTR_rAcc = cute.make_rmem_tensor( 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_fragment(tTR_rAcc.shape, self.c_dtype) + tTR_rC = cute.make_rmem_tensor(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 e769566218fe..dee10d230030 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,6 +192,7 @@ 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 94db5d143085..1516eafc22ec 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_fragment(4, cutlass.Int32) + values = cute.make_rmem_tensor(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_fragment(NUM_W_IN_REG * next_n, self.epi_dtype) + w_cache = cute.make_rmem_tensor(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_fragment(next_n, self.output_dtype) + result_arr = cute.make_rmem_tensor(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_fragment(NUM_W_IN_REG * next_n, self.epi_dtype) + w_cache = cute.make_rmem_tensor(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_fragment(next_n, self.output_dtype) + result_arr = cute.make_rmem_tensor(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 4367f06f442a..722ac73cebd0 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_fragment(NUM_W_IN_REG * next_n, self.epi_dtype) + w_cache = cute.make_rmem_tensor(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_fragment(NUM_W_IN_REG * next_n, self.epi_dtype) + w_cache = cute.make_rmem_tensor(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 159d44cb3445..5ce3ab29cdcc 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_fragment((vecsize_out, nvec_per_thread), self.dtype) - topk_indices = cute.make_fragment((vecsize_out, nvec_per_thread), cutlass.Int32) + 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) 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_fragment((vec_size,), self.dtype) + scan_frag = cute.make_rmem_tensor((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 44275ab832c3..b3725ef4e6c9 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_fragment((vec_w,), self.dtype) + rng_frag = cute.make_rmem_tensor((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_fragment((vec_w,), self.dtype) + tail_frag = cute.make_rmem_tensor((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_fragment((SEG,), cutlass.Int32) + seg_frag = cute.make_rmem_tensor((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_fragment((SEG,), cutlass.Int32) + seg_frag = cute.make_rmem_tensor((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_fragment((vec_w,), self.dtype) + rng_frag = cute.make_rmem_tensor((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_fragment((vec_w,), self.dtype) + tail_frag = cute.make_rmem_tensor((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_fragment((M,), cutlass.Float32) - cnt_frag = cute.make_fragment((M,), cutlass.Int32) + thr_frag = cute.make_rmem_tensor((M,), cutlass.Float32) + cnt_frag = cute.make_rmem_tensor((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_fragment((vec_w,), self.dtype) + rng_frag = cute.make_rmem_tensor((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_fragment((vec_w,), self.dtype) + tail_frag = cute.make_rmem_tensor((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_fragment((vec_w,), self.dtype) + rng_frag = cute.make_rmem_tensor((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_fragment((vec_w,), self.dtype) + tail_frag = cute.make_rmem_tensor((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_fragment((ppl,), cutlass.Int32) + cnt_frag = cute.make_rmem_tensor((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 5dbe90ced305..66165a286729 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_fragment((per,), cutlass.Int32) + h = cute.make_rmem_tensor((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_fragment((4,), cutlass.Float32) for _ in range(n_batch) + cute.make_rmem_tensor((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 472edfb0ef3d..4ee9beb8e745 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_fragment((R,), cutlass.Float32) - cnt = cute.make_fragment((R,), cutlass.Int32) + tr = cute.make_rmem_tensor((R,), cutlass.Float32) + cnt = cute.make_rmem_tensor((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_fragment((MAXV * 4,), cutlass.Float32) - frag4 = cute.make_fragment((4,), cutlass.Float32) + a = cute.make_rmem_tensor((MAXV * 4,), cutlass.Float32) + frag4 = cute.make_rmem_tensor((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_fragment((8,), cutlass.Int32) + h = cute.make_rmem_tensor((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 9c194a556546..9bdb5e5a82cb 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_fragment((R,), cutlass.Float32) - cnt = cute.make_fragment((R,), cutlass.Int32) + tr = cute.make_rmem_tensor((R,), cutlass.Float32) + cnt = cute.make_rmem_tensor((R,), cutlass.Int32) for r in cutlass.range_constexpr(R): tr[r] = s_rungs[r] cnt[r] = cutlass.Int32(0) frags = [ - cute.make_fragment((4,), cutlass.Float32) for _ in range(U) + cute.make_rmem_tensor((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_fragment((R,), cutlass.Float32) - cnt = cute.make_fragment((R,), cutlass.Int32) + tr = cute.make_rmem_tensor((R,), cutlass.Float32) + cnt = cute.make_rmem_tensor((R,), cutlass.Int32) for r in cutlass.range_constexpr(R): tr[r] = s_rungs[r] cnt[r] = cutlass.Int32(0) - frag = cute.make_fragment((4,), cutlass.Float32) + frag = cute.make_rmem_tensor((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_fragment((4,), cutlass.Float32) for _ in range(4)] + frags = [cute.make_rmem_tensor((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_fragment((4,), cutlass.Float32) - hok = cute.make_fragment((4,), cutlass.Int32) + hv = cute.make_rmem_tensor((4,), cutlass.Float32) + hok = cute.make_rmem_tensor((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_fragment((R,), cutlass.Float32) - cnt = cute.make_fragment((R,), cutlass.Int32) + tr = cute.make_rmem_tensor((R,), cutlass.Float32) + cnt = cute.make_rmem_tensor((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_fragment((4,), cutlass.Float32) for _ in range(U)] + frags = [cute.make_rmem_tensor((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_fragment((4,), cutlass.Float32) for _ in range(4)] + frags = [cute.make_rmem_tensor((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_fragment((4,), cutlass.Float32) + frag = cute.make_rmem_tensor((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_fragment((8,), cutlass.Int32) + h = cute.make_rmem_tensor((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 a277bf6439e5..e7b5262c7732 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_fragment((vec_size,), self.dtype) + frag = cute.make_rmem_tensor((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_fragment((vec_size,), self.ordered_type) + ordered_frag = cute.make_rmem_tensor((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 c6f4784b8110..6a65635867e9 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.5.0`` which is the first release that - ships all of them; older wheels return ``(False, reason)``. + ``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)``. 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 532e86e91e08..b65eb1e9a645 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py @@ -74,7 +74,8 @@ def _fwd( seqused_k: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Calls _flash_attn_fwd with torch.compile disabled. Returns (output, lse).""" - output, lse = _flash_attn_fwd( + # FA4's private forward API may append diagnostics that this backend does not consume. + 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 2ced031ae28c..1e882733b529 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,14 +13,10 @@ # 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 a2d814b2dbc8..10473119ed69 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.core.ThrMma, + thr_mma_qk: cute.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_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) + 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) 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.core.ThrMma, + thr_mma_pv: cute.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_fragment( + corr_tCrO_ld = cute.make_rmem_tensor( 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_fragment( + wb_tCrO_ld = cute.make_rmem_tensor( cute.select(wb_tCcO_ld.shape, mode=[0, 1, 2]), self.acc_dtype ) - wb_tCrO_ld_half = cute.make_fragment(wb_tCrO_ld.layout, sO.element_type) + wb_tCrO_ld_half = cute.make_rmem_tensor(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_fragment(wb_tCrO_ld.layout, self.acc_dtype) + wb_tCrO_reduction = cute.make_rmem_tensor(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 093803572e59..4bee4dad352a 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,22 +28,24 @@ 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: # cuda-bindings / cutlass-dsl not installed +except (ImportError, OSError) as error: _cuda = None cute = None from_dlpack = None - VideoSparseAttentionForward = None + _VideoSparseAttentionForward = None CUTE_AVAILABLE = False + _cute_import_error = error __all__ = [ @@ -106,14 +108,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 " - "cutlass-dsl is not importable." - ) + f"cutlass-dsl is not importable: {_cute_import_error}" + ) from _cute_import_error 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." ) @@ -143,7 +145,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 12fa15bc5930..9a1cd7b7dc55 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,12 +12,8 @@ # 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 -import cutlass.cute as cute -from cutlass._mlir.dialects import llvm, nvvm +from cutlass._mlir.dialects import llvm from cutlass.cute import typing as cutlass_typing from cutlass.cutlass_dsl import dsl_user_op @@ -255,123 +251,3 @@ 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 bfaa0c637875..e7f9f01dd932 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -227,6 +227,7 @@ 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 6ed163bfcc9c..978441f849ed 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 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 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 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause # Redistribution and use in source and binary forms, with or without @@ -56,9 +56,11 @@ 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 @@ -216,31 +218,19 @@ 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_tensor, a_torch = cutlass_torch.cute_tensor_like( + _, 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 ) # Configure gemm kernel @@ -277,8 +267,8 @@ def run( n, k, batch, - a_tensor, - b_tensor, + a_ptr, + b_ptr, c_tensor, a_stride_m, a_stride_batch, @@ -298,7 +288,6 @@ def run_kernel(a_t, b_t, c_t, stream): c_t, a_stride_m, a_stride_batch, - max_active_clusters, stream, ) else: @@ -309,8 +298,8 @@ def run_kernel(a_t, b_t, c_t, stream): n, k, batch, - a_tensor, - b_tensor, + a_ptr, + b_ptr, c_tensor, max_active_clusters, current_stream, @@ -326,14 +315,13 @@ 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_tensor, b_tensor, c_tensor, current_stream) + run_kernel(a_ptr, b_ptr, c_tensor, current_stream) print("Verifying results...") # Reference: C = einsum("mkl,nkl->mnl", A, B) @@ -351,82 +339,86 @@ 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(): - a_tensor_new, _ = cutlass_torch.cute_tensor_like( + def generate_tensors() -> cute.testing.JitArguments: + _, a_torch_new = cutlass_torch.cute_tensor_like( a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - b_tensor_new, _ = cutlass_torch.cute_tensor_like( + _, b_torch_new = cutlass_torch.cute_tensor_like( b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - c_tensor_new, _ = cutlass_torch.cute_tensor_like( + _, c_torch_new = cutlass_torch.cute_tensor_like( c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 ) - a_tensor_new.mark_compact_shape_dynamic( - mode=1, - stride_order=(2, 0, 1), - divisibility=1, + a_ptr_new = make_ptr( + ab_dtype, + a_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, + b_ptr_new = make_ptr( + ab_dtype, + b_torch_new.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, ) - c_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 ) + backing_tensors.append((a_torch_new, b_torch_new, c_torch_new)) return cute.testing.JitArguments( m, n, k, batch, - a_tensor_new, - b_tensor_new, + a_ptr_new, + b_ptr_new, c_tensor_new, a_stride_m, a_stride_batch, - max_active_clusters, current_stream, ) else: - def generate_tensors(): - a_tensor_new, _ = cutlass_torch.cute_tensor_like( + def generate_tensors() -> cute.testing.JitArguments: + _, a_torch_new = cutlass_torch.cute_tensor_like( a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - b_tensor_new, _ = cutlass_torch.cute_tensor_like( + _, b_torch_new = cutlass_torch.cute_tensor_like( b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - c_tensor_new, _ = cutlass_torch.cute_tensor_like( + _, c_torch_new = cutlass_torch.cute_tensor_like( c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 ) - a_tensor_new.mark_compact_shape_dynamic( - mode=1, - stride_order=(2, 0, 1), - divisibility=1, + a_ptr_new = make_ptr( + ab_dtype, + a_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, + b_ptr_new = make_ptr( + ab_dtype, + b_torch_new.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, ) - c_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 ) + backing_tensors.append((a_torch_new, b_torch_new, c_torch_new)) return cute.testing.JitArguments( m, n, k, batch, - a_tensor_new, - b_tensor_new, + a_ptr_new, + b_ptr_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 9f1a9366c870..9fab7b82ef79 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,6 +7,8 @@ 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 @@ -18,6 +20,30 @@ 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