tilelang-rocm 0.1.4.post5__cp310-cp310-manylinux1_x86_64.whl → 0.1.4.post9__cp310-cp310-manylinux1_x86_64.whl

tilelang/quantize/quantization.py

@@ -0,0 +1,234 @@
+# Copyright 2018 The apache/tvm Authors. All Rights Reserved.
+# 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.
+#
+# Copyright (c) Tile-AI Corporation.
+# Licensed under the MIT License.
+# The code below is mostly copied from mlc.ai quantization.py in mlc-llm.
+# pylint: disable=invalid-name,missing-function-docstring,unused-variable
+"""TIR computation utilities for quantization."""
+
+from tilelang import tvm as tvm
+from tvm import tir
+
+
+# fmt: off
+def _tir_f32x2_to_bf16x2_to_u32(v0: tir.PrimExpr, v1: tir.PrimExpr, round_to_even: bool = True):
+    mask = tir.const((1 << 16) - 1, "uint32")
+    res = []
+    for data in [v0, v1]:
+        u32_val = tir.reinterpret("uint32", data)
+        if round_to_even:
+            rounding_bias = ((u32_val >> tir.const(16, "uint32"))
+                             & tir.const(1, "uint32")) + tir.const(0x7FFF, "uint32")
+            u32_val += rounding_bias
+        res.append((u32_val >> tir.const(16, "uint32")) & mask)
+    return res[0] | (res[1] << tir.const(16, "uint32"))
+
+
+def _tir_u32_to_bf16x2_to_f32x2(x: tir.PrimExpr):
+    mask = tir.const((1 << 16) - 1, "uint32")
+    x0 = x & mask
+    x1 = (x >> 16) & mask
+    return (tir.reinterpret("float32", x << tir.const(16, "uint32")) for x in [x0, x1])
+
+
+def _tir_u32_to_int_to_float(nbit: int, val: tir.PrimExpr, pos: tir.PrimExpr, dtype: str):
+    assert val.dtype == "uint32"
+    mask = tvm.tir.const((1 << nbit) - 1, "uint32")
+    return tir.Cast(dtype, (val >> (pos * nbit).astype("uint32")) & mask)
+
+
+def _tir_packed_uint_to_uint_to_float(storage_nbit: int):
+    storage_dtype = "uint" + str(storage_nbit)
+
+    def f_convert(nbit: int, val: tir.PrimExpr, pos: tir.PrimExpr, dtype: str):
+        assert val.dtype == storage_dtype, f"{val.dtype} != {storage_dtype}"
+        max_int_value = (1 << (nbit - 1)) - 1
+        return ((val >> (pos.astype("uint32") * tir.const(nbit, "uint32"))) & tir.const(
+            (1 << nbit) - 1, "uint32")).astype(dtype) - tir.const(max_int_value, dtype)
+
+    return f_convert
+
+
+def _tir_packed_int_to_int_to_float(storage_nbit: int):
+    storage_dtype = "int" + str(storage_nbit)
+
+    def f_convert(nbit: int, val: tir.PrimExpr, pos: tir.PrimExpr, dtype: str):
+        assert val.dtype == storage_dtype, f"{val.dtype} != {storage_dtype}"
+        mask = tir.const((1 << nbit) - 1, "int32")
+        unextended = (val >> (pos.astype("int32") * tir.const(nbit, "int32"))) & mask
+        return tir.Cast(
+            dtype, (unextended << tir.const(32 - nbit, "int32")) >> tir.const(32 - nbit, "int32"))
+
+    return f_convert
+
+
+def _tir_f32_to_uint_to_f4(val: tir.PrimExpr):
+    assert val.dtype == "float32"
+    val_u32 = tir.reinterpret("uint32", val)
+    # e_f32 >  120 -> e_f4 = min(e_f32 - 120 + M_h, 7)
+    # e_f32 == 120 -> e_f4 = 1
+    # e_f32 < 120 -> e_f4 = 0
+    m_h = (val_u32 >> tir.const(22, "uint32")) & tir.const(1, "uint32")
+    e_f32 = (val_u32 >> tir.const(23, "uint32")) & tir.const(255, "uint32")
+    s = (val_u32 >> tir.const(31, "uint32"))
+    e_f4 = tir.Select(
+        e_f32 > tir.const(120, "uint32"),
+        tir.Min(e_f32 - tir.const(120, "uint32") + m_h, tir.const(7, "uint32")),
+        tir.Select(e_f32 == tir.const(120, "uint32"), tir.const(1, "uint32"),
+                   tir.const(0, "uint32")))
+    return (s << tir.const(3, "uint32")) | e_f4
+
+
+def _tir_f16_to_uint_to_f4(val: tir.PrimExpr):
+    assert val.dtype == "float16"
+    val_u32 = tir.Cast("uint32", tir.reinterpret("uint16", val))
+    m_h = (val_u32 >> tir.const(9, "uint32")) & tir.const(1, "uint32")
+    e_f16 = (val_u32 >> tir.const(10, "uint32")) & tir.const(31, "uint32")
+    s = (val_u32 >> tir.const(15, "uint32"))
+    e_f4 = tir.Select(
+        e_f16 > tir.const(8, "uint32"),
+        tir.Min(e_f16 - tir.const(8, "uint32") + m_h, tir.const(7, "uint32")),
+        tir.Select(e_f16 == tir.const(8, "uint32"), tir.const(1, "uint32"), tir.const(0, "uint32")))
+    return (s << tir.const(3, "uint32")) | e_f4
+
+
+def _tir_u32_to_f4_to_f32(nbit: int, val: tir.PrimExpr, pos: tir.PrimExpr, dtype: str):
+    assert nbit == 4
+    assert dtype == "float32"
+    assert val.dtype == "uint32"
+    # e_f4 == 0 -> e_f32 = 0
+    # e_f4 != 0 -> e_f32 = e_f4 + 120 = e_f4 | (1111000)_2
+    mask = tvm.tir.const((1 << nbit) - 1, "uint32")
+    f4 = (val >> (pos.astype("uint32") * tir.const(nbit, "uint32"))) & mask
+    s = f4 >> tir.const(3, "uint32")
+    e_f4 = f4 & tir.const(7, "uint32")
+    e_f32 = e_f4 | tir.const(120, "uint32")
+    val_f32 = tir.reinterpret("float32",
+                              (e_f32 | (s << tir.const(8, "uint32"))) << tir.const(23, "uint32"))
+    return tir.Select(e_f4 == tir.const(0, "uint32"), tir.const(0, "float32"), val_f32)
+
+
+def _tir_packed_to_fp4_to_f16(nbit: int, val: tir.PrimExpr, pos: tir.PrimExpr, dtype: str):
+    assert nbit == 4
+    assert dtype == "float16"
+    assert val.dtype == "uint32"
+    # e_f4 == 0 -> e_f16 = 0
+    # e_f4 != 0 -> e_f16 = e_f4 + 8 = e_f4 | (1000)_2
+    mask = tvm.tir.const((1 << nbit) - 1, "uint16")
+    f4 = (val >> (pos.astype("uint16") * tir.const(nbit, "uint16"))) & mask
+    s = f4 >> tir.const(3, "uint16")
+    e_f4 = f4 & tir.const(7, "uint16")
+    e_f16 = e_f4 | tir.const(8, "uint16")
+    val_f16 = tir.reinterpret("float16",
+                              ((e_f16 | (s << tir.const(5, "uint16"))) << tir.const(10, "uint16")).astype("uint16"))
+    return tir.Select(e_f4 == tir.const(0, "uint16"), tir.const(0, "float16"), val_f16)
+
+def _tir_packed_to_fp4_to_f16(storage_type="uint", storage_nbit=8):
+    storage_dtype = storage_type + str(storage_nbit)
+
+    def f_convert(nbit: int, val: tvm.tir.PrimExpr, pos: tvm.tir.PrimExpr, dtype: str):
+        assert val.dtype == storage_dtype, f"{val.dtype} != {storage_dtype}"
+        mask = tvm.tir.const((1 << nbit) - 1, storage_dtype)
+        f4 = ((val >> (pos * nbit).astype(storage_dtype)) & mask).astype(storage_dtype)
+        f4 = (val >> (pos.astype(storage_dtype) * tir.const(nbit, storage_dtype))) & mask
+        s = f4 >> tir.const(3, storage_dtype)
+        e_f4 = f4 & tir.const(7, storage_dtype)
+        e_f16 = e_f4 | tir.const(8, storage_dtype)
+        val_f16 = tir.reinterpret("float16",
+                                ((e_f16 | (s << tir.const(5, storage_dtype))) << tir.const(10, storage_dtype)).astype("uint16"))
+        return tir.Select(e_f4 == tir.const(0, storage_dtype), tir.const(0, "float16"), val_f16)
+
+    return f_convert
+
+def _tir_u8_to_f8_e4m3_to_f16_naive(nbit: int, val: tir.PrimExpr, dtype: str):
+    assert nbit == 8
+    assert dtype == "float16"
+    s_f16 = (val >> tir.const(7, "uint16")) << tir.const(15, "uint16")
+    e4 = val & tir.const(0x40, "uint16")
+    prefix = tir.Select(e4 == tir.const(0, "uint16"), tir.const(0x2000, "uint16"),
+                        tir.const(0x4000, "uint16"))
+    e_f16 = (((val & tir.const(63, "uint16")) << tir.const(7, "uint16"))) | prefix
+    return tir.reinterpret("float16", s_f16 | e_f16)
+
+
+def _tir_u8_to_f8_e4m3_to_f16(nbit: int, val: tir.PrimExpr, dtype: str):
+    assert nbit == 8
+    assert dtype == "float16"
+    s_f16 = (val >> tir.const(7, "uint16")) << tir.const(15, "uint16")
+    e4 = val & tir.const(0x40, "uint16")
+    e_f16 = (((val & tir.const(63, "uint16")) << tir.const(7, "uint16"))) | (e4 << tir.const(8, "uint16")) | (e4 << tir.const(7, "uint16"))
+    e_f16 = e_f16 ^ tir.const(0x2000, "uint16")
+    return tir.reinterpret("float16", s_f16 | e_f16)
+
+
+def _tir_u8_to_f8_e5m2_to_f16(nbit: int, val: tir.PrimExpr, dtype: str):
+    assert nbit == 8
+    assert dtype == "float16"
+    return tir.reinterpret("e5m2_float8", val).astype("float16")
+
+
+def _tir_packed_to_signed_convert(storage_type="uint", storage_nbit=8):
+    storage_dtype = storage_type + str(storage_nbit)
+
+    def f_convert(nbit: int, val: tir.PrimExpr, pos: tir.PrimExpr, dtype: str):
+        assert val.dtype == storage_dtype, f"{val.dtype} != {storage_dtype}"
+        max_int_value = (1 << (nbit - 1))
+        return ((val >> (pos.astype("uint32") * tir.const(nbit, "uint32"))) & tir.const(
+            (1 << nbit) - 1, "uint32")).astype(dtype) - tir.const(max_int_value, dtype)
+
+    return f_convert
+
+
+def _tir_packed_to_unsigned_convert(storage_type="uint", storage_nbit=8):
+    storage_dtype = storage_type + str(storage_nbit)
+
+    def f_convert(nbit: int, val: tvm.tir.PrimExpr, pos: tvm.tir.PrimExpr, dtype: str):
+        assert val.dtype == storage_dtype, f"{val.dtype} != {storage_dtype}"
+        mask = tvm.tir.const((1 << nbit) - 1, storage_dtype)
+        return ((val >> (pos * nbit).astype(storage_dtype)) & mask).astype(dtype)
+
+    return f_convert
+
+
+def _tir_packed_to_unsigned_convert_with_zeros(storage_type="uint", storage_nbit=8):
+    storage_dtype = storage_type + str(storage_nbit)
+
+    def f_convert(nbit: int, val: tvm.tir.PrimExpr, pos: tvm.tir.PrimExpr, zero: tvm.tir.PrimExpr,
+                  dtype: str):
+        assert val.dtype == storage_dtype, f"{val.dtype} != {storage_dtype}"
+        mask = tvm.tir.const((1 << nbit) - 1, storage_dtype)
+        return (((val >> (pos * nbit).astype(storage_dtype)) & mask) - zero).astype(dtype)
+
+    return f_convert
+
+
+def _tir_packed_int_to_int_convert(storage_type="uint", storage_nbit=8):
+    storage_dtype = storage_type + str(storage_nbit)
+
+    def f_convert(nbit: int, val: tir.PrimExpr, pos: tir.PrimExpr, dtype: str):
+        assert val.dtype == storage_dtype, f"{val.dtype} != {storage_dtype}"
+        mask = tir.const((1 << nbit) - 1, "int32")
+        unextended = (val >> (pos.astype("int32") * tir.const(nbit, "int32"))) & mask
+        return tir.Cast(
+            dtype, (unextended << tir.const(32 - nbit, "int32")) >> tir.const(32 - nbit, "int32"))
+
+    return f_convert
+
+
+# fmt: on

tilelang/quantize/utils.py

@@ -0,0 +1,126 @@
+# Copyright (c) Tile-AI Corporation.
+# Licensed under the MIT License.
+
+
+def gen_quant4(k, n, groupsize=-1):
+    import torch
+    import torch.nn as nn
+    maxq = 2**4
+    w = torch.randn((k, n), dtype=torch.half, device="cpu")
+
+    original_w = w.clone()
+
+    if groupsize == -1:
+        groupsize = k
+
+    if groupsize != -1:
+        w = w.reshape((-1, groupsize, n))
+        w = w.permute(1, 0, 2)
+        w = w.reshape((groupsize, -1))
+
+    s = torch.max(torch.abs(w), 0, keepdim=True)[0]
+    s *= 2 / maxq
+
+    # Quantize.
+    w = torch.round(w / s).int()
+
+    # Unsigned storage.
+    w += (maxq) // 2
+
+    w = torch.clamp(w, 0, maxq)
+
+    # Dequantize.
+    ref = (w - (maxq) // 2).half() * s
+
+    if groupsize != -1:
+
+        def reshape(w):
+            w = w.reshape((groupsize, -1, n))
+            w = w.permute(1, 0, 2)
+            w = w.reshape((k, n)).contiguous()
+            return w
+
+        ref = reshape(ref)
+        w = reshape(w)
+
+    s = s.reshape((-1, n)).contiguous()
+    linear = nn.Linear(k, n, bias=False)
+    linear.weight.data = ref.t()
+
+    return original_w, linear, s, (w - (maxq) // 2)
+
+
+def general_compress(lowprecision_weight, source_bits=4, storage_dtype=None):
+    import torch
+    if storage_dtype is None:
+        storage_dtype = torch.int8
+    elems_per_byte = 8 // source_bits
+    if lowprecision_weight.dtype == torch.float16:
+        lowprecision_weight = lowprecision_weight.to(torch.int8)
+    int8_weight = torch.zeros(
+        (*lowprecision_weight.shape[:-1], lowprecision_weight.shape[-1] // elems_per_byte),
+        dtype=torch.int8,
+        device=lowprecision_weight.device)
+    for j in range(lowprecision_weight.shape[-1] // elems_per_byte):
+        for k in range(elems_per_byte):
+            int8_weight[..., j] |= (lowprecision_weight[..., j * elems_per_byte + k] <<
+                                    (source_bits * k)).to(torch.int8)
+
+    return int8_weight.to(storage_dtype)
+
+
+# interleave weight numpy implementation
+def interleave_weight(qweight, nbits=4, target_dtype="float16"):
+    """Interleave the weight to the target data type.
+
+    Args:
+        qweight (_type_): _description_
+        nbits (int, optional): _description_. Defaults to 4.
+        target_dtype (str, optional): _description_. Defaults to "float16".
+
+    Returns:
+        _type_: _description_
+    
+    Example:
+        qweight = torch.randint(0, 127, (10, 10), dtype=torch.int8).cuda()
+        interleave_weight(qweight, 4, "float16")
+    """
+    import torch
+    assert target_dtype in ["float16", "int8"]
+    # reinterpret the data type of qweight to int32
+    qweight = qweight.view(torch.int32)
+    new_qweight = torch.zeros_like(qweight)
+    bits_stride = 8 if target_dtype == "int8" else 16
+    mask = (1 << nbits) - 1  # for 4bit the val is 0x0000000f
+    num_groups = 32 // bits_stride
+    elems_per_group = bits_stride // nbits
+    for i in range(num_groups):
+        for j in range(elems_per_group):
+            offset = i * elems_per_group + j
+            shift = (offset % num_groups) * bits_stride + (offset // num_groups) * nbits
+            new_qweight |= ((qweight >> (nbits * offset)) & mask) << shift
+
+    if nbits == 1 and target_dtype == "int8":
+        # special handling for 1b interleave
+        n16_weight = new_qweight & torch.int32(0xF0F00F0F)
+        n16_weight |= ((new_qweight & torch.int32(0x000000F0)) >> 4) << 16
+        n16_weight |= ((new_qweight & torch.int32(0x0000F000)) >> 12) << 24
+        n16_weight |= ((new_qweight & torch.int32(0x000F0000)) >> 16) << 4
+        n16_weight |= ((new_qweight & torch.int32(0x0F000000)) >> 24) << 12
+        return n16_weight.view(torch.int8)
+    elif nbits == 2 and target_dtype == "float16":
+        n8_weight = new_qweight & torch.int32(0xFF0000FF)
+        n8_weight |= ((new_qweight & torch.int32(0x0000FF00)) >> 8) << 16
+        n8_weight |= ((new_qweight & torch.int32(0x00FF0000)) >> 16) << 8
+        return n8_weight.view(torch.int8)
+    elif nbits == 1 and target_dtype == "float16":
+        n8_weight = new_qweight & torch.int32(0xF000000F)
+        n8_weight |= ((new_qweight & torch.int32(0x000000F0)) >> 4) << 8
+        n8_weight |= ((new_qweight & torch.int32(0x00000F00)) >> 8) << 16
+        n8_weight |= ((new_qweight & torch.int32(0x0000F000)) >> 12) << 24
+        n8_weight |= ((new_qweight & torch.int32(0x000F0000)) >> 16) << 4
+        n8_weight |= ((new_qweight & torch.int32(0x00F00000)) >> 20) << 12
+        n8_weight |= ((new_qweight & torch.int32(0x0F000000)) >> 24) << 20
+        return n8_weight.view(torch.int8)
+
+    return new_qweight.view(torch.int8)

tilelang/src/tl_templates/cuda/common.h

@@ -181,6 +181,19 @@ TL_DEVICE void AtomicAddx2(bfloat16_t *address, bfloat16_t *val) {
 }
 #endif
 
+#if (defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ >= 900))
+// AtomicAdd Functions for FLOAT16x2
+TL_DEVICE void AtomicAddx2(float *address, float *val) {
+  atomicAdd(reinterpret_cast<float2 *>(address),
+            static_cast<float2>(*reinterpret_cast<float2 *>(val)));
+}
+// AtomicAdd Functions for FLOAT16x4
+TL_DEVICE void AtomicAddx4(float *address, float *val) {
+  atomicAdd(reinterpret_cast<float4 *>(address),
+            static_cast<float4>(*reinterpret_cast<float4 *>(val)));
+}
+#endif
+
 // DP4A
 template <typename InDatatype, typename OutDatatype>
 TL_DEVICE void DP4A(InDatatype *a, InDatatype *b, OutDatatype *c) {
@@ -210,4 +223,14 @@ template <typename T> TL_DEVICE bool All(T *a, int size) {
   }
   return true;
 }
+
+// Pow of int
+template <int y = 1, typename T> TL_DEVICE T pow_of_int(T x) {
+  T result = x;
+  for (int i = 1; i < y; i++) {
+    result *= x;
+  }
+  return result;
+}
+
 } // namespace tl

tilelang/src/tl_templates/cuda/cuda_fp8.h

@@ -3,21 +3,50 @@
 #pragma once
 
 #include <cute/numeric/numeric_types.hpp>
+
 using fp8_e4_t = cute::float_e4m3_t;
-using fp8_e4_2_t = __nv_fp8x2_e4m3;
-using fp8_e4_4_t = __nv_fp8x4_e4m3;
-struct fp8_e4_8_t {
-  fp8_e4_t data[8];
+using fp8_e5_t = cute::float_e5m2_t;
+
+struct __CUDA_ALIGN__(2) fp8_e4_2_t {
+  fp8_e4_t x;
+  fp8_e4_t y;
+};
+
+struct __CUDA_ALIGN__(4) fp8_e4_4_t {
+  fp8_e4_t x;
+  fp8_e4_t y;
+  fp8_e4_t z;
+  fp8_e4_t w;
 };
-struct fp8_e4_16_t {
-  fp8_e4_t data[16];
+
+struct __CUDA_ALIGN__(8) fp8_e4_8_t {
+  fp8_e4_4_t x;
+  fp8_e4_4_t y;
 };
-using fp8_e5_t = cute::float_e5m2_t;
-using fp8_e5_2_t = __nv_fp8x2_e5m2;
-using fp8_e5_4_t = __nv_fp8x4_e5m2;
-struct fp8_e5_8_t {
-  fp8_e5_t data[8];
+
+struct __CUDA_ALIGN__(16) fp8_e4_16_t {
+  fp8_e4_8_t x;
+  fp8_e4_8_t y;
 };
-struct fp8_e5_16_t {
-  fp8_e5_t data[16];
+
+struct __CUDA_ALIGN__(2) fp8_e5_2_t {
+  fp8_e5_t x;
+  fp8_e5_t y;
+};
+
+struct __CUDA_ALIGN__(4) fp8_e5_4_t {
+  fp8_e5_t x;
+  fp8_e5_t y;
+  fp8_e5_t z;
+  fp8_e5_t w;
+};
+
+struct __CUDA_ALIGN__(8) fp8_e5_8_t {
+  fp8_e5_4_t x;
+  fp8_e5_4_t y;
+};
+
+struct __CUDA_ALIGN__(16) fp8_e5_16_t {
+  fp8_e5_8_t x;
+  fp8_e5_8_t y;
 };

tilelang/src/tl_templates/cuda/debug.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "./cuda_fp8.h"
 #include "common.h"
 #include <stdio.h>
 
@@ -78,10 +79,25 @@ __device__ void debug_print_var<double>(const char *msg, double var) {
          threadIdx.z, var);
 }
 
-#pragma once
+// Specialization for fp8_e4_t type
+template <>
+__device__ void debug_print_var<fp8_e4_t>(const char *msg, fp8_e4_t var) {
+  printf(
+      "msg='%s' BlockIdx=(%d, %d, %d), ThreadIdx=(%d, %d, %d): dtype=fp8_e4_t "
+      "value=%f\n",
+      msg, blockIdx.x, blockIdx.y, blockIdx.z, threadIdx.x, threadIdx.y,
+      threadIdx.z, (float)var);
+}
 
-#include "common.h"
-#include <stdio.h>
+// Specialization for fp8_e5_t type
+template <>
+__device__ void debug_print_var<fp8_e5_t>(const char *msg, fp8_e5_t var) {
+  printf(
+      "msg='%s' BlockIdx=(%d, %d, %d), ThreadIdx=(%d, %d, %d): dtype=fp8_e5_t "
+      "value=%f\n",
+      msg, blockIdx.x, blockIdx.y, blockIdx.z, threadIdx.x, threadIdx.y,
+      threadIdx.z, (float)var);
+}
 
 // Template declaration for device-side debug printing (buffer only)
 template <typename T>
@@ -175,3 +191,25 @@ __device__ void debug_print_buffer_value<double>(const char *msg,
          msg, blockIdx.x, blockIdx.y, blockIdx.z, threadIdx.x, threadIdx.y,
          threadIdx.z, buf_name, index, var);
 }
+
+// Specialization for fp8_e4_t type
+template <>
+__device__ void debug_print_buffer_value<fp8_e4_t>(const char *msg,
+                                                   const char *buf_name,
+                                                   int index, fp8_e4_t var) {
+  printf("msg='%s' BlockIdx=(%d, %d, %d), ThreadIdx=(%d, %d, %d): buffer=%s, "
+         "index=%d, dtype=fp8_e4_t value=%f\n",
+         msg, blockIdx.x, blockIdx.y, blockIdx.z, threadIdx.x, threadIdx.y,
+         threadIdx.z, buf_name, index, (float)var);
+}
+
+// Specialization for fp8_e5_t type
+template <>
+__device__ void debug_print_buffer_value<fp8_e5_t>(const char *msg,
+                                                   const char *buf_name,
+                                                   int index, fp8_e5_t var) {
+  printf("msg='%s' BlockIdx=(%d, %d, %d), ThreadIdx=(%d, %d, %d): buffer=%s, "
+         "index=%d, dtype=fp8_e5_t value=%f\n",
+         msg, blockIdx.x, blockIdx.y, blockIdx.z, threadIdx.x, threadIdx.y,
+         threadIdx.z, buf_name, index, (float)var);
+}

tilelang/src/tl_templates/cuda/gemm_sm80.h

@@ -58,6 +58,23 @@ struct DispatchInstruction<half_t, half_t, float, num_warp_m, num_warp_n, N> {
 };
 #endif
 
+template <int N, int num_warp_n, bool transpose> struct SelectCopy {
+  static constexpr int remainder = (N / num_warp_n) % 16;
+  using type = std::conditional_t<
+      remainder == 4 || remainder == 8 || remainder == 0,
+      std::conditional_t<
+          transpose,
+          std::conditional_t<
+              remainder == 4, SM75_U32x1_LDSM_N,
+              std::conditional_t<remainder == 8, SM75_U32x2_LDSM_N,
+                                 SM75_U32x4_LDSM_N>>,
+          std::conditional_t<
+              remainder == 4, SM75_U16x2_LDSM_T,
+              std::conditional_t<remainder == 8, SM75_U16x4_LDSM_T,
+                                 SM75_U16x8_LDSM_T>>>,
+      DefaultCopy>;
+};
+
 template <int Bits, int N, int K, bool K_inner, int num_warp_n,
           typename Enable = void>
 struct OperandTraits {
@@ -77,8 +94,7 @@ struct OperandTraits<16, N, K, true, num_warp_n,
   using LayoutAtom = decltype(composition(
       Swizzle<2, 3, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{}));
   using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}));
-  using Copy = typename std::conditional<N == 8 * num_warp_n, SM75_U32x2_LDSM_N,
-                                         SM75_U32x4_LDSM_N>::type;
+  using Copy = typename SelectCopy<N, num_warp_n, true>::type;
 };
 
 template <int N, int K, int num_warp_n>
@@ -87,8 +103,7 @@ struct OperandTraits<16, N, K, true, num_warp_n,
   using LayoutAtom = decltype(composition(
       Swizzle<3, 3, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{}));
   using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}));
-  using Copy = typename std::conditional<N == 8 * num_warp_n, SM75_U32x2_LDSM_N,
-                                         SM75_U32x4_LDSM_N>::type;
+  using Copy = typename SelectCopy<N, num_warp_n, true>::type;
 };
 
 template <int N, int K, int num_warp_n>
@@ -98,7 +113,7 @@ struct OperandTraits<16, N, K, false, num_warp_n,
       Swizzle<2, 3, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{}));
   using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{},
                                         Step<_2, _1>{}));
-  using Copy = SM75_U16x8_LDSM_T;
+  using Copy = typename SelectCopy<N, num_warp_n, false>::type;
 };
 
 template <int N, int K, int num_warp_n>
@@ -108,7 +123,7 @@ struct OperandTraits<16, N, K, false, num_warp_n,
       Swizzle<3, 3, 3>{}, Layout<Shape<_64, _8>, Stride<_1, _64>>{}));
   using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{},
                                         Step<_2, _1>{}));
-  using Copy = SM75_U16x8_LDSM_T;
+  using Copy = typename SelectCopy<N, num_warp_n, false>::type;
 };
 
 template <int N, int K, int num_warp_n>
@@ -117,8 +132,7 @@ struct OperandTraits<32, N, K, true, num_warp_n,
   using LayoutAtom = decltype(composition(
       Swizzle<3, 2, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{}));
   using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}));
-  using Copy = typename std::conditional<N == 8 * num_warp_n, SM75_U32x2_LDSM_N,
-                                         SM75_U32x4_LDSM_N>::type;
+  using Copy = typename SelectCopy<N, num_warp_n, true>::type;
 };
 
 template <int N, int K, int num_warp_n>
@@ -127,8 +141,7 @@ struct OperandTraits<32, N, K, true, num_warp_n,
   using LayoutAtom = decltype(composition(
       Swizzle<2, 2, 3>{}, Layout<Shape<_8, _16>, Stride<_16, _1>>{}));
   using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}));
-  using Copy = typename std::conditional<N == 8 * num_warp_n, SM75_U32x2_LDSM_N,
-                                         SM75_U32x4_LDSM_N>::type;
+  using Copy = typename SelectCopy<N, num_warp_n, true>::type;
 };
 
 template <int N, int K, int num_warp_n>
@@ -157,7 +170,7 @@ struct OperandTraits<8, N, K, true, num_warp_n,
   using LayoutAtom = decltype(composition(
       Swizzle<2, 4, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{}));
   using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}));
-  using Copy = SM75_U32x4_LDSM_N;
+  using Copy = typename SelectCopy<N, num_warp_n, true>::type;
 };
 
 template <int N, int K, int num_warp_n>
@@ -166,8 +179,7 @@ struct OperandTraits<8, N, K, true, num_warp_n,
   using LayoutAtom = decltype(composition(
       Swizzle<3, 4, 3>{}, Layout<Shape<_8, _128>, Stride<_128, _1>>{}));
   using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}));
-  using Copy = typename std::conditional<N == 8 * num_warp_n, SM75_U32x2_LDSM_N,
-                                         SM75_U32x4_LDSM_N>::type;
+  using Copy = typename SelectCopy<N, num_warp_n, true>::type;
 };
 
 template <int N, int K, int num_warp_n>