chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,478 @@
+# 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.
+# pylint: disable=invalid-name, unused-wildcard-import, wildcard-import, pointless-exception-statement
+# ruff: noqa: E501, F403, F405
+"""Generator for CUTLASS GEMM kernels."""
+
+from .library import *
+
+
+class GemmOperation:
+    """Describes various attributes for instantiating GEMM kernels."""
+
+    def __init__(
+        self,
+        arch,
+        tile_description,
+        A,
+        B,
+        C,
+        element_epilogue,
+        epilogue_functor=EpilogueFunctor.LinearCombination,
+        swizzling_functor=SwizzlingFunctor.Identity8,
+    ):
+        self.operation_kind = OperationKind.Gemm
+        self.arch = arch
+        self.tile_description = tile_description
+        self.A = A
+        self.B = B
+        self.C = C
+        self.element_epilogue = element_epilogue
+        self.epilogue_functor = epilogue_functor
+        self.swizzling_functor = swizzling_functor
+
+    def accumulator_type(self):
+        return self.tile_description.math_instruction.element_accumulator
+
+    def short_math_name(self):
+        return ShortDataTypeNames[self.accumulator_type()]
+
+    def core_name(self):
+        """The basic operation kind is prefixed with a letter indicating the accumulation type."""
+        inst_shape = ""
+        intermediate_type = ""
+
+        if (
+            self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp
+            or self.tile_description.math_instruction.opcode_class == OpcodeClass.WmmaTensorOp
+        ):
+            inst_shape = "{}{}{}".format(*self.tile_description.math_instruction.instruction_shape)
+            if (
+                self.tile_description.math_instruction.element_a != self.A.element
+                and self.tile_description.math_instruction.element_a
+                != self.tile_description.math_instruction.element_accumulator
+            ):
+                intermediate_type = DataTypeNames[self.tile_description.math_instruction.element_a]
+
+        return f"{self.short_math_name()}{inst_shape}{intermediate_type}gemm"
+
+    def extended_name(self):
+        """Append data types if they differ from compute type."""
+        if (
+            self.C.element != self.tile_description.math_instruction.element_accumulator
+            and self.A.element != self.tile_description.math_instruction.element_accumulator
+        ):
+            extended_name = "${element_c}_${core_name}_${element_a}"
+        elif (
+            self.C.element == self.tile_description.math_instruction.element_accumulator
+            and self.A.element != self.tile_description.math_instruction.element_accumulator
+        ):
+            extended_name = "${core_name}_${element_a}"
+        else:
+            extended_name = "${core_name}"
+
+        extended_name = substitute_template(
+            extended_name,
+            {
+                "element_a": DataTypeNames[self.A.element],
+                "element_c": DataTypeNames[self.C.element],
+                "core_name": self.core_name(),
+            },
+        )
+
+        return extended_name
+
+    def layout_name(self):
+        return f"{ShortLayoutTypeNames[self.A.layout]}{ShortLayoutTypeNames[self.B.layout]}"
+
+    def procedural_name(self):
+        """The full procedural name indicates architecture, extended name, tile size,
+        and layout.
+        """
+        threadblock = self.tile_description.procedural_name()
+        opcode_class_name = OpcodeClassNames[self.tile_description.math_instruction.opcode_class]
+
+        return substitute_template(
+            "cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}_align${alignment}",
+            {
+                "opcode_class": opcode_class_name,
+                "extended_name": self.extended_name(),
+                "threadblock": threadblock,
+                "layout": self.layout_name(),
+                "alignment": f"{self.A.alignment}",
+            },
+        )
+
+    def leading_dim(self):
+        """lda, ldb, ldc, according to the leading dimension."""
+        if self.A.layout == LayoutType.RowMajor:
+            lda = "K"
+        elif self.A.layout == LayoutType.ColumnMajor:
+            lda = "M"
+        else:
+            ValueError("The layout of A is not implemented.")
+
+        if self.B.layout == LayoutType.RowMajor:
+            ldb = "N"
+        elif self.B.layout == LayoutType.ColumnMajor:
+            ldb = "K"
+        else:
+            ValueError("The layout of B is not implemented.")
+
+        if self.C.layout == LayoutType.RowMajor:
+            ldc = "N"
+        elif self.C.layout == LayoutType.ColumnMajor:
+            ldc = "M"
+        else:
+            ValueError("The layout of B is not implemented.")
+
+        return substitute_template(
+            "int lda = ${lda_val};\n\tint ldb = ${ldb_val};\n\tint ldc = ${ldc_val};\n",
+            {"lda_val": lda, "ldb_val": ldb, "ldc_val": ldc},
+        )
+
+
+class EmitGemmInstance:
+    """Responsible for emitting a CUTLASS template definition."""
+
+    def __init__(self):
+        self.epilogue_default = """
+    ${epilogue_functor}<
+      ${element_c},
+      ${epilogue_vector_length},
+      ${element_accumulator},
+      ${element_epilogue}
+    >"""
+
+        self.epilogue_no_beta_scaling = """
+    ${epilogue_functor}<
+      ${element_c},
+      ${epilogue_vector_length},
+      ${element_accumulator},
+      ${element_epilogue},
+      cutlass::epilogue::thread::ScaleType::NoBetaScaling
+    >"""
+
+        self.epilogue_residual_block = """
+    ${epilogue_functor}<
+      ${element_c},
+      ${element_accumulator},
+      ${element_epilogue},
+      ${element_c},
+      ${epilogue_vector_length},
+      ${activation},
+      ${binary_op},
+      ${unary_op}
+    >"""
+
+        self.gemm_template = """
+  // Gemm operator ${operation_name}
+  using Operation_${operation_name} = cutlass::gemm::device::${kernel_name}<
+    ${element_a}, ${layout_a},
+    ${element_b}, ${layout_b},
+    ${element_c}, ${layout_c},
+    ${element_accumulator},
+    ${opcode_class},
+    ${arch},
+    cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
+    cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,
+    cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
+    ${epilogue},
+    ${swizzling_functor},
+    ${stages},
+    ${align_a},
+    ${align_b}
+  >;
+"""
+
+    def emit(self, operation, no_beta_scaling=False, batched=False, residual_block_info=False):
+        """Instantiate a GEMM kernel from given `operation`."""
+        warp_shape = [
+            operation.tile_description.threadblock_shape[idx]
+            // operation.tile_description.warp_count[idx]
+            for idx in range(3)
+        ]
+        epilogue_vector_length = (
+            min(operation.C.alignment * DataTypeSize[operation.C.element], 128)
+            // DataTypeSize[operation.C.element]
+        )
+        values = {
+            "operation_name": operation.procedural_name(),
+            "element_a": DataTypeTag[operation.A.element],
+            "layout_a": LayoutTag[operation.A.layout],
+            "element_b": DataTypeTag[operation.B.element],
+            "layout_b": LayoutTag[operation.B.layout],
+            "element_c": DataTypeTag[operation.C.element],
+            "layout_c": LayoutTag[operation.C.layout],
+            "element_accumulator": DataTypeTag[operation.accumulator_type()],
+            "opcode_class": OpcodeClassTag[
+                operation.tile_description.math_instruction.opcode_class
+            ],
+            "arch": f"cutlass::arch::Sm{operation.arch}",
+            "threadblock_shape_m": str(operation.tile_description.threadblock_shape[0]),
+            "threadblock_shape_n": str(operation.tile_description.threadblock_shape[1]),
+            "threadblock_shape_k": str(operation.tile_description.threadblock_shape[2]),
+            "warp_shape_m": str(warp_shape[0]),
+            "warp_shape_n": str(warp_shape[1]),
+            "warp_shape_k": str(warp_shape[2]),
+            "instruction_shape_m": str(
+                operation.tile_description.math_instruction.instruction_shape[0]
+            ),
+            "instruction_shape_n": str(
+                operation.tile_description.math_instruction.instruction_shape[1]
+            ),
+            "instruction_shape_k": str(
+                operation.tile_description.math_instruction.instruction_shape[2]
+            ),
+            "epilogue_vector_length": str(epilogue_vector_length),
+            "element_epilogue": str(DataTypeTag[operation.element_epilogue]),
+            "epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],
+            "swizzling_functor": SwizzlingFunctorTag[operation.swizzling_functor],
+            "stages": str(operation.tile_description.stages),
+            "align_a": str(operation.A.alignment),
+            "align_b": str(operation.B.alignment),
+            "math_operation": MathOperationTag[
+                operation.tile_description.math_instruction.math_operation
+            ],
+        }
+
+        values["kernel_name"] = "GemmBatched" if batched else "Gemm"
+
+        if residual_block_info:
+            values["kernel_name"] = "GemmUniversalWithBroadcast"
+            template = substitute_template(
+                self.gemm_template, {"epilogue": self.epilogue_residual_block}
+            )
+            values.update(
+                {
+                    "unary_op": residual_block_info["unary_op"],
+                    "binary_op": residual_block_info["binary_op"],
+                    "activation": residual_block_info["activation"],
+                }
+            )
+        elif no_beta_scaling:
+            template = substitute_template(
+                self.gemm_template, {"epilogue": self.epilogue_no_beta_scaling}
+            )
+        else:
+            template = substitute_template(self.gemm_template, {"epilogue": self.epilogue_default})
+
+        return substitute_template(template, values)
+
+
+def instantiate_gemm_template(attrs):
+    """Return CUTLASS host code for GEMM based on a template and the provided attribute map."""
+
+    argument_template_default = """
+  typename ${kernel}::Arguments arguments{
+   problem_size,
+   {static_cast<ElementInputA*>(ptr_a), ${lda}}, ${batch_stride_A}
+   {static_cast<ElementInputB*>(ptr_b), ${ldb}}, ${batch_stride_B}
+   {static_cast<ElementOutput*>(${ptr_c}), ${c_stride}}, ${batch_stride_C}
+   {static_cast<ElementOutput*>(ptr_out), ${ldc}}, ${batch_stride_D}
+   {${alpha_beta}},
+   ${split_k_slices_or_batch}
+  };
+    """
+
+    # See cutlass/gemm/kernel/gemm_with_fused_epilogue.h
+    argument_template_residual = """
+  typename ${kernel}::Arguments arguments{
+    cutlass::gemm::GemmUniversalMode::${gemm_universal_mode},
+    problem_size,
+    ${split_k_slices_or_batch}, // batch_count
+    {${alpha_beta}},
+    static_cast<ElementInputA*>(ptr_a),
+    static_cast<ElementInputB*>(ptr_b),
+    static_cast<ElementOutput*>(ptr_residual),
+    static_cast<ElementOutput*>(ptr_out),
+    static_cast<ElementOutput*>(ptr_bias),
+    nullptr, // ptr_Tensor
+    ${batch_stride_A}
+    ${batch_stride_B}
+    ${batch_stride_C}
+    ${batch_stride_D}
+    0, // batch_stride_Vector,
+    0, // batch_stride_Tensor,
+    ${lda},
+    ${ldb},
+    ${ldc},
+    ${ldc},
+    0, // ldv, the stride for bias
+    0, // ldt
+  };
+    """
+
+    template = """
+  using ElementInputA = ${ElementInputA};
+  using ElementInputB = ${ElementInputB};
+  using ElementOutput = ${ElementOutput};
+  using ElementComputeEpilogue = ${ElementOutput};
+
+  ${cutlass_op_def}
+
+  using ${kernel} = Operation_${cutlass_op_name};
+  int M = ${M};
+  int N = ${N};
+  int K = ${K};
+  cutlass::gemm::GemmCoord problem_size(M, N, K);
+  ElementComputeEpilogue alpha = ElementComputeEpilogue(1);
+  ElementComputeEpilogue beta = ElementComputeEpilogue(${beta});
+  void* ptr_a = (void*)(${lhs_arg}->data);
+  void* ptr_b = (void*)(${rhs_arg}->data);
+  ${bias_decl}
+  ${residual_decl}
+  void* ptr_out = (void*)(out0->data);
+
+  ${argument}
+  size_t workspace_size = ${kernel}::get_workspace_size(arguments);
+  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
+  ${kernel} gemm_op;
+  cutlass::Status status = gemm_op.can_implement(arguments);
+  TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
+  status = gemm_op.initialize(arguments, workspace.get());
+  TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
+
+  cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${A_arg}->device.device_id));
+
+  status = gemm_op(stream);
+  TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
+"""
+    op_type = attrs["op_type"]
+    has_bias = "bias" in op_type
+    is_gelu = "gelu" in op_type
+    batched = "batch" in attrs
+    has_residual_block = "residual" in op_type
+    aux_map = {"kernel": "Gemm"}
+
+    if has_bias:
+        aux_map.update(
+            {
+                "bias_decl": "void* ptr_bias = (void*)(${bias_arg}->data);\n",
+                "ptr_c": "ptr_bias",
+                "c_stride": (
+                    "(${bias_arg}->ndim == 1 ||"
+                    " ${bias_arg}->shape[${bias_arg}->ndim - 2] == 1) ? 0 : " + attrs["ldc"]
+                ),
+            }
+        )
+    else:
+        aux_map.update({"bias_decl": "", "ptr_c": "ptr_out", "c_stride": attrs["ldc"]})
+
+    if is_gelu or has_residual_block:
+        # GeLU epilogue does not compile with NoBetaScaling, so we explicitly specify the scale.
+        aux_map["beta"] = 1
+    else:
+        aux_map["beta"] = 0
+
+    if has_bias and not is_gelu and not has_residual_block:
+        aux_map["alpha_beta"] = "alpha"
+    else:
+        aux_map["alpha_beta"] = "alpha, beta"
+
+    for key in ["batch_stride_A", "batch_stride_B", "batch_stride_C"]:
+        if not batched and not has_residual_block:
+            aux_map[key] = ""
+        else:
+            aux_map[key] = attrs.get(key, "0") + ","
+
+    aux_map["batch_stride_D"] = aux_map["batch_stride_C"]
+    if has_bias and batched and not has_residual_block:
+        aux_map["batch_stride_C"] = "0,"
+
+    if batched:
+        attrs["split_k_slices_or_batch"] = attrs["batch"]
+    else:
+        attrs["split_k_slices_or_batch"] = 1
+
+    if has_residual_block:
+        template = substitute_template(template, {"argument": argument_template_residual})
+        aux_map["residual_decl"] = "void* ptr_residual = (void*)(${residual_arg}->data);\n"
+        aux_map["gemm_universal_mode"] = "kBatched" if batched else "kGemm"
+    else:
+        template = substitute_template(template, {"argument": argument_template_default})
+        aux_map["residual_decl"] = ""
+
+    template = substitute_template(template, aux_map)
+
+    return substitute_template(template, attrs)
+
+
+def emit_fp16A_intB_matmul(attrs):
+    """Return CUTLASS host code for fp16 A and int4 or int8 B GEMM."""
+    if attrs["group_size"] > 0:
+        attrs["quant_op"] = "cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY"
+    else:
+        attrs["quant_op"] = "cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY"
+        attrs["group_size"] = "k"
+
+    attrs["template_common"] = substitute_template(
+        """
+  using namespace fastertransformer;
+  constexpr auto QuantOp = ${quant_op};
+
+  int m = ${M};
+  int n = ${B_arg}->shape[1] * ${float_per_int};
+  int k = ${B_arg}->shape[0];
+
+  cudaStream_t stream = static_cast<cudaStream_t>(
+    TVMFFIEnvGetStream(kDLCUDA, ${A_arg}->device.device_id));
+    """,
+        attrs,
+    )
+
+    template = """
+  ${template_common}
+  gemm_fp16_int_bias_act<${weight_dtype}, QuantOp>(static_cast<cutlass::half_t*>(${A_arg}->data),
+                static_cast<${weight_dtype}*>(${B_arg}->data),
+                static_cast<cutlass::half_t*>(${scales_arg}->data),
+                ${bias},
+                static_cast<cutlass::half_t*>(out0->data),
+                "${activation}",
+                m, n, k, ${group_size}, ${bias_stride}, nullptr, 0, stream);
+"""
+
+    template_residual = """
+  ${template_common}
+  gemm_fp16_int_bias_act_residual<${weight_dtype}, QuantOp>(
+                static_cast<cutlass::half_t*>(${A_arg}->data),
+                static_cast<${weight_dtype}*>(${B_arg}->data),
+                static_cast<cutlass::half_t*>(${scales_arg}->data),
+                ${bias},
+                static_cast<cutlass::half_t*>(${residual_arg}->data),
+                static_cast<cutlass::half_t*>(out0->data),
+                "${activation}", "${binary_op}", "${unary_op}",
+                m, n, k, ${group_size}, nullptr, 0, stream);
+"""
+
+    if "residual_arg" in attrs:
+        if "bias_arg" in attrs:
+            bias = "static_cast<cutlass::half_t*>(${bias_arg}->data)"
+        else:
+            bias = "nullptr"
+
+        template_residual = substitute_template(template_residual, {"bias": bias})
+        return substitute_template(template_residual, attrs)
+
+    if "bias_arg" in attrs:
+        template = substitute_template(
+            template, {"bias": "static_cast<cutlass::half_t*>(${bias_arg}->data)"}
+        )
+    else:
+        template = substitute_template(template, {"bias": "nullptr"})
+
+    return substitute_template(template, attrs)