wehub-resource-sync
909 lines
37 KiB
Python
909 lines
37 KiB
Python
# 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
|
|
# ruff: noqa: F821
|
|
"""Common functions and classes for CUTLASS GEMM and Conv2d geneator."""
|
|
|
|
import logging
|
|
import math
|
|
import multiprocessing
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import tempfile
|
|
|
|
import tvm_ffi
|
|
|
|
from tvm.runtime import Object
|
|
from tvm.tirx import IntImm
|
|
|
|
from . import _ffi_api as ffi
|
|
from .attention_operation import (
|
|
instantiate_attention_template,
|
|
instantiate_flash_attention_template,
|
|
instantiate_flash_attention_var_len_template,
|
|
)
|
|
from .conv2d_operation import instantiate_conv2d_template
|
|
from .gemm_operation import emit_fp16A_intB_matmul, instantiate_gemm_template
|
|
from .layer_norm_operation import instantiate_layer_norm_template
|
|
from .library import (
|
|
DataType,
|
|
DataTypeSize,
|
|
DataTypeTag,
|
|
EpilogueFunctor,
|
|
MathInstruction,
|
|
MathOperation,
|
|
OpcodeClass,
|
|
TileDescription,
|
|
)
|
|
from .rms_norm_operation import instantiate_rms_norm_template
|
|
|
|
logger = logging.getLogger("cutlass")
|
|
|
|
|
|
dtype_map = {
|
|
"int8": DataType.s8,
|
|
"uint8": DataType.u8,
|
|
"int32": DataType.s32,
|
|
"float32": DataType.f32,
|
|
"float16": DataType.f16,
|
|
}
|
|
|
|
|
|
def generate_tensor_op_common(
|
|
math_instructions, alignment_constraints, get_tile_descriptions, op_creator
|
|
):
|
|
"""Common kernel generator to be used by archtecture specific generators."""
|
|
ops = []
|
|
for math_inst in math_instructions:
|
|
tile_descriptions = get_tile_descriptions(math_inst)
|
|
data_type = [
|
|
math_inst.element_a,
|
|
math_inst.element_b,
|
|
math_inst.element_c,
|
|
math_inst.element_accumulator,
|
|
]
|
|
|
|
out = op_creator(tile_descriptions, data_type, alignment_constraints)
|
|
|
|
ops.extend(out)
|
|
|
|
return ops
|
|
|
|
|
|
def generate_sm50_simt(out_dtype, arg0_dtype, arg1_dtype, op_creator, accumulator_dtype="float32"):
|
|
"""Gemerate GEMM or Conv2D SIMT kernels"""
|
|
# pylint: disable=unused-argument
|
|
min_cc = 50
|
|
max_cc = 1024
|
|
if arg0_dtype == "float32" and arg1_dtype == "float32":
|
|
assert out_dtype == "float32" and accumulator_dtype == "float32"
|
|
math_instructions = [
|
|
MathInstruction(
|
|
[1, 1, 1],
|
|
DataType.f32,
|
|
DataType.f32,
|
|
DataType.f32,
|
|
DataType.f32,
|
|
OpcodeClass.Simt,
|
|
MathOperation.multiply_add,
|
|
)
|
|
]
|
|
alignment_constraints = [1]
|
|
tile_descriptions = [
|
|
([128, 128, 8], 2, [4, 2, 1], min_cc, max_cc),
|
|
([128, 64, 8], 2, [2, 2, 1], min_cc, max_cc),
|
|
([64, 128, 8], 2, [2, 2, 1], min_cc, max_cc),
|
|
([64, 64, 8], 2, [2, 1, 1], min_cc, max_cc),
|
|
([128, 32, 8], 2, [2, 1, 1], min_cc, max_cc),
|
|
([32, 128, 8], 2, [1, 2, 1], min_cc, max_cc),
|
|
]
|
|
|
|
def get_tile_descriptions(math_inst):
|
|
return [
|
|
TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
|
|
for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
|
|
]
|
|
|
|
return generate_tensor_op_common(
|
|
math_instructions, alignment_constraints, get_tile_descriptions, op_creator
|
|
)
|
|
else:
|
|
raise NotImplementedError()
|
|
|
|
|
|
def generate_sm75_tensor_op_1688(
|
|
out_dtype,
|
|
arg0_dtype,
|
|
arg1_dtype,
|
|
op_creator,
|
|
check_align,
|
|
_,
|
|
profile_all_alignments=False,
|
|
accumlator_dtype="float32",
|
|
):
|
|
"""Generate GEMM or Conv2D kernels for Turing."""
|
|
assert out_dtype in ["float32", "float16", "int32"]
|
|
min_cc = 75
|
|
max_cc = 1024
|
|
|
|
if arg0_dtype == "float16" and arg1_dtype == "float16":
|
|
math_instructions = [
|
|
MathInstruction(
|
|
[16, 8, 8],
|
|
DataType.f16,
|
|
DataType.f16,
|
|
dtype_map[out_dtype],
|
|
dtype_map[accumlator_dtype],
|
|
OpcodeClass.TensorOp,
|
|
MathOperation.multiply_add,
|
|
)
|
|
]
|
|
alignment_constraints = [8, 4, 2, 1]
|
|
tile_descriptions = [
|
|
([256, 128, 32], 2, [4, 2, 1], min_cc, max_cc),
|
|
([128, 256, 32], 2, [2, 4, 1], min_cc, max_cc),
|
|
([128, 128, 32], 2, [2, 2, 1], min_cc, max_cc),
|
|
([64, 128, 32], 2, [2, 2, 1], min_cc, max_cc),
|
|
([128, 64, 32], 2, [2, 2, 1], min_cc, max_cc),
|
|
([64, 64, 32], 2, [2, 2, 1], min_cc, max_cc),
|
|
([64, 128, 64], 2, [1, 2, 2], min_cc, max_cc),
|
|
]
|
|
|
|
elif "int8" in arg0_dtype and "int8" in arg1_dtype:
|
|
assert out_dtype == "int32"
|
|
math_instructions = [
|
|
MathInstruction(
|
|
[8, 8, 16],
|
|
dtype_map[arg0_dtype],
|
|
dtype_map[arg1_dtype],
|
|
DataType.s32,
|
|
DataType.s32,
|
|
OpcodeClass.TensorOp,
|
|
MathOperation.multiply_add_saturate,
|
|
)
|
|
]
|
|
alignment_constraints = [16, 8, 4, 2, 1]
|
|
tile_descriptions = [
|
|
([256, 128, 64], 2, [4, 2, 1], min_cc, max_cc),
|
|
([128, 256, 64], 2, [2, 4, 1], min_cc, max_cc),
|
|
([128, 128, 64], 2, [2, 2, 1], min_cc, max_cc),
|
|
([64, 256, 64], 2, [1, 4, 1], min_cc, max_cc),
|
|
([256, 64, 64], 2, [4, 1, 1], min_cc, max_cc),
|
|
([64, 128, 64], 2, [2, 2, 1], min_cc, max_cc),
|
|
([128, 64, 64], 2, [2, 2, 1], min_cc, max_cc),
|
|
([64, 64, 64], 2, [2, 2, 1], min_cc, max_cc),
|
|
]
|
|
elif arg0_dtype == "float32" and arg1_dtype == "float32" and out_dtype == "float32":
|
|
return generate_sm50_simt(out_dtype, arg0_dtype, arg1_dtype, op_creator, accumlator_dtype)
|
|
else:
|
|
raise NotImplementedError()
|
|
|
|
alignment_constraints = [align for align in alignment_constraints if check_align(align)]
|
|
assert len(alignment_constraints) > 0
|
|
|
|
if not profile_all_alignments:
|
|
alignment_constraints = [alignment_constraints[0]]
|
|
|
|
def get_tile_descriptions(math_inst):
|
|
return [
|
|
TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
|
|
for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
|
|
]
|
|
|
|
return generate_tensor_op_common(
|
|
math_instructions, alignment_constraints, get_tile_descriptions, op_creator
|
|
)
|
|
|
|
|
|
def generate_sm80_tensor_op_16816(
|
|
out_dtype,
|
|
arg0_dtype,
|
|
arg1_dtype,
|
|
op_creator,
|
|
check_align,
|
|
use_3xtf32=True,
|
|
profile_all_alignments=False,
|
|
accumlator_dtype="float32",
|
|
):
|
|
"""Generate GEMM or Conv2D kernels for Ampere."""
|
|
min_cc = 80
|
|
max_cc = 1024
|
|
max_cc_smem_limited = 80
|
|
|
|
def get_default_tile_descriptions(block_k_factor):
|
|
return [
|
|
([128, 256, int(32 * block_k_factor)], 3, [2, 4, 1], min_cc, max_cc),
|
|
([256, 128, int(32 * block_k_factor)], 3, [4, 2, 1], min_cc, max_cc),
|
|
([256, 64, int(32 * block_k_factor)], 3, [4, 1, 1], min_cc, max_cc),
|
|
([256, 64, int(32 * block_k_factor)], 4, [4, 1, 1], min_cc, max_cc),
|
|
([64, 256, int(32 * block_k_factor)], 4, [1, 4, 1], min_cc, max_cc),
|
|
([128, 128, int(32 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
|
|
([128, 128, int(32 * block_k_factor)], 4, [2, 2, 1], min_cc, max_cc),
|
|
([128, 128, int(32 * block_k_factor)], 5, [2, 2, 1], min_cc, max_cc),
|
|
([128, 64, int(32 * block_k_factor)], 6, [2, 2, 1], min_cc, max_cc),
|
|
([64, 128, int(32 * block_k_factor)], 6, [2, 2, 1], min_cc, max_cc),
|
|
([64, 64, int(32 * block_k_factor)], 10, [2, 2, 1], min_cc, max_cc),
|
|
([256, 128, int(64 * block_k_factor)], 3, [4, 2, 1], min_cc, max_cc_smem_limited),
|
|
([128, 256, int(64 * block_k_factor)], 3, [2, 4, 1], min_cc, max_cc_smem_limited),
|
|
([256, 64, int(64 * block_k_factor)], 4, [4, 1, 1], min_cc, max_cc_smem_limited),
|
|
([64, 256, int(64 * block_k_factor)], 4, [1, 4, 1], min_cc, max_cc_smem_limited),
|
|
([128, 128, int(64 * block_k_factor)], 4, [2, 2, 1], min_cc, max_cc),
|
|
([256, 64, int(64 * block_k_factor)], 3, [4, 1, 1], min_cc, max_cc),
|
|
([64, 256, int(64 * block_k_factor)], 3, [1, 4, 1], min_cc, max_cc),
|
|
([128, 128, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
|
|
([128, 64, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
|
|
([64, 128, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
|
|
([64, 64, int(64 * block_k_factor)], 5, [2, 2, 1], min_cc, max_cc),
|
|
]
|
|
|
|
if arg0_dtype == "float16" and arg1_dtype == "float16":
|
|
math_instructions = [
|
|
MathInstruction(
|
|
[16, 8, 16],
|
|
DataType.f16,
|
|
DataType.f16,
|
|
dtype_map[out_dtype],
|
|
dtype_map[accumlator_dtype],
|
|
OpcodeClass.TensorOp,
|
|
MathOperation.multiply_add,
|
|
)
|
|
]
|
|
alignment_constraints = [8, 4, 2]
|
|
tile_descriptions = get_default_tile_descriptions(1)
|
|
elif arg0_dtype == "float32" and arg1_dtype == "float32":
|
|
math_instructions = [
|
|
MathInstruction(
|
|
[16, 8, 8],
|
|
DataType.f32,
|
|
DataType.f32,
|
|
DataType.f32,
|
|
DataType.f32,
|
|
OpcodeClass.TensorOp,
|
|
MathOperation.multiply_add_fast_f32 if use_3xtf32 else MathOperation.multiply_add,
|
|
)
|
|
]
|
|
alignment_constraints = [4, 2, 1]
|
|
|
|
if use_3xtf32:
|
|
# tf32
|
|
tile_descriptions = [
|
|
([128, 128, 16], 4, [4, 2, 1], min_cc, max_cc),
|
|
([128, 128, 16], 3, [4, 2, 1], min_cc, max_cc),
|
|
([256, 64, 16], 3, [4, 2, 1], min_cc, max_cc),
|
|
([64, 256, 16], 3, [2, 4, 1], min_cc, max_cc),
|
|
([128, 64, 16], 4, [2, 2, 1], min_cc, max_cc),
|
|
([64, 128, 16], 4, [2, 2, 1], min_cc, max_cc),
|
|
([64, 64, 16], 3, [2, 2, 1], min_cc, max_cc),
|
|
([128, 128, 32], 3, [4, 2, 1], min_cc, max_cc),
|
|
([256, 64, 32], 3, [4, 2, 1], min_cc, max_cc_smem_limited),
|
|
([64, 256, 32], 3, [2, 4, 1], min_cc, max_cc_smem_limited),
|
|
([128, 64, 32], 3, [2, 2, 1], min_cc, max_cc),
|
|
([64, 128, 32], 3, [2, 2, 1], min_cc, max_cc),
|
|
([64, 64, 32], 3, [2, 2, 1], min_cc, max_cc),
|
|
]
|
|
else:
|
|
tile_descriptions = get_default_tile_descriptions(0.5)
|
|
else:
|
|
assert out_dtype == "int32"
|
|
math_instructions = [
|
|
MathInstruction(
|
|
[16, 8, 32],
|
|
dtype_map[arg0_dtype],
|
|
dtype_map[arg1_dtype],
|
|
DataType.s32,
|
|
DataType.s32,
|
|
OpcodeClass.TensorOp,
|
|
MathOperation.multiply_add_saturate,
|
|
)
|
|
]
|
|
alignment_constraints = [16, 8, 4]
|
|
tile_descriptions = get_default_tile_descriptions(2)
|
|
|
|
def get_tile_descriptions(math_inst):
|
|
return [
|
|
TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
|
|
for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
|
|
]
|
|
|
|
alignment_constraints = [align for align in alignment_constraints if check_align(align)]
|
|
|
|
if len(alignment_constraints) > 0 and not profile_all_alignments:
|
|
alignment_constraints = [alignment_constraints[0]]
|
|
|
|
if arg0_dtype != "float32" and arg1_dtype != "float32":
|
|
sm75_kernels = generate_sm75_tensor_op_1688(
|
|
out_dtype,
|
|
arg0_dtype,
|
|
arg1_dtype,
|
|
op_creator,
|
|
check_align,
|
|
False,
|
|
profile_all_alignments,
|
|
accumlator_dtype=accumlator_dtype,
|
|
)
|
|
else:
|
|
# TF32 (float32 + float32 case) is only supported on sm80
|
|
sm75_kernels = []
|
|
|
|
if len(alignment_constraints) > 0:
|
|
sm80_kernels = generate_tensor_op_common(
|
|
math_instructions, alignment_constraints, get_tile_descriptions, op_creator
|
|
)
|
|
else:
|
|
sm80_kernels = []
|
|
|
|
# TODO(masahi): For int8 kernels, The CUTLASS generator modifies the output tensor alignment
|
|
# after ops are created. Revisit how important this modification is.
|
|
# for op in operations:
|
|
# if op.tile_description.threadblock_shape[1] >= 128:
|
|
# op.C.alignment = 16
|
|
# else:
|
|
# op.C.alignment = 8
|
|
|
|
return sm75_kernels + sm80_kernels
|
|
|
|
|
|
GENERATOR_FUNC_TABLE = {75: generate_sm75_tensor_op_1688, 80: generate_sm80_tensor_op_16816}
|
|
|
|
|
|
# (Epilogue functor name, no_beta_scaling)
|
|
EPILOGUE_MAP = {
|
|
"cutlass.dense": (EpilogueFunctor.LinearCombination, False),
|
|
"cutlass.dense_bias": (EpilogueFunctor.LinearCombinationBias, True),
|
|
"cutlass.dense_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
|
|
"cutlass.dense_bias_gelu_fp16": (EpilogueFunctor.LinearCombinationGelu, False),
|
|
"cutlass.dense_bias_gelu_fp32": (EpilogueFunctor.LinearCombinationGelu, False),
|
|
"cutlass.matmul": (EpilogueFunctor.LinearCombination, False),
|
|
"cutlass.matmul_bias": (EpilogueFunctor.LinearCombinationBias, True),
|
|
"cutlass.matmul_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
|
|
"cutlass.matmul_bias_gelu": (EpilogueFunctor.LinearCombinationGelu, False),
|
|
"cutlass.matmul_transposed": (EpilogueFunctor.LinearCombination, False),
|
|
"cutlass.matmul_transposed_bias": (EpilogueFunctor.LinearCombinationBias, True),
|
|
"cutlass.matmul_transposed_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
|
|
"cutlass.matmul_transposed_bias_gelu": (EpilogueFunctor.LinearCombinationGelu, False),
|
|
"cutlass.batch_matmul": (EpilogueFunctor.LinearCombination, False),
|
|
"cutlass.conv2d_bias_hardswish": (EpilogueFunctor.LinearCombinationHardSwish, False),
|
|
"cutlass.conv2d_bias_silu": (EpilogueFunctor.LinearCombinationSilu, False),
|
|
"cutlass.conv2d_bias_sigmoid": (EpilogueFunctor.LinearCombinationSigmoid, False),
|
|
"cutlass.conv2d_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
|
|
"cutlass.conv2d_bias": (EpilogueFunctor.LinearCombinationBias, True),
|
|
"cutlass.conv2d": (EpilogueFunctor.LinearCombination, False),
|
|
"cutlass.conv2d_transpose": (EpilogueFunctor.LinearCombination, False),
|
|
"cutlass.conv2d_backward_weight": (EpilogueFunctor.LinearCombination, False),
|
|
}
|
|
|
|
|
|
class ProfilerEngine:
|
|
"""Compile and run a given profiler executable."""
|
|
|
|
def __init__(self, cuda_arch, cutlass_path, binary_prefix):
|
|
self.cuda_arch = cuda_arch
|
|
self.binary_prefix = binary_prefix
|
|
self.cutlass = cutlass_path
|
|
self.cflags = f"-I{cutlass_path}/include -I{cutlass_path}/tools/util/include -O3 -std=c++17"
|
|
self.cflags += " -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1"
|
|
self.cflags += (
|
|
f" -gencode=arch=compute_{cuda_arch},code=[sm_{cuda_arch},compute_{cuda_arch}]"
|
|
)
|
|
self.cflags += " -Xcompiler=-Wconversion -Xcompiler=-fno-strict-aliasing"
|
|
self.cmd = "nvcc {cflags} {src} -o {output}"
|
|
|
|
def _compile(self, op):
|
|
os.makedirs(self.binary_prefix, exist_ok=True)
|
|
opath = os.path.join(self.binary_prefix, op["name"])
|
|
if os.path.exists(opath):
|
|
return
|
|
fi = tempfile.NamedTemporaryFile("w", delete=False, prefix=self.binary_prefix, suffix=".cu")
|
|
fi.write(op["src"])
|
|
fi.close()
|
|
cmd = self.cmd.format(cflags=self.cflags, src=fi.name, output=opath)
|
|
logger.info("invoking compilation %s", cmd)
|
|
os.system(cmd)
|
|
os.unlink(fi.name)
|
|
|
|
def compile_all(self, ops, use_multiprocessing=False):
|
|
"""Compile all profiler executables."""
|
|
if use_multiprocessing:
|
|
pool = multiprocessing.Pool(multiprocessing.cpu_count())
|
|
pool.map(self._compile, ops)
|
|
else:
|
|
for op in ops:
|
|
self._compile(op)
|
|
|
|
def evaluate(self, op, args):
|
|
"""Run the profiler executable corresponding to op_name with args."""
|
|
op_name = op["name"]
|
|
opath = os.path.join(self.binary_prefix, op_name)
|
|
if not os.path.exists(opath):
|
|
self._compile(op)
|
|
if not os.path.exists(opath):
|
|
# Bail out if compilation fails for a whatever reason (e.g. static assert failure)
|
|
return float("inf")
|
|
cmd = [opath]
|
|
for arg in args:
|
|
cmd.append(str(arg))
|
|
try:
|
|
logger.info("invoking evaluation %s", cmd)
|
|
sp = subprocess.run(cmd, capture_output=True, check=True)
|
|
rt = float(sp.stdout)
|
|
if rt == 0.0:
|
|
# This seems to happen with split-k using invalid split-k-slices
|
|
rt = float("inf")
|
|
logger.info("%s, %f", op_name, rt)
|
|
except subprocess.CalledProcessError:
|
|
rt = float("inf")
|
|
return rt
|
|
|
|
|
|
class CodegenResult(Object):
|
|
"""The holder for the generated code and required headers."""
|
|
|
|
def __init__(self, code, headers):
|
|
self.__init_handle_by_constructor__(ffi.CodegenResult, code, headers)
|
|
|
|
|
|
def _get_optional_int_annotation(annotations, key, default=None):
|
|
value = annotations.get(key, default)
|
|
if value is None:
|
|
return default
|
|
return int(value)
|
|
|
|
|
|
@tvm_ffi.register_global_func("contrib.cutlass.instantiate_template")
|
|
def instantiate_template(func_name, annotations, func_args):
|
|
"""Return CUTLASS host code based on a template and the provided annotations.
|
|
|
|
Parameters
|
|
----------
|
|
func_name: str
|
|
A string to identify the type of the kernel (dense/matmul, batched_matmul, or conv2d).
|
|
|
|
annotations: tvm_ffi.Map
|
|
Key and value pairs annotated during kernel selection.
|
|
|
|
func_args: list
|
|
Names of the function arguments.
|
|
|
|
Returns
|
|
-------
|
|
codegen_result : CodegenResult
|
|
Generated CUTLASS host code and required header-file names.
|
|
"""
|
|
attrs = {}
|
|
|
|
for k in ["lda", "ldb", "ldc", "cutlass_op_def", "cutlass_op_name", "op_type"]:
|
|
if k in annotations:
|
|
attrs[k] = annotations[k]
|
|
|
|
headers = ["tvm/ffi/function.h", "tvm/ffi/extra/c_env_api.h"]
|
|
|
|
if "relu" in func_name:
|
|
headers.append("cutlass/epilogue/thread/linear_combination_bias_relu.h")
|
|
elif "gelu" in func_name:
|
|
headers.append("cutlass/epilogue/thread/linear_combination_gelu.h")
|
|
elif "sigmoid" in func_name:
|
|
headers.append("cutlass/epilogue/thread/linear_combination_sigmoid.h")
|
|
elif "silu" in func_name:
|
|
headers.append("cutlass/epilogue/thread/linear_combination_silu.h")
|
|
elif "hardswish" in func_name:
|
|
headers.append("cutlass/epilogue/thread/linear_combination_hardswish.h")
|
|
else:
|
|
headers.append("cutlass/epilogue/thread/linear_combination.h")
|
|
|
|
if "residual" in func_name:
|
|
headers.append("cutlass/epilogue/thread/linear_combination_residual_block.h")
|
|
|
|
def get_dim(shape_annot, var_name, axis_idx, batched_offset=0):
|
|
if isinstance(shape_annot, IntImm):
|
|
return str(int(shape_annot))
|
|
return f"{var_name}->shape[{batched_offset + axis_idx}]"
|
|
|
|
def get_batch_stride(stride_annot, arg0_idx, arg1_idx, arg0_axis_idx, arg1_axis_idx):
|
|
if isinstance(stride_annot, IntImm):
|
|
return str(int(stride_annot))
|
|
dim1 = func_args[arg0_idx] + f"->shape[{arg0_axis_idx}]"
|
|
dim2 = func_args[arg1_idx] + f"->shape[{arg1_axis_idx}]"
|
|
return dim1 + " * " + dim2
|
|
|
|
def get_flattened_batch_dim(arg_name, batch_rank):
|
|
return " * ".join([f"{arg_name}->shape[{i}]" for i in range(batch_rank)])
|
|
|
|
if "decode_matmul" in func_name:
|
|
headers.append("cutlass_kernels/fpA_intB_gemm.h")
|
|
lhs_arg_idx = _get_optional_int_annotation(annotations, "lhs_arg_idx", 0)
|
|
rhs_arg_idx = _get_optional_int_annotation(annotations, "rhs_arg_idx", 1)
|
|
scales_arg_idx = _get_optional_int_annotation(annotations, "scales_arg_idx", 2)
|
|
bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
|
|
residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)
|
|
|
|
attrs["A_arg"] = func_args[lhs_arg_idx]
|
|
attrs["B_arg"] = func_args[rhs_arg_idx]
|
|
attrs["scales_arg"] = func_args[scales_arg_idx]
|
|
attrs["activation"] = annotations.get("activation", "identity")
|
|
attrs["bias_stride"] = annotations["bias_stride"]
|
|
attrs["M"] = annotations["M"]
|
|
attrs["group_size"] = annotations["group_size"]
|
|
|
|
if not isinstance(attrs["M"], tvm.tirx.IntImm):
|
|
attrs["M"] = get_flattened_batch_dim(
|
|
func_args[lhs_arg_idx], int(annotations["batch_rank"])
|
|
)
|
|
|
|
if bias_arg_idx is not None:
|
|
attrs["bias_arg"] = func_args[bias_arg_idx]
|
|
|
|
if residual_arg_idx is not None:
|
|
attrs["residual_arg"] = func_args[residual_arg_idx]
|
|
attrs["binary_op"] = annotations["binary_op"]
|
|
attrs["unary_op"] = annotations["unary_op"]
|
|
|
|
if annotations["weight_nbit"] == 4:
|
|
attrs["weight_dtype"] = "cutlass::uint4b_t"
|
|
attrs["float_per_int"] = 2
|
|
else:
|
|
assert annotations["weight_nbit"] == 8
|
|
attrs["weight_dtype"] = "uint8_t"
|
|
attrs["float_per_int"] = 1
|
|
|
|
code = emit_fp16A_intB_matmul(attrs)
|
|
return CodegenResult(code, headers)
|
|
|
|
elif "dense" in func_name or "matmul" in func_name:
|
|
batched = "batch" in annotations
|
|
# dense is equal to transposed_matmul
|
|
transposed = "transposed" in func_name or "dense" in func_name
|
|
lhs_arg_idx = _get_optional_int_annotation(annotations, "lhs_arg_idx", 0)
|
|
rhs_arg_idx = _get_optional_int_annotation(annotations, "rhs_arg_idx", 1)
|
|
if "bias" in func_name:
|
|
bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", 2)
|
|
else:
|
|
bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
|
|
residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)
|
|
|
|
lhs_arg = func_args[lhs_arg_idx]
|
|
rhs_arg = func_args[rhs_arg_idx]
|
|
lhs_shape = annotations[f"arg{lhs_arg_idx}_shape"]
|
|
rhs_shape = annotations[f"arg{rhs_arg_idx}_shape"]
|
|
lhs_batched_offset = len(lhs_shape) - 2
|
|
rhs_batched_offset = len(rhs_shape) - 2
|
|
|
|
attrs["lhs_arg"] = lhs_arg
|
|
attrs["rhs_arg"] = rhs_arg
|
|
|
|
if bias_arg_idx is not None:
|
|
attrs["bias_arg"] = func_args[bias_arg_idx]
|
|
if residual_arg_idx is not None:
|
|
attrs["residual_arg"] = func_args[residual_arg_idx]
|
|
|
|
attrs["ElementInputA"] = DataTypeTag[dtype_map[annotations[f"arg{lhs_arg_idx}_dtype"]]]
|
|
attrs["ElementInputB"] = DataTypeTag[dtype_map[annotations[f"arg{rhs_arg_idx}_dtype"]]]
|
|
attrs["ElementOutput"] = DataTypeTag[dtype_map[annotations["ret_dtype"]]]
|
|
|
|
attrs["K"] = lhs_shape[lhs_batched_offset + 1]
|
|
attrs["M"] = get_dim(lhs_shape[lhs_batched_offset], lhs_arg, 0, lhs_batched_offset)
|
|
|
|
if transposed:
|
|
attrs["N"] = get_dim(rhs_shape[rhs_batched_offset], rhs_arg, 0, rhs_batched_offset)
|
|
else:
|
|
attrs["N"] = get_dim(rhs_shape[rhs_batched_offset + 1], rhs_arg, 1, rhs_batched_offset)
|
|
|
|
if batched:
|
|
headers.append("cutlass/gemm/device/gemm_batched.h")
|
|
|
|
def get_batch_on_arg(arg_name, arg_shape):
|
|
return " * ".join(f"{arg_name}->shape[{i}]" for i in range(len(arg_shape) - 2))
|
|
|
|
if isinstance(annotations["batch"], IntImm):
|
|
attrs["batch"] = str(int(annotations["batch"]))
|
|
elif annotations["batch_stride_A"] == 0:
|
|
# 2D x ND
|
|
attrs["batch"] = get_batch_on_arg(rhs_arg, rhs_shape)
|
|
else:
|
|
# ND x 2D or ND x ND
|
|
attrs["batch"] = get_batch_on_arg(lhs_arg, lhs_shape)
|
|
|
|
attrs["batch_stride_A"] = get_batch_stride(
|
|
annotations["batch_stride_A"],
|
|
lhs_arg_idx,
|
|
lhs_arg_idx,
|
|
lhs_batched_offset,
|
|
lhs_batched_offset + 1,
|
|
)
|
|
attrs["batch_stride_B"] = get_batch_stride(
|
|
annotations["batch_stride_B"],
|
|
rhs_arg_idx,
|
|
rhs_arg_idx,
|
|
rhs_batched_offset,
|
|
rhs_batched_offset + 1,
|
|
)
|
|
|
|
if transposed:
|
|
attrs["batch_stride_C"] = get_batch_stride(
|
|
annotations["batch_stride_C"],
|
|
lhs_arg_idx,
|
|
rhs_arg_idx,
|
|
lhs_batched_offset,
|
|
rhs_batched_offset,
|
|
)
|
|
else:
|
|
attrs["batch_stride_C"] = get_batch_stride(
|
|
annotations["batch_stride_C"],
|
|
lhs_arg_idx,
|
|
rhs_arg_idx,
|
|
lhs_batched_offset,
|
|
rhs_batched_offset + 1,
|
|
)
|
|
else:
|
|
headers.append("cutlass/gemm/device/gemm.h")
|
|
|
|
if "residual" in func_name:
|
|
headers.append("cutlass/gemm/device/gemm_universal_with_broadcast.h")
|
|
|
|
code = instantiate_gemm_template(attrs)
|
|
return CodegenResult(code, headers)
|
|
|
|
elif "conv2d" in func_name:
|
|
data_arg_idx = _get_optional_int_annotation(annotations, "data_arg_idx", 0)
|
|
weight_arg_idx = _get_optional_int_annotation(annotations, "weight_arg_idx", 1)
|
|
bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
|
|
residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)
|
|
|
|
attrs["data_arg"] = func_args[data_arg_idx]
|
|
attrs["weight_arg"] = func_args[weight_arg_idx]
|
|
|
|
if bias_arg_idx is not None:
|
|
attrs["bias_arg"] = func_args[bias_arg_idx]
|
|
if residual_arg_idx is not None:
|
|
attrs["residual_arg"] = func_args[residual_arg_idx]
|
|
|
|
activation_shape = annotations[f"arg{data_arg_idx}_shape"]
|
|
weight_shape = annotations[f"arg{weight_arg_idx}_shape"]
|
|
output_shape = annotations["ret_shape"]
|
|
|
|
if "conv2d_transpose" in func_name:
|
|
headers.append("cutlass/conv/kernel/default_conv2d_dgrad.h")
|
|
activation_shape = output_shape
|
|
output_shape = annotations["arg0_shape"]
|
|
elif "backward" in func_name:
|
|
headers.append("cutlass/conv/kernel/default_conv2d_wgrad.h")
|
|
activation_shape = annotations["arg1_shape"]
|
|
weight_shape = output_shape
|
|
output_shape = annotations["arg0_shape"]
|
|
elif "residual" in func_name:
|
|
headers.append("cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h")
|
|
else:
|
|
headers.append("cutlass/conv/kernel/default_conv2d_fprop.h")
|
|
|
|
headers.append("cutlass/conv/device/implicit_gemm_convolution.h")
|
|
|
|
op_name = attrs["cutlass_op_name"]
|
|
|
|
if "splitk" in op_name:
|
|
headers += [
|
|
"cutlass/reduction/device/reduce_split_k.h",
|
|
"cutlass/reduction/thread/reduction_operators.h",
|
|
]
|
|
|
|
data_arg = attrs["data_arg"]
|
|
attrs["N"] = get_dim(activation_shape[0], data_arg, 0)
|
|
attrs["H"] = get_dim(activation_shape[1], data_arg, 1)
|
|
attrs["W"] = get_dim(activation_shape[2], data_arg, 2)
|
|
attrs["C"] = activation_shape[3]
|
|
attrs["P"] = get_dim(output_shape[1], "out0", 1)
|
|
attrs["Q"] = get_dim(output_shape[2], "out0", 2)
|
|
attrs["K"] = output_shape[3]
|
|
attrs["R"] = weight_shape[1]
|
|
attrs["S"] = weight_shape[2]
|
|
attrs["pad_h"] = annotations["padding"][0]
|
|
attrs["pad_w"] = annotations["padding"][1]
|
|
attrs["stride_h"] = annotations["strides"][0]
|
|
attrs["stride_w"] = annotations["strides"][1]
|
|
attrs["dilation_h"] = annotations["dilation"][0]
|
|
attrs["dilation_w"] = annotations["dilation"][1]
|
|
|
|
if "splitk" in op_name:
|
|
attrs["split_k_mode"] = "kParallel"
|
|
attrs["split_k_slices"] = str(re.search(r"splitk(\d+)", op_name).group(1))
|
|
else:
|
|
attrs["split_k_mode"] = "kSerial"
|
|
attrs["split_k_slices"] = 1
|
|
|
|
if "residual_shape" in annotations:
|
|
attrs["residual_shape"] = annotations["residual_shape"]
|
|
|
|
code = instantiate_conv2d_template(attrs)
|
|
return CodegenResult(code, headers)
|
|
|
|
elif "attention" in func_name:
|
|
is_var_len = "var_len" in func_name
|
|
data_type = dtype_map[annotations["arg0_dtype"]]
|
|
|
|
attrs["qkv_layout"] = annotations["qkv_layout"]
|
|
if attrs["qkv_layout"] == "default":
|
|
attrs["query"] = func_args[0]
|
|
attrs["key"] = func_args[1]
|
|
attrs["value"] = func_args[2]
|
|
attrs["num_queries"] = s = get_dim(annotations["num_queries"], func_args[0], 1)
|
|
attrs["num_keys"] = get_dim(annotations["num_keys"], func_args[1], 1)
|
|
if len(func_args) > 4 and not is_var_len: # +1 for workspace, the last arg
|
|
attrs["bias"] = func_args[3]
|
|
elif attrs["qkv_layout"] == "qkv_stacked":
|
|
attrs["qkv"] = func_args[0]
|
|
attrs["num_queries"] = s = annotations["num_queries"]
|
|
attrs["num_keys"] = annotations["num_keys"]
|
|
if len(func_args) > 2 and not is_var_len: # +1 for workspace, the last arg
|
|
attrs["bias"] = func_args[1]
|
|
else:
|
|
raise NotImplementedError()
|
|
|
|
attrs["data_type"] = DataTypeTag[data_type]
|
|
attrs["num_batches"] = b = annotations["num_batches"]
|
|
attrs["head_dim"] = h = annotations["head_dim"]
|
|
attrs["head_dim_value"] = h_v = annotations["head_dim_value"]
|
|
attrs["kMaxK"] = max(int(attrs["head_dim"]), int(attrs["head_dim_value"]))
|
|
attrs["scale"] = (
|
|
float(1 / math.sqrt(h.value)) if annotations["scale"] is None else annotations["scale"]
|
|
)
|
|
|
|
if is_var_len:
|
|
attrs["seqstart_q"] = func_args[int(annotations["seqstart_q_idx"])]
|
|
attrs["seqstart_k"] = func_args[int(annotations["seqstart_k_idx"])]
|
|
attrs["max_seqlen_q"] = func_args[int(annotations["max_seqlen_q_idx"])]
|
|
attrs["max_seqlen_k"] = func_args[int(annotations["max_seqlen_k_idx"])]
|
|
|
|
is_mqa = annotations["num_q_heads"] != annotations["num_kv_heads"]
|
|
|
|
use_flash = (
|
|
annotations["ret_dtype"] == "float16"
|
|
and "bias" not in attrs
|
|
and int(attrs["head_dim"]) <= 256
|
|
and int(attrs["head_dim"]) % 8 == 0
|
|
and int(attrs["head_dim"]) == int(attrs["head_dim_value"])
|
|
# For the causal case (custom mask = "BottomRight"), only use flash for multi-query
|
|
# attention workloads. Otherwise, CUTLASS fMHA seems faster for causal attention
|
|
# with a single query.
|
|
# In addition, sliding-window attention is only supported by flash.
|
|
and (
|
|
int(annotations["custom_mask_type"]) == 0
|
|
or (int(annotations["custom_mask_type"]) == 2 and is_mqa)
|
|
or (int(annotations["custom_mask_type"]) == 2 and "window_size" in annotations)
|
|
)
|
|
# Flash v2 is currently not supported for sm < 80
|
|
and int(annotations["arch"]) >= 80
|
|
)
|
|
|
|
# See https://github.com/Dao-AILab/flash-attention/blob/
|
|
# 92dd5703ecdb99aa4a4aee9817f28557907403a2/csrc/flash_attn/flash_api.cpp#L111-L116
|
|
if "window_size" in annotations:
|
|
assert use_flash, "Sliding-window attention is supported only by Flash Attention."
|
|
assert int(annotations["custom_mask_type"]) == 2, (
|
|
"Sliding-window attention is only supported for causal with bottom right mask."
|
|
)
|
|
attrs["window_size_left"] = int(annotations["window_size"]) - 1
|
|
attrs["window_size_right"] = 0
|
|
attrs["is_causal"] = False
|
|
else:
|
|
if int(annotations["custom_mask_type"]) == 2:
|
|
attrs["window_size_left"] = attrs["num_keys"]
|
|
attrs["window_size_right"] = 0
|
|
attrs["is_causal"] = True
|
|
else:
|
|
attrs["window_size_left"] = -1
|
|
attrs["window_size_right"] = -1
|
|
attrs["is_causal"] = False
|
|
|
|
if use_flash:
|
|
headers.append("flash.h")
|
|
attrs["num_q_heads"] = annotations["num_q_heads"]
|
|
attrs["num_kv_heads"] = annotations["num_kv_heads"]
|
|
|
|
if is_var_len:
|
|
code = instantiate_flash_attention_var_len_template(attrs)
|
|
else:
|
|
code = instantiate_flash_attention_template(attrs)
|
|
else:
|
|
headers.append("kernel_forward.h")
|
|
|
|
assert not is_mqa, (
|
|
"The number of query and KV heads need to be the same for CUTLASS fMHA."
|
|
)
|
|
|
|
attrs["num_heads"] = n = annotations["num_q_heads"]
|
|
|
|
data_type_size = DataTypeSize[data_type]
|
|
if (data_type_size * h // 8) % 16 == 0 and (data_type_size * h_v // 8) % 16 == 0:
|
|
attrs["kIsAligned"] = True
|
|
elif (h % 4 == 0) and (h_v % 4 == 0):
|
|
attrs["kIsAligned"] = False
|
|
else:
|
|
raise NotImplementedError()
|
|
if h_v > 64:
|
|
attrs["kQueriesPerBlock"] = 32
|
|
attrs["kKeysPerBlock"] = 128
|
|
attrs["kSingleValueIteration"] = h_v <= 128
|
|
else:
|
|
attrs["kQueriesPerBlock"] = 64
|
|
attrs["kKeysPerBlock"] = 64
|
|
attrs["kSingleValueIteration"] = True
|
|
|
|
assert attrs["scale"] > 0 or attrs["scale"] < 0, (
|
|
"Cutlass may generate nan occasionally when scale == 0.0"
|
|
)
|
|
attrs["arch"] = "cutlass::arch::Sm{}".format(annotations["arch"])
|
|
attrs["kSupportsDropout"] = False
|
|
|
|
attrs["output_size"] = f"{b} * {s} * {n} * {h_v}"
|
|
|
|
attrs["custom_mask_type"] = annotations["custom_mask_type"]
|
|
|
|
for arg in func_args:
|
|
if "workspace" in arg:
|
|
attrs["workspace"] = arg
|
|
if "bias" in attrs:
|
|
attrs["kSupportsBias"] = True
|
|
if len(annotations["bias_shape"]) == 4:
|
|
strides = "p.num_keys"
|
|
if annotations["bias_shape"][2] == 1:
|
|
attrs["bias_strideM"] = 0
|
|
else:
|
|
attrs["bias_strideM"] = strides
|
|
strides = f"p.num_queries * {strides}"
|
|
if annotations["bias_shape"][1] == 1:
|
|
attrs["bias_strideH"] = 0
|
|
else:
|
|
attrs["bias_strideH"] = strides
|
|
strides = f"p.num_heads * {strides}"
|
|
if annotations["bias_shape"][0] == 1:
|
|
attrs["bias_strideB"] = 0
|
|
else:
|
|
attrs["bias_strideB"] = strides
|
|
else:
|
|
raise NotImplementedError()
|
|
else:
|
|
# To support negative scale in current Cutlass implementation,
|
|
# kSupportsBias should be set true, or there are nan's as result.
|
|
attrs["kSupportsBias"] = attrs["scale"] < 0
|
|
|
|
code = instantiate_attention_template(attrs)
|
|
|
|
return CodegenResult(code, headers)
|
|
elif "layer_norm" in func_name:
|
|
headers.append("cutlass/util/device_layernorm.h")
|
|
headers.append("cutlass/layout/matrix.h")
|
|
attrs = {"input": func_args[0], "gamma": func_args[1], "beta": func_args[2]}
|
|
attrs.update(dict(annotations))
|
|
|
|
if not isinstance(attrs["M"], tvm.tirx.IntImm):
|
|
attrs["M"] = get_flattened_batch_dim(func_args[0], int(attrs["batch_rank"]))
|
|
|
|
code = instantiate_layer_norm_template(attrs)
|
|
return CodegenResult(code, headers)
|
|
elif "rms_norm" in func_name:
|
|
headers.append("cutlass/util/device_rmsnorm.h")
|
|
headers.append("cutlass/layout/matrix.h")
|
|
attrs = {"input": func_args[0], "weight": func_args[1]}
|
|
attrs.update(dict(annotations))
|
|
|
|
if not isinstance(attrs["M"], tvm.tirx.IntImm):
|
|
attrs["M"] = get_flattened_batch_dim(func_args[0], int(attrs["batch_rank"]))
|
|
|
|
code = instantiate_rms_norm_template(attrs)
|
|
return CodegenResult(code, headers)
|
|
|
|
raise ValueError(f"Do not have a template for {func_name}")
|