From 26cddfa9d9f216ccd255042e61ad8ccde3014716 Mon Sep 17 00:00:00 2001 From: cchung100m Date: Sun, 28 Jun 2026 17:49:01 +0800 Subject: [PATCH 1/2] [TOPI][CUDA] Fix topk/sort gridDim.y overflow by moving the width-block dim to blockIdx.z --- python/tvm/topi/gpu/sort.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py index 317a3c57e3d3..461a90ca2980 100644 --- a/python/tvm/topi/gpu/sort.py +++ b/python/tvm/topi/gpu/sort.py @@ -499,17 +499,17 @@ def dual_mergepath( nbx = tvm.tirx.generic.cast(ceil_div(width, max_threads * thread_work), "int32") nbz = tvm.tirx.generic.cast(ceil_div(size, width), "int32") - tx, bx, by, _, _, _ = _get_threads(ntx, nbx, nthread_by * nbz) + tx, bx, by, _, _, _ = _get_threads(ntx, nbx, nthread_by) + bz = te.thread_axis("blockIdx.z") with T.frame_scope( [ T.attr(tx, "thread_extent", ntx), T.attr(bx, "thread_extent", nbx), - T.attr(by, "thread_extent", nthread_by * nbz), + T.attr(by, "thread_extent", nthread_by), + T.attr(bz, "thread_extent", nbz), ] ): - by_val = by % nthread_by - bz = by // nthread_by - base_idx = by_val * size + base_idx = by * size # calculate the start, mid, and end points of this section start_pos = width * bz From b84daff5b1a22b8b047c9e9db5b1578d74643b54 Mon Sep 17 00:00:00 2001 From: cchung100m Date: Mon, 29 Jun 2026 21:06:16 +0800 Subject: [PATCH 2/2] [TOPI][CUDA] Fix topk/sort gridDim overflow by remapping grid axes in sort_ir --- python/tvm/topi/gpu/sort.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py index 461a90ca2980..4368590dd3ea 100644 --- a/python/tvm/topi/gpu/sort.py +++ b/python/tvm/topi/gpu/sort.py @@ -499,8 +499,10 @@ def dual_mergepath( nbx = tvm.tirx.generic.cast(ceil_div(width, max_threads * thread_work), "int32") nbz = tvm.tirx.generic.cast(ceil_div(size, width), "int32") - tx, bx, by, _, _, _ = _get_threads(ntx, nbx, nthread_by) - bz = te.thread_axis("blockIdx.z") + tx = te.thread_axis("threadIdx.x") + bx = te.thread_axis("blockIdx.z") # nbx + by = te.thread_axis("blockIdx.y") # batch + bz = te.thread_axis("blockIdx.x") # nbz (largest extent) with T.frame_scope( [ T.attr(tx, "thread_extent", ntx),