Merge remote-tracking branch 'origin/main' into pr-2364-main-merge

# Conflicts: # examples/dequantize_gemm/quantize/nvfp4.py
LeiWang1999 · Jul 30, 2026 · Jul 3, 2026 · Jul 8, 2026 · Jul 8, 2026 · Jul 8, 2026
commit 97abc1f400ecb85df7eaa67b889348066aebd3e1
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -47,7 +47,7 @@ jobs:
 
       - name: Setup Python 3.10
         id: setup-pylowest
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@v7
         with:
           python-version: "3.10"
           update-environment: true

diff --git a/.github/workflows/publish-docs.yml b/.github/workflows/publish-docs.yml
@@ -31,7 +31,7 @@ jobs:
           submodules: recursive
 
       - name: Setup Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@v7
         with:
           python-version: "3.10"
 

diff --git a/3rdparty/tvm b/3rdparty/tvm
diff --git a/benchmark/matmul/benchmark_matmul_sp.py b/benchmark/matmul/benchmark_matmul_sp.py
@@ -1,13 +1,18 @@
 import argparse
 import itertools
 import logging
+from pathlib import Path
+import sys
+
 import torch
 from triton.testing import do_bench
 
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "examples" / "gemm_sp"))
+
 import tilelang.language as T
 from tilelang.autotuner import autotune
 from tilelang import jit
-from tilelang.utils.sparse import get_e_factor
+from sparse_utils import get_e_factor
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)

diff --git a/benchmark/matmul/benchmark_matmul_sp_compress.py b/benchmark/matmul/benchmark_matmul_sp_compress.py
@@ -1,8 +1,13 @@
 import argparse
+from pathlib import Path
+import sys
 
 import torch
 from tilelang.profiler import do_bench
-from tilelang.utils.sparse import compress, randn_semi_sparse, randint_semi_sparse, torch_compress
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "examples" / "gemm_sp"))
+
+from sparse_utils import compress, randn_semi_sparse, randint_semi_sparse, torch_compress
 
 SUPPORTED_DTYPE_NAMES = ["float16", "bfloat16", "float32", "int8"]
 SUPPORTED_META_DTYPE_NAMES = ["int8", "int16", "int32"]

diff --git a/docs/deeplearning_operators/matmul_sparse.md b/docs/deeplearning_operators/matmul_sparse.md
@@ -38,10 +38,10 @@ To utilize sparse Tensor Cores, a dense tensor must first be **compressed** into
 
 Both `PyTorch` and `vLLM` use `CUTLASS` as their computation backend (see references [here](https://github.com/pytorch/pytorch/blob/a8d6afb511a69687bbb2b7e88a3cf67917e1697e/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredOps.cu#L47) and [here](https://github.com/vllm-project/vllm/blob/a5dd03c1ebc5e4f56f3c9d3dc0436e9c582c978f/csrc/sparse/cutlass/sparse_scaled_mm_c3x.cuh#L116)), leveraging `CUTLASS`’s built-in compressor (or reimplementing it in `PyTorch`).
 
-A compressor is provided in `tilelang.utils.sparse`. Pass in a dense 2:4-sparse tensor and optionally a metadata dtype to get back the compressed values and metadata:
+A compressor is provided with the sparse GEMM example in `examples/gemm_sp/sparse_utils.py`. Pass in a dense 2:4-sparse tensor and optionally a metadata dtype to get back the compressed values and metadata:
 
 ```python
-from tilelang.utils.sparse import compress
+from examples.gemm_sp.sparse_utils import compress
 A_sparse, E = compress(A)                        # default: int16 metadata for fp16/bf16
 A_sparse, E = compress(A.t().contiguous())       # compress the transposed layout
 ```
@@ -56,7 +56,7 @@ The default metadata dtype for fp16/bf16 is `int16` with an E-factor of 16 (one
 
 ```python
 import tilelang.language as T
-from tilelang.utils.sparse import get_e_factor
+from examples.gemm_sp.sparse_utils import get_e_factor
 
 def matmul_sp(
     M, N, K,

diff --git a/examples/dequantize_gemm/example_dequant_gemm_bf16_fp4_hopper.py b/examples/dequantize_gemm/example_dequant_gemm_bf16_fp4_hopper.py
@@ -76,7 +76,7 @@ def matmul(
     - num_bits (default 4) is the bit-width of the quantized elements; storage_dtype is uint8 and num_elems_per_byte = 8 // num_bits.
     - QK = K // num_elems_per_byte and Block_QK = block_K // num_elems_per_byte determine B and shared-buffer shapes.
     - Asserts that K % (block_K * split) == 0; K must be divisible by block_K * split for the tiling to be valid.
-    - When fast_dequant is True, a valid mxfp intrinsic group (C source and function name) must be available via tilelang.quantize.get_mxfp_intrin_group.
+    - When fast_dequant is True, a valid mxfp intrinsic group (C source and function name) must be available via quantize.get_mxfp_intrin_group.
     - The kernel launches a 2D grid over ceildiv(N, block_N) and ceildiv(M, block_M) and uses `threads` threads per block with `num_stages` pipeline stages.
 
     Parameters that alter kernel layout/behavior (brief):
@@ -101,7 +101,7 @@ def matmul(
     B_dequantize_shared_shape = (block_N, block_K)
     assert K % (block_K * split) == 0
 
-    from tilelang.quantize import get_mxfp_intrin_group
+    from quantize import get_mxfp_intrin_group
 
     # fast_dequant_bf16_fp4_twiddling
     # It requires that the 2 consecutive uint8 elements (16bits) contains 4 fp4 elements in a bit-twiddling way.

diff --git a/examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_cdna4.py b/examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_cdna4.py
@@ -19,7 +19,7 @@
 from tvm import tirx
 import torch
 
-from tilelang.quantize import get_mxfp_intrin_group
+from quantize import get_mxfp_intrin_group
 from dequantize_utils import torch_convert_bit_twiddling, torch_convert
 
 

diff --git a/examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper.py b/examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper.py
@@ -154,7 +154,7 @@ def matmul(
     B_dequantize_shared_shape = (block_N, block_K)
     assert K % (block_K * split) == 0
 
-    from tilelang.quantize import get_mxfp_intrin_group
+    from quantize import get_mxfp_intrin_group
 
     # fast_dequant_bf16_fp4_twiddling
     mxfp_intrin_info = get_mxfp_intrin_group(

diff --git a/examples/dequantize_gemm/example_dequant_gemm_fine_grained.py b/examples/dequantize_gemm/example_dequant_gemm_fine_grained.py
@@ -28,7 +28,7 @@ def matmul(
     threads,
     num_bits=4,
 ):
-    from tilelang.quantize import _tir_packed_to_unsigned_convert
+    from quantize import _tir_packed_to_unsigned_convert
 
     num_elems_per_byte = 8 // num_bits
     storage_dtype = T.int8

diff --git a/examples/dequantize_gemm/example_dequant_gemv_fp16xint4.py b/examples/dequantize_gemm/example_dequant_gemv_fp16xint4.py
@@ -3,7 +3,7 @@
 from typing import Optional, Callable, Any
 import torch
 from tilelang import DataType
-from tilelang.quantize import (
+from quantize import (
     _tir_packed_int_to_int_convert,
 )
 
@@ -58,7 +58,7 @@ def dequantize_gemv(
     if fast_decoding is True:
         # Lazy import to decrease the startup time
         # as intrin registry may take a while to load
-        from tilelang.quantize import get_lop3_intrin_group
+        from quantize import get_lop3_intrin_group
 
         lop3_intrin_info = get_lop3_intrin_group(
             out_dtype=in_dtype,
@@ -199,7 +199,7 @@ def main() -> None:
     C = torch.zeros(M, N, dtype=getattr(torch, accum_dtype)).cuda()
 
     if fast_decoding:
-        from tilelang.quantize.utils import interleave_weight
+        from quantize.utils import interleave_weight
 
         qB = interleave_weight(qB, num_bits, in_dtype)
     kernel(A, qB, C)
@@ -261,7 +261,7 @@ def run_regression_perf():
     C = torch.zeros(M, N, dtype=getattr(torch, accum_dtype)).cuda()
 
     if fast_decoding:
-        from tilelang.quantize.utils import interleave_weight
+        from quantize.utils import interleave_weight
 
         qB = interleave_weight(qB, num_bits, in_dtype)
     kernel(A, qB, C)

diff --git a/tilelang/quantize/__init__.py → ...ples/dequantize_gemm/quantize/__init__.py b/tilelang/quantize/__init__.py → ...ples/dequantize_gemm/quantize/__init__.py
diff --git a/tilelang/quantize/lop3.py → examples/dequantize_gemm/quantize/lop3.py b/tilelang/quantize/lop3.py → examples/dequantize_gemm/quantize/lop3.py
diff --git a/tilelang/quantize/mxfp.py → examples/dequantize_gemm/quantize/mxfp.py b/tilelang/quantize/mxfp.py → examples/dequantize_gemm/quantize/mxfp.py
diff --git a/tilelang/quantize/nvfp4.py → examples/dequantize_gemm/quantize/nvfp4.py b/tilelang/quantize/nvfp4.py → examples/dequantize_gemm/quantize/nvfp4.py
diff --git a/tilelang/quantize/quantization.py → .../dequantize_gemm/quantize/quantization.py b/tilelang/quantize/quantization.py → .../dequantize_gemm/quantize/quantization.py
diff --git a/tilelang/quantize/utils.py → examples/dequantize_gemm/quantize/utils.py b/tilelang/quantize/utils.py → examples/dequantize_gemm/quantize/utils.py
diff --git a/examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.py b/examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.py
@@ -25,7 +25,7 @@
 import tilelang
 import tilelang.language as T
 from tilelang.profiler import do_bench
-from tilelang.quantize import swizzle_blockscaled_chunk_kmajor_scale_words
+from examples.dequantize_gemm.quantize import swizzle_blockscaled_chunk_kmajor_scale_words
 
 
 _SM120_SCALE_LAYOUT = "blockscaled_chunk_kmajor"
@@ -50,7 +50,7 @@ def sm120_nvfp4_blockscaled_gemm(
     # Tail tiles: TMA loads zero-fill out-of-bounds rows and the C store is
     # predicated, so M (or N) may be arbitrary as long as the other dimension
     # is a multiple of its tile. The scale source is padded to full 128-row
-    # tiles (the packers in tilelang.quantize do this automatically).
+    # tiles (the helpers in examples.dequantize_gemm.quantize do this automatically).
     # N must keep 16-byte aligned bf16 rows, the same contiguous-dim rule as
     # CuTeDSL's is_valid_tensor_alignment.
     assert N % 8 == 0, "N must be a multiple of 8 (16-byte aligned output rows)"

diff --git a/examples/gemm_sp/example_gemm_sp.py b/examples/gemm_sp/example_gemm_sp.py
@@ -3,7 +3,7 @@
 import tilelang
 import tilelang.language as T
 
-from tilelang.utils.sparse import compress, randn_semi_sparse, get_e_factor
+from sparse_utils import compress, randn_semi_sparse, get_e_factor
 from tilelang.profiler import do_bench
 
 import torch

diff --git a/tilelang/utils/sparse.py → examples/gemm_sp/sparse_utils.py b/tilelang/utils/sparse.py → examples/gemm_sp/sparse_utils.py
@@ -5,7 +5,7 @@
 
 # 2:4 sparsity layout metadata lives in the leaf module so the mma_sp code
 # generator can import it without dragging this torch/@tilelang.jit recipe into
-# the bootstrap closure. Re-exported here for backward compatibility.
+# the bootstrap closure. Import it here so this example uses the same contract.
 from tilelang.cuda.intrinsics.sparse_layout import (  # noqa: F401
     GROUP_CONFIG,
     get_e_factor,

diff --git a/maint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.py b/maint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.py
@@ -29,14 +29,14 @@
 import tilelang.language as T
 from tilelang.carver.arch import driver
 from tilelang.profiler import do_bench
-from tilelang.quantize import (
+from examples.dequantize_gemm.quantize import (
     swizzle_blockscaled_chunk_kmajor_scale_words,
     unswizzle_blockscaled_chunk_kmajor_scale_words,
 )
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
 from tilelang_nvfp4_quantizer import tilelang_quantize_bf16_to_nvfp4_blockscaled  # noqa: E402
-from tilelang.quantize.nvfp4 import (
+from examples.dequantize_gemm.quantize.nvfp4 import (
     blockscaled_chunk_kmajor_tile_source_coords,
     blockscaled_chunk_kmajor_word_offset,
     decode_ue4m3_scale_bytes,
@@ -159,7 +159,7 @@ def tilelang_nvfp4_blockscaled_gemm(
         def copy_blockscaled_chunk_kmajor_scale_tile(SF, SF_shared, tile_row, block_rows, ko, stage, tx):
             # Producer-warp-group staging: 128 producer lanes stride the tile,
             # with the layout addressing shared via
-            # tilelang.quantize.nvfp4.blockscaled_chunk_kmajor_tile_source_coords.
+            # examples.dequantize_gemm.quantize.nvfp4.blockscaled_chunk_kmajor_tile_source_coords.
             scale_tile_words = block_rows * sf_words_per_block_k
             scale_lane = tx - 256
             for scale_iter in T.serial((scale_tile_words + 127) // 128):

diff --git a/maint/gemm/gemm_sm120/tilelang_nvfp4_quantizer.py b/maint/gemm/gemm_sm120/tilelang_nvfp4_quantizer.py
@@ -2,13 +2,13 @@
 
 This is the performance quantizer used by the SM120 NVFP4 GEMM maint
 benchmark. It builds on the device-side format helpers and layout contract in
-``tilelang.quantize.nvfp4``; the library itself only ships those helpers plus
-a torch reference implementation (``quantize_bf16_to_nvfp4_blockscaled``).
+``examples.dequantize_gemm.quantize.nvfp4`` alongside the torch reference
+implementation (``quantize_bf16_to_nvfp4_blockscaled``).
 """
 
 import tilelang
 import tilelang.language as T
-from tilelang.quantize.nvfp4 import (
+from examples.dequantize_gemm.quantize.nvfp4 import (
     _BLOCKSCALED_CHUNK_ROWS,
     _BLOCKSCALED_CHUNK_WORDS,
     _FP4_E2M1_MAX,

diff --git a/pyproject.toml b/pyproject.toml
@@ -58,7 +58,7 @@ nvcc = [
     # On Windows, ship nvrtc through the runtime deps so `pip install tilelang`
     # works without a host CUDA install (TileLang JITs CUDA kernels via nvrtc).
     "nvidia-cuda-nvrtc>=13; platform_system == 'Windows'",
-    # cuRAND is required when kernels use tilelang.language.random
+    # cuRAND is required when kernels use tilelang.cuda.language.random
     "nvidia-curand>=10.4; platform_system == 'Windows'",
 ]
 

diff --git a/reproducer_metal_address_space.py b/reproducer_metal_address_space.py
diff --git a/reproducer_vec_type.py b/reproducer_vec_type.py
@@ -0,0 +1,37 @@
+"""
+Minimal reproducer for the missing vec_type arithmetic operators in common.h.
+
+Without the fix, compiling a CPU kernel with `c` target fails:
+  error: no match for 'operator-' (operand types are 'float4' and 'float4')
+
+Requires: tilelang, torch
+"""
+
+import tilelang
+import tilelang.language as T
+import torch
+
+
+@T.prim_func
+def vec_add(
+    A: T.Tensor((256,), "float32"),
+    B: T.Tensor((256,), "float32"),
+    C: T.Tensor((256,), "float32"),
+):
+    for i in T.Parallel(256):
+        C[i] = A[i] + B[i]
+
+
+def main():
+    f = tilelang.compile(vec_add,>
+    A = torch.randn(256)
+    B = torch.randn(256)
+    C = torch.zeros(256)
+    f(A, B, C)
+    ref = A + B
+    assert (C - ref).abs().max().item() < 1e-5
+    print("OK: vec_add compiled and ran successfully")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/backend/common/op/finalize_reducer.h b/src/backend/common/op/finalize_reducer.h
@@ -6,6 +6,7 @@
 #ifndef TVM_TL_BACKEND_COMMON_OP_FINALIZE_REDUCER_H_
 #define TVM_TL_BACKEND_COMMON_OP_FINALIZE_REDUCER_H_
 
+#include "backend/common/op/reduce.h"
 #include "op/finalize_reducer.h"
 #include "support/check.h"
 
@@ -52,6 +53,7 @@ template <typename Impl> struct FinalizeReducerLowerer {
     auto op_str = op_names[static_cast<int>(op.op)];
 
     int reducing_threads = extent;
+    reduce::CheckAllReduceWidth(reducing_threads, 1, "tl.finalize_reducer");
     auto thread_offset = lower_args.thread_bounds->min;
 
     int64_t layout_batch_size = 1;

diff --git a/src/backend/common/op/reduce.h b/src/backend/common/op/reduce.h
@@ -104,6 +104,25 @@ inline int GetPreferedVectorizedSize(DataType dt,
   return 1;
 }
 
+inline void CheckAllReduceWidth(int reducing_threads, int scale,
+                                const char *op_name) {
+  ICHECK_GT(reducing_threads, 0)
+      << op_name << ": AllReduce threads must be positive, got "
+      << reducing_threads;
+  ICHECK_GT(scale, 0) << op_name << ": AllReduce scale must be positive, got "
+                      << scale;
+  ICHECK_EQ(reducing_threads % scale, 0)
+      << op_name << ": AllReduce threads (" << reducing_threads
+      << ") must be divisible by scale (" << scale << ")";
+  int logical_width = reducing_threads / scale;
+  int shift = 0;
+  ICHECK(tirx::is_const_power_of_two_integer(Integer(logical_width), &shift))
+      << op_name << ": XOR-butterfly AllReduce requires logical_width "
+      << "(threads / scale) to be a positive power of two, got "
+      << logical_width << " (threads=" << reducing_threads
+      << ", scale=" << scale << ")";
+}
+
 inline PrimExpr MakeInitValue(const ReduceOpNode &op, int vsize = 1) {
   auto dst_dtype = op.dst->dtype;
   auto is_int = dst_dtype.is_int();
@@ -909,6 +928,8 @@ template <typename Impl> struct ReduceLowerer {
 
         for (const auto &thread_step : reduce_plan.thread_steps) {
           int reducing_threads = thread_step.ReducingThreads();
+          reduce::CheckAllReduceWidth(reducing_threads, thread_step.scale,
+                                      "tl.reduce");
           int block_threads =
               static_cast<int>(*as_const_int(lower_args.thread_bounds->extent));
           auto thread_offset = lower_args.thread_bounds->min;
@@ -1101,6 +1122,8 @@ template <typename Impl> struct ReduceLowerer {
 
       for (const auto &thread_step : reduce_plan.thread_steps) {
         int reducing_threads = thread_step.ReducingThreads();
+        reduce::CheckAllReduceWidth(reducing_threads, thread_step.scale,
+                                    "tl.reduce");
         auto thread_offset = lower_args.thread_bounds->min;
         std::string allreduce = Impl::MakeScalarAllReduce(
             reduce::MakeCodegenReducer(op).value(), reducing_threads,