This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new f0ea2d43b8 [TIRx][CUDA] Preserve explicit single-CTA cluster launches
(#20180)
f0ea2d43b8 is described below
commit f0ea2d43b8d6b2b6b14bcbdf725fc0165eb7c541
Author: Bohan Hou <[email protected]>
AuthorDate: Tue Aug 25 16:57:37 2026 -0400
[TIRx][CUDA] Preserve explicit single-CTA cluster launches (#20180)
## Motivation
Calling `K.cta_id_in_cluster(...)` is an explicit request for CUDA
cluster-launch semantics, including when every cluster extent is one.
Lowering previously removed extent-one `clusterCtaIdx.*` bindings, and
the CUDA runtime inferred cluster presence only from dimensions greater
than one. As a result, an explicit `(1, 1, 1)` cluster launch became
indistinguishable from an ordinary CTA launch.
## Changes
- Preserve `clusterCtaIdx.*` launch bindings whenever cluster CTA scope
is present, including extent-one bindings.
- Derive a `use_cluster_launch` flag from regular cluster launch-tag
presence in `LaunchParamConfig`.
- Emit `CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION` whenever that flag is
set, including `(1, 1, 1)`.
- Keep the missing-`clusterCtaIdx` fallback to constant zero for
compatibility with pre-existing IR.
- Add lowering, `SplitHostDevice`, and C++ launch-parameter coverage for
ordinary, singleton, multi-CTA, and preferred-cluster cases.
This does not change serialized module metadata. Preferred-cluster tags
alone do not request a regular cluster launch.
## Launch behavior
NCU 2026.1 measurements on NVIDIA B200:
| Probe | Cluster dimensions | Cluster Size | Uses Blocks as Clusters |
|---|---:|---:|---:|
| Ordinary CTA | `(0, 0, 0)` | 0 | 0 |
| Explicit singleton | `(1, 1, 1)` | 1 | 0 |
| Explicit multi-CTA | `(2, 1, 1)` | 2 | 0 |
## Validation
- TVM build completed successfully.
- LowerTIRx cluster/preferred-cluster tests: 6 passed.
- SplitHostDevice launch tests: 6 passed.
- `LaunchParamConfig` C++ test: 1 passed.
- Relevant CUDA cluster codegen tests: 2 passed.
- Downstream affected correctness matrix: 278/278 passed without skips:
- amax: 59
- fast top-k: 34
- dense SwiGLU: 75
- interleaved quant: 58
- grouped dGLU: 52
- Staged-file pre-commit checks passed.
Downstream B200 benchmarks used Proton timing, five independent samples
per implementation, one-second cooldowns, reference implementations
enabled, interference monitoring, and no sample splicing:
| Matrix | Passing rows | Minimum `mean(reference) / mean(tirx)` |
|---|---:|---:|
| Targeted affected-family gate | 15/15 | 0.996504 |
| Full amax matrix | 88/88 | 0.997440 |
| Full fast top-k matrix | 34/34 | 1.006541 |
The maximum patched-versus-main TIRx slowdown in the targeted comparison
was 0.911%, below the 1% investigation threshold. Six externally
interfered amax attempts were recorded and replaced only by complete
five-round row reruns.
A full TIRx suite run was also performed. The strict kernel registry
passed. The broad parallel run encountered unrelated
environment/reference failures from cross-device DLPack state, aggregate
GPU memory pressure, a missing optional `cutlass.experimental` package,
and reference PTX incompatibility. An isolated core rerun produced 2631
passed, 73 skipped, 3 xpassed, and one pre-existing stale packed-dtype
error-message assertion; no affected cluster correctness configuration
failed.
---
src/backend/cuda/runtime/cuda_module.cc | 2 +-
src/runtime/thread_storage_scope.h | 8 ++
src/tirx/ir/exec_scope.cc | 7 +-
src/tirx/transform/tile_primitive_dispatch.cc | 2 +-
tests/cpp/runtime/thread_storage_scope_test.cc | 42 +++++++++
.../test_tir_transform_split_host_device.py | 29 ++++++
.../tirx/transform/test_transform_lower_tirx.py | 103 +++++++++++++++++++++
7 files changed, 187 insertions(+), 6 deletions(-)
diff --git a/src/backend/cuda/runtime/cuda_module.cc
b/src/backend/cuda/runtime/cuda_module.cc
index 0d210ac0da..604a7de33c 100644
--- a/src/backend/cuda/runtime/cuda_module.cc
+++ b/src/backend/cuda/runtime/cuda_module.cc
@@ -246,7 +246,7 @@ class CUDAWrappedFunc {
unsigned int num_attrs = 0;
// 1) Cluster
- if (wl.cluster_dim(0) != 1 || wl.cluster_dim(1) != 1 || wl.cluster_dim(2)
!= 1) {
+ if (launch_param_config_.use_cluster_launch()) {
CUlaunchAttribute attr{};
attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
attr.value.clusterDim.x = wl.cluster_dim(0);
diff --git a/src/runtime/thread_storage_scope.h
b/src/runtime/thread_storage_scope.h
index c4c3b50cfb..0443fc68a0 100644
--- a/src/runtime/thread_storage_scope.h
+++ b/src/runtime/thread_storage_scope.h
@@ -284,6 +284,7 @@ class LaunchParamConfig {
public:
void Init(size_t base, const ffi::Array<ffi::String>& launch_param_tags) {
base_ = base;
+ use_cluster_launch_ = false;
std::vector<bool> filled(12, false);
for (size_t i = 0; i < launch_param_tags.size(); ++i) {
std::string tag(launch_param_tags[i]);
@@ -297,6 +298,9 @@ class LaunchParamConfig {
use_cooperative_launch_ = true;
} else {
ThreadScope ts = ThreadScope::Create(tag);
+ if (ts.IsClusterCtaIdx()) {
+ use_cluster_launch_ = true;
+ }
arg_index_map_.push_back(ts.rank * 3 + ts.dim_index);
filled[ts.rank * 3 + ts.dim_index] = true;
}
@@ -333,6 +337,8 @@ class LaunchParamConfig {
bool use_cooperative_launch() const { return use_cooperative_launch_; }
+ bool use_cluster_launch() const { return use_cluster_launch_; }
+
private:
/*! \brief base axis */
size_t base_;
@@ -346,6 +352,8 @@ class LaunchParamConfig {
bool use_programmatic_dependent_launch_{false};
/*! \brief Whether or not use cooperative launch. */
bool use_cooperative_launch_{false};
+ /*! \brief Whether the kernel declares a cluster-to-CTA scope. */
+ bool use_cluster_launch_{false};
};
} // namespace runtime
diff --git a/src/tirx/ir/exec_scope.cc b/src/tirx/ir/exec_scope.cc
index b3a5b8d4ba..47a2198fdc 100644
--- a/src/tirx/ir/exec_scope.cc
+++ b/src/tirx/ir/exec_scope.cc
@@ -381,10 +381,9 @@ ffi::Array<PrimExpr> ResolveCuda(ScopeBinding binding,
case ScopeBinding::kKernelCta:
return Trivial3DResolve(params, "blockIdx.", out_dim);
case ScopeBinding::kClusterCta:
- // A launch whose cluster is a single CTA binds no clusterCtaIdx var; the
- // coordinate is then the constant 0, which is what GetThread returns
for a
- // missing tag. blockIdx and threadIdx are always bound, so they keep the
- // strict lookup.
+ // Keep the missing-tag fallback for compatibility with pre-existing IR.
A missing
+ // clusterCtaIdx coordinate resolves to the constant zero; blockIdx and
threadIdx are
+ // always bound, so they keep the strict lookup.
return Trivial3DResolve(params, "clusterCtaIdx.", out_dim,
/*allow_missing=*/true);
case ScopeBinding::kCtaThread:
return Trivial3DResolve(params, "threadIdx.", out_dim);
diff --git a/src/tirx/transform/tile_primitive_dispatch.cc
b/src/tirx/transform/tile_primitive_dispatch.cc
index 940c0d5004..c36b7d9f8d 100644
--- a/src/tirx/transform/tile_primitive_dispatch.cc
+++ b/src/tirx/transform/tile_primitive_dispatch.cc
@@ -776,7 +776,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator {
}
};
auto cluster_cta_it = id_set.find(ScopeBinding::kClusterCta);
- if (cluster_cta_it == id_set.end() ||
is_one((*cluster_cta_it).second.fused_extent())) {
+ if (cluster_cta_it == id_set.end()) {
// no cluster
add_launch_param(ScopeBinding::kKernelCta, "blockIdx.");
} else {
diff --git a/tests/cpp/runtime/thread_storage_scope_test.cc
b/tests/cpp/runtime/thread_storage_scope_test.cc
new file mode 100644
index 0000000000..50579cf286
--- /dev/null
+++ b/tests/cpp/runtime/thread_storage_scope_test.cc
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "../../../src/runtime/thread_storage_scope.h"
+
+#include <gtest/gtest.h>
+
+namespace tvm {
+namespace runtime {
+namespace {
+
+TEST(LaunchParamConfigTest, DerivesClusterLaunchFromTagPresence) {
+ LaunchParamConfig config;
+ config.Init(0, {"blockIdx.x", "threadIdx.x"});
+ EXPECT_FALSE(config.use_cluster_launch());
+
+ config.Init(0, {"blockIdx.x", "clusterCtaIdx.x", "threadIdx.x"});
+ EXPECT_TRUE(config.use_cluster_launch());
+
+ config.Init(0, {"blockIdx.x", "preferredClusterCtaIdx.x", "threadIdx.x"});
+ EXPECT_FALSE(config.use_cluster_launch());
+}
+
+} // namespace
+} // namespace runtime
+} // namespace tvm
diff --git
a/tests/python/tirx-transform/test_tir_transform_split_host_device.py
b/tests/python/tirx-transform/test_tir_transform_split_host_device.py
index 5379b2e3b5..d04f2338a1 100644
--- a/tests/python/tirx-transform/test_tir_transform_split_host_device.py
+++ b/tests/python/tirx-transform/test_tir_transform_split_host_device.py
@@ -437,6 +437,35 @@ def test_cuda_launch_preserves_flag_metadata():
assert int(launch.args[-1]) == 16
+def test_cuda_launch_preserves_singleton_cluster_dimensions():
+ @I.ir_module
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer(1, "float32")):
+ T.func_attr({"target": T.target("cuda", host="llvm")})
+ with T.attr(T.target("cuda"), "target", 0):
+ T.launch_thread("blockIdx.x", 4)
+ T.launch_thread("clusterCtaIdx.x", 1)
+ T.launch_thread("clusterCtaIdx.y", 1)
+ T.launch_thread("clusterCtaIdx.z", 1)
+ T.launch_thread("threadIdx.x", 32)
+ A[0] = 0.0
+
+ after = tvm.tirx.transform.SplitHostDevice()(Before)
+ kernel = after["main_kernel"]
+ assert list(kernel.attrs["tirx.kernel_launch_params"]) == [
+ "blockIdx.x",
+ "clusterCtaIdx.x",
+ "clusterCtaIdx.y",
+ "clusterCtaIdx.z",
+ "threadIdx.x",
+ ]
+
+ launch = after["main"].body.value
+ assert isinstance(launch, tvm.ir.Call)
+ assert [int(arg) for arg in launch.args[-5:]] == [4, 1, 1, 1, 32]
+
+
def test_device_scope_region_extracted_as_device_kernel():
"""A bare device_scope is annotated and extracted as a device kernel."""
diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py
b/tests/python/tirx/transform/test_transform_lower_tirx.py
index 86f6fb0b29..9f9c254853 100644
--- a/tests/python/tirx/transform/test_transform_lower_tirx.py
+++ b/tests/python/tirx/transform/test_transform_lower_tirx.py
@@ -48,6 +48,17 @@ def _int_triple(side, axis):
return tuple(int(x) for x in side[axis])
+def _launch_thread_extents(func):
+ extents = {}
+
+ def collect(node):
+ if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key ==
"thread_extent":
+ extents[str(node.node.thread_tag)] = int(node.value)
+
+ tvm.tirx.stmt_functor.post_order_visit(func.body, collect)
+ return extents
+
+
L_LANE = T.TileLayout(T.S[32 : 1 @ laneid])
@@ -414,6 +425,98 @@ def test_lower_scope_id():
compare(before3, after3, LowerTIRx)
+def test_lower_ordinary_cta_has_no_cluster_launch_tags():
+ @T.prim_func(private=True)
+ def before() -> None:
+ T.device_entry()
+ T.cta_id([1])
+ T.thread_id([32])
+
+ with tvm.target.Target("cuda"):
+ after = LowerTIRx()(tvm.IRModule({"main": before}))["main"]
+
+ launch_extents = _launch_thread_extents(after)
+ assert launch_extents == {"blockIdx.x": 1, "threadIdx.x": 32}
+
+
+def test_lower_explicit_singleton_cluster_launch_tags_survive_when_unused():
+ @T.prim_func(private=True)
+ def cluster_2d() -> None:
+ T.device_entry()
+ unused_cbx, unused_cby = T.cta_id_in_cluster([1, 1])
+ unused_bx, unused_by = T.cta_id([1, 1])
+ T.thread_id([32])
+
+ @T.prim_func(private=True)
+ def cluster_3d() -> None:
+ T.device_entry()
+ unused_cbx, unused_cby, unused_cbz = T.cta_id_in_cluster([1, 1, 1])
+ unused_bx, unused_by, unused_bz = T.cta_id([1, 1, 1])
+ T.thread_id([32])
+
+ with tvm.target.Target("cuda"):
+ after = LowerTIRx()(tvm.IRModule({"cluster_2d": cluster_2d,
"cluster_3d": cluster_3d}))
+
+ assert _launch_thread_extents(after["cluster_2d"]) == {
+ "blockIdx.x": 1,
+ "blockIdx.y": 1,
+ "clusterCtaIdx.x": 1,
+ "clusterCtaIdx.y": 1,
+ "threadIdx.x": 32,
+ }
+ assert _launch_thread_extents(after["cluster_3d"]) == {
+ "blockIdx.x": 1,
+ "blockIdx.y": 1,
+ "blockIdx.z": 1,
+ "clusterCtaIdx.x": 1,
+ "clusterCtaIdx.y": 1,
+ "clusterCtaIdx.z": 1,
+ "threadIdx.x": 32,
+ }
+
+
+def test_lower_multi_cta_cluster_launch_tags_remain_unchanged():
+ @T.prim_func(private=True)
+ def before() -> None:
+ T.device_entry()
+ unused_cbx, unused_cby = T.cta_id_in_cluster([2, 1])
+ unused_bx, unused_by = T.cta_id([2, 1])
+ T.thread_id([32])
+
+ with tvm.target.Target("cuda"):
+ after = LowerTIRx()(tvm.IRModule({"main": before}))["main"]
+
+ assert _launch_thread_extents(after) == {
+ "blockIdx.x": 2,
+ "blockIdx.y": 1,
+ "clusterCtaIdx.x": 2,
+ "clusterCtaIdx.y": 1,
+ "threadIdx.x": 32,
+ }
+
+
+def test_lower_singleton_cluster_preserves_preferred_cluster_tags():
+ @T.prim_func(private=True)
+ def before() -> None:
+ T.device_entry()
+ unused_cbx, unused_cby = T.cta_id_in_cluster([1, 1], preferred=[2, 2])
+ unused_bx, unused_by = T.cta_id([1, 1])
+ T.thread_id([32])
+
+ with tvm.target.Target("cuda"):
+ after = LowerTIRx()(tvm.IRModule({"main": before}))["main"]
+
+ assert _launch_thread_extents(after) == {
+ "blockIdx.x": 1,
+ "blockIdx.y": 1,
+ "clusterCtaIdx.x": 1,
+ "clusterCtaIdx.y": 1,
+ "preferredClusterCtaIdx.x": 2,
+ "preferredClusterCtaIdx.y": 2,
+ "threadIdx.x": 32,
+ }
+
+
def test_lower_scope_id2():
@T.inline
def func(warp_id, tx):
