diff --git a/aphrodite/quantization/__init__.py b/aphrodite/quantization/__init__.py index 494f45926d..f58f73ad22 100644 --- a/aphrodite/quantization/__init__.py +++ b/aphrodite/quantization/__init__.py @@ -11,6 +11,7 @@ from aphrodite.quantization.eetq import EETQConfig from aphrodite.quantization.experts_int8 import ExpertsInt8Config from aphrodite.quantization.fbgemm_fp8 import FBGEMMFp8Config +from aphrodite.quantization.flute import FluteConfig from aphrodite.quantization.fp6 import QuantLLMFPConfig from aphrodite.quantization.fp8 import Fp8Config from aphrodite.quantization.gguf import GGUFConfig @@ -35,6 +36,7 @@ "fp8": Fp8Config, "quant_llm": QuantLLMFPConfig, "fbgemm_fp8": FBGEMMFp8Config, + "flute": FluteConfig, "modelopt": ModelOptFp8Config, "gguf": GGUFConfig, # The order of gptq methods is important for config.py iteration over diff --git a/aphrodite/quantization/flute.py b/aphrodite/quantization/flute.py new file mode 100644 index 0000000000..f75288d1d4 --- /dev/null +++ b/aphrodite/quantization/flute.py @@ -0,0 +1,365 @@ +from typing import Any, Dict, List, Optional + +import torch +from torch.nn.parameter import Parameter + +from aphrodite.distributed import (divide, get_tp_group, + tensor_model_parallel_all_gather) +from aphrodite.modeling.layers.linear import (LinearBase, LinearMethodBase, + set_weight_attrs) +from aphrodite.quantization.base_config import QuantizationConfig + + +class PackFactor(object): + + def __init__(self, pack_bits: int, num_bits: int) -> None: + if num_bits not in [2, 3, 4]: + raise ValueError + self.pack_bits = pack_bits + self.num_bits = num_bits + + def __rfloordiv__(self, other: int) -> int: + if not isinstance(other, int): + raise TypeError + # the sole purpose of this class is to change the ordering + # of the operands, so that it works with 3-bits. That is, + # instead of using `other // (pack_bits // num_bits)`, + # we use `(other // pack_bits) * num_bits`. The former + # does not work with 3-bits, while the latter does. + return divide(other, self.pack_bits) * self.num_bits + + +class FluteConfig(QuantizationConfig): + """Config class for FLUTE Quantization.""" + + def __init__( + self, + num_bits: int, + group_size: int, + num_sms_packed: int, + ) -> None: + if num_bits not in [2, 3, 4]: + raise ValueError + + self.num_bits = num_bits + self.group_size = group_size + self.pack_factor = PackFactor(pack_bits=16, num_bits=num_bits) + self.num_sms_packed = num_sms_packed + + def __repr__(self) -> str: + return (f"FluteConfig(" + f"num_bits={self.num_bits}, " + f"group_size={self.group_size}, " + f"num_sms_packed={self.num_sms_packed})") + + @classmethod + def get_name(cls) -> str: + return "flute" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.float16, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["quantize_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "FluteConfig": + num_bits = cls.get_from_keys(config, ["num_bits"]) + group_size = cls.get_from_keys(config, ["group_size"]) + num_sms_packed = cls.get_from_keys(config, ["num_sms"]) + + return cls(num_bits=num_bits, + group_size=group_size, + num_sms_packed=num_sms_packed) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional["FluteLinearMethod"]: + if isinstance(layer, LinearBase): + return FluteLinearMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class FluteLinearMethod(LinearMethodBase): + """Linear method for Flute. + Args: + quant_config: The Flute quantization config. + """ + + def __init__(self, quant_config: FluteConfig) -> None: + + try: + import flute + if flute.__version__ < "0.0.6": + raise ImportError("flute version is wrong. Please " + "install flute>=0.0.6.") + except ImportError as err: + raise ImportError("Please install flute>=0.0.6 via " + "`pip install flute-kernel>=0.0.6` to use " + "flute quantizer.") from err + + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + + import flute + import flute.utils + + if params_dtype not in [torch.float16, torch.bfloat16]: + raise TypeError + + K = input_size_per_partition + N = sum(output_partition_sizes) + P = int(N / 16 * self.quant_config.num_bits) + G = int(K / self.quant_config.group_size) + device = "cuda" + + weight = Parameter( + torch.empty( + (P, K), + dtype=torch.int16, + device=device, + ), + requires_grad=False, + ) + set_weight_attrs( + weight, + { + **extra_weight_attrs, + "input_dim": 1, + "output_dim": 0, + "packed_dim": 0, + "pack_factor": self.quant_config.pack_factor, + }, + ) + + scales = Parameter( + torch.empty( + (N, G), + dtype=params_dtype, + device=device, + ), + requires_grad=False, + ) + set_weight_attrs( + scales, + { + **extra_weight_attrs, + "input_dim": 1, + "output_dim": 0, + # it's unclear if we need to specify `packed_dim`, but looks + # like this is only useful if `pack_factor == output_dim` + # "packed_dim": 1, + }, + ) + + tables = Parameter( + torch.arange( + 2**self.quant_config.num_bits, + dtype=params_dtype, + device=device, + ), + requires_grad=False, + ) + set_weight_attrs( + tables, + { + **extra_weight_attrs, + "input_dim": None, + "output_dim": None, + "ignore_warning": True, + }, + ) + + tables2 = Parameter( + flute.utils.make_qmap2_from_qmap(tables), + requires_grad=False, + ) + set_weight_attrs( + tables2, + { + **extra_weight_attrs, + "input_dim": None, + "output_dim": None, + "ignore_warning": True, + }, + ) + + layer.num_bits = self.quant_config.num_bits + layer.group_size = self.quant_config.group_size + layer.num_sms_packed = self.quant_config.num_sms_packed + layer.workspace = flute.utils.get_workspace_streamk(weight.device) + + layer.register_parameter("weight", weight) + layer.register_parameter("scales", scales) + layer.register_parameter("tables", tables) + layer.register_parameter("tables2", tables2) + + layer.needs_repacking = True + layer.flute_input_size = input_size + layer.flute_output_size = output_size + layer.flute_output_partition_sizes = output_partition_sizes + layer.flute_input_size_per_partition = input_size_per_partition + layer.flute_is_K_partitioned = (input_size_per_partition != input_size) + layer.flute_is_N_partitioned = (sum(output_partition_sizes) != + output_size) + + def _maybe_tensor_all_gather( + self, + tensor: torch.Tensor, + shard_dim: Optional[int], + ) -> torch.Tensor: + + if shard_dim is None: + return tensor + + # NCCL does not support int16 + if tensor.dtype == torch.int16: + tensor_dtype = tensor.dtype + tensor_casted = tensor.to(dtype=torch.int32) + if not (tensor_casted == tensor).all(): + raise ValueError + else: + tensor_dtype = None + tensor_casted = tensor + + tensor_gathered = tensor_model_parallel_all_gather(tensor_casted, + dim=shard_dim) + + if tensor_dtype is not None: + tensor_gathered_casted = tensor_gathered.to(dtype=tensor_dtype) + if not (tensor_gathered_casted == tensor_gathered).all(): + raise ValueError + return tensor_gathered_casted + else: + return tensor_gathered + + def _maybe_tensor_shard( + self, + tensor: torch.Tensor, + shard_dim: Optional[int], + ) -> torch.Tensor: + + if shard_dim is None: + return tensor + + tp_group = get_tp_group() + shard_dim_size = divide(tensor.shape[shard_dim], tp_group.world_size) + tensor_shards = torch.split(tensor, shard_dim_size, dim=shard_dim) + # NOTE: torch.split does not create contiguous tensors by default. + tensor_shard = tensor_shards[tp_group.rank].contiguous() + return tensor_shard + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + + import flute + import flute.utils + + # the weights packing are possibly shape-specialized, but as vLLM + # potentially fuses parameters and/or partitions the weights (TP), + # we need to potentially re-pack the weights. + if ((not hasattr(layer, "needs_repacking") + or not layer.needs_repacking)): + return + + if ((layer.flute_is_K_partitioned is False) + and (layer.flute_is_N_partitioned is False)): + shard_dim = None + if ((layer.flute_is_K_partitioned is True) + and (layer.flute_is_N_partitioned is False)): + shard_dim = 1 + if ((layer.flute_is_K_partitioned is False) + and (layer.flute_is_N_partitioned is True)): + shard_dim = 0 + if ((layer.flute_is_K_partitioned is True) + and (layer.flute_is_N_partitioned is True)): + raise NotImplementedError + + # split the combined tensors into individual tensors + # weight: [P, K] + # scales: [N, G] + Ns = layer.flute_output_partition_sizes + Ps = [int(N / 16 * layer.num_bits) for N in Ns] + Qs = torch.split(layer.weight, Ps, dim=0) + Ss = torch.split(layer.scales, Ns, dim=0) + + Qs_unpacked = [] + for Q, S in zip(Qs, Ss): + # when the tensors are sharded, gather them before unpacking + Q_gathered = self._maybe_tensor_all_gather(Q, shard_dim=shard_dim) + S_gathered = self._maybe_tensor_all_gather(S, shard_dim=shard_dim) + + # unpack + Q_gathered_unpacked = flute.utils.unpack( + weight=Q_gathered, + scales=S_gathered, + workspace=layer.workspace, + num_bits=layer.num_bits, + group_size=layer.group_size, + num_sms_packed=layer.num_sms_packed) + + # re-shard + Qs_unpacked.append( + self._maybe_tensor_shard(Q_gathered_unpacked, + shard_dim=shard_dim)) + + # reconstruct the unpacked tensor + Q_unpacked = torch.cat(Qs_unpacked, dim=0) + + # re-pack the tensors + Q_repacked = flute.utils.pack( + Q_unpacked.T.contiguous().to(device="cpu"), + num_bits=layer.num_bits, + group_size=layer.group_size).to(device=layer.weight.device) + + if not all([ + Q_repacked.shape == layer.weight.shape, Q_repacked.dtype + == layer.weight.dtype, Q_repacked.device == layer.weight.device + ]): + raise ValueError + layer.weight = Parameter(Q_repacked, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + import flute + + output = flute.qgemm_simple( + x, + layer.weight, + layer.scales, + layer.tables, + layer.tables2, + layer.workspace, + layer.num_bits, + layer.group_size, + ) + + if bias is not None: + output.add_(bias) # In-place add + + return output diff --git a/kernels/quantization/flute/config.hpp b/kernels/quantization/flute/config.hpp new file mode 100644 index 0000000000..69984331ca --- /dev/null +++ b/kernels/quantization/flute/config.hpp @@ -0,0 +1,561 @@ +#pragma once + +#include +#include +#include + + +namespace config { + +using namespace cute; + +// --- Type Traits --- +template +struct G2SCopyOpTraits; + + +template <> +struct G2SCopyOpTraits<_16, _8> { + using type = cute::uint128_t; + using op = SM80_CP_ASYNC_CACHEGLOBAL; +}; + + +template <> +struct G2SCopyOpTraits<_16, _4> { + using type = cute::uint64_t; + using op = SM80_CP_ASYNC_CACHEALWAYS; +}; + + +template <> +struct G2SCopyOpTraits<_16, _2> { + using type = cute::uint32_t; + using op = SM80_CP_ASYNC_CACHEALWAYS; +}; + + +template <> +struct G2SCopyOpTraits<_32, _2> { + using type = cute::uint64_t; + using op = SM80_CP_ASYNC_CACHEALWAYS; +}; + + +template <> +struct G2SCopyOpTraits<_32, _1> { + using type = cute::uint32_t; + using op = SM80_CP_ASYNC_CACHEALWAYS; +}; + + +template +struct G2SCopyTraits { + using TypeSize = Int>; + using Op = typename G2SCopyOpTraits::op; + using Traits = Copy_Traits; + using Atom = Copy_Atom; +}; + + +// Strategies for decomposing the problem +enum class DecompositionModeEnum { + // Split-K and Slice-K decomposition + SplitK, + + // Stream-K decomposition + StreamK +}; + + +// Strategies for computing reductions between CTAs computing portions of a given output tile +enum class ReductionModeEnum { + // Participating CTAs perform reduction in a turnstile fashion in order of the K extent + // covered by each CTA. This requires a lock to be held exclusively be the CTA that is + // currently accumulating. + Deterministic, + + // Participating CTAs perform reduction atomically to the same workspace (mostly) without locking. + // Locks are used only to wait for the first CTA to write its partial values (to initialize the + // workspace), and for all but the final CTA to have accumulated (so that the final CTA can load + // the accumulated value and accumulate it into registers on top of which the epilogue will + // be performed). + Nondeterministic +}; + + +enum class QuantMapModeEnum { + // look-up one entry at a time + Basic, + + // look-up two entries at a time + Vectorized, + + // vectorized and duplicated 32 times + Vectorized_32, + + // vectorized and duplicated 16 times + Vectorized_16, + + // vectorized and duplicated 8 times + Vectorized_8, + + // look-up one entry at time using warp shuffle (deprecated) + WarpShuffle, + + // Marlin-style integer dequantization + Marlin +}; + + +enum class AccumulationModeEnum { + // accumulate in FP16/BF16 + Low, + + // accumulate in FP32 + High, + + // accumulate in FP32, but reduce in FP16/BF16 + Mixed +}; + + +template +struct VectorizedQuantMapTraits { + using Duplicates = _1; + using IsVectorized = false_type; + using IsDuplicated = false_type; +}; + + +template +struct VectorizedQuantMapTraits { + using Duplicates = _1; + using IsVectorized = true_type; + using IsDuplicated = false_type; +}; + + +template +struct VectorizedQuantMapTraits { + CUTE_STATIC_ASSERT_V(NumBits{} == _4{}); + using Duplicates = _32; + using IsVectorized = true_type; + using IsDuplicated = true_type; +}; + + +template +struct VectorizedQuantMapTraits { + CUTE_STATIC_ASSERT_V(NumBits{} == _4{}); + using Duplicates = _16; + using IsVectorized = true_type; + using IsDuplicated = true_type; +}; + + +template +struct VectorizedQuantMapTraits { + CUTE_STATIC_ASSERT_V(NumBits{} == _4{}); + using Duplicates = _8; + using IsVectorized = true_type; + using IsDuplicated = true_type; +}; + + +template < + // types + typename T_, + typename TQ_, + // threads and tile sizes + typename Threads_, + typename TileM_, + typename TileK_, + typename TileP_, + typename Stages_, + // quantization + typename NumBits_, + typename GroupSize_, + // enums + QuantMapModeEnum QuantMapMode_, + AccumulationModeEnum AccumulationMode_, + DecompositionModeEnum DecompositionMode_, + // misc + typename G2STiledCopySizeS_, + typename MmaPrmK_ +> +struct GemmConfig { + + // enums + static constexpr QuantMapModeEnum QuantMapMode = QuantMapMode_; + static constexpr AccumulationModeEnum AccumulationMode = AccumulationMode_; + static constexpr DecompositionModeEnum DecompositionMode = DecompositionMode_; + + // type configuration + using T = T_; + using TQ = TQ_; // type for quantization + using TC = conditional_t; // type for accumulation + using TR = conditional_t; // type for reduction + using T2 = conditional_t, __half2, __nv_bfloat162>; + + CUTE_STATIC_ASSERT(sizeof(T ) == 2); + CUTE_STATIC_ASSERT(sizeof(TQ) == 2); + CUTE_STATIC_ASSERT(is_same_v == true || is_same_v == true); + + // threads configuration + using Warps = _32; // using compile-time constant instead of runtime `warpSize` + using Threads = Threads_; + CUTE_STATIC_ASSERT_V(Threads{} % _128{} == _0{}); + + // quantization configuration + using NumBits = NumBits_; + using NumPacked = decltype(_8{} * Int{} / NumBits{}); + using GroupSize = GroupSize_; + + // 2-bit: 1 x 16-bit -> 8 x 2-bit + // 4-bit: 1 x 16-bit -> 4 x 4-bit + // 8-bit: 3 x 16-bit -> 16 x 3-bit + using UnpackSourceSize = conditional_t, _3, _3>; // workaround: use size = 3 when NumBits != 3 as well + using UnpackTargetSize = conditional_t, _16, NumPacked>; + + // tile configuration + using TileM = TileM_; + using TileK = TileK_; + using TileP = TileP_; + using Stages = Stages_; + + using G2STiledCopySizeA = decltype(cute::min(TileM{} * TileK{} / Threads{}, _8{})); + using G2STiledCopySizeQ = _8; + using G2STiledCopySizeQ2 = _8; + using G2STiledCopySizeS = G2STiledCopySizeS_; + using EpilogueCopySizeC = _8; + CUTE_STATIC_ASSERT_V(G2STiledCopySizeS{} == _2{} || G2STiledCopySizeS{} == _8{}); + + // derived tile configuration + using TileP2 = decltype(TileP{} * _2{}); + using TileN = decltype(TileP{} * UnpackTargetSize{}); + // we want `TileN x TileG = Threads x G2STiledCopySizeS` and `TileG >= G2STiledCopySizeS` + using TileG = decltype(Threads{} * G2STiledCopySizeS{} / cute::min(TileN{}, Threads{})); + // note that for scales, each `TileG` can service multiple `TileK`, hence we only need + // `Stages / (GroupSize * TileG / TileK)` stages of scales, with a minimum of 2 stages. + using TileKsPerTileG = decltype(GroupSize{} * TileG{} / TileK{}); + // in the prefetching stage, we load `Stages{} - 1` stages of A and Q. We thus need to add + // an additional stage to the `scales` similarly. + using StagesGRaw = decltype(ceil_div(Stages{} - _1{}, TileKsPerTileG{}) + _1{}); + using StagesG = decltype(ceil_div(StagesGRaw{}, _2{}) * _2{}); // round up to the nearest multiple of 2 + + CUTE_STATIC_ASSERT_V(TileM{} == _16{} || TileM{} == _32{} || TileM{} == _64{}); + CUTE_STATIC_ASSERT_V((TileP {} * UnpackTargetSize{}) == TileN{}); // TileP * NumPacked must be equal to TileN + CUTE_STATIC_ASSERT_V((TileM {} * TileK{}) % (Threads{} * G2STiledCopySizeA {}) == _0{}); // TileA is too small for G2S copy + CUTE_STATIC_ASSERT_V((TileP {} * TileK{}) % (Threads{} * G2STiledCopySizeQ {}) == _0{}); // TileQ is too small for G2S copy + CUTE_STATIC_ASSERT_V((TileP2{} * TileK{}) % (Threads{} * G2STiledCopySizeQ2{}) == _0{}); // TileQ2 is too small for G2S copy + CUTE_STATIC_ASSERT_V((TileN {} * TileG{}) % (Threads{} * G2STiledCopySizeS {}) == _0{}); // TileS is too small for G2S copy + // CUTE_STATIC_ASSERT_V(Stages{} >= StagesG{}); // do we really need this? + // CUTE_STATIC_ASSERT_V(Stages{} % TileKsPerTileG{} == _0{}); + CUTE_STATIC_ASSERT_V((GroupSize{} * TileG{}) % TileK{} == _0{}); + CUTE_STATIC_ASSERT_V((GroupSize{} * TileG{}) == (TileK{} * TileKsPerTileG{})); + CUTE_STATIC_ASSERT_V(((GroupSize{} >= TileK{}) && (GroupSize{} % TileK{} == _0{})) || // GroupSize must be a multiple of TileK + ((GroupSize{} < TileK{}) && (TileK{} % GroupSize{} == _0{}))); // TileK must be a multiple of GroupSize + + // + // --- SMEM Layouts --- + // + + // https://github.com/NVIDIA/cutlass/blob/main/test/unit/gemm/device/default_gemm_configuration.hpp#L84 + using SmemLayoutAtomK = decltype(composition (Swizzle<3, 3, 3>{}, make_layout(make_shape(_8{}, _64{}), make_stride(_64{}, _1{})))); + using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtomK{}, make_shape(TileM{}, TileK{}, Stages{}))); + using SmemLayoutB = decltype(tile_to_shape(SmemLayoutAtomK{}, make_shape(TileN{}, TileK{}, Stages{}))); + using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomK{}, make_shape(TileP{}, TileK{}, Stages{}))); + using SmemLayoutQ2 = decltype(tile_to_shape(SmemLayoutAtomK{}, make_shape(TileP2{},TileK{}, Stages{}))); + using SmemLayoutAtomG = decltype(composition (Swizzle<3, 3, 3>{}, make_layout(make_shape(_8{}, TileG{}), make_stride(TileG{}, _1{})))); + using SmemLayoutS = decltype(tile_to_shape(SmemLayoutAtomG{}, make_shape(TileN{}, TileG{}, StagesG{}))); + + // `SmemLayoutSView` is `SmemLayoutS` broadcasted to `TileK` dimension + // 1. `GroupSize > TileK`: (TileN, TileK, (GroupSize / TileK, TileG, StagesG)) + // (TileG, 0 , (0 , 1 , TileN * TileG)) + // 2. `TileK > GroupSize`: (TileN, (GroupSize, TileK / GroupSize), (TileG / (TileK / GroupSize), StagesG)) + // (TileG, (0 , 1 ), (TileK / GroupSize , TileN * TileG)) + using TileKsPerGroup = decltype(ceil_div (GroupSize{}, TileK{})); // used when `GroupSize > TileK` + using GroupsPerTileK = decltype(ceil_div (TileK {}, GroupSize{})); // used when `TileK > GroupSize` + using TileGView = decltype(ceil_div (TileG {}, GroupsPerTileK{})); // ceil(TileG / ceil(TileK / GroupSize)) + using SmemLayoutSViewShapeCase1 = decltype(make_shape (TileN {}, TileK {}, TileKsPerGroup{}, TileG {}, StagesG{})); + using SmemLayoutSViewShapeCase2 = decltype(make_shape (TileN {}, GroupSize{}, GroupsPerTileK{}, TileGView {}, StagesG{})); + using SmemLayoutSViewStrideCase1 = decltype(make_stride(TileG {}, _0{}, _0{}, _1{}, TileN{} * TileG{})); + using SmemLayoutSViewStrideCase2 = decltype(make_stride(TileG {}, _0{}, _1{}, GroupsPerTileK{}, TileN{} * TileG{})); + using SmemLayoutSViewCaseRaw1 = decltype(composition(Swizzle<3, 3, 3>{}, make_layout(SmemLayoutSViewShapeCase1{}, SmemLayoutSViewStrideCase1{}))); + using SmemLayoutSViewCaseRaw2 = decltype(composition(Swizzle<3, 3, 3>{}, make_layout(SmemLayoutSViewShapeCase2{}, SmemLayoutSViewStrideCase2{}))); + using SmemLayoutSViewCase1 = decltype( group<2, 5>(SmemLayoutSViewCaseRaw1{})); + using SmemLayoutSViewCase2 = decltype(group<2, 4>(group<1, 3>(SmemLayoutSViewCaseRaw2{}))); + using SmemLayoutSView = conditional_t<(TileK{} <= GroupSize{}), SmemLayoutSViewCase1, SmemLayoutSViewCase2>; + using StagesGView = decltype(size<2> (SmemLayoutSView{})); + + // `SmemLayoutAtomK` requirements + CUTE_STATIC_ASSERT_V(TileM{} >= size<0>(SmemLayoutAtomK{})); + CUTE_STATIC_ASSERT_V(TileN{} >= size<0>(SmemLayoutAtomK{})); + CUTE_STATIC_ASSERT_V(TileP{} >= size<0>(SmemLayoutAtomK{})); + CUTE_STATIC_ASSERT_V(TileP2{} >= size<0>(SmemLayoutAtomK{})); + CUTE_STATIC_ASSERT_V(TileK{} >= size<1>(SmemLayoutAtomK{})); + + // + // --- Quantization Maps --- + // + + using QuantMapVecTraits = VectorizedQuantMapTraits; + using QuantMapDuplicates = typename QuantMapVecTraits::Duplicates; + + // 2^NumBits + using QuantMapSize = decltype(_1{} << (NumBits{})); + using QuantMapSize2 = decltype(_1{} << (NumBits{} * _2{})); + using SmemLayoutQM = decltype(make_layout(make_shape(QuantMapSize{}))); + using SmemLayoutQM2 = decltype(make_layout(make_shape(QuantMapSize2{}))); + using SmemLayoutQM3 = decltype(make_layout(make_shape(QuantMapSize2{}, QuantMapDuplicates{}), LayoutRight{})); // TODO: ablate layout left + using SmemLayoutQMView = decltype(make_layout(make_shape(_1{}, QuantMapSize{}))); + using SmemLayoutQM2View = decltype(make_layout(make_shape(QuantMapSize2{}, QuantMapDuplicates{}), make_stride(_1{}, _0{}))); + + // + // --- MMA --- + // + // https://github.com/NVIDIA/cutlass/discussions/1142 + // https://zhuanlan.zhihu.com/p/663092747 + // https://github.com/NVIDIA/cutlass/issues/1028#issuecomment-1668088899 + // It is recommended to tile TiledMma to cover the 4 sub-partitions of the SM (128 thread large tiled TiledMma, or some multiple of it). + + using MmaOpFP16 = conditional_t; + using MmaOpBF16 = conditional_t; + using MmaOp = conditional_t == true, MmaOpFP16, MmaOpBF16>; + using MmaTraits = MMA_Traits; + using MmaAtom = MMA_Atom; + CUTE_STATIC_ASSERT(is_same_v == true || AccumulationMode != AccumulationModeEnum::Low); // BF16 does not low-precision accumulation + + // We can increase the size of the computation via + // 1. adding more threads (`kMmaThr*`) + using LargeM = decltype(TileM {} > _32{}); + using NumMmas = decltype(Threads{} / _32{}); // Each `MmaOp` takes 32 threads. + using MmaThrM = conditional_t; + using MmaThrN = decltype(NumMmas{} / MmaThrM{}); + using MmaThrK = _1; + + // 2. adding more works to threads (`kMmaPrm*`) + using MmaPrmM = decltype(TileM{} / (get<0>(typename MmaTraits::Shape_MNK{}) * MmaThrM{})); + using MmaPrmN = decltype(TileP{} / (get<1>(typename MmaTraits::Shape_MNK{}) * MmaThrN{})); + using MmaPrmK = MmaPrmK_; + + using MmaPermutations = decltype(make_tile( + get<0>(typename MmaTraits::Shape_MNK{}) * MmaThrM{} * MmaPrmM{}, + get<1>(typename MmaTraits::Shape_MNK{}) * MmaThrN{} * MmaPrmN{} * UnpackTargetSize{}, + get<2>(typename MmaTraits::Shape_MNK{}) * MmaThrK{} * MmaPrmK{})); + using MmaPermutationsQ = decltype(make_tile( + get<0>(typename MmaTraits::Shape_MNK{}) * MmaThrM{} * MmaPrmM{}, + get<1>(typename MmaTraits::Shape_MNK{}) * MmaThrN{} * MmaPrmN{}, + get<2>(typename MmaTraits::Shape_MNK{}) * MmaThrK{} * MmaPrmK{})); + using MmaPermutationsQ2 = decltype(make_tile( + get<0>(typename MmaTraits::Shape_MNK{}) * MmaThrM{} * MmaPrmM{}, + get<1>(typename MmaTraits::Shape_MNK{}) * MmaThrN{} * MmaPrmN{} * (UnpackSourceSize{} - _1{}), // `-1` because we already have one in `MmaPermutationsQ` + get<2>(typename MmaTraits::Shape_MNK{}) * MmaThrK{} * MmaPrmK{})); + using MmaThrLayout = decltype(make_layout(make_shape(MmaThrM{}, MmaThrN{}, MmaThrK{}))); + using TiledMma = decltype(make_tiled_mma(MmaAtom{}, MmaThrLayout{}, MmaPermutations{})); + using TiledMmaQ = decltype(make_tiled_mma(MmaAtom{}, MmaThrLayout{}, MmaPermutationsQ{})); + using TiledMmaQ2 = decltype(make_tiled_mma(MmaAtom{}, MmaThrLayout{}, MmaPermutationsQ2{})); + + CUTE_STATIC_ASSERT_V(MmaThrK{} == _1{}); // We don't know how to handle MmaThrK > 1, yet + CUTE_STATIC_ASSERT_V(MmaPrmK{} == _1{} || MmaPrmK{} == _2{}); + CUTE_STATIC_ASSERT_V(Threads{} == size(TiledMma{})); + CUTE_STATIC_ASSERT_V(Threads{} == size(TiledMmaQ{})); + CUTE_STATIC_ASSERT_V(Threads{} == size(TiledMmaQ2{})); + // we cannot compute with more data than we will load + CUTE_STATIC_ASSERT_V(tile_size<0>(TiledMma{}) == TileM{}); + CUTE_STATIC_ASSERT_V(tile_size<1>(TiledMma{}) == TileN{}); + CUTE_STATIC_ASSERT_V(tile_size<2>(TiledMma{}) <= TileK{}); + CUTE_STATIC_ASSERT_V(tile_size<0>(TiledMmaQ{}) == TileM{}); + CUTE_STATIC_ASSERT_V(tile_size<1>(TiledMmaQ{}) == TileP{}); + CUTE_STATIC_ASSERT_V(tile_size<2>(TiledMmaQ{}) <= TileK{}); + CUTE_STATIC_ASSERT_V(tile_size<0>(TiledMmaQ2{}) == TileM{}); + CUTE_STATIC_ASSERT_V(tile_size<1>(TiledMmaQ2{}) == TileP2{}); + CUTE_STATIC_ASSERT_V(tile_size<2>(TiledMmaQ2{}) <= TileK{}); + CUTE_STATIC_ASSERT(is_same_v == true); // Mma and TC must have the same accumulation type + CUTE_STATIC_ASSERT_V(Threads{} == (MmaThrM{} * MmaThrN{} * MmaThrK{} * _32{})); // TiledMma should cover all threads + + // just to make sure we allocate enough scratch space for StreamK + CUTE_STATIC_ASSERT_V((MmaPrmM{}) <= _2{}); + CUTE_STATIC_ASSERT_V((MmaPrmN{} * UnpackTargetSize{}) <= _32{}); + + // + // --- Global to Shared Memory Copy --- + // + + using G2STiledCopySizeQM = decltype(ceil_div(QuantMapSize {}, Threads{})); + using G2STiledCopySizeQM2 = decltype(ceil_div(QuantMapSize2{}, Threads{})); + + using G2SA_M = TileM; + using G2SQ_P = TileP; + using G2SQ2_P2 = TileP2; + using G2SS_N = decltype(cute::min(TileN{}, Threads{})); // `G2SS_N` cannot exceed `Threads`, otherwise `G2SS_G` will be zero. + using G2SA_K = decltype(Threads{} / G2SA_M{}); + using G2SQ_K = decltype(Threads{} / G2SQ_P{}); + using G2SQ2_K = decltype(Threads{} / G2SQ2_P2{}); + using G2SS_G = decltype(Threads{} / G2SS_N{}); + + using G2SCopyAtomA = typename G2SCopyTraits::Atom; + using G2SCopyAtomQ = typename G2SCopyTraits::Atom; + using G2SCopyAtomQ2 = typename G2SCopyTraits::Atom; + using G2SCopyAtomS = typename G2SCopyTraits::Atom; + using G2SCopyAtomQM = Copy_Atom; // the copy size is too small for `cp.async` + using G2SCopyAtomQM2 = typename G2SCopyTraits::Atom; + + // https://zhuanlan.zhihu.com/p/664671157 (Section 2.2.2) + // Similarly, we choose many atoms needed to cover the 4 sub-partitions of the SM (some multiple of 128 thread). + // https://github.com/NVIDIA/cutlass/blob/v3.4.0/test/unit/gemm/device/default_gemm_configuration.hpp#L90 + // This operation uses Threads threads in a G2SA_M x G2SA_K shape. Each thread loads 1x8 elements + // of 16-bits to cover, in total, a [G2SA_M, G2SA_K x 8] shape. + + using G2STiledCopyA = decltype(make_tiled_copy( + G2SCopyAtomA{}, + make_layout(make_shape(G2SA_M{}, G2SA_K{}), make_stride(G2SA_K{}, _1{})), + make_layout(make_shape(_1{}, G2STiledCopySizeA{})))); + + using G2STiledCopyQ = decltype(make_tiled_copy( + G2SCopyAtomQ{}, + make_layout(make_shape(G2SQ_P{}, G2SQ_K{}), make_stride(G2SQ_K{}, _1{})), + make_layout(make_shape(_1{}, G2STiledCopySizeQ{})))); + + using G2STiledCopyQ2 = decltype(make_tiled_copy( + G2SCopyAtomQ2{}, + make_layout(make_shape(G2SQ2_P2{}, G2SQ2_K{}), make_stride(G2SQ2_K{}, _1{})), + make_layout(make_shape(_1{}, G2STiledCopySizeQ2{})))); + + using G2STiledCopyS = decltype(make_tiled_copy( + G2SCopyAtomS{}, + make_layout(make_shape(G2SS_N{}, G2SS_G{}), make_stride(G2SS_G{}, _1{})), + make_layout(make_shape(_1{}, G2STiledCopySizeS{})))); + + // each thread load one element, with predication. + using G2STiledCopyQM = decltype(make_tiled_copy( + G2SCopyAtomQM{}, + make_layout(make_shape(Threads{})), + make_layout(make_shape(G2STiledCopySizeQM{})))); + + using G2STiledCopyQM2 = decltype(make_tiled_copy( + G2SCopyAtomQM2{}, + make_layout(make_shape(Threads{})), + make_layout(make_shape(G2STiledCopySizeQM2{})))); + + using G2STiledCopyShapeA = decltype(shape(typename G2STiledCopyA ::Tiler_MN{})); + using G2STiledCopyShapeQ = decltype(shape(typename G2STiledCopyQ ::Tiler_MN{})); + using G2STiledCopyShapeQ2 = decltype(shape(typename G2STiledCopyQ2 ::Tiler_MN{})); + using G2STiledCopyShapeS = decltype(shape(typename G2STiledCopyS ::Tiler_MN{})); + using G2STiledCopyShapeQM = decltype(shape(typename G2STiledCopyQM ::Tiler_MN{})); + using G2STiledCopyShapeQM2 = decltype(shape(typename G2STiledCopyQM2::Tiler_MN{})); + + // CUTE_STATIC_ASSERT_V(G2SQ_P{} == _16{}); // Not supported yet + CUTE_STATIC_ASSERT_V(Threads{} == size(G2STiledCopyA{})); + CUTE_STATIC_ASSERT_V(Threads{} == size(G2STiledCopyQ{})); + CUTE_STATIC_ASSERT_V(Threads{} == size(G2STiledCopyQ2{})); + CUTE_STATIC_ASSERT_V(Threads{} == size(G2STiledCopyS{})); + CUTE_STATIC_ASSERT_V(Threads{} == size(G2STiledCopyQM{})); + CUTE_STATIC_ASSERT_V(Threads{} == size(G2STiledCopyQM2{})); + CUTE_STATIC_ASSERT_V(TileM {} % size<0>(G2STiledCopyShapeA {}) == _0{}); + CUTE_STATIC_ASSERT_V(TileK {} % size<1>(G2STiledCopyShapeA {}) == _0{}); + CUTE_STATIC_ASSERT_V(TileP {} % size<0>(G2STiledCopyShapeQ {}) == _0{}); + CUTE_STATIC_ASSERT_V(TileK {} % size<1>(G2STiledCopyShapeQ {}) == _0{}); + CUTE_STATIC_ASSERT_V(TileP2{} % size<0>(G2STiledCopyShapeQ2 {}) == _0{}); + CUTE_STATIC_ASSERT_V(TileK {} % size<1>(G2STiledCopyShapeQ2 {}) == _0{}); + CUTE_STATIC_ASSERT_V(TileN {} % size<0>(G2STiledCopyShapeS {}) == _0{}); + CUTE_STATIC_ASSERT_V(TileG {} % size<1>(G2STiledCopyShapeS {}) == _0{}); + CUTE_STATIC_ASSERT_V(QuantMapSize {} <= size (G2STiledCopyShapeQM {})); + CUTE_STATIC_ASSERT_V(QuantMapSize2{} <= size (G2STiledCopyShapeQM2{})); + CUTE_STATIC_ASSERT_V(QuantMapSize {} <= Warps{}); + CUTE_STATIC_ASSERT_V(QuantMapSize {} <= Threads{}); + + // + // --- Shared Memory to Shared Memory Copy --- + // + + using S2STiledCopySizeQM3 = _4; // 128b = 4 x 32b + using S2SQM3_1 = decltype(ceil_div(QuantMapDuplicates{}, S2STiledCopySizeQM3{})); + using S2SQM3_0 = decltype(ceil_div(Threads {}, S2SQM3_1 {})); + using S2SCopyAtomQM3 = Copy_Atom; + using S2STiledCopyQM3 = decltype(make_tiled_copy( + S2SCopyAtomQM3{}, + make_layout(make_shape(S2SQM3_0{}, S2SQM3_1{}), make_stride(S2SQM3_1{}, _1{})), + make_layout(make_shape(_1{}, S2STiledCopySizeQM3{})))); + + using S2STiledCopyShapeQM3 = decltype(shape(typename S2STiledCopyQM3::Tiler_MN{})); + CUTE_STATIC_ASSERT (sizeof(T2) == 4); + CUTE_STATIC_ASSERT_V((Threads {} == size (S2STiledCopyQM3{}))); + CUTE_STATIC_ASSERT_V((QuantMapSize2 {} % size<0>(S2STiledCopyShapeQM3{}) == _0{}) || (typename QuantMapVecTraits::IsDuplicated{} == false_type{})); + CUTE_STATIC_ASSERT_V((QuantMapDuplicates{} % size<1>(S2STiledCopyShapeQM3{}) == _0{}) || (typename QuantMapVecTraits::IsDuplicated{} == false_type{})); + + // + // --- Shared Memory to Registers Copy --- + // + using S2RCopyOpA = SM75_U32x4_LDSM_N; + using S2RCopyOpQ = conditional_t, SM75_U32x2_LDSM_N, SM75_U32x4_LDSM_N>; + using S2RCopyOpQ2 = SM75_U32x4_LDSM_N; + + using S2RCopyTraitsA = Copy_Traits; + using S2RCopyTraitsQ = Copy_Traits; + using S2RCopyTraitsQ2 = Copy_Traits; + + using S2RCopyAtomA = Copy_Atom; + using S2RCopyAtomQ = Copy_Atom; + using S2RCopyAtomQ2 = Copy_Atom; + using S2RCopyAtomSView = Copy_Atom; + using S2RCopyAtomQM = Copy_Atom; + + // + // --- Epilogue (Register to Global via Shared Memory) --- + // + using TiledMmaM = Int(TiledMma{})>; + using TiledMmaN = Int(TiledMma{})>; + using SmemLayoutAtomC = decltype(composition(Swizzle<2, 3, 3>{}, make_layout(make_shape(TiledMmaM{}, TiledMmaN{}), make_stride(TiledMmaN{}, _1{})))); + using SmemLayoutC = decltype(tile_to_shape(SmemLayoutAtomC{}, make_shape(TileM{}, TileN{}))); + // for some reason, the compiler will complain if we use `cosize_t` directly + using SmemLayoutCSize = cosize_t; + + // https://github.com/NVIDIA/cutlass/blob/v3.4.0/examples/50_hopper_gemm_with_epilogue_swizzle/50_hopper_gemm_with_epilogue_swizzle.cu + // note that all of these copy atoms are of type `T` instead of `TC` because we will convert + // the accumulated values to `T` before the epilogue. Performing epilogue type conversions before + // these copy operations can both save memory bandwidth and size of epilogue shared memory buffer. + using R2SCopyAtomC = Copy_Atom; + using S2RCopyAtomC = Copy_Atom, T>; + using R2GCopyAtomC = Copy_Atom, T>; + + using S2RC_M = TileM; + using S2RC_N = decltype(Threads{} / S2RC_M{}); + using S2RTiledCopyC = decltype(make_tiled_copy( + S2RCopyAtomC{}, + make_layout(make_shape(S2RC_M{}, S2RC_N{}), make_stride(S2RC_N{}, _1{})), + make_layout(make_shape(_1{}, EpilogueCopySizeC{})))); + + using S2RTiledCopyCShape = decltype(shape(typename S2RTiledCopyC::Tiler_MN{})); + CUTE_STATIC_ASSERT_V(TileM{} % size<0>(S2RTiledCopyCShape{}) == _0{}); + CUTE_STATIC_ASSERT_V(TileN{} % size<1>(S2RTiledCopyCShape{}) == _0{}); + + // + // --- Shared Memory Size --- + // + + struct SharedStorage + { + + array_aligned> smem_A; + array_aligned> smem_Q; + array_aligned> smem_S; + array_aligned> smem_QM; + + // optional + static constexpr int kSmemLayoutQ2Size = conditional_t , cosize_t, _0>::value; + static constexpr int kSmemLayoutQM2Size = conditional_t, cosize_t, _0>::value; + static constexpr int kSmemLayoutQM3Size = conditional_t, cosize_t, _0>::value; + static constexpr int kSmemLayoutCSize = conditional_t::value; + array_aligned smem_Q2; + array_aligned smem_QM2; + array_aligned smem_QM3; + array_aligned smem_C; + }; + + static constexpr int kSmemSize = int(sizeof(SharedStorage)); + +}; + +} // namespace config \ No newline at end of file diff --git a/kernels/quantization/flute/conversion_utils.hpp b/kernels/quantization/flute/conversion_utils.hpp new file mode 100644 index 0000000000..3115cf0c2b --- /dev/null +++ b/kernels/quantization/flute/conversion_utils.hpp @@ -0,0 +1,57 @@ +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/numeric_conversion.h" + + +namespace conversion_utils { +using namespace cute; + + +// https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp#L1260 +template < + class EngineSource, + class EngineTarget, + class TensorLayout, + int ConversionVectorWidth = cosize_v +> +CUTLASS_DEVICE void +convert_tensor( + Tensor const& source, + Tensor & target, + cute::Int width = {}) +{ + + /// This is an element-wise conversion where we expect both tensors to have the same layout. + /// As a result, we can cast as a cutlass array to use the fast numeric converters without + /// worrying about indexing into the layout. + constexpr int N = cosize_v; + + /// The inputs must be backed by registers & be statically sized. + static_assert(is_rmem::value, "Input tensor for A conversion must come from registers"); + static_assert(is_rmem::value, "Output tensor for A conversion must come from registers"); + static_assert(is_static_v, "Tensor layout for the conversion must be static"); + static_assert(cosize_v == size(TensorLayout{}), "Cosize and size of the layout must be equal."); + static_assert(N % ConversionVectorWidth == 0, "Conversion vector width must divide cosize of the tensor layout."); + + using SrcType = typename EngineSource::value_type; + using DstType = typename EngineTarget::value_type; + + using SrcArray = cutlass::Array; + using DstArray = cutlass::Array; + + constexpr cutlass::FloatRoundStyle RoundStyle = cutlass::FloatRoundStyle::round_to_nearest; + using Converter = cutlass::NumericArrayConverter; + + constexpr int NumIterations = N / ConversionVectorWidth; + + for (int ii = 0; ii < NumIterations; ++ii) + { + SrcArray const* src_array_ptr = reinterpret_cast(raw_pointer_cast(source.data())) + ii; + DstArray* dst_array_ptr = reinterpret_cast(raw_pointer_cast(target.data())) + ii; + *dst_array_ptr = Converter::convert(*src_array_ptr); + } +} + + +} // namespace conversion_utils \ No newline at end of file diff --git a/kernels/quantization/flute/cutlass_extensions_bf16.h b/kernels/quantization/flute/cutlass_extensions_bf16.h new file mode 100644 index 0000000000..a294c38eb2 --- /dev/null +++ b/kernels/quantization/flute/cutlass_extensions_bf16.h @@ -0,0 +1,42 @@ +#pragma once + +#include "cutlass/block_striped.h" + + +namespace cutlass { + + +/// Utility for performing block-striped access (load, store, reduce) of trivially-copyable, +/// statically-sized array types to global memory. +/// (Specialization for bfloat16_t. Uses nv_bfloat162 vectorized-reduction.) +template < + int BlockThreads, + typename ArrayT> +struct BlockStripedReduce : + BlockStriped< + BlockThreads, + ArrayT, + nv_bfloat162> +{ + static_assert(BlockStripedReduce::kStripes % 2 == 0, "Array of half must be even number in length"); + + /// Reduce + CUTLASS_DEVICE + static void reduce(ArrayT *ptr, const ArrayT &data, int thread_idx) + { + // This operation is natively supported by devices of compute + // capability 9.x and higher, older devices use emulation path + cutlass::atomic_add reduce; + nv_bfloat162 *access_output = reinterpret_cast(ptr); + const nv_bfloat162 *access_data = reinterpret_cast(&data); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < BlockStripedReduce::kStripes; ++i) + { + reduce(access_output + (BlockThreads * i) + thread_idx, access_data[i]); + } + } +}; + + +} // namespace cutlass \ No newline at end of file diff --git a/kernels/quantization/flute/hadamard_transform.cpp b/kernels/quantization/flute/hadamard_transform.cpp new file mode 100644 index 0000000000..88f2a0aeaa --- /dev/null +++ b/kernels/quantization/flute/hadamard_transform.cpp @@ -0,0 +1,57 @@ +// borrowed from https://github.com/pytorch-labs/applied-ai/tree/main/kernels/cuda/inference/hadamard_transform + +#include +#include +#include +#include + +using namespace torch::indexing; + +template +void run_fht(void* a, void* out, uint32_t numel, uint32_t had_size, cudaStream_t stream); + +constexpr bool is_power_of_two(uint32_t x) { + return x && !(x & (x - 1)); +} + +torch::Tensor hadamard_transform(at::Tensor& in, bool inplace) { + auto dtype = in.scalar_type(); + TORCH_CHECK(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently"); + TORCH_CHECK(in.is_cuda()); + + const int had_size = in.size(-1); + TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)), + "Only power of two Hadamard sizes up to 2^15 are supported, got ", had_size); + + const auto res_shape = in.sizes(); + torch::Tensor x = in.reshape({-1, had_size}); + + auto numel = in.numel(); + if (numel % 256 != 0) { + x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 0, 0, (256 - numel % 256) / had_size})); + } + + if (x.stride(-1) != 1) { + x = x.contiguous(); + } + torch::Tensor out = inplace ? x : torch::empty_like(x); + + at::cuda::CUDAGuard device_guard{(char)x.get_device()}; + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + if (dtype == torch::ScalarType::Half) { + run_fht(x.data_ptr(), out.data_ptr(), x.numel(), had_size, stream); + } else { + run_fht(x.data_ptr(), out.data_ptr(), x.numel(), had_size, stream); + } + + if (numel % 256 != 0) { + out = out.index({Slice(0, numel / had_size)}); + } + + if (inplace && out.data_ptr() != in.data_ptr()) { + in.copy_(out.view(res_shape)); + return in; + } + return out.reshape(res_shape); +} \ No newline at end of file diff --git a/kernels/quantization/flute/hadamard_transform_cuda.cu b/kernels/quantization/flute/hadamard_transform_cuda.cu new file mode 100644 index 0000000000..37ba47fbc6 --- /dev/null +++ b/kernels/quantization/flute/hadamard_transform_cuda.cu @@ -0,0 +1,751 @@ +// borrowed from https://github.com/pytorch-labs/applied-ai/tree/main/kernels/cuda/inference/hadamard_transform + +#include +#include +#include +#include +#include +#include + +#ifndef __CUDACC__ +#define __launch_bounds__(x,y) +#endif + +#define MAX_WARPS_PER_SM 48 + +#define MIN(a, b) ((a) < (b) ? (a) : (b)) + +typedef uint32_t b32; +typedef uint16_t b16; + +constexpr int launch_configs_big[7][3] = { + // default + {2, 1, 24}, + {2, 2, 16}, + {2, 4, 8}, + {2, 8, 4}, + {2, 16, 3}, + {4, 16, 2}, + {8, 16, 1} + // // extra coalescing + // {2, 1, 24}, + // {2, 2, 16}, + // {2, 4, 8}, + // {2, 8, 4}, + // {4, 8, 3}, + // {8, 8, 2}, + // {16, 8, 1} + // // less coalescing + // {2, 1, 24}, + // {2, 2, 16}, + // {2, 4, 8}, + // {2, 8, 4}, + // {1, 32, 1}, + // {2, 32, 1}, + // {4, 32, 1} +}; + +// a 4x2, b 2x2, c 2x2 +template +__device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32& c0, b32& c1){ + static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16); + // d, a, b, c + b32 zero = 0; + if constexpr(dtype == torch::ScalarType::Half) { + asm ( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};\n\t" + : "=r"(c0), "=r"(c1) : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "r"(zero), "r"(zero) + ); + } else { + b32 temp0, temp1, temp2, temp3; + asm ( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n\t" + : "=r"(temp0), "=r"(temp1), "=r"(temp2), "=r"(temp3) : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "r"(zero), "r"(zero), "r"(zero), "r"(zero) + ); + asm ("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c0) : "r"(temp1), "r"(temp0)); + asm ("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c1) : "r"(temp3), "r"(temp2)); + } +} + +// a 4x2, b 4x2, c 4x2 +template +__device__ __forceinline__ void mma_m16_n16_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32 b2, b32 b3, b32& c0, b32& c1, b32& c2, b32& c3){ + mma_m16_n8_k16_b16_b16_b16_noacc(a0, a1, a2, a3, b0, b1, c0, c1); + mma_m16_n8_k16_b16_b16_b16_noacc(a0, a1, a2, a3, b2, b3, c2, c3); +} + +__device__ __forceinline__ void matrix_transpose_m8_n8_b16_inplace(b32& a0) { + asm ( + "movmatrix.sync.aligned.m8n8.trans.b16 " + "%0, %1;\n\t" + : "=r"(a0) : "r"(a0) + ); +} + +#define p_p(i) ((val_1p[i] & 0x0000FFFF) | val_1p[i] << 16) +#define p_n(i) ((val_1p[i] & 0x0000FFFF) | val_1n[i] << 16) +#define n_p(i) ((val_1n[i] & 0x0000FFFF) | val_1p[i] << 16) +#define n_n(i) ((val_1n[i] & 0x0000FFFF) | val_1n[i] << 16) + +template +__global__ void __launch_bounds__(32 * warps_per_block, blocks_per_sm) +// a is column major, b is row major +hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) { + static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently"); + + b32 b_frag_all[num_chunks][4]; // for all chunks, holds matrix fragment (which takes 4 regs of b16x2 * 32 threads) + + uint blockid = blockIdx.x * warps_per_block + threadIdx.x / 32; + uint threadid = threadIdx.x % 32; + extern __shared__ b32 bfrag_arr[]; // num_chunks * warps_per_block * 128 + int real_num_chunks = ((blockid + 1) * num_chunks) > total_num_chunks ? (total_num_chunks - (blockid * num_chunks)) : num_chunks; + int diff_num_chunks = real_num_chunks - num_chunks; + + b32* a_start_ptr = (b32*) (a + blockid * num_chunks * 256); // offset a to where this warp starts + b32* out_start_ptr = (b32*) (out + blockid * num_chunks * 256); + b32* a_ptr = a_start_ptr + threadid * 4; + b32* b_frag_ptr = bfrag_arr + (blockid % warps_per_block) * num_chunks * 128 + threadid * 4; + + #if (__CUDA_ARCH__ < 900) // SM80, SM89 + uint64_t cache_policy; + asm volatile( + "createpolicy.fractional.L2::evict_first.b64 %0, 1.0;\n" + : "=l"(cache_policy) + ); + #endif + + #pragma unroll + for (int k = 0; k < num_chunks; k++) { + size_t shared_ptr = __cvta_generic_to_shared(b_frag_ptr); + #if (__CUDA_ARCH__ >= 900) // SM90 + asm volatile( + "cp.async.cg.shared.global [%0], [%1], 16;\n" + "cp.async.commit_group;\n" + :: "l"(shared_ptr), "l"(a_ptr) + ); + #else // SM80, SM89 + asm volatile( + "cp.async.cg.shared.global.L2::cache_hint.L2::256B [%0], [%1], 16, %2;\n" + "cp.async.commit_group;\n" + :: "l"(shared_ptr), "l"(a_ptr), "l"(cache_policy) + ); + #endif + + a_ptr += 128; + b_frag_ptr += 128; + } + + // generate hadamard 16x16 (up to 2 of them) + constexpr b16 fp16_1p[4] = {0b0011100110101000, 0b0011100000000000, 0b0011010110101000, 0b0011010000000000}; + constexpr b16 fp16_1n[4] = {0b1011100110101000, 0b1011100000000000, 0b1011010110101000, 0b1011010000000000}; + constexpr b16 bf16_1p[4] = {0b0011111100110101, 0b0011111100000000, 0b0011111010110101, 0b0011111010000000}; + constexpr b16 bf16_1n[4] = {0b1011111100110101, 0b1011111100000000, 0b1011111010110101, 0b1011111010000000}; + + #define val_type_1p(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i])) + #define val_type_1n(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i])) + constexpr b16 val_1p[4] = {val_type_1p(0), val_type_1p(1), val_type_1p(2), val_type_1p(3)}; + constexpr b16 val_1n[4] = {val_type_1n(0), val_type_1n(1), val_type_1n(2), val_type_1n(3)}; + + constexpr b32 p_p[4] = {p_p(0), p_p(1), p_p(2), p_p(3)}; + constexpr b32 p_n[4] = {p_n(0), p_n(1), p_n(2), p_n(3)}; + constexpr b32 n_p[4] = {n_p(0), n_p(1), n_p(2), n_p(3)}; + constexpr b32 n_n[4] = {n_n(0), n_n(1), n_n(2), n_n(3)}; + const b32 had_16_p1[4][4] = { + { + 0b10001000010001000010001000010001, + 0b00000000000000000000000000000000, + 0b00000000000000000000000000000000, + 0b10001000010001000010001000010001 + }, + { + 0b11001100100010000011001100100010, + 0b00000000000000000000000000000000, + 0b00000000000000000000000000000000, + 0b11001100100010000011001100100010 + }, + { + 0b11111111101010101100110010011001, + 0b00000000000000000000000000000000, + 0b00000000000000000000000000000000, + 0b11111111101010101100110010011001 + }, + { + 0b11111111101010101100110010011001, + 0b11111111101010101100110010011001, + 0b11111111101010101100110010011001, + 0b00000000010101010011001101100110 + } + }; + const b32 had_16_p2[4][4] = { + { + 0b10000000010000000010000000010000, + 0b00000000000000000000000000000000, + 0b00000000000000000000000000000000, + 0b10000000010000000010000000010000 + }, + { + 0b11000000100001000011000000100001, + 0b00000000000000000000000000000000, + 0b00000000000000000000000000000000, + 0b11000000100001000011000000100001 + }, + { + 0b11110000101001011100001110010110, + 0b00000000000000000000000000000000, + 0b00000000000000000000000000000000, + 0b11110000101001011100001110010110 + }, + { + 0b11110000101001011100001110010110, + 0b11110000101001011100001110010110, + 0b11110000101001011100001110010110, + 0b00001111010110100011110001101001 + } + }; + const b32 had_16_mask[3][4] = { + { + 0b10001000010001000010001000010001, + 0b00000000000000000000000000000000, + 0b00000000000000000000000000000000, + 0b10001000010001000010001000010001 + }, + { + 0b11001100110011000011001100110011, + 0b00000000000000000000000000000000, + 0b00000000000000000000000000000000, + 0b11001100110011000011001100110011 + }, + { + 0b11111111111111111111111111111111, + 0b00000000000000000000000000000000, + 0b00000000000000000000000000000000, + 0b11111111111111111111111111111111 + } + }; + b32 had_frag[8]; + #pragma unroll + for (int i = 0; i < 2; i++) { + int c_log_h = (i == 0) ? MIN(4, log_had_size) : log_had_size % 4; + #pragma unroll + for (int j = 0; j < 4; j++) { + if (c_log_h < 4) { + bool mask = had_16_mask[c_log_h - 1][j] & (1 << (31 - threadid)); + if (!mask) { + had_frag[i * 4 + j] = 0; + continue; + } + } + bool pred1 = had_16_p1[c_log_h - 1][j] & (1 << (31 - threadid)); + bool pred2 = had_16_p2[c_log_h - 1][j] & (1 << (31 - threadid)); + b32 val = pred1 ? (pred2 ? p_p[c_log_h - 1] : p_n[c_log_h - 1]) : (pred2 ? n_p[c_log_h - 1] : n_n[c_log_h - 1]); + had_frag[i * 4 + j] = val; + } + if constexpr(log_had_size <= 4 || log_had_size % 4 == 0) break; + } + + // log had size above 8, only used for above 2^8 = 256 size + constexpr int part8_log_had_size = log_had_size - 8; + + b32* a_chunk_ptr = a_start_ptr; // first chunk starts at this warp's data starts + b32* out_chunk_ptr = out_start_ptr; + + #pragma unroll + for (int l = 0; l < 2; l++) { + if constexpr(log_had_size <= 8) { // l == 0 guaranteed, redundant simplified version of else body, to help compiler warnings + b_frag_ptr = bfrag_arr + (blockid % warps_per_block) * num_chunks * 128; + } else { + b_frag_ptr = bfrag_arr + (blockid % warps_per_block) * num_chunks * (l == 0 ? 128 : (128 >> part8_log_had_size)); + } + + if (l == 1) { + if constexpr(log_had_size > 8) { + __syncthreads(); // sync between first and second iterations if above size 256 + + if constexpr(log_had_size >= 12) { + // sizes 4k and above + + // a + threadblock offset + warp offset + // can then index into all chunks owned by this warp + b32* store = bfrag_arr + (128 >> part8_log_had_size) * (num_chunks * (blockid % warps_per_block)); + + #pragma unroll + for (int j = 0; j < 4; j++) { + #pragma unroll + for (int k = 0; k < num_chunks; k++) { + // here, j represents register, and k represents 8-offset/chunk + int real_chunk_num = (num_chunks - (threadid % num_chunks) + k) % num_chunks; // chunk at which you have target thread #'s data + + int real_thread_id = (threadid / num_chunks) * num_chunks + k; // target thread # + int chunk_idx = 128 * real_chunk_num; // index due to fetching from another chunk (chunk in which this thread has the target thread's original data) + int thread_group_idx = (real_thread_id / 4) * 16; // index due to fetching from another group of num_chunk threads (since shuffle is between num_chunk threads) + int thread_idx = (real_thread_id % 4) * 2; // index due to original thread's position within the group of num_chunk threads + int reg_idx = (j / 2) * 8 + (j % 2); // index due to target register + int idx = chunk_idx + thread_group_idx + thread_idx + reg_idx; // final index + + // fix idx for majorness + int rowidx = idx % (1 << part8_log_had_size); + int colidx = idx >> part8_log_had_size; + + // store[rowidx * 128 + colidx] = data; + b32 data = store[rowidx * 128 + colidx]; + + // compiler generates excessive instructions, so we manually do the if statement + #pragma unroll + for (int i = 0; i < num_chunks; i++) { + asm volatile ( + "{\n\t" + " .reg .pred p0;\n\t" + " setp.eq.u32 p0, %1, %2;\n\t" + " @p0 mov.b32 %0, %3;\n\t" + "}\n\t" + : "+r"(b_frag_all[i][j]) // Output operand %0 + : "r"(real_chunk_num), "r"(i), "r"(data) // Input operands %1, %2, %3 + ); + } + } + } + + #pragma unroll + for (int j = 0; j < 4; j++) { + #pragma unroll + for (int k = 1; k < num_chunks; k++) { + int threadid_contig = threadid % num_chunks; + int threadid_mul = threadid / num_chunks; + int threadid2 = (threadid_contig + num_chunks - k) % num_chunks + threadid_mul * num_chunks; // thread to give your data to + b_frag_all[k][j] = __shfl_sync(0xFFFFFFFF, b_frag_all[k][j], threadid2); + } + } + } + } + } + + #pragma unroll + for (int k = 0; k < num_chunks; k++) { + if constexpr(enable_mask) { + if (k >= real_num_chunks) + break; + } + if (l == 0) { + // bad fix for k not being recognized as a constexpr by compiler + // asm("cp.async.wait_group %0;\n" :: "n"(num_chunks - k - 1)); + #define SWITCH_WAIT_ASYNC_LOAD_GROUP(i) case i: asm volatile("cp.async.wait_group %0;\n" :: "n"(num_chunks - i - 1)); break; + if constexpr(enable_mask) { + switch(k + diff_num_chunks) { + SWITCH_WAIT_ASYNC_LOAD_GROUP(0) + SWITCH_WAIT_ASYNC_LOAD_GROUP(1) + SWITCH_WAIT_ASYNC_LOAD_GROUP(2) + SWITCH_WAIT_ASYNC_LOAD_GROUP(3) + SWITCH_WAIT_ASYNC_LOAD_GROUP(4) + SWITCH_WAIT_ASYNC_LOAD_GROUP(5) + SWITCH_WAIT_ASYNC_LOAD_GROUP(6) + SWITCH_WAIT_ASYNC_LOAD_GROUP(7) + SWITCH_WAIT_ASYNC_LOAD_GROUP(8) + SWITCH_WAIT_ASYNC_LOAD_GROUP(9) + SWITCH_WAIT_ASYNC_LOAD_GROUP(10) + SWITCH_WAIT_ASYNC_LOAD_GROUP(11) + SWITCH_WAIT_ASYNC_LOAD_GROUP(12) + SWITCH_WAIT_ASYNC_LOAD_GROUP(13) + SWITCH_WAIT_ASYNC_LOAD_GROUP(14) + SWITCH_WAIT_ASYNC_LOAD_GROUP(15) + SWITCH_WAIT_ASYNC_LOAD_GROUP(16) + SWITCH_WAIT_ASYNC_LOAD_GROUP(17) + SWITCH_WAIT_ASYNC_LOAD_GROUP(18) + SWITCH_WAIT_ASYNC_LOAD_GROUP(19) + SWITCH_WAIT_ASYNC_LOAD_GROUP(20) + SWITCH_WAIT_ASYNC_LOAD_GROUP(21) + SWITCH_WAIT_ASYNC_LOAD_GROUP(22) + SWITCH_WAIT_ASYNC_LOAD_GROUP(23) + SWITCH_WAIT_ASYNC_LOAD_GROUP(24) + SWITCH_WAIT_ASYNC_LOAD_GROUP(25) + SWITCH_WAIT_ASYNC_LOAD_GROUP(26) + SWITCH_WAIT_ASYNC_LOAD_GROUP(27) + SWITCH_WAIT_ASYNC_LOAD_GROUP(28) + SWITCH_WAIT_ASYNC_LOAD_GROUP(29) + SWITCH_WAIT_ASYNC_LOAD_GROUP(30) + SWITCH_WAIT_ASYNC_LOAD_GROUP(31) + } + } else { + switch(k) { + SWITCH_WAIT_ASYNC_LOAD_GROUP(0) + SWITCH_WAIT_ASYNC_LOAD_GROUP(1) + SWITCH_WAIT_ASYNC_LOAD_GROUP(2) + SWITCH_WAIT_ASYNC_LOAD_GROUP(3) + SWITCH_WAIT_ASYNC_LOAD_GROUP(4) + SWITCH_WAIT_ASYNC_LOAD_GROUP(5) + SWITCH_WAIT_ASYNC_LOAD_GROUP(6) + SWITCH_WAIT_ASYNC_LOAD_GROUP(7) + SWITCH_WAIT_ASYNC_LOAD_GROUP(8) + SWITCH_WAIT_ASYNC_LOAD_GROUP(9) + SWITCH_WAIT_ASYNC_LOAD_GROUP(10) + SWITCH_WAIT_ASYNC_LOAD_GROUP(11) + SWITCH_WAIT_ASYNC_LOAD_GROUP(12) + SWITCH_WAIT_ASYNC_LOAD_GROUP(13) + SWITCH_WAIT_ASYNC_LOAD_GROUP(14) + SWITCH_WAIT_ASYNC_LOAD_GROUP(15) + SWITCH_WAIT_ASYNC_LOAD_GROUP(16) + SWITCH_WAIT_ASYNC_LOAD_GROUP(17) + SWITCH_WAIT_ASYNC_LOAD_GROUP(18) + SWITCH_WAIT_ASYNC_LOAD_GROUP(19) + SWITCH_WAIT_ASYNC_LOAD_GROUP(20) + SWITCH_WAIT_ASYNC_LOAD_GROUP(21) + SWITCH_WAIT_ASYNC_LOAD_GROUP(22) + SWITCH_WAIT_ASYNC_LOAD_GROUP(23) + SWITCH_WAIT_ASYNC_LOAD_GROUP(24) + SWITCH_WAIT_ASYNC_LOAD_GROUP(25) + SWITCH_WAIT_ASYNC_LOAD_GROUP(26) + SWITCH_WAIT_ASYNC_LOAD_GROUP(27) + SWITCH_WAIT_ASYNC_LOAD_GROUP(28) + SWITCH_WAIT_ASYNC_LOAD_GROUP(29) + SWITCH_WAIT_ASYNC_LOAD_GROUP(30) + SWITCH_WAIT_ASYNC_LOAD_GROUP(31) + } + } + } + + if (l == 0) { + // loading for the first iteration + + // thread 0 loads [t0r0, t16r1, t0r2, t16r3] + // thread 16 loads [t0r1, t16r0, t0r3, t16r2] + // allows full coalescing, same for t1/t17, t2/t18, etc. + #pragma unroll + for (int j = 0; j < 4; j++) { + int reg = ((threadid & 16) == 0) ? j : (j / 2 * 2 + (1 - j % 2)); + int real_thread_id = (reg == 0 || reg == 2) ? threadid : (threadid ^ 16); + int real_row = real_thread_id % 4; + int real_col = real_thread_id / 4; + b_frag_all[k][j] = b_frag_ptr[(real_row + (reg % 2) * 4) + (real_col + (j / 2) * 8) * 8]; + } + + // for t16 swap r0/r1 and r2/r3 to have [t16r0, t0r1, t16r2, t0r3] + // so registers are in right order, same for t17, t18, etc. + if ((threadid & 16) != 0) { + b32 temp = b_frag_all[k][0]; + b_frag_all[k][0] = b_frag_all[k][1]; + b_frag_all[k][1] = temp; + + temp = b_frag_all[k][2]; + b_frag_all[k][2] = b_frag_all[k][3]; + b_frag_all[k][3] = temp; + } + + // t0 and t16 swap r1 and r3 to have their own data, + // same for t1/t17, t2/18, etc. + #pragma unroll + for (int j = 1; j < 4; j += 2) { + b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], 16); + } + } else if constexpr(log_had_size > 8) { // condition is redundant to help compiler warnings + if constexpr(log_had_size < 12) { + // sizes 512, 1k, and 2k + + // for 512: + // thread 0 loads [t0r0, t0r1, t16r2, t16r3] + // thread 16 loads [t0r2, t0r3, t16r0, t16r1] + // same for t1/t17, t2/t18, etc. + // for 1k and 2k: + // thread 0 loads [t0r0, t0r1, t1r2, t1r3] + // thread 1 loads [t0r2, t0r3, t1r0, t1r1] + // same for t2/t3, t4/t5, etc. + // allows full coalescing for 512 and 1k, 16x coalescing for 2k + constexpr int xor_val = log_had_size == 9 ? 16 : 1; + + #pragma unroll + for (int j = 0; j < 4; j++) { + int reg = ((threadid & xor_val) == 0) ? j : (j + 2) % 4; + int real_thread_id = reg < 2 ? threadid : (threadid ^ xor_val); + int idx = (real_thread_id / 4 * 16) + (real_thread_id % 4 * 2) + (reg / 2 * 8) + (reg % 2); + int rowidx = idx % (1 << part8_log_had_size); + int colidx = idx >> part8_log_had_size; + b_frag_all[k][j] = b_frag_ptr[rowidx * 128 + colidx]; + } + + if ((threadid & xor_val) != 0) { + b32 temp = b_frag_all[k][0]; + b_frag_all[k][0] = b_frag_all[k][2]; + b_frag_all[k][2] = temp; + + temp = b_frag_all[k][1]; + b_frag_all[k][1] = b_frag_all[k][3]; + b_frag_all[k][3] = temp; + } + + #pragma unroll + for (int j = 2; j < 4; j++) { + b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], xor_val); + } + } + } + + if (l == 1) { + // for second iteration, we load 2 consecutive b16s (1 b32) per register, + // but tensor core register layout requires 2 b16s that are in the + // same column/consecutive rows to be in the same register, so do the swap + b32 f0 = ((b_frag_all[k][1] & 0xFFFF) << 16) | (b_frag_all[k][0] & 0xFFFF); + b32 f1 = ((b_frag_all[k][3] & 0xFFFF) << 16) | (b_frag_all[k][2] & 0xFFFF); + b32 f2 = (b_frag_all[k][1] & 0xFFFF0000) | (b_frag_all[k][0] >> 16); + b32 f3 = (b_frag_all[k][3] & 0xFFFF0000) | (b_frag_all[k][2] >> 16); + b_frag_all[k][0] = f0; + b_frag_all[k][1] = f1; + b_frag_all[k][2] = f2; + b_frag_all[k][3] = f3; + } + + #pragma unroll + for(int i = 0, remaining_log_had_size = log_had_size - l * 8; i < 2 && remaining_log_had_size > 0; i++) { + int had_off = ((remaining_log_had_size < 4) && !(log_had_size <= 4 || log_had_size % 4 == 0)) ? 4 : 0; + mma_m16_n16_k16_b16_b16_b16_noacc(had_frag[had_off + 0], had_frag[had_off + 1], had_frag[had_off + 2], had_frag[had_off + 3], b_frag_all[k][0], b_frag_all[k][1], b_frag_all[k][2], b_frag_all[k][3], b_frag_all[k][0], b_frag_all[k][1], b_frag_all[k][2], b_frag_all[k][3]); + + remaining_log_had_size -= 4; + if (remaining_log_had_size <= 0 && i == 0) { + // TODO: consider different storing so no need for transpose + matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][0]); + matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][1]); + matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][2]); + matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][3]); + } else { + // swap and use output directly as b_frag for next iteration as an actually free transpose + b32 temp = b_frag_all[k][1]; + b_frag_all[k][1] = b_frag_all[k][2]; + b_frag_all[k][2] = temp; + } + } + + if (l == 1) { + // invert swap from above for second iteration + b32 f0 = ((b_frag_all[k][2] & 0xFFFF) << 16) | (b_frag_all[k][0] & 0xFFFF); + b32 f1 = (b_frag_all[k][2] & 0xFFFF0000) | (b_frag_all[k][0] >> 16); + b32 f2 = ((b_frag_all[k][3] & 0xFFFF) << 16) | (b_frag_all[k][1] & 0xFFFF); + b32 f3 = (b_frag_all[k][3] & 0xFFFF0000) | (b_frag_all[k][1] >> 16); + b_frag_all[k][0] = f0; + b_frag_all[k][1] = f1; + b_frag_all[k][2] = f2; + b_frag_all[k][3] = f3; + } + + if (l == 0) { + // inverse of coalesced load for first iteration to store result + #pragma unroll + for (int j = 1; j < 4; j += 2) { + b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], 16); + } + + if ((threadid & 16) != 0) { + b32 temp = b_frag_all[k][0]; + b_frag_all[k][0] = b_frag_all[k][1]; + b_frag_all[k][1] = temp; + + temp = b_frag_all[k][2]; + b_frag_all[k][2] = b_frag_all[k][3]; + b_frag_all[k][3] = temp; + } + + // if only going up to 256 size, store directly back to global memory, + // otherwise store back to shared memory for next iteration + b32* store = (log_had_size <= 8) ? out_chunk_ptr : b_frag_ptr; + + #pragma unroll + for (int j = 0; j < 4; j++) { + int reg = ((threadid & 16) == 0) ? j : (j / 2 * 2 + (1 - j % 2)); + int real_thread_id = (reg == 0 || reg == 2) ? threadid : (threadid ^ 16); + int real_row = real_thread_id % 4; + int real_col = real_thread_id / 4; + store[(real_row + (reg % 2) * 4) + (real_col + (reg / 2) * 8) * 8] = b_frag_all[k][j]; + } + } else if constexpr(log_had_size > 8) { // condition is redundant to help compiler warnings + if (log_had_size < 12) { + // inverse of coalesced load for sizes 512, 1k and 2k to store result + constexpr int xor_val = log_had_size == 9 ? 16 : 1; + #pragma unroll + for (int j = 2; j < 4; j++) { + b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], xor_val); + } + + if ((threadid & xor_val) != 0) { + b32 temp = b_frag_all[k][0]; + b_frag_all[k][0] = b_frag_all[k][2]; + b_frag_all[k][2] = temp; + + temp = b_frag_all[k][1]; + b_frag_all[k][1] = b_frag_all[k][3]; + b_frag_all[k][3] = temp; + } + + b32* store = (b32*)(out + (blockid / warps_per_block) * (num_chunks * warps_per_block) * 256 + (256 >> part8_log_had_size) * (num_chunks * (blockid % warps_per_block) + k)); + #pragma unroll + for (int j = 0; j < 4; j++) { + int reg = ((threadid & xor_val) == 0) ? j : (j + 2) % 4; + b32 data = b_frag_all[k][j]; + int real_thread_id = reg < 2 ? threadid : (threadid ^ xor_val); + int idx = (real_thread_id / 4 * 16) + (real_thread_id % 4 * 2) + (reg / 2 * 8) + (reg % 2); + int rowidx = idx % (1 << part8_log_had_size); + int colidx = idx >> part8_log_had_size; + store[rowidx * 128 + colidx] = data; + } + } + // for size 4k and above, wait to process all chunks so a final store can be performed coalesced + } + + a_chunk_ptr += 128; // (only affects first 256 size) move on to next chunk by skipping 256 elements in b16 (= 128 in b32) + out_chunk_ptr += 128; + if constexpr(log_had_size > 8) { + b_frag_ptr += (l == 0 ? 128 : (128 >> part8_log_had_size)); + } else { // else is redundant, simplified version of if body, to help compiler warnings + b_frag_ptr += 128; + } + } + if (log_had_size <= 8) + break; + } + + if constexpr(log_had_size >= 12) { + // for sizes 4k and above, perform final coalesced store after processing all chunks + #pragma unroll + for (int j = 0; j < 4; j++) { + #pragma unroll + for (int k = 1; k < num_chunks; k++) { + int threadid_contig = threadid % num_chunks; + int threadid_mul = threadid / num_chunks; + int threadid2 = (threadid_contig + k) % num_chunks + threadid_mul * num_chunks; // thread to give your data to + b_frag_all[k][j] = __shfl_sync(0xFFFFFFFF, b_frag_all[k][j], threadid2); + } + } + + // a + threadblock offset + warp offset + // can then index into all chunks owned by this warp + b32* store = bfrag_arr + (128 >> part8_log_had_size) * (num_chunks * (blockid % warps_per_block)); + + #pragma unroll + for (int j = 0; j < 4; j++) { + #pragma unroll + for (int k = 0; k < num_chunks; k++) { + // here, j represents register, and k represents 8-offset/chunk + int real_chunk_num = (num_chunks - (threadid % num_chunks) + k) % num_chunks; // chunk at which you have target thread #'s data + + // b32 data = b_frag_all[real_chunk_num][j]; // target thread data + b32 data; + #pragma unroll + for (int i = 0; i < num_chunks; i++) { + if (real_chunk_num == i) data = b_frag_all[i][j]; + } + + int real_thread_id = (threadid / num_chunks) * num_chunks + k; // target thread # + int chunk_idx = 128 * real_chunk_num; // index due to fetching from another chunk (chunk in which this thread has the target thread's original data) + int thread_group_idx = (real_thread_id / 4) * 16; // index due to fetching from another group of num_chunk threads (since shuffle is between num_chunk threads) + int thread_idx = (real_thread_id % 4) * 2; // index due to original thread's position within the group of num_chunk threads + int reg_idx = (j / 2) * 8 + (j % 2); // index due to target register + int idx = chunk_idx + thread_group_idx + thread_idx + reg_idx; // final index + + // fix idx for majorness + int rowidx = idx % (1 << part8_log_had_size); + int colidx = idx >> part8_log_had_size; + + store[rowidx * 128 + colidx] = data; + } + } + + __syncthreads(); + store = ((b32*) out) + (blockid / warps_per_block) * (num_chunks * warps_per_block) * 128; + int4* store4 = (int4*) store; + int4* bfrag_arr4 = (int4*) bfrag_arr; + // flush smem, simply linearly write to store + // always divisible by 128*32b, so (32*4)*32b is ok + #pragma unroll + for (int warp_off = 0; warp_off < (num_chunks * warps_per_block * 128 / 4); warp_off += 32 * warps_per_block) { + int total_off = warp_off + threadid + (blockid % warps_per_block) * 32; + store4[total_off] = bfrag_arr4[total_off]; + } + } + +} + +constexpr int ceil_div(int a, int b) { + return (a + b - 1) / b; +} + +template +void __forceinline__ run_kernel(b16* a_mat, b16* out, int num_chunks, cudaStream_t stream) { + int shared_size = chunks_per_warp * warps_per_block * 128 * 4; + dim3 block_size = 32 * warps_per_block; + + #define CHECK_SHARED_LIM() { \ + if (shared_size > 48 * 1024) { \ + C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \ + } \ + } \ + + if constexpr(check_masking) { + if (num_chunks % (chunks_per_warp * warps_per_block) != 0) { + dim3 grid_size = ceil_div(ceil_div(num_chunks, chunks_per_warp), warps_per_block); + auto kernel = hadamard_transform_kernel; + CHECK_SHARED_LIM(); + kernel<<>>(a_mat, out, num_chunks); + } else { + dim3 grid_size = num_chunks / chunks_per_warp / warps_per_block; + auto kernel = hadamard_transform_kernel; + CHECK_SHARED_LIM(); + kernel<<>>(a_mat, out, num_chunks); + } + } else { + dim3 grid_size = num_chunks / chunks_per_warp / warps_per_block; + auto kernel = hadamard_transform_kernel; + CHECK_SHARED_LIM(); + kernel<<>>(a_mat, out, num_chunks); + } + + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void run_fht(void* a_mat_ptr, void* out_ptr, uint32_t numel, uint32_t had_size, cudaStream_t stream) { + uint32_t num_chunks = numel / 256; // caller required to ensure divisible by 256 + // for size 256, use (2, 1) + // for size 32k use (8, 16) + constexpr int chunks_per_warp_small = 1;// 8; + constexpr int warps_per_block_small = 1;//2;//16; + constexpr int blocks_per_sm_small = 24; + constexpr int chunks_per_warp_large = 2; + constexpr int warps_per_block_large = 1; + constexpr int blocks_per_sm_large = 24; + + // constexpr torch::ScalarType dtype = torch::ScalarType::Half; + + b16* a_mat = (b16*) a_mat_ptr; + b16* out = (b16*) out_ptr; + + if (numel <= 256) { + switch (had_size) { + case (1<<1): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<2): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<3): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<4): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<5): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<6): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<7): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<8): run_kernel(a_mat, out, num_chunks, stream); break; + } + } else { + switch (had_size) { + case (1<<1): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<2): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<3): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<4): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<5): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<6): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<7): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<8): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<9): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<10): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<11): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<12): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<13): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<14): run_kernel(a_mat, out, num_chunks, stream); break; + case (1<<15): run_kernel(a_mat, out, num_chunks, stream); break; + } + } +} + +template void run_fht(void* a_mat_ptr, void* out_ptr, uint32_t numel, uint32_t had_size, cudaStream_t stream); +template void run_fht(void* a_mat_ptr, void* out_ptr, uint32_t numel, uint32_t had_size, cudaStream_t stream); \ No newline at end of file diff --git a/kernels/quantization/flute/marlin_utils.hpp b/kernels/quantization/flute/marlin_utils.hpp new file mode 100644 index 0000000000..da02f887e1 --- /dev/null +++ b/kernels/quantization/flute/marlin_utils.hpp @@ -0,0 +1,95 @@ +/* + * Copyright (C) Marlin.2024 Elias Frantar (elias.frantar@ist.ac.at) + * + * Licensed 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. + */ + + +#ifndef MARLIN_CUDA_KERNEL_CUH +#define MARLIN_CUDA_KERNEL_CUH + + +#include +#include +#include +#include + + +namespace marlin_utils { + + +// Instances of `Vec` are used to organize groups of >>registers<<, as needed for instance as inputs to tensor core +// operations. Consequently, all corresponding index accesses must be compile-time constants, which is why we +// extensively use `#pragma unroll` throughout the kernel code to guarantee this. +template +struct Vec { + T elems[n]; + __device__ T& operator[](int i) { + return elems[i]; + } +}; + + +using I4 = Vec; + +// Matrix fragments for tensor core instructions; their precise layout is documented here: +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type +using FragA = Vec; +using FragB = Vec; +using FragC = Vec; +using FragS = Vec; // quantization scales + + +// Lookup-table based 3-input logical operation; explicitly used for dequantization as the compiler does not seem to +// automatically recognize it in all cases. +template +__device__ inline int lop3(int a, int b, int c) { + int res; + asm volatile( + "lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(res) : "r"(a), "r"(b), "r"(c), "n"(lut) + ); + return res; +} + +// Efficiently dequantize an int32 value into a full B-fragment of 4 fp16 values. +// We mostly follow the strategy in the link below, with some small changes: +// https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h +__device__ inline FragB dequant(int q) { + const int LO = 0x000f000f; + const int HI = 0x00f000f0; + const int EX = 0x64006400; + // Guarantee that the `(a & b) | c` operations are LOP3s. + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); + // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point directly into `SUB` and `ADD`. + const int SUB = 0x64086408; + const int MUL = 0x2c002c00; + const int ADD = 0xd480d480; + FragB frag_b; + frag_b[0] = __hsub2( + *reinterpret_cast(&lo), + *reinterpret_cast(&SUB) + ); + frag_b[1] = __hfma2( + *reinterpret_cast(&hi), + *reinterpret_cast(&MUL), *reinterpret_cast(&ADD) + ); + return frag_b; +} + + +} // namespace marlin_utils + + +#endif \ No newline at end of file diff --git a/kernels/quantization/flute/packbits_utils.hpp b/kernels/quantization/flute/packbits_utils.hpp new file mode 100644 index 0000000000..d705ae3ec3 --- /dev/null +++ b/kernels/quantization/flute/packbits_utils.hpp @@ -0,0 +1,429 @@ +#pragma once + +#include +#include +#include +#include + +#include "config.hpp" +#include "marlin_utils.hpp" + + +namespace packbits_utils { + + +template +struct DequantizationTraits +{ + + CUTE_DEVICE static + void + apply( + cute::Tensor const& source, + cute::Tensor const& source2, + cute::Tensor & target, + cute::Tensor const& scale, + cute::Tensor const& qmap, + cute::Tensor const& qmap2, + cute::Tensor const& qmap3) + { + + using TQ = cute::uint16_t; + using TQ2 = cute::uint32_t; + using T = typename TargetEngine::value_type; + using TI = cute::conditional_t, __half , __nv_bfloat16 >; + using T2 = cute::conditional_t, __half2, __nv_bfloat162>; + CUTE_STATIC_ASSERT(cute::is_same_v == true || + cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + + static constexpr int kNumBits = NumBits::value; + static constexpr int kNumPacked2 = NumBits::value == 4 ? 8 : 16; + static constexpr cute::uint16_t kMask = NumBits::value == 4 ? 0x000f : 0x0003; + static constexpr cute::uint32_t kMask2 = NumBits::value == 4 ? 0x000000ff : 0x0000000f; + static constexpr cute::uint32_t kMaskSync = 0xffffffff; + + // vectorize the source and target + auto source_vec = cute::recast(source); + auto source2_vec = cute::recast(source2); // unused + auto target_vec = cute::recast(target); + auto scale_vec = cute::recast(scale); + auto qmap_view = cute::recast(qmap); + auto qmap2_view = cute::recast(qmap2); + auto qmap3_view = cute::recast(qmap3); + + CUTE_STATIC_ASSERT_V(NumBits{} == cute::_4{} || NumBits{} == cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(source ) == cute::size<0>(source_vec ) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(source2) == cute::size<0>(source2_vec) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(target ) == cute::size<0>(target_vec ) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(scale ) == cute::size<0>(scale_vec ) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<1>(source ) == cute::size<1>(source_vec )); + CUTE_STATIC_ASSERT_V(cute::size<1>(source2) == cute::size<1>(source2_vec)); + CUTE_STATIC_ASSERT_V(cute::size<1>(target ) == cute::size<1>(target_vec )); + CUTE_STATIC_ASSERT_V(cute::size<1>(scale ) == cute::size<1>(scale_vec )); + CUTE_STATIC_ASSERT_V(cute::size (qmap ) == cute::size (qmap_view )); + CUTE_STATIC_ASSERT_V(cute::size (qmap2 ) == cute::size (qmap2_view )); + CUTE_STATIC_ASSERT_V(cute::size (qmap3 ) == cute::size (qmap3_view )); + + CUTE_UNROLL + for (int i = 0; i < cute::size<0>(source_vec); ++i) + { + + CUTE_UNROLL + for (int p = 0; p < cute::size<1>(source_vec); ++p) + { + auto src_crd = cute::make_coord(i, p); + + CUTE_UNROLL + for (int k2 = 0; k2 < kNumPacked2; k2 += 2) + { + auto k = k2 / 2; + auto tgt_crd = cute::make_coord(i, k * cute::size<1>(source_vec) + p); + auto src_raw = source_vec(src_crd) >> (k2 * kNumBits); + T2 src_val; + + if constexpr ((QuantMapMode == config::QuantMapModeEnum::Vectorized ) || + (QuantMapMode == config::QuantMapModeEnum::Vectorized_32) || + (QuantMapMode == config::QuantMapModeEnum::Vectorized_16) || + (QuantMapMode == config::QuantMapModeEnum::Vectorized_8)) + { + // vectorized table lookup + src_val = qmap2_view[src_raw & kMask2]; + + } + else + { + + TI src_val_0; + TI src_val_1; + + if constexpr (QuantMapMode == config::QuantMapModeEnum::WarpShuffle) + { + // in-register table lookup + src_val_0 = __shfl_sync(kMaskSync, qmap3_view(0), (src_raw >> kNumBits) & kMask); + src_val_1 = __shfl_sync(kMaskSync, qmap3_view(0), (src_raw ) & kMask); + } + else + { + // normal table lookup + src_val_0 = qmap_view[(src_raw >> kNumBits) & kMask]; + src_val_1 = qmap_view[(src_raw ) & kMask]; + } + + if constexpr (cute::is_same_v) + { + src_val = __halves2half2 (src_val_0, src_val_1); + } + else + { + src_val = __halves2bfloat162(src_val_0, src_val_1); + } + + } + + // vectorized scaling + target_vec(tgt_crd) = __hmul2(src_val, scale_vec(tgt_crd)); + } + } + } + } +}; + + +template +struct DequantizationTraits, + config::QuantMapModeEnum::Marlin> +{ + + CUTE_DEVICE static + void + apply( + cute::Tensor const& source, + cute::Tensor const& source2, + cute::Tensor & target, + cute::Tensor const& scale, + cute::Tensor const& qmap, + cute::Tensor const& qmap2, + cute::Tensor const& qmap3) + { + + using TQ = cute::uint16_t; + using TQ2 = cute::uint32_t; + using T = typename TargetEngine::value_type; + using TI = cute::conditional_t, __half , __nv_bfloat16 >; + using T2 = cute::conditional_t, __half2, __nv_bfloat162>; + CUTE_STATIC_ASSERT(cute::is_same_v == true || + cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + + static constexpr int kNumBits = 4; + static constexpr int kNumPacked2 = 8; + + // vectorize the source and target + auto source_vec = cute::recast(source); + auto source2_vec = cute::recast(source2); // unused + auto target_vec = cute::recast(target); + auto scale_vec = cute::recast(scale); + auto qmap_view = cute::recast(qmap); + auto qmap2_view = cute::recast(qmap2); + auto qmap3_view = cute::recast(qmap3); + + CUTE_STATIC_ASSERT_V(cute::size<0>(source ) == cute::size<0>(source_vec ) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(source2) == cute::size<0>(source2_vec) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(target ) == cute::size<0>(target_vec ) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(scale ) == cute::size<0>(scale_vec ) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<1>(source ) == cute::size<1>(source_vec )); + CUTE_STATIC_ASSERT_V(cute::size<1>(source2) == cute::size<1>(source2_vec)); + CUTE_STATIC_ASSERT_V(cute::size<1>(target ) == cute::size<1>(target_vec )); + CUTE_STATIC_ASSERT_V(cute::size<1>(scale ) == cute::size<1>(scale_vec )); + CUTE_STATIC_ASSERT_V(cute::size (qmap ) == cute::size (qmap_view )); + CUTE_STATIC_ASSERT_V(cute::size (qmap2 ) == cute::size (qmap2_view )); + CUTE_STATIC_ASSERT_V(cute::size (qmap3 ) == cute::size (qmap3_view )); + + CUTE_UNROLL + for (int i = 0; i < cute::size<0>(source_vec); ++i) + { + + CUTE_UNROLL + for (int p = 0; p < cute::size<1>(source_vec); ++p) + { + auto src_crd = cute::make_coord(i, p); + + CUTE_UNROLL + for (int k4 = 0; k4 < kNumPacked2; k4 += 4) + { + auto k = k4 / 4; + auto src_raw = source_vec(src_crd) >> (k * 8); + auto src_val = marlin_utils::dequant(src_raw); + + auto tgt0_crd = cute::make_coord(i, (k * 2 + 0) * cute::size<1>(source_vec) + p); + auto tgt1_crd = cute::make_coord(i, (k * 2 + 1) * cute::size<1>(source_vec) + p); + target_vec(tgt0_crd) = __hmul2(src_val[0], scale_vec(tgt0_crd)); + target_vec(tgt1_crd) = __hmul2(src_val[1], scale_vec(tgt1_crd)); + } + } + } + } +}; + + +template +struct DequantizationTraits, + QuantMapMode> +{ + + CUTE_DEVICE static + void + apply( + cute::Tensor const& source, + cute::Tensor const& source2, + cute::Tensor & target, + cute::Tensor const& scale, + cute::Tensor const& qmap, + cute::Tensor const& qmap2, + cute::Tensor const& qmap3) + { + + using TQ = cute::uint16_t; + using TQ2 = cute::uint32_t; + using T = typename TargetEngine::value_type; + using TI = cute::conditional_t, __half , __nv_bfloat16 >; + using T2 = cute::conditional_t, __half2, __nv_bfloat162>; + CUTE_STATIC_ASSERT(cute::is_same_v == true || + cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + CUTE_STATIC_ASSERT(cute::is_same_v == true); + + CUTE_STATIC_ASSERT (QuantMapMode == config::QuantMapModeEnum::Vectorized); + CUTE_STATIC_ASSERT_V(cute::size<1>(source2) == cute::size<1>(source) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(source2) == cute::size<0>(source)); + + static constexpr int kNumBits = 3; + static constexpr int kNumPacked2 = 10; + static constexpr cute::uint32_t kMask2 = 0x0000003f; + static constexpr cute::uint32_t kMaskF2 = 0x00000003; + + // vectorize the source and target + auto source0_vec = cute::recast(source); + auto source1_vec = cute::recast(source2); + auto source2_vec = cute::recast(source2); // the same as source2 + auto target_vec = cute::recast(target); + auto scale_vec = cute::recast(scale); + auto qmap_view = cute::recast(qmap); + auto qmap2_view = cute::recast(qmap2); + auto qmap3_view = cute::recast(qmap3); + + CUTE_STATIC_ASSERT_V(cute::size<0>(source ) == cute::size<0>(source0_vec) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(source2) == cute::size<0>(source1_vec) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(source2) == cute::size<0>(source2_vec) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(target ) == cute::size<0>(target_vec ) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<0>(scale ) == cute::size<0>(scale_vec ) * cute::_2{}); + CUTE_STATIC_ASSERT_V(cute::size<1>(source ) == cute::size<1>(source0_vec)); + CUTE_STATIC_ASSERT_V(cute::size<1>(source2) == cute::size<1>(source1_vec)); + CUTE_STATIC_ASSERT_V(cute::size<1>(source2) == cute::size<1>(source2_vec)); + CUTE_STATIC_ASSERT_V(cute::size<1>(target ) == cute::size<1>(target_vec )); + CUTE_STATIC_ASSERT_V(cute::size<1>(scale ) == cute::size<1>(scale_vec )); + CUTE_STATIC_ASSERT_V(cute::size (qmap ) == cute::size (qmap_view )); + CUTE_STATIC_ASSERT_V(cute::size (qmap2 ) == cute::size (qmap2_view )); + CUTE_STATIC_ASSERT_V(cute::size (qmap3 ) == cute::size (qmap3_view )); + + CUTE_UNROLL + for (int i = 0; i < cute::size<0>(source0_vec); ++i) + { + + CUTE_UNROLL + for (int p = 0; p < cute::size<1>(source0_vec); ++p) + { + auto src0_crd = cute::make_coord(i, p); + auto src1_crd = cute::make_coord(i, p * 2); + auto src2_crd = cute::make_coord(i, p * 2 + 1); + + CUTE_UNROLL + for (int k2 = 0; k2 < kNumPacked2; k2 += 2) + { + auto k = k2 / 2; + // using `source0_vec` for the stride, since others sahre the same value + auto tgt0_crd = cute::make_coord(i, (k * 3 + 0) * cute::size<1>(source0_vec) + p); + auto tgt1_crd = cute::make_coord(i, (k * 3 + 1) * cute::size<1>(source0_vec) + p); + auto tgt2_crd = cute::make_coord(i, (k * 3 + 2) * cute::size<1>(source0_vec) + p); + + auto src0_raw = source0_vec(src0_crd) >> (k2 * kNumBits); + auto src0_val = qmap2_view [src0_raw & kMask2]; + target_vec(tgt0_crd) = __hmul2 (src0_val , scale_vec(tgt0_crd)); + + auto src1_raw = source1_vec(src1_crd) >> (k2 * kNumBits); + auto src1_val = qmap2_view [src1_raw & kMask2]; + target_vec(tgt1_crd) = __hmul2 (src1_val , scale_vec(tgt1_crd)); + + auto src2_raw = source2_vec(src2_crd) >> (k2 * kNumBits); + auto src2_val = qmap2_view [src2_raw & kMask2]; + target_vec(tgt2_crd) = __hmul2 (src2_val , scale_vec(tgt2_crd)); + } + + // handle the last element + auto tgt3_crd = cute::make_coord(i, ((kNumPacked2 / 2) * 3) * cute::size<1>(source0_vec) + p); + auto src3_raw = ((((source0_vec(src0_crd) >> (kNumPacked2 * kNumBits)) & kMaskF2) ) | + (((source1_vec(src1_crd) >> (kNumPacked2 * kNumBits)) & kMaskF2) << 2) | + (((source2_vec(src2_crd) >> (kNumPacked2 * kNumBits)) & kMaskF2) << 4)); + auto src3_val = qmap2_view[src3_raw & kMask2]; + target_vec(tgt3_crd) = __hmul2(src3_val, scale_vec(tgt3_crd)); + } + } + } +}; + + +template +CUTE_DEVICE +void +dequantize( + cute::Tensor const& source, + cute::Tensor const& source2, + cute::Tensor && target, + cute::Tensor const& scale, + cute::Tensor const& qmap, + cute::Tensor const& qmap2, + cute::Tensor const& qmap3, + NumBits) +{ + + CUTE_STATIC_ASSERT_V(cute::rank (source ) == cute::_2{}); // ((dim0, dim1), Mma_P) + CUTE_STATIC_ASSERT_V(cute::rank (source2) == cute::_2{}); // ((dim0, dim1), Mma_P * 2) + CUTE_STATIC_ASSERT_V(cute::rank (target ) == cute::_2{}); // ((dim0, dim1), Mma) + CUTE_STATIC_ASSERT_V(cute::rank (scale ) == cute::_2{}); // ((dim0, dim1), Mma) + CUTE_STATIC_ASSERT_V(cute::rank (qmap ) == cute::_1{}); // (2 ** (NumBits),) + CUTE_STATIC_ASSERT_V(cute::rank (qmap2 ) == cute::_1{}); // (2 ** (NumBits * 2),) + CUTE_STATIC_ASSERT_V(cute::rank (qmap3 ) == cute::_1{}); // (1,) + CUTE_STATIC_ASSERT_V(cute::size<0>(target) == cute::size<0>(source)); + CUTE_STATIC_ASSERT_V(cute::size<0>(target) == cute::size<0>(source2)); + CUTE_STATIC_ASSERT_V(cute::size<0>(target) == cute::size<0>(scale)); + CUTE_STATIC_ASSERT_V(cute::size<1>(target) == cute::size<1>(scale)); + CUTE_STATIC_ASSERT_V(cute::size (qmap3) == cute::_1{}); + CUTE_STATIC_ASSERT (cute::is_same_v == true); + CUTE_STATIC_ASSERT (cute::is_same_v == true); + CUTE_STATIC_ASSERT (cute::is_same_v == true || cute::is_same_v == true); + CUTE_STATIC_ASSERT (cute::is_same_v == true || cute::is_same_v == true); + CUTE_STATIC_ASSERT (cute::is_same_v == true || cute::is_same_v == true); + CUTE_STATIC_ASSERT (cute::is_same_v == true || cute::is_same_v == true); + CUTE_STATIC_ASSERT (cute::is_same_v == true || cute::is_same_v == true); + + DequantizationTraits< + SourceEngine , SourceLayout , + SourceEngine2 , SourceLayout2 , + TargetEngine , TargetLayout , + ScaleEngine , ScaleLayout , + QuantMapEngine , QuantMapLayout , + QuantMapEngine2, QuantMapLayout2, + QuantMapEngine3, QuantMapLayout3, + NumBits, + QuantMapMode>::apply( + source, + source2, + target, + scale, + qmap, + qmap2, + qmap3); +} + +} // namespace packbits_utils \ No newline at end of file diff --git a/kernels/quantization/flute/qgemm.cpp b/kernels/quantization/flute/qgemm.cpp new file mode 100644 index 0000000000..74251c26f0 --- /dev/null +++ b/kernels/quantization/flute/qgemm.cpp @@ -0,0 +1,260 @@ +#include +#include +#include +#include +#include +#include +#include "cute/numeric/integral_constant.hpp" + + +torch::Tensor +hadamard_transform(at::Tensor& in, + bool inplace); + + +template < + typename T, + typename TQ, + typename T2, + typename NumBits, + typename GroupSize +> +void +_qgemm_raw(int64_t M, + int64_t N, + int64_t K, + int64_t P, + const T * const __restrict__ A, + const TQ* const __restrict__ Q, + T * __restrict__ D, + const T * const __restrict__ S, + const T * const __restrict__ QM, + const T2* const __restrict__ QM2, + void* __restrict__ workspace, + const int64_t template_id, + const int64_t num_sms, + const cudaStream_t stream); + + +template < + typename T, + typename NumBits, + typename GroupSize +> +void +qgemm_raw(const at::Tensor& input, + const at::Tensor& weight, + at::Tensor& output, + const at::Tensor& scales, + const at::Tensor& table, + const at::Tensor& table2, + at::Tensor& workspace, + const int64_t template_id, + const int64_t num_sms, + const cudaStream_t stream) +{ + using namespace cute; + using TQ = cute::uint16_t; + using T2 = conditional_t, __half2, __nv_bfloat162>; + + _qgemm_raw< + T, + TQ, + T2, + NumBits, + GroupSize + > ( + output.size(0), // M + output.size(1), // N + input .size(1), // K + weight.size(0), // P + reinterpret_cast(input .data_ptr()), + reinterpret_cast(weight .data_ptr()), + reinterpret_cast< T *>(output .data_ptr()), + reinterpret_cast(scales .data_ptr()), + reinterpret_cast(table .data_ptr()), + reinterpret_cast(table2 .data_ptr()), + reinterpret_cast< void*>(workspace.data_ptr()), + template_id, + num_sms, + stream); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +at::Tensor +qgemm_raw_simple(const at::Tensor& input, + const at::Tensor& weight, + const at::Tensor& scales, + const at::Tensor& table, + const at::Tensor& table2, + at::Tensor& workspace, + const cute::int64_t num_bits, + const cute::int64_t group_size, + const cute::int64_t template_id, + const cute::int64_t num_sms) +{ + + // Set the device of this function, primarily used when + // we have multiple devices in the same process. + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + + // Get the current CUDA stream, primarily used + // to make CUDA Graphs work. + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + // Squash the batch dimensions of the input tensor with its + // next-to-last dimensions. + const auto input_sizes = input.sizes().vec(); + const auto input_2d = input.reshape({-1, input_sizes.back()}); + auto output = at::empty( + { + input_2d.size(0), + scales.size(0) + }, + at::TensorOptions() + .dtype(input_2d.dtype()) + .device(input_2d.device())); + +#define RUN_QGEMM_RAW(T, NUM_BITS, GROUP_SIZE) \ + do { \ + qgemm_raw< \ + T, \ + cute::Int, \ + cute::Int \ + > ( \ + input_2d, \ + weight, \ + output, \ + scales, \ + table, \ + table2, \ + workspace, \ + template_id, \ + num_sms, \ + stream); \ + } while (false) + +#define RUN_QGEMM_RAW_SWITCH_GROUP_SIZE(T, NUM_BITS) \ + do { \ + switch (group_size) \ + { \ + case 64: \ + RUN_QGEMM_RAW(T, NUM_BITS, 64); \ + break; \ + case 128: \ + RUN_QGEMM_RAW(T, NUM_BITS, 128); \ + break; \ + case 256: \ + RUN_QGEMM_RAW(T, NUM_BITS, 256); \ + break; \ + default: \ + AT_ERROR("Unsupported `group_size`"); \ + } \ + } while (false) + +#define RUN_QGEMM_RAW_SWITCH_NUM_BITS_AND_GROUP_SIZE(T) \ + do { \ + switch (num_bits) \ + { \ + case 2: \ + RUN_QGEMM_RAW_SWITCH_GROUP_SIZE(T, 2); \ + break; \ + case 3: \ + RUN_QGEMM_RAW_SWITCH_GROUP_SIZE(T, 3); \ + break; \ + case 4: \ + RUN_QGEMM_RAW_SWITCH_GROUP_SIZE(T, 4); \ + break; \ + default: \ + AT_ERROR("Unsupported `num_bits`"); \ + } \ + } while (false) + + + AT_DISPATCH_SWITCH( + input.scalar_type(), + "qgemm_raw_simple", + AT_DISPATCH_CASE( + at::ScalarType::Half, + [&]() { + RUN_QGEMM_RAW_SWITCH_NUM_BITS_AND_GROUP_SIZE(cute::half_t); + return; + } + ) + AT_DISPATCH_CASE( + at::ScalarType::BFloat16, + [&]() { + RUN_QGEMM_RAW_SWITCH_NUM_BITS_AND_GROUP_SIZE(cute::bfloat16_t); + return; + } + ) + ); + + auto output_sizes = input_sizes; + output_sizes.back() = scales.size(0); + return output.reshape(output_sizes); +} + + +at::Tensor +apply_hadamard(const at::Tensor& input, + const cute::int64_t hadamard_size) +{ + auto input_sizes = input.sizes(); + auto flat_input = input.reshape({-1, hadamard_size}); + auto had_input = hadamard_transform( + flat_input, false + ); + return had_input.reshape(input_sizes); +} + + +at::Tensor +qgemm_raw_simple_hadamard(const at::Tensor& input, + const at::Tensor& weight, + const at::Tensor& scales, + const at::Tensor& table, + const at::Tensor& table2, + at::Tensor& workspace, + const cute::int64_t num_bits, + const cute::int64_t group_size, + const cute::int64_t hadamard_size, + const cute::int64_t template_id, + const cute::int64_t num_sms) +{ + auto had_input = apply_hadamard( + input, + hadamard_size + ); + + return qgemm_raw_simple( + had_input, + weight, + scales, + table, + table2, + workspace, + num_bits, + group_size, + template_id, + num_sms + ); +} + + +// Registers _C as an extension module. +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {} + +// Defines the operators +TORCH_LIBRARY(flute, m) { + m.def("qgemm_raw_simple(Tensor input, Tensor weight, Tensor scales, Tensor table, Tensor table2, Tensor(a!) workspace, int num_bits, int group_size, int template_id, int num_sms) -> Tensor"); + m.def("qgemm_raw_simple_hadamard(Tensor input, Tensor weight, Tensor scales, Tensor table, Tensor table2, Tensor(a!) workspace, int num_bits, int group_size, int hadamard_size, int template_id, int num_sms) -> Tensor"); +} + + +TORCH_LIBRARY_IMPL(flute, CUDA, m) { + m.impl("qgemm_raw_simple", &qgemm_raw_simple); + m.impl("qgemm_raw_simple_hadamard", &qgemm_raw_simple_hadamard); +} \ No newline at end of file diff --git a/kernels/quantization/flute/qgemm_kernel.hpp b/kernels/quantization/flute/qgemm_kernel.hpp new file mode 100644 index 0000000000..63941c29f2 --- /dev/null +++ b/kernels/quantization/flute/qgemm_kernel.hpp @@ -0,0 +1,939 @@ +#pragma once + +#include +#include +#include +#include "cutlass/workspace.h" + +#include "config.hpp" +#include "packbits_utils.hpp" +#include "conversion_utils.hpp" +#include "tile_scheduler_utils.hpp" +#define DEBUG 0 + + +// Essentially implements the following +// https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/gemm/collective/sm80_mma_multistage.hpp +template < + class Config, + class TileScheduler, + class T, + class TQ, + class T2 +> +__global__ /* __launch_bounds__(128, 1) */ +void +qgemm_device(const T * const __restrict__ A, + const TQ* const __restrict__ Q, + T * __restrict__ D, + const T * const __restrict__ S, + const T * const __restrict__ QM, + const T2* const __restrict__ QM2, + void* __restrict__ workspace, + TileScheduler scheduler) +{ + using namespace cute; + using X = Underscore; + + static constexpr config::QuantMapModeEnum QuantMapMode = Config::QuantMapMode; + static constexpr config::AccumulationModeEnum AccumulationMode = Config::AccumulationMode; + static constexpr config::DecompositionModeEnum DecompositionMode = Config::DecompositionMode; + using SharedStorage = typename Config::SharedStorage; + using TC = typename Config::TC; // accumulator type + using TR = typename Config::TR; // reduction type + CUTE_STATIC_ASSERT(is_same_v == true); + CUTE_STATIC_ASSERT(is_same_v == true); + CUTE_STATIC_ASSERT(is_same_v == true); + + using Warps = typename Config::Warps; + using Threads = typename Config::Threads; + using TileM = typename Config::TileM; + using TileN = typename Config::TileN; + using TileK = typename Config::TileK; + using TileP = typename Config::TileP; + using TileP2 = typename Config::TileP2; + using TileG = typename Config::TileG; + using Stages = typename Config::Stages; + using StagesG = typename Config::StagesG; + using StagesGView = typename Config::StagesGView; + using NumBits = typename Config::NumBits; + using GroupSize = typename Config::GroupSize; + using QuantMapSize = typename Config::QuantMapSize; + using QuantMapSize2 = typename Config::QuantMapSize2; + using QuantMapVecTraits = typename Config::QuantMapVecTraits; + using QuantMapDuplicates = typename Config::QuantMapDuplicates; + using TileKsPerTileG = typename Config::TileKsPerTileG; + + // shared memory layout + using SmemLayoutA = typename Config::SmemLayoutA; + using SmemLayoutB = typename Config::SmemLayoutB; + using SmemLayoutC = typename Config::SmemLayoutC; + using SmemLayoutQ = typename Config::SmemLayoutQ; + using SmemLayoutQ2 = typename Config::SmemLayoutQ2; + using SmemLayoutS = typename Config::SmemLayoutS; + using SmemLayoutSView = typename Config::SmemLayoutSView; + using SmemLayoutQM = typename Config::SmemLayoutQM; + using SmemLayoutQM2 = typename Config::SmemLayoutQM2; + using SmemLayoutQM3 = typename Config::SmemLayoutQM3; + using SmemLayoutQMView = typename Config::SmemLayoutQMView; + using SmemLayoutQM2View = typename Config::SmemLayoutQM2View; + + // global to shared memory copy + using G2STiledCopyA = typename Config::G2STiledCopyA; + using G2STiledCopyQ = typename Config::G2STiledCopyQ; + using G2STiledCopyQ2 = typename Config::G2STiledCopyQ2; + using G2STiledCopyS = typename Config::G2STiledCopyS; + using G2STiledCopyQM = typename Config::G2STiledCopyQM; + using G2STiledCopyQM2 = typename Config::G2STiledCopyQM2; + + // shared to shared memory copy + using S2STiledCopyQM3 = typename Config::S2STiledCopyQM3; + + // shared to register copy + using S2RCopyAtomA = typename Config::S2RCopyAtomA; + using S2RCopyAtomQ = typename Config::S2RCopyAtomQ; + using S2RCopyAtomQ2 = typename Config::S2RCopyAtomQ2; + using S2RCopyAtomSView = typename Config::S2RCopyAtomSView; + using S2RCopyAtomQM = typename Config::S2RCopyAtomQM; + + // mma + using TiledMma = typename Config::TiledMma; + using TiledMmaQ = typename Config::TiledMmaQ; + using TiledMmaQ2 = typename Config::TiledMmaQ2; + using MmaPrmM = typename Config::MmaPrmM; + using MmaPrmN = typename Config::MmaPrmN; + using UnpackTargetSize = typename Config::UnpackTargetSize; + + // register to shared copy + using R2SCopyAtomC = typename Config::R2SCopyAtomC; + + // shared to register copy + using S2RTiledCopyC = typename Config::S2RTiledCopyC; + + // register to global copy + using R2GCopyAtomC = typename Config::R2GCopyAtomC; + + // thread and lane index + int thr_index = threadIdx.x; + int lane_index = threadIdx.x % Warps{}; + +#if DEBUG + if(thread0()){ + print("\n\n"); + print("TileM = "); print(TileM{}) ; print("\n"); + print("TileN = "); print(TileN{}) ; print("\n"); + print("TileK = "); print(TileK{}) ; print("\n"); + print("TileP = "); print(TileP{}) ; print("\n"); + print("TileP2 = "); print(TileP2{}); print("\n"); + print("TileG = "); print(TileG{}) ; print("\n"); + print("Stages = "); print(Stages{}); print("\n"); + } +#endif + + // use Tensor notation to represent device pointer + dimension + Tensor mD = make_tensor(make_gmem_ptr(D) , make_shape(scheduler.M(), scheduler.N()), make_stride(scheduler.N(), _1{})); // (M, N) + Tensor mA = make_tensor(make_gmem_ptr(A) , make_shape(scheduler.M(), scheduler.K()), make_stride(scheduler.K(), _1{})); // (M, K) + Tensor mQ = make_tensor(make_gmem_ptr(Q) , make_shape(scheduler.P(), scheduler.K()), make_stride(scheduler.K(), _1{})); // (P, K) + Tensor mS = make_tensor(make_gmem_ptr(S) , make_shape(scheduler.N(), scheduler.G()), make_stride(scheduler.G(), _1{})); // (N, G) + Tensor mQM = make_tensor(make_gmem_ptr(QM) , make_shape(QuantMapSize{})); + Tensor mQM2 = make_tensor(make_gmem_ptr(QM2), make_shape(QuantMapSize2{})); + // this implicitly assumes that Q is K-major. + Tensor mQ2 = make_tensor(make_gmem_ptr(Q + scheduler.P() * scheduler.K()), make_shape(scheduler.P2(), scheduler.K()), make_stride(scheduler.K(), _1{})); // (P2, K) + + // slice the tensor to small one which is used for current thread block. + Tensor gD = local_tile(mD, make_tile(TileM{}, TileN{}), make_coord(_, _)); // (TileM , TileN, Num_Tiles_M , Num_Tiles_N) + Tensor gA = local_tile(mA, make_tile(TileM{}, TileK{}), make_coord(_, _)); // (TileM , TileK, Num_Tiles_M , Num_Tiles_K) + Tensor gQ = local_tile(mQ, make_tile(TileP{}, TileK{}), make_coord(_, _)); // (TileP , TileK, Num_Tiles_P , Num_Tiles_K) + Tensor gQ2 = local_tile(mQ2, make_tile(TileP2{}, TileK{}), make_coord(_, _)); // (TileP2, TileK, Num_Tiles_P2, Num_Tiles_K) + Tensor gS = local_tile(mS, make_tile(TileN{}, TileG{}), make_coord(_, _)); // (TileN , TileG, Num_Tiles_N , Num_Tiles_G) + + // shared memory + extern __shared__ char smem_buf[]; + SharedStorage& shared_storage = *reinterpret_cast(smem_buf); + auto smem_C_data = shared_storage.smem_C.data(); + if constexpr (DecompositionMode == config::DecompositionModeEnum::SplitK) + { + // In Stream-K mode, CTA might continue processing tiles after an epilogue, so we + // cannot reuse the buffer for C. Instead, we allocate a separate buffer for C. + // However, for Split-K mode, since CTA has processed all tiles after an epilogue, + // we will reuse the same buffer for C. + smem_C_data = reinterpret_cast(smem_buf); + } + + auto sA = make_tensor(make_smem_ptr(shared_storage.smem_A .data()), SmemLayoutA{}); // (TileM, TileK, Stages) + auto sB = make_tensor(make_smem_ptr(shared_storage.smem_A .data()), SmemLayoutB{}); // (TileN, TileK, Stages), using `smemA` as this is mostly as a placeholder to infer the shapes + auto sQ = make_tensor(make_smem_ptr(shared_storage.smem_Q .data()), SmemLayoutQ{}); // (TileP, TileK, Stages) + auto sQ2 = make_tensor(make_smem_ptr(shared_storage.smem_Q2 .data()), SmemLayoutQ2{}); // (TileP2, TileK, Stages) + auto sS = make_tensor(make_smem_ptr(shared_storage.smem_S .data()), SmemLayoutS{}); // (TileN, TileG, StagesG) + auto sQM = make_tensor(make_smem_ptr(shared_storage.smem_QM .data()), SmemLayoutQM{}); + auto sQM2 = make_tensor(make_smem_ptr(shared_storage.smem_QM2.data()), SmemLayoutQM2{}); + auto sQM3 = make_tensor(make_smem_ptr(shared_storage.smem_QM3.data()), SmemLayoutQM3{}); + auto sC = make_tensor(make_smem_ptr(smem_C_data ), SmemLayoutC{}); + // this is a view of `sS` broadcasted with `TileK` + auto sSv = make_tensor(make_smem_ptr(shared_storage.smem_S .data()), SmemLayoutSView{}); // (TileN, TileK, (...), StagesG) + // this is a view of `sQM` with an extra leading dimension of `1` + auto sQMv = make_tensor(make_smem_ptr(shared_storage.smem_QM .data()), SmemLayoutQMView{}); + // this is a view of `sQM2` with broadcasted to an extra dimension of `QuantMapDuplicates` + auto sQM2v = make_tensor(make_smem_ptr(shared_storage.smem_QM2.data()), SmemLayoutQM2View{}); + + CUTE_STATIC_ASSERT_V(size<0>(gA) == size<0>(sA)); // TileM + CUTE_STATIC_ASSERT_V(size<1>(gA) == size<1>(sA)); // TileK + CUTE_STATIC_ASSERT_V(size<0>(gQ) == size<0>(sQ)); // TileP + CUTE_STATIC_ASSERT_V(size<1>(gQ) == size<1>(sQ)); // TileK + CUTE_STATIC_ASSERT_V(size<0>(gQ2) == size<0>(sQ2)); // TileP2 + CUTE_STATIC_ASSERT_V(size<1>(gQ2) == size<1>(sQ2)); // TileK + CUTE_STATIC_ASSERT_V(size<0>(gS) == size<0>(sS)); // TileN + CUTE_STATIC_ASSERT_V(size<1>(gS) == size<1>(sS)); // TileG + CUTE_STATIC_ASSERT_V(size<0>(gS) == size<0>(sSv)); // TileN + CUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sQ)); // TileK + CUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sSv)); // TileK + CUTE_STATIC_ASSERT_V(size<0>(sB) == size<0>(sSv)); // TileN + CUTE_STATIC_ASSERT_V(size<1>(sB) == size<1>(sSv)); // TileK + CUTE_STATIC_ASSERT_V(Stages{} == size<2>(sA)); // Stages + CUTE_STATIC_ASSERT_V(Stages{} == size<2>(sQ)); // Stages + CUTE_STATIC_ASSERT_V(Stages{} == size<2>(sQ2)); // Stages + CUTE_STATIC_ASSERT_V(StagesG{} == size<2>(sS)); // StagesG + + // + // MMA Atom partitioning + // + + // dispatch TileA/TileB/TileC mma tensor into thread fragment via partition method + TiledMma tiled_mma; + TiledMmaQ tiled_mma_Q; // primarily used to get the right shape for loading `Q` + TiledMmaQ2 tiled_mma_Q2; // primarily used to get the right shape for loading `Q2` + auto thr_mma = tiled_mma .get_slice(thr_index); + auto thr_mma_Q = tiled_mma_Q .get_slice(thr_index); + auto thr_mma_Q2 = tiled_mma_Q2.get_slice(thr_index); + auto accum = thr_mma .partition_fragment_C(gD (_, _, _0{}, _0{})); // (Mma, Mma_M, Mma_N) + auto tCrA = thr_mma .partition_fragment_A(sA (_, _, _0{} )); // (Mma, Mma_M, Mma_K) + auto tCrB = thr_mma .partition_fragment_B(sB (_, _, _0{} )); // (Mma, Mma_N, Mma_K) + auto tCrSv = thr_mma .partition_fragment_B(sSv(_, _, _0{} )); // (Mma, Mma_N, Mma_K) + auto tCrQ_tmp = thr_mma_Q .partition_fragment_B(sQ (_, _, _0{} )); // (Mma, Mma_P, Mma_K) + auto tCrQ2_tmp = thr_mma_Q2 .partition_fragment_B(sQ2(_, _, _0{} )); // (Mma, Mma_P2, Mma_K) + auto accum_epilogue = make_fragment_like(accum); // output types is the same as input types + auto accum_reduction = make_fragment_like(accum); // reduction types could be different from compute types + auto tCrQ = make_fragment_like(tCrQ_tmp); // (Mma, Mma_P, Mma_K), tCrQ_tmp is of type `T` + auto tCrQ2 = make_fragment_like(tCrQ2_tmp); // (Mma, Mma_P2, Mma_K), tCrQ2_tmp is of type `T` + + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum)); // Mma_M + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(accum)); // Mma_N + CUTE_STATIC_ASSERT_V(size<1>(tCrSv) == size<2>(accum)); // Mma_N + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // Mma_K + CUTE_STATIC_ASSERT_V(size<2>(tCrSv) == size<2>(tCrB)); // Mma_K + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrQ)); // Mma_K + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrQ2)); // Mma_K + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrSv)); // Mma_K + CUTE_STATIC_ASSERT_V(size<0>(accum) == size<0>(accum_epilogue)); + CUTE_STATIC_ASSERT_V(size<1>(accum) == size<1>(accum_epilogue)); + CUTE_STATIC_ASSERT_V(size<2>(accum) == size<2>(accum_epilogue)); + CUTE_STATIC_ASSERT_V(size<0>(accum) == size<0>(accum_reduction)); + CUTE_STATIC_ASSERT_V(size<1>(accum) == size<1>(accum_reduction)); + CUTE_STATIC_ASSERT_V(size<2>(accum) == size<2>(accum_reduction)); + CUTE_STATIC_ASSERT_V(size<1>(accum) == (MmaPrmM{})); + CUTE_STATIC_ASSERT_V(size<2>(accum) == (MmaPrmN{} * UnpackTargetSize{})); + + // + // Copy Atom + // + + // global to shared memory copy, partition the copying of A and B tiles across the threads + G2STiledCopyA g2s_tiled_copy_A; + G2STiledCopyQ g2s_tiled_copy_Q; + G2STiledCopyQ2 g2s_tiled_copy_Q2; + G2STiledCopyS g2s_tiled_copy_S; + auto g2s_thr_copy_A = g2s_tiled_copy_A .get_slice(thr_index); + auto g2s_thr_copy_Q = g2s_tiled_copy_Q .get_slice(thr_index); + auto g2s_thr_copy_Q2 = g2s_tiled_copy_Q2.get_slice(thr_index); + auto g2s_thr_copy_S = g2s_tiled_copy_S .get_slice(thr_index); + auto tAgA = g2s_thr_copy_A .partition_S(gA); // (G2S_CPY, G2S_CPY_M , G2S_CPY_K, Num_Tiles_M , Num_Tiles_K) + auto tBgQ = g2s_thr_copy_Q .partition_S(gQ); // (G2S_CPY, G2S_CPY_P , G2S_CPY_K, Num_Tiles_P , Num_Tiles_K) + auto tBgQ2 = g2s_thr_copy_Q2.partition_S(gQ2); // (G2S_CPY, G2S_CPY_P2, G2S_CPY_K, Num_Tiles_P2, Num_Tiles_K) + auto tSgS = g2s_thr_copy_S .partition_S(gS); // (G2S_CPY, G2S_CPY_N , G2S_CPY_G, Num_Tiles_N , Num_Tiles_G) + auto tAsA = g2s_thr_copy_A .partition_D(sA); // (G2S_CPY, G2S_CPY_M , G2S_CPY_K, Stages) + auto tBsQ = g2s_thr_copy_Q .partition_D(sQ); // (G2S_CPY, G2S_CPY_P , G2S_CPY_K, Stages) + auto tBsQ2 = g2s_thr_copy_Q2.partition_D(sQ2); // (G2S_CPY, G2S_CPY_P2, G2S_CPY_K, Stages) + auto tSsS = g2s_thr_copy_S .partition_D(sS); // (G2S_CPY, G2S_CPY_N , G2S_CPY_G, Stages) + + CUTE_STATIC_ASSERT_V((size(tAsA) * Threads{}) == size(sA)); // sA is too small for Threads each with tAsA + CUTE_STATIC_ASSERT_V((size(tBsQ) * Threads{}) == size(sQ)); // sQ is too small for Threads each with tBsQ + CUTE_STATIC_ASSERT_V((size(tBsQ2)* Threads{}) == size(sQ2)); // sQ2 is too small for Threads each with tBsQ2 + CUTE_STATIC_ASSERT_V((size(tSsS) * Threads{}) == size(sS)); // sS is too small for Threads each with tSsS + + // shared to register copy + auto s2r_tiled_copy_A = make_tiled_copy_A(S2RCopyAtomA{} , tiled_mma); + auto s2r_tiled_copy_Q = make_tiled_copy_B(S2RCopyAtomQ{} , tiled_mma_Q); + auto s2r_tiled_copy_Q2 = make_tiled_copy_B(S2RCopyAtomQ2{} , tiled_mma_Q2); + auto s2r_tiled_copy_Sv = make_tiled_copy_B(S2RCopyAtomSView{}, tiled_mma); + auto s2r_thr_copy_A = s2r_tiled_copy_A .get_slice(thr_index); + auto s2r_thr_copy_Q = s2r_tiled_copy_Q .get_slice(thr_index); + auto s2r_thr_copy_Q2 = s2r_tiled_copy_Q2.get_slice(thr_index); + auto s2r_thr_copy_Sv = s2r_tiled_copy_Sv.get_slice(thr_index); + auto tCsA = s2r_thr_copy_A .partition_S(sA); // ? (S2R_CPY, S2R_CPY_M, S2R_CPY_K, Stages) + auto tCsQ = s2r_thr_copy_Q .partition_S(sQ); // ? (S2R_CPY, S2R_CPY_P, S2R_CPY_K, Stages) + auto tCsQ2 = s2r_thr_copy_Q2 .partition_S(sQ2); // ? (S2R_CPY, S2R_CPY_P2, S2R_CPY_K, Stages) + auto tCsSv = s2r_thr_copy_Sv .partition_S(sSv); // ? (S2R_CPY, S2R_CPY_N, S2R_CPY_G, StagesG) + auto tCrA_view = s2r_thr_copy_A .retile_D(tCrA); // ? (S2R_CPY, S2R_CPY_M, S2R_CPY_K) + auto tCrQ_view = s2r_thr_copy_Q .retile_D(tCrQ); // ? (S2R_CPY, S2R_CPY_P, S2R_CPY_K) + auto tCrQ2_view = s2r_thr_copy_Q2 .retile_D(tCrQ2); // ? (S2R_CPY, S2R_CPY_P2, S2R_CPY_K) + auto tCrSv_view = s2r_thr_copy_Sv .retile_D(tCrSv); // ? (S2R_CPY, S2R_CPY_N, S2R_CPY_G) + + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_view)); // CPY_M + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_view)); // CPY_K + CUTE_STATIC_ASSERT_V(size<1>(tCsQ) == size<1>(tCrQ_view)); // CPY_P + CUTE_STATIC_ASSERT_V(size<2>(tCsQ) == size<2>(tCrQ_view)); // CPY_K + CUTE_STATIC_ASSERT_V(size<1>(tCsQ2) == size<1>(tCrQ2_view)); // CPY_P2 + CUTE_STATIC_ASSERT_V(size<2>(tCsQ2) == size<2>(tCrQ2_view)); // CPY_K + CUTE_STATIC_ASSERT_V(size<1>(tCsSv) == size<1>(tCrSv_view)); // CPY_N + CUTE_STATIC_ASSERT_V(size<2>(tCsSv) == size<2>(tCrSv_view)); // CPY_G + + // quant map, note that every thread block loads the entire QM + G2STiledCopyQM g2s_tiled_copy_QM; + G2STiledCopyQM2 g2s_tiled_copy_QM2; + auto g2s_thr_copy_QM = g2s_tiled_copy_QM .get_slice(thr_index); + auto g2s_thr_copy_QM2 = g2s_tiled_copy_QM2.get_slice(thr_index); + auto tQMgQM = g2s_thr_copy_QM .partition_S(mQM); // ((1, 1), (1,)) + auto tQM2gQM2 = g2s_thr_copy_QM2 .partition_S(mQM2); // ((1, ?), (?,)) + auto tQMsQM = g2s_thr_copy_QM .partition_D(sQM); // ((1, 1), (1,)) + auto tQM2sQM2 = g2s_thr_copy_QM2 .partition_D(sQM2); // ((1, ?), (?,)) + auto tQMrQM = make_tensor(Shape<_1>{}); + + S2STiledCopyQM3 s2s_tiled_copy_QM3; + auto s2s_thr_copy_QM3 = s2s_tiled_copy_QM3.get_slice(thr_index); + auto tQM3sQM2 = s2s_thr_copy_QM3 .partition_S(sQM2v); // ((1, ?), (?,)) + auto tQM3sQM3 = s2s_thr_copy_QM3 .partition_D(sQM3); // ((1, ?), (?,)) + + CUTE_STATIC_ASSERT_V(size(tQMgQM) == _1{}); // QM is too large for Threads + CUTE_STATIC_ASSERT_V(size(tQMsQM) == _1{}); // QM is too large for Threads + // CUTE_STATIC_ASSERT_V(size(tQM2gQM2) == _1{}); // QM is too large for Threads + // CUTE_STATIC_ASSERT_V(size(tQM2sQM2) == _1{}); // QM is too large for Threads + CUTE_STATIC_ASSERT_V(size(tQMrQM) == _1{}); // QM is too large for Threads + + // + // PREDICATES + // + + // Allocate predicate tensors for m and n + auto tApA = make_tensor(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1, _0>{}); + auto tBpQ = make_tensor(make_shape(size<1>(tBsQ), size<2>(tBsQ)), Stride<_1, _0>{}); + auto tBpQ2= make_tensor(make_shape(size<1>(tBsQ2),size<2>(tBsQ2)),Stride<_1, _0>{}); + auto tSpS = make_tensor(make_shape(size<1>(tSsS), size<2>(tSsS)), Stride<_1, _0>{}); + + // Construct identity layout for sA and sB + auto cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (TileM, TileK) -> (tile_m, tile_k) + auto cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (TileP, TileK) -> (tile_p, tile_k) + auto cQ2= make_identity_tensor(make_shape(size<0>(sQ2),size<1>(sQ2))); // (TileP2,TileK) -> (tile_p2,tile_k) + auto cS = make_identity_tensor(make_shape(size<0>(sS), size<1>(sS))); // (TileN, TileG) -> (tile_n, tile_g) + + // Repeat the partitioning with identity layouts + auto tAcA = g2s_thr_copy_A .partition_S(cA); // (ACPY,ACPY_M,ACPY_K) -> (tile_m,tile_k) + auto tBcQ = g2s_thr_copy_Q .partition_S(cQ); // (BCPY,BCPY_N,BCPY_K) -> (tile_n,tile_k) + auto tBcQ2 = g2s_thr_copy_Q2.partition_S(cQ2); // (BCPY,BCPY_N,BCPY_K) -> (tile_n,tile_k) + auto tScS = g2s_thr_copy_S .partition_S(cS); // (SCPY,SCPY_N,SCPY_G) -> (tile_n,tile_g) + + CUTE_STATIC_ASSERT_V(TileM {} == size<0>(gA )); + CUTE_STATIC_ASSERT_V(TileN {} == size<1>(gD )); + CUTE_STATIC_ASSERT_V(TileP {} == size<0>(gQ )); + CUTE_STATIC_ASSERT_V(TileP2{} == size<0>(gQ2)); + + // + // Epilogue + // https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp + // + + // Partition sC to match the accumulator partitioning + // 1. note that we will write `accum_epilogue` instead of `accum` + // 2. note that `tiled_mma` has accumulator of type `TC`, but we want to copy in type `T`, + // but it seems like `make_tiled_copy_C` uses just the layout, not the type info + auto r2s_tiled_copy_C = make_tiled_copy_C(R2SCopyAtomC{}, tiled_mma); + auto r2s_thr_copy_C = r2s_tiled_copy_C.get_slice(thr_index); + auto tCaC = r2s_thr_copy_C.retile_S(accum_epilogue); // (R2S_CPY=(Atom, AtomNum), Mma_M, Mma_N) + auto tCsC = r2s_thr_copy_C.partition_D(sC); // (R2S_CPY=(Atom, AtomNum), PIPE_M, PIPE_N) + + // Tile gD and gC by the shape of SmemLayout first + auto gD_tile = make_shape(size<0>(sC), size<1>(sC)); + auto gDt = flat_divide(gD, gD_tile); // (SMEM_M, SMEM_N, TileM, TileN, Num_Tiles_M, Num_Tiles_N) + + // Partition sC, and gD for the output + S2RTiledCopyC s2r_tiled_copy_C; + auto s2r_thr_copy_C = s2r_tiled_copy_C.get_slice(thr_index); + auto tDsC = s2r_thr_copy_C.partition_S(sC); // (S2R_COPY=(Atom, AtomNum), ATOM_M, ATOM_N) + auto tDgD = s2r_thr_copy_C.partition_D(gDt); // (S2R_COPY=(Atom, AtomNum), ATOM_M, ATOM_N, TileM, TileN, Num_Tiles_M, Num_Tiles_N) + + // Allocate intermediate registers on the dst tensors + // note that `tDrC` is of type `T` instead of `TC` because we will write `accum_epilogue` instead of `accum` + auto tDrC = make_tensor(take<0, 3>(shape(tDgD))); // ((Atom, AtomNum), ATOM_M, ATOM_N) + // auto tDrD = make_tensor(shape(tDrC)); // ((Atom, AtomNum), ATOM_M, ATOM_N) + + // Repeat the D-partitioning for coordinates and predication + auto cD = make_identity_tensor(make_shape(size<0>(gD), size<1>(gD))); // (TileM, TileN) -> (tile_m, tile_n) + auto cDt = flat_divide(cD, gD_tile); // (TileM, TileN, Num_Tiles_M=1, Num_Tiles_N=1) + auto tDcD = s2r_thr_copy_C.partition_D(cDt); // (S2R_COPY=(Atom, AtomNum), ATOM_M, ATOM_N, Num_Tiles_M=1, Num_Tiles_N=1) + + CUTE_STATIC_ASSERT(size<1>(tCaC) % size<3>(tDgD) == 0); // TileM divides Mma_M + CUTE_STATIC_ASSERT(size<2>(tCaC) % size<4>(tDgD) == 0); // TileN divides Mma_N + // CUTE_STATIC_ASSERT(typename S2RTiledCopyC::TiledNumThr{} == size<0>(typename TiledMma::AtomLayoutC_TV{})); + + +#if DEBUG + +#define PPRINT_LAYOUT(name, tensor) \ + do { \ + print("\n\n"); \ + print(name); \ + print("\n"); \ + print(tensor.layout()); \ + } while(0) + +#define PPRINT_DATA(name, data) \ + do { \ + print("\n\n"); \ + print(name); \ + print("\n"); \ + print(data); \ + } while(0) + + if(thread0()) + { + PPRINT_LAYOUT("mD" , mD); + PPRINT_LAYOUT("mA" , mA); + PPRINT_LAYOUT("mQ" , mQ); + PPRINT_LAYOUT("mQ2", mQ2); + PPRINT_LAYOUT("mS" , mS); + + PPRINT_LAYOUT("gD" , gD); + PPRINT_LAYOUT("gA" , gA); + PPRINT_LAYOUT("gQ" , gQ); + PPRINT_LAYOUT("gQ2", gQ2); + PPRINT_LAYOUT("gS" , gS); + + PPRINT_LAYOUT("sA" , sA); + PPRINT_LAYOUT("sB" , sB); + PPRINT_LAYOUT("sQ" , sQ); + PPRINT_LAYOUT("sQ2", sQ2); + PPRINT_LAYOUT("sS" , sS); + PPRINT_LAYOUT("sSv", sSv); + PPRINT_LAYOUT("sC" , sC); + print("\n\n"); + + print("--------------------- Global to Shared ---------------------"); + PPRINT_LAYOUT("tAgA" , tAgA); + PPRINT_LAYOUT("tBgQ" , tBgQ); + PPRINT_LAYOUT("tBgQ2", tBgQ2); + PPRINT_LAYOUT("tSgS" , tSgS); + PPRINT_LAYOUT("tAsA" , tAsA); + PPRINT_LAYOUT("tBsQ" , tBsQ); + PPRINT_LAYOUT("tBsQ2", tBsQ2); + PPRINT_LAYOUT("tSsS" , tSsS); + print("\n\n"); + + print("--------------------- Shared to Registers ---------------------"); + PPRINT_LAYOUT("tCsA" , tCsA); + PPRINT_LAYOUT("tCsQ" , tCsQ); + PPRINT_LAYOUT("tCsQ2", tCsQ2); + PPRINT_LAYOUT("tCsSv", tCsSv); + PPRINT_LAYOUT("tCrA_view" , tCrA_view); + PPRINT_LAYOUT("tCrQ_view" , tCrQ_view); + PPRINT_LAYOUT("tCrQ2_view", tCrQ2_view); + PPRINT_LAYOUT("tCrSv_view", tCrSv_view); + print("\n\n"); + + print("--------------------- QM ---------------------"); + PPRINT_LAYOUT("mQM" , mQM); + PPRINT_LAYOUT("mQM2" , mQM2); + PPRINT_LAYOUT("sQM" , sQM); + PPRINT_LAYOUT("sQM2" , sQM2); + PPRINT_LAYOUT("sQM3" , sQM3); + PPRINT_LAYOUT("sQMv" , sQMv); + PPRINT_LAYOUT("sQM2v", sQM2v); + PPRINT_LAYOUT("tQMgQM" , tQMgQM); + PPRINT_LAYOUT("tQM2gQM2", tQM2gQM2); + PPRINT_LAYOUT("tQM3sQM2", tQM3sQM2); + PPRINT_LAYOUT("tQMsQM" , tQMsQM); + PPRINT_LAYOUT("tQM2sQM2", tQM2sQM2); + PPRINT_LAYOUT("tQM3sQM3", tQM3sQM3); + PPRINT_LAYOUT("tQMrQM" , tQMrQM); + print("\n\n"); + + print("--------------------- TiledMma ---------------------"); + print("\n\ntile_shape(tiled_mma)\n"); + print(tile_shape(tiled_mma)); + print("\n\ntile_shape(tiled_mma_Q)\n"); + print(tile_shape(tiled_mma_Q)); + print("\n\ntile_shape(tiled_mma_Q2)\n"); + print(tile_shape(tiled_mma_Q2)); + PPRINT_LAYOUT("accum" , accum); + PPRINT_LAYOUT("accum_epilogue" , accum_epilogue); + PPRINT_LAYOUT("tCrA" , tCrA); + PPRINT_LAYOUT("tCrB" , tCrB); + PPRINT_LAYOUT("tCrQ" , tCrQ); + PPRINT_LAYOUT("tCrQ2" , tCrQ2); + PPRINT_LAYOUT("tCrSv" , tCrSv); + print("\n\n"); + + print("--------------------- Predicates ---------------------"); + PPRINT_LAYOUT("tApA", tApA); + PPRINT_LAYOUT("tBpQ", tBpQ); + PPRINT_LAYOUT("tBpQ2",tBpQ2); + PPRINT_LAYOUT("tSpS", tSpS); + PPRINT_LAYOUT("cA" , cA); + PPRINT_LAYOUT("cQ" , cQ); + PPRINT_LAYOUT("cQ2" , cQ2); + PPRINT_LAYOUT("cS" , cS); + PPRINT_LAYOUT("tAcA", tAcA); + PPRINT_LAYOUT("tBcQ", tBcQ); + PPRINT_LAYOUT("tBcQ2",tBcQ2); + PPRINT_LAYOUT("tScS", tScS); + print("\n\n"); + + print("--------------------- Registers to Shared ---------------------"); + PPRINT_LAYOUT("tCaC", tCaC); + PPRINT_LAYOUT("tCsC", tCsC); + print("\n\n"); + + print("--------------------- Shared to Registers ---------------------"); + PPRINT_LAYOUT("tDrC", tDrC); + // PPRINT_LAYOUT("tDrD", tDrD); + print("\n\n"); + + print("--------------------- Registers to Global ---------------------"); + PPRINT_LAYOUT("gDt" , gDt); + PPRINT_LAYOUT("tDsC", tDsC); + PPRINT_LAYOUT("tDgD", tDgD); + print("\n\n"); + + print("--------------------- Epilogue Predicates ---------------------"); + PPRINT_LAYOUT("cD" , cD); + PPRINT_LAYOUT("cDt" , cDt); + PPRINT_LAYOUT("tDcD", tDcD); + print("\n\n"); + } +#endif + + // + // PIPELINED MAIN LOOP + // + + // initialize the tile scheduler + scheduler.initialize(tApA, tBpQ, tBpQ2, tSpS, tAcA, tBcQ, tBcQ2, tScS); + + int smem_pipe_read = 0; + int smem_pipe_read_G = scheduler.smem_pipe_read_G_offset() % StagesGView{}; // the starting K tile index might not be aligned wth the G tile + int smem_pipe_read_raw = 0; + int smem_pipe_write = 0; + int smem_pipe_write_G = 0; + + // size of the register pipeline + auto num_mma_K = size<2>(tCrA); + + // partition the workspace + void* workspace_barriers = workspace; + void* workspace_partials = static_cast(workspace) + scheduler.workspace_size_barriers(); + + // Clear the smem tiles to account for predicated off loads + clear(tAsA); + clear(tBsQ); + // fill zero for accumulator + clear(accum); + + // + // ------- Prefetching ------- + // + + // start async gmem -> shm loads for all pipes but the last + + // the copy size is too small for `cp.async` so this is sync + if(thr_index < QuantMapSize{}) { + cute::copy(g2s_tiled_copy_QM, tQMgQM, tQMsQM); + } + + // prefetch quant map, relying on the next `cp_async_fence`. + if constexpr (is_same_v) + { + if(thr_index < QuantMapSize2{}) + { + cute::copy(g2s_tiled_copy_QM2, tQM2gQM2, tQM2sQM2); + } + } + + CUTLASS_PRAGMA_UNROLL + for (int stage_index = 0; stage_index < Stages{} - 1; ++stage_index) + { + auto tile_coord_A = scheduler.get_tile_coord_A(); + auto tile_coord_Q = scheduler.get_tile_coord_Q(); + cute::copy_if(g2s_tiled_copy_A, tApA, tAgA(tile_coord_A), tAsA(_, _, _, stage_index)); + cute::copy_if(g2s_tiled_copy_Q, tBpQ, tBgQ(tile_coord_Q), tBsQ(_, _, _, stage_index)); + + if (scheduler.start_of_group()) + { + auto tile_coord_S = scheduler.get_tile_coord_S(); + cute::copy_if(g2s_tiled_copy_S, tSpS, tSgS(tile_coord_S), tSsS(_, _, _, smem_pipe_write_G)); + smem_pipe_write_G = (smem_pipe_write_G + 1) % StagesG{}; + } + + if constexpr (is_same_v) + { + auto tile_coord_Q2 = scheduler.get_tile_coord_Q2(); + cute::copy_if(g2s_tiled_copy_Q2, tBpQ2, tBgQ2(tile_coord_Q2), tBsQ2(_, _, _, stage_index)); + } + + cp_async_fence(); + + // increment index + ++smem_pipe_write; + scheduler.step_read(tApA, tBpQ, tBpQ2, tSpS, tAcA, tBcQ, tBcQ2, tScS); + + } + // smem_pipe_read == Stages % Stages (== 0) + // smem_pipe_write == Stages - 1 + + // prefetch register pipeline, wait one submitted gmem->smem done + cp_async_wait(); + __syncthreads(); + + // prefer the first rmem from the first k-tile + // smem -> reg + cute::copy(S2RCopyAtomQM{} , sQMv (_, lane_index % QuantMapSize{}), tQMrQM); + cute::copy(s2r_tiled_copy_A , tCsA (_, _, _0{}, smem_pipe_read ), tCrA_view (_, _, _0{})); + cute::copy(s2r_tiled_copy_Q , tCsQ (_, _, _0{}, smem_pipe_read ), tCrQ_view (_, _, _0{})); + cute::copy(s2r_tiled_copy_Sv, tCsSv(_, _, _0{}, smem_pipe_read_G ), tCrSv_view(_, _, _0{})); + + if constexpr (is_same_v) + { + cute::copy(s2r_tiled_copy_Q2, tCsQ2(_, _, _0{}, smem_pipe_read), tCrQ2_view(_, _, _0{})); + } + + // duplicate the quant map for each lane + if constexpr (is_same_v) + { + cute::copy(s2s_tiled_copy_QM3, tQM3sQM2, tQM3sQM3); + } + + // + // ------- Main Loop ------- + // + + // loop over K: i. load tile, ii. mma + CUTLASS_PRAGMA_NO_UNROLL + for ( ; scheduler.tile_is_in_bound(); scheduler.step()) + { + + // Note, the for_each() function is required here to ensure `mma_index` is of type Int. + for_each(make_int_sequence{}, [&] (auto mma_index) + { + + if (mma_index == num_mma_K - 1) { + // increment tile index + ++smem_pipe_read_raw; + smem_pipe_read = (smem_pipe_read_raw) % Stages{}; + smem_pipe_read_G = (smem_pipe_read_raw + scheduler.smem_pipe_read_G_offset()) % StagesGView{}; + + // wait one submitted gmem->smem done + cp_async_wait(); + __syncthreads(); + } + + // shm -> reg s[tile_index][mma_index + 1] -> r[mma_index + 1] + auto mma_index_next = (mma_index + _1{}) % num_mma_K; + cute::copy(s2r_tiled_copy_A , tCsA (_, _, mma_index_next, smem_pipe_read ), tCrA_view (_, _, mma_index_next)); + cute::copy(s2r_tiled_copy_Q , tCsQ (_, _, mma_index_next, smem_pipe_read ), tCrQ_view (_, _, mma_index_next)); + cute::copy(s2r_tiled_copy_Sv, tCsSv(_, _, mma_index_next, smem_pipe_read_G), tCrSv_view(_, _, mma_index_next)); + + if constexpr (is_same_v) { + cute::copy(s2r_tiled_copy_Q2, tCsQ2(_, _, mma_index_next, smem_pipe_read), tCrQ2_view(_, _, mma_index_next)); + } + + // copy gmem to smem before computing gemm on each k-pipe + if (mma_index == 0) + { + if (scheduler.tile_read_is_in_bound()) + { + auto tile_coord_A = scheduler.get_tile_coord_A(); + auto tile_coord_Q = scheduler.get_tile_coord_Q(); + cute::copy_if(g2s_tiled_copy_A, tApA, tAgA(tile_coord_A), tAsA(_, _, _, smem_pipe_write)); + cute::copy_if(g2s_tiled_copy_Q, tBpQ, tBgQ(tile_coord_Q), tBsQ(_, _, _, smem_pipe_write)); + + if (scheduler.start_of_group()) + { + auto tile_coord_S = scheduler.get_tile_coord_S(); + cute::copy_if(g2s_tiled_copy_S, tSpS, tSgS(tile_coord_S), tSsS(_, _, _, smem_pipe_write_G)); + smem_pipe_write_G = (smem_pipe_write_G + 1) % StagesG{}; + } + + if constexpr (is_same_v) + { + auto tile_coord_Q2 = scheduler.get_tile_coord_Q2(); + cute::copy_if(g2s_tiled_copy_Q2, tBpQ2, tBgQ2(tile_coord_Q2), tBsQ2(_, _, _, smem_pipe_write)); + } + + smem_pipe_write = (smem_pipe_write + 1) % Stages{}; + scheduler.step_read(tApA, tBpQ, tBpQ2, tSpS, tAcA, tBcQ, tBcQ2, tScS); + + } + + cp_async_fence(); + } + + // dequantize + if constexpr (is_same_v) + { + packbits_utils::dequantize( + tCrQ (_, _, mma_index), + tCrQ2(_, _, mma_index), + tCrB (_, _, mma_index), + tCrSv(_, _, mma_index), + sQM, + sQM2, + tQMrQM, + NumBits{}); + } + else + { + packbits_utils::dequantize( + tCrQ (_, _, mma_index), + tCrQ2(_, _, mma_index), + tCrB (_, _, mma_index), + tCrSv(_, _, mma_index), + sQM, + sQM3 (_, lane_index % QuantMapDuplicates{}), + tQMrQM, + NumBits{}); + } + + + // mma + cute::gemm( + tiled_mma, + accum, + tCrA(_, _, mma_index), + tCrB(_, _, mma_index), + accum); + + }); // for mma_index + + + // + // ------- Epilogue ------- + // + + if (scheduler.needs_fixup()) + { + if constexpr (AccumulationMode != config::AccumulationModeEnum::Mixed) + { + scheduler.maybe_fixup(accum, thr_index, workspace_partials, workspace_barriers); + } + else + { + conversion_utils::convert_tensor(accum, accum_reduction); + scheduler.maybe_fixup(accum_reduction, thr_index, workspace_partials, workspace_barriers); + } + + } + + if (scheduler.needs_epilogue()) + { + int epilogue_tile_M_index; + int epilogue_tile_N_index; + int epilogue_residue_M; + int epilogue_residue_N; + scheduler.prepare_epilogue(epilogue_tile_M_index, epilogue_tile_N_index, epilogue_residue_M, epilogue_residue_N); + + // output the possibly lower-precision type + // note that `tCaC` is a view of `accum_epilogue` + if constexpr (AccumulationMode != config::AccumulationModeEnum::Mixed) + { + conversion_utils::convert_tensor(accum, accum_epilogue); + } + else + { + // this implicitly assumes that `needs_epilogue` always follows `needs_fixup` + // otherwise, we need to convert `accum` to `accum_reduction` here + conversion_utils::convert_tensor(accum_reduction, accum_epilogue); + } + + // For each tiling needed for SmemLayout to cover shape(gD) + CUTLASS_PRAGMA_UNROLL + for (int step_m = 0; step_m < size<2>(cDt); ++step_m) // Num_Tiles_M + { + CUTLASS_PRAGMA_UNROLL + for (int step_n = 0; step_n < size<3>(cDt); ++step_n) // Num_Tiles_N + { + // Step 1. Copy to SMEM + CUTLASS_PRAGMA_UNROLL + for (int pipe_m = 0; pipe_m < size<1>(tCsC); ++pipe_m) // PIPE_M + { + CUTLASS_PRAGMA_UNROLL + for (int pipe_n = 0; pipe_n < size<2>(tCsC); ++pipe_n) // PIPE_N + { + int mma_m = step_m * size<1>(tCsC) + pipe_m; + int mma_n = step_n * size<2>(tCsC) + pipe_n; + copy(r2s_tiled_copy_C, tCaC(_, mma_m, mma_n), tCsC(_, pipe_m, pipe_n)); + } + } + + // Step 2. Wait for SMEM writes to complete + __syncthreads(); + + // Step 3. Copy from SMEM into a fragment + copy(s2r_tiled_copy_C, tDsC, tDrC); + + // Step 4. Wait for SMEM reads to complete + __syncthreads(); + + auto tDgDmn = tDgD(_, _, _, step_m, step_n, epilogue_tile_M_index, epilogue_tile_N_index); + auto tDcDmn = tDcD(_, _, _, step_m, step_n); + + // Step 5. Elementwise operation with conversion + // CUTLASS_PRAGMA_UNROLL + // for (int i = 0; i < size(tDrC); ++i) + // { + // tDrD(i) = epilogue_op(tDrC(i)); + // } + + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < size<1>(tDgDmn); ++m) + { + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < size<2>(tDgDmn); ++n) + { + // Predication + if (get<0>(tDcDmn(0, m, n)) < epilogue_residue_M && + get<1>(tDcDmn(0, m, n)) < epilogue_residue_N) + { + // Step 6. Copy to GMEM + // copy(R2GCopyAtomC{}, tDrD(_, m, n), tDgDmn(_, m, n)); + copy(R2GCopyAtomC{}, tDrC(_, m, n), tDgDmn(_, m, n)); + } + } + } + } // for step_n + } // for step_m + } + + if (scheduler.needs_to_clear_accum()) + { + // Clear the accumulator for the next output tile + clear(accum); + } + + + } // tile_index +} + + +template < + typename T, + typename TQ, + typename T2, + typename Threads, + typename TileM, + typename TileK, + typename TileP, + typename Stages, + typename NumBits, + typename GroupSize, + config::QuantMapModeEnum QuantMapMode, + config::AccumulationModeEnum AccumulationMode, + config::DecompositionModeEnum DecompositionMode, + typename G2STiledCopySizeS, + typename MmaPrmK +> +void +qgemm_host(int M, + int N, + int K, + int P, + const T * const __restrict__ A, + const TQ* const __restrict__ Q, + T * __restrict__ D, + const T * const __restrict__ S, + const T * const __restrict__ QM, + const T2* const __restrict__ QM2, + void* __restrict__ workspace, + const int blocks, + const cudaStream_t stream) +{ + using namespace cute; + + CUTE_STATIC_ASSERT_V(Threads{} % _128{} == _0{}); + CUTE_STATIC_ASSERT_V(NumBits{} == _4{} || NumBits{} == _3{} || NumBits{} == _2{}); + + using Config = config::GemmConfig; + using TileScheduler = config::TileScheduler; + auto qgemm_device_func = qgemm_device; + + // assume `slices = 0` since it is deprecated + TileScheduler scheduler(M, N, K, P, 0, blocks); + dim3 grid = scheduler.grid(); + dim3 block = scheduler.block(); + int smem_size = scheduler.smem_size(); + +#if DEBUG + +#define CUDA_CHECK(call) \ + do { \ + cudaError_t status = call; \ + if(status != cudaSuccess) { \ + printf("FAIL: call='%s'. Reason:%s\n", #call, \ + cudaGetErrorString(status)); \ + } \ + } while (0) + + int devId; + int numProcs; + CUDA_CHECK(cudaGetDevice(&devId)); + CUDA_CHECK(cudaDeviceGetAttribute( + &numProcs, + cudaDevAttrMultiProcessorCount, + devId)); + + print("TiledMma\n"); + print(typename Config::TiledMma{}); + print("\n"); + print("TiledMmaQ\n"); + print(typename Config::TiledMmaQ{}); + print("\n"); + print("G2STiledCopyA\n"); + print(typename Config::G2STiledCopyA{}); + print("\n"); + print("G2STiledCopyQ\n"); + print(typename Config::G2STiledCopyQ{}); + print("\n"); + print("G2STiledCopyS\n"); + print(typename Config::G2STiledCopyS{}); + print("\n"); + print("G2STiledCopyQM\n"); + print(typename Config::G2STiledCopyQM{}); + print("\n"); + print("S2RTiledCopyC\n"); + print(typename Config::S2RTiledCopyC{}); + print("\n"); + print("Grid = \t (%d, %d, %d)\n", grid.x, grid.y, grid.z); + print("Block = \t (%d, %d, %d)\n", block.x, block.y, block.z); + print("numProcs = \t %d\n", numProcs); +#endif + + cudaFuncSetAttribute( + qgemm_device_func, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + qgemm_device_func + <<>> + (A, Q, D, S, QM, QM2, + workspace, + scheduler); +} \ No newline at end of file diff --git a/kernels/quantization/flute/qgemm_kernel_raw_generated.cu b/kernels/quantization/flute/qgemm_kernel_raw_generated.cu new file mode 100644 index 0000000000..8a0d12bf4b --- /dev/null +++ b/kernels/quantization/flute/qgemm_kernel_raw_generated.cu @@ -0,0 +1,825 @@ +#include +#include +#include +#include +#include "qgemm_kernel.hpp" + + +template < + typename T, + typename TQ, + typename T2, + typename NumBits, + typename GroupSize +> +void +_qgemm_raw(int M, + int N, + int K, + int P, + const T * const __restrict__ A, + const TQ* const __restrict__ Q, + T * __restrict__ D, + const T * const __restrict__ S, + const T * const __restrict__ QM, + const T2* const __restrict__ QM2, + void* __restrict__ workspace, + const int template_id, + const int num_sms, + const cudaStream_t stream) +{ + + using namespace cute; + static constexpr config::QuantMapModeEnum kVectorized = config::QuantMapModeEnum ::Vectorized; + static constexpr config::QuantMapModeEnum kVectorized_32 = config::QuantMapModeEnum ::Vectorized_32; + static constexpr config::QuantMapModeEnum kVectorized_16 = config::QuantMapModeEnum ::Vectorized_16; + static constexpr config::QuantMapModeEnum kVectorized_8 = config::QuantMapModeEnum ::Vectorized_8; + static constexpr config::AccumulationModeEnum kLow = config::AccumulationModeEnum ::Low; + static constexpr config::AccumulationModeEnum kHigh = config::AccumulationModeEnum ::High; + static constexpr config::AccumulationModeEnum kMixed = config::AccumulationModeEnum ::Mixed; + static constexpr config::DecompositionModeEnum kStreamK = config::DecompositionModeEnum::StreamK; + +#define RUN_QGEMM(T, \ + TQ, \ + T2, \ + SMS_MULTIPLE, \ + THREADS, \ + TILE_M, \ + TILE_K, \ + TILE_P, \ + STAGES, \ + NUM_BITS, \ + GROUP_SIZE, \ + QUANT_MAP_MODE, \ + ACCUMULATION_MODE, \ + DECOMPOSITION_MODE, \ + G2S_TILED_COPY_SIZE_S, \ + MMA_PRM_K) \ + do { \ + qgemm_host< \ + T, \ + TQ, \ + T2, \ + cute::Int, \ + cute::Int, \ + cute::Int, \ + cute::Int, \ + cute::Int, \ + cute::Int, \ + cute::Int, \ + QUANT_MAP_MODE, \ + ACCUMULATION_MODE, \ + DECOMPOSITION_MODE, \ + cute::Int, \ + cute::Int \ + > ( \ + M, \ + N, \ + K, \ + P, \ + A, \ + Q, \ + D, \ + S, \ + QM, \ + QM2, \ + workspace, \ + num_sms * SMS_MULTIPLE, \ + stream); \ + } while (false) + + // Generated Code Below + if constexpr (cute::is_same_v>) + { + switch (template_id) + { + case 0: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 1: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 2: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 3: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 4: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 5: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 6: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 7: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 8: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 9: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 10: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 11: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 12: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 13: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 14: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 15: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 16: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 17: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 18: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 19: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 20: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 21: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 22: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 23: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 24: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 25: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 26: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 27: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 28: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 29: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 30: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 31: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 32: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 33: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 34: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 35: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + default: + AT_ERROR("Unsupported template_id value"); + } + } + else if constexpr (cute::is_same_v>) + { + switch (template_id) + { + case 0: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 1: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 2: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 3: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 4: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 5: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 6: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 7: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 8: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 9: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 10: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 11: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 12: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 13: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 14: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 15: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 16: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 17: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 18: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 19: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 20: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 21: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 22: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 23: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 24: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 25: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 26: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 27: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 28: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 29: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 30: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 31: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 32: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 33: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 34: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 35: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + default: + AT_ERROR("Unsupported template_id value"); + } + } + else if constexpr (cute::is_same_v>) + { + switch (template_id) + { + case 0: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 1: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 2: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 3: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 4: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 5: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 6: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 7: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 8: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 9: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 10: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 11: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 12: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 13: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 14: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 15: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 16: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 17: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 18: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 19: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 20: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 21: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 22: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 23: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 24: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 25: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 26: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 27: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 28: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 29: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 30: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 31: + RUN_QGEMM(T, TQ, T2, 1, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 32: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 33: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 34: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 35: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 36: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 37: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 38: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 39: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 40: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 41: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 42: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 43: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 44: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 45: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 46: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 47: + RUN_QGEMM(T, TQ, T2, 1, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 48: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 49: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 50: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 51: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 52: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 53: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 54: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 55: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 56: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 57: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 58: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 59: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 60: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 61: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 62: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 63: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 64: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 65: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 66: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 67: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 68: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 69: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 70: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 71: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 72: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 73: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 74: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 75: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 76: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 77: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 78: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 79: + RUN_QGEMM(T, TQ, T2, 2, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 80: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 81: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 82: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 83: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 84: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 85: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 86: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 87: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 88: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 89: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 90: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 91: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 92: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 93: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 94: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 95: + RUN_QGEMM(T, TQ, T2, 2, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 96: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 97: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 98: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 99: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 2, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 100: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 101: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 102: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 103: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 3, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 104: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 105: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 106: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 107: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 4, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 108: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 109: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 110: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 111: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 64, 5, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 112: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 113: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 114: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 115: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 116: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 117: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 118: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 119: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 120: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 121: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 122: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 123: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 124: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 125: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 126: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 127: + RUN_QGEMM(T, TQ, T2, 4, 256, 32, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 128: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 129: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 130: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 131: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 2, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 132: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 133: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 134: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 135: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 3, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 136: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 137: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 138: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 139: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 4, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + case 140: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized , kMixed, kStreamK, 2, 1); + break; + case 141: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_32, kMixed, kStreamK, 2, 1); + break; + case 142: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_16, kMixed, kStreamK, 2, 1); + break; + case 143: + RUN_QGEMM(T, TQ, T2, 4, 128, 16, 64, 32, 5, NumBits::value, GroupSize::value, kVectorized_8 , kMixed, kStreamK, 2, 1); + break; + default: + AT_ERROR("Unsupported template_id value"); + } + } + else + { + AT_ERROR("Unsupported NumBits value"); + } +} + + +#define INSTANTIATE_TEMPLATE(T, \ + TQ, \ + T2, \ + NUM_BITS, \ + GROUP_SIZE) \ + template \ + void \ + _qgemm_raw< \ + T, \ + TQ, \ + T2, \ + cute::Int, \ + cute::Int \ + > ( \ + int M, \ + int N, \ + int K, \ + int P, \ + const T * const __restrict__ A, \ + const TQ* const __restrict__ Q, \ + T * __restrict__ D, \ + const T * const __restrict__ S, \ + const T * const __restrict__ QM, \ + const T2* const __restrict__ QM2, \ + void* __restrict__ workspace, \ + const int template_id, \ + const int num_sms, \ + const cudaStream_t stream) + + +// INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 2, 32); +INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 2, 64); +INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 2, 128); +INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 2, 256); +// INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 3, 32); +INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 3, 64); +INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 3, 128); +INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 3, 256); +// INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 4, 32); +INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 4, 64); +INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 4, 128); +INSTANTIATE_TEMPLATE(cute::half_t , cute::uint16_t, __half2 , 4, 256); + +// INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 2, 32); +INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 2, 64); +INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 2, 128); +INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 2, 256); +// INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 3, 32); +INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 3, 64); +INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 3, 128); +INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 3, 256); +// INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 4, 32); +INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 4, 64); +INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 4, 128); +INSTANTIATE_TEMPLATE(cute::bfloat16_t, cute::uint16_t, __nv_bfloat162, 4, 256); \ No newline at end of file diff --git a/kernels/quantization/flute/tile_scheduler_utils.hpp b/kernels/quantization/flute/tile_scheduler_utils.hpp new file mode 100644 index 0000000000..c16e6900b8 --- /dev/null +++ b/kernels/quantization/flute/tile_scheduler_utils.hpp @@ -0,0 +1,1060 @@ +#pragma once + +#include +#include +#include +#include "cutlass/array.h" +#include "cutlass/barrier.h" +#include "cutlass/fast_math.h" +#include "cutlass/block_striped.h" + +// custom extensions to make `BlockStripedReduce` work with `nv_bfloat162` +#include "cutlass_extensions_bf16.h" + +// debugging helpers +#define TS_DEBUG 0 // 1 +#define TS_DEBUG_THR 0 // 255 +#define TS_DEBUG_BLK 0 // 131071 +#define PPRINT_HEADER() do { print("[thread: %d, block: %d]\t", TS_DEBUG_THR, TS_DEBUG_BLK); } while(0) +#define BACKWARDS 1 + + +// 2/4 +namespace config { + +using namespace cute; + + +// Pad the given allocation size up to the nearest cache line +CUTE_HOST_DEVICE static +size_t +cacheline_align_up(size_t size) +{ + static const int CACHELINE_SIZE = 128; + return (size + CACHELINE_SIZE - 1) / CACHELINE_SIZE * CACHELINE_SIZE; +} + + +// Get the workspace size needed for intermediate partial sums +CUTE_HOST_DEVICE +size_t +get_workspace_size_partials(int sk_blocks, int threads, int accum_size) +{ + return cacheline_align_up(sk_blocks * threads * accum_size); +} + + +// Get the workspace size needed for barrier +CUTE_HOST_DEVICE +size_t +get_workspace_size_barriers(int sk_blocks) +{ + // For atomic reduction, each SK-block needs a synchronization flag. For parallel reduction, + // each reduction block needs its own synchronization flag. + return cacheline_align_up(sizeof(typename cutlass::Barrier::T) * sk_blocks); +} + + +template +struct FixupHelper +{ + + using Config = Config_; + using Threads = typename Config::Threads; + static constexpr ReductionModeEnum ReductionMode = ReductionModeEnum::Nondeterministic; + + // Share accumulators with peers + template + CUTE_DEVICE static + void + initialize_or_accumulate( + Tensor& accum, + int thread_index, + int block_index, + int block_index_first, + void* workspace_partials, + void* workspace_barriers) + { + + using AccumulatorArrayT = cutlass::Array; + using BlockStripedReduceT = cutlass::BlockStripedReduce; + + auto accum_array = reinterpret_cast(&accum); + auto workspace_accum = reinterpret_cast(workspace_partials); + auto workspace_index = block_index_first * Threads{}; + +#if TS_DEBUG + + using BarrierT = typename cutlass::Barrier::T; + BarrierT* flag_ptr = reinterpret_cast(workspace_barriers) + block_index_first; + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + print("\n--------------- Fixup ---------------\n"); + PPRINT_HEADER(); print("type: %s\n", (block_index == block_index_first) ? "initialization" : "accumulation"); + PPRINT_HEADER(); print("thread_index: %d\n", thread_index); + PPRINT_HEADER(); print("block_index: %d\n", block_index); + PPRINT_HEADER(); print("block_index_first: %d\n", block_index_first); + PPRINT_HEADER(); print("workspace_index: %d\n", workspace_index); + PPRINT_HEADER(); print("blocks_to_wait: %d\n", BACKWARDS ? block_index_first - block_index : block_index - block_index_first); + PPRINT_HEADER(); print("flag_index: %d\n", block_index_first); + PPRINT_HEADER(); print("flag_value (old): %d\n", *flag_ptr); + } +#endif + + if (block_index == block_index_first) + { + // First peer initializes the workspace partials + BlockStripedReduceT::store(workspace_accum + workspace_index, *accum_array, thread_index); + } + else + { + // Subsequent peers atomically accumulate into the workspace partials + if constexpr (ReductionMode == ReductionModeEnum::Nondeterministic) + { + // Non-deterministic reduction order: wait for the first peer to have initialized the partials before we add to them + cutlass::Barrier::wait_lt(workspace_barriers, thread_index, block_index_first, 1); + } + else + { + // Turnstile reduction order: wait until the previous peer has written +#if BACKWARDS + // in the BACKWARDS case, we define the `block_index_first` as the last + // logical block, hence all non-first blocks will have lower logical index + auto blocks_to_wait = block_index_first - block_index; +#else + auto blocks_to_wait = block_index - block_index_first; +#endif + cutlass::Barrier::wait_eq(workspace_barriers, thread_index, block_index_first, blocks_to_wait); + } + + // Perform reduction in workspace + BlockStripedReduceT::reduce(workspace_accum + workspace_index, *accum_array, thread_index); + } + + // Signal our arrival + cutlass::Barrier::arrive_inc(workspace_barriers, thread_index, block_index_first); + +#if TS_DEBUG + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + PPRINT_HEADER(); print("flag_value (new): %d\n", *flag_ptr); + } +#endif + + } + + // Acquire accumulators from peers + template + CUTE_DEVICE static + void + acquire( + Tensor& accum, + int thread_index, + int block_index, + int block_index_first, + void* workspace_partials, + void* workspace_barriers) + { + + using AccumulatorArrayT = cutlass::Array; + using BlockStripedReduceT = cutlass::BlockStripedReduce; + + auto accum_array = reinterpret_cast(&accum); + auto workspace_accum = reinterpret_cast(workspace_partials); + auto workspace_index = block_index_first * Threads{}; +#if BACKWARDS + // in the BACKWARDS case, we define the `block_index_first` as the last + // logical block, hence all non-first blocks will have lower logical index + auto blocks_to_wait = block_index_first - block_index; +#else + auto blocks_to_wait = block_index - block_index_first; +#endif + + +#if TS_DEBUG + + using BarrierT = typename cutlass::Barrier::T; + BarrierT* flag_ptr = reinterpret_cast(workspace_barriers) + block_index_first; + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + print("\n--------------- Fixup ---------------\n"); + PPRINT_HEADER(); print("type: %s\n", "acquire"); + PPRINT_HEADER(); print("thread_index: %d\n", thread_index); + PPRINT_HEADER(); print("block_index: %d\n", block_index); + PPRINT_HEADER(); print("block_index_first: %d\n", block_index_first); + PPRINT_HEADER(); print("workspace_index: %d\n", workspace_index); + PPRINT_HEADER(); print("blocks_to_wait: %d\n", blocks_to_wait); + PPRINT_HEADER(); print("flag_index: %d\n", block_index_first); + PPRINT_HEADER(); print("flag_value (old): %d\n", *flag_ptr); + } +#endif + + // Wait for arrival + cutlass::Barrier::wait_eq_reset(workspace_barriers, thread_index, block_index_first, blocks_to_wait); + + // Load and add peer-partials accumulator tile to local accumulator tile + BlockStripedReduceT::load_add(*accum_array, workspace_accum + workspace_index, thread_index); + +#if TS_DEBUG + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + PPRINT_HEADER(); print("flag_value (new): %d\n", *flag_ptr); + } +#endif + + } + +}; + + +template +class TileScheduler +{ + +private: + + using Config = Config_; + using FixupHelperT = FixupHelper; + using Threads = typename Config::Threads; + using TileM = typename Config::TileM; + using TileN = typename Config::TileN; + using TileK = typename Config::TileK; + using TileP = typename Config::TileP; + using TileP2 = typename Config::TileP2; + using TileG = typename Config::TileG; + using NumBits = typename Config::NumBits; + using GroupSize = typename Config::GroupSize; + using TileKsPerTileG = typename Config::TileKsPerTileG; + static constexpr int kSmemSize = Config::kSmemSize; + static constexpr DecompositionModeEnum DecompositionMode = Config::DecompositionMode; + // CUTE_STATIC_ASSERT(((DecompositionMode == DecompositionModeEnum::SplitK ) && decltype(Blocks{} == _0{})::value) || + // ((DecompositionMode == DecompositionModeEnum::StreamK) && decltype(Slices{} == _0{})::value)); + + // + // Member state + // + + int m_slices; + int m_blocks; + + int m_M; + int m_N; + int m_K; + int m_P; + int m_P2; + int m_G; + int m_tiles; + int m_tiles_M; + int m_tiles_N; + int m_tiles_K; + int m_tiles_P; + int m_tile_index; + int m_tile_index_read; + int m_tiles_per_block; + int m_tiles_this_block; + int m_tiles_typical_streamk; + int m_tiles_special_streamk; + int m_blocks_typical_streamk; + int m_blocks_special_streamk; + int m_smem_pipe_read_G_offset; + + CUTE_DEVICE + auto + get_block_index_streamk() const + { +#if BACKWARDS + return m_blocks - blockIdx.x - 1; +#else + return blockIdx.x; +#endif + } + + CUTE_DEVICE + auto + get_global_tile_index_streamk(int tile_index) const + { + // when the CTA is at `block_index`, it means it is at `block_index + 1` block, + // and there are `block_index` blocks before it. We then need to figure out + // how many of the previous `block_index` blocks are typical blocks and how many + // are special blocks. + auto block_index = get_block_index_streamk(); + auto blocks_typical = cute::min(block_index, m_blocks_typical_streamk); + auto blocks_special = cute::max(block_index, m_blocks_typical_streamk) - m_blocks_typical_streamk; + return tile_index + + blocks_typical * (m_tiles_per_block) + + blocks_special * (m_tiles_per_block + 1); + } + + CUTE_DEVICE + auto + get_tiles_this_block() const + { + if constexpr (DecompositionMode == DecompositionModeEnum::SplitK) + { + // we assume that partitioning is even in Split-K + return m_tiles_per_block; + } + else + { + // we assume that the special blocks are at the end + return (get_block_index_streamk() < m_blocks_typical_streamk) + ? m_tiles_per_block + : m_tiles_per_block + 1; + } + } + + CUTE_DEVICE + auto + get_tile_coord(int tile_index) const + { + if constexpr (DecompositionMode == DecompositionModeEnum::SplitK) + { + auto tile_M_index = blockIdx.y; + auto tile_N_index = blockIdx.x; // == tile_P_index == tile_P2_index (3-bit case) + auto slice_index = blockIdx.z; + auto tile_K_index = slice_index * m_tiles_per_block + tile_index; + return make_coord(tile_M_index, tile_N_index, tile_K_index); + } + else + { + auto tiles_shape = make_shape(m_tiles_M, m_tiles_N, m_tiles_K); + auto tiles_layout = make_layout(tiles_shape, LayoutRight{}); + auto global_tile_index = get_global_tile_index_streamk(tile_index); + return tiles_layout.get_hier_coord(global_tile_index); + } + } + + // Compute tile residues for predication + // https://github.com/NVIDIA/cutlass/blob/v3.4.0/include/cutlass/gemm/kernel/sm70_gemm.hpp#L227 + + CUTE_DEVICE + auto + get_residue_M(int tile_index) const + { + auto tile_coord = get_tile_coord(tile_index); + auto tile_M_index = get<0>(tile_coord); + return m_M - TileM{} * tile_M_index; + } + + CUTE_DEVICE + auto + get_residue_N(int tile_index) const + { + auto tile_coord = get_tile_coord(tile_index); + auto tile_N_index = get<1>(tile_coord); + return m_N - TileN{} * tile_N_index; + } + + CUTE_DEVICE + auto + get_residue_P(int tile_index) const + { + auto tile_coord = get_tile_coord(tile_index); + auto tile_P_index = get<1>(tile_coord); + return m_P - TileP{} * tile_P_index; + } + + CUTE_DEVICE + auto + get_residue_P2(int tile_index) const + { + auto tile_coord = get_tile_coord(tile_index); + auto tile_P2_index = get<1>(tile_coord); + return m_P2 - TileP2{} * tile_P2_index; + } + + template < + typename PrdEngine, typename PrdLayout, + typename CrdEngine, typename CrdLayout + > + CUTE_DEVICE + auto + set_predicates( + cute::Tensor & pred, + cute::Tensor const& coord, + int residue) const + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(pred); ++i) { + // tile coord < residue + pred(i, 0) = get<0>(coord(0, i, 0)) < residue; + } + } + + CUTE_DEVICE + bool + finished_output_tile(int tile_index) const + { + auto tile_coord = get_tile_coord(tile_index); + auto tile_K_index = get<2>(tile_coord); + return (tile_K_index == m_tiles_K - 1); + } + + CUTE_DEVICE + bool + is_last_tile(int tile_index) const + { + return (tile_index == m_tiles_this_block - 1); + } + + CUTE_DEVICE + bool + started_output_tile(int tile_index) const + { + // We assume that CTA processes tiles in order. If we are at `tile_K_index` + // and have processed `tile_index > tile_K_index`, then we have started the + // output tile. Similarly, if we are at `tile_K_index` and have processed + // `tile_index < tile_K_index`, then we have not started the output tile. + auto tile_coord = get_tile_coord(tile_index); + auto tile_K_index = get<2>(tile_coord); + return tile_K_index <= tile_index; + } + +public: + + // Default Constructor + TileScheduler() = default; + + // Constructor + TileScheduler( + int const M_, + int const N_, + int const K_, + int const P_, + int const slices_, + int const blocks_) + : + m_slices(slices_), + m_blocks(blocks_), + m_M(M_), + m_N(N_), + m_K(K_), + m_P(P_), + m_P2(0), + m_tile_index(0), + m_tile_index_read(0), + m_tiles_typical_streamk(0), + m_tiles_special_streamk(0), + m_blocks_typical_streamk(0), + m_blocks_special_streamk(0) + { + + // in the 3-bit case, we split Q into two sub-matrices + // Q : [N / (sizeof(TQ) / 1), K], the first 1 bits + // Q2: [N / (sizeof(TQ) / 2), K], the remaining 2 bits + if constexpr (is_same_v) { + m_P = ceil_div(m_N, 16); // N / (sizeof(TQ) / 1) == N / (P / 3 * 1) + m_P2 = ceil_div(m_N, 8); // N / (sizeof(TQ) / 2) == N / (P / 3 * 2) + + // in an older version, we set + // `tiles_P = P_ / SuperTileP = (N x 3 / 16) / (TileP + 2 x TileP)` + // here, we instead set + // `tiles_P = P / TileP = N / 16 / TileP` + } + + // number of groups + m_G = ceil_div(m_K, GroupSize{}); + // Note that `tiles_N == tiles_P == tiles_P2 (3-bit case)` + m_tiles_M = ceil_div(m_M, TileM{}); + m_tiles_N = ceil_div(m_N, TileN{}); + m_tiles_K = ceil_div(m_K, TileK{}); + m_tiles_P = ceil_div(m_P, TileP{}); + m_tiles = m_tiles_M * m_tiles_N * m_tiles_K; + + if constexpr (DecompositionMode == DecompositionModeEnum::SplitK) + { + // we assume that K is divisible by TileK * Slices + m_tiles_per_block = m_tiles_K / m_slices; + } + else + { + // the last `m_tiles_remaining` logical blocks will have one extra tile, + // while the rest will have `m_tiles_per_block` tiles + m_tiles_per_block = m_tiles / m_blocks; + m_blocks_special_streamk = m_tiles - m_tiles_per_block * m_blocks; + m_blocks_typical_streamk = m_blocks - m_blocks_special_streamk; + m_tiles_typical_streamk = m_blocks_typical_streamk * (m_tiles_per_block); + m_tiles_special_streamk = m_blocks_special_streamk * (m_tiles_per_block + 1); + } + +#if TS_DEBUG + dim3 grid_dim = grid(); + print("\n--------------- TileScheduler ---------------\n"); + print("M : %5d \t TileM : %5d \t tiles_M: %5d \n", m_M , TileM ::value, m_tiles_M); + print("N : %5d \t TileN : %5d \t tiles_N: %5d \n", m_N , TileN ::value, m_tiles_N); + print("K : %5d \t TileK : %5d \t tiles_K: %5d \n", m_K , TileK ::value, m_tiles_K); + print("P : %5d \t TileP : %5d \t tiles_P: %5d \n", m_P , TileP ::value, m_tiles_P); + print("P2: %5d \t TileP2: %5d \n", m_P2, TileP2::value); + print("G: %5d \t TileG : %5d \n", m_G , TileG ::value); + print("tiles: %5d (%5d per block) \n", m_tiles, m_tiles_per_block); + print("tiles_typical_streamk: %5d \n", m_tiles_typical_streamk); + print("tiles_special_streamk: %5d \n", m_tiles_special_streamk); + print("blocks_typical_streamk: %5d \n", m_blocks_typical_streamk); + print("blocks_special_streamk: %5d \n", m_blocks_special_streamk); + print("grid_dim: (%d, %d, %d)\n", grid_dim.x, grid_dim.y, grid_dim.z); +#endif + + } + + template < + typename PrdEngineA , typename PrdLayoutA , + typename PrdEngineQ , typename PrdLayoutQ , + typename PrdEngineQ2, typename PrdLayoutQ2, + typename PrdEngineS , typename PrdLayoutS , + typename CrdEngineA , typename CrdLayoutA , + typename CrdEngineQ , typename CrdLayoutQ , + typename CrdEngineQ2, typename CrdLayoutQ2, + typename CrdEngineS , typename CrdLayoutS + > + CUTE_DEVICE + void + initialize( + cute::Tensor & pred_A, + cute::Tensor & pred_Q, + cute::Tensor & pred_Q2, + cute::Tensor & pred_S, + cute::Tensor const& coord_A, + cute::Tensor const& coord_Q, + cute::Tensor const& coord_Q2, + cute::Tensor const& coord_S) + { + m_tile_index = 0; + m_tile_index_read = 0; + m_tiles_this_block = get_tiles_this_block(); + +#if TS_DEBUG + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + print("\n--------------- Initialization ---------------\n"); + PPRINT_HEADER(); print("tile_index: %d (ended = %d)\n", m_tile_index , !tile_is_in_bound()); + PPRINT_HEADER(); print("tile_index_read: %d (ended = %d)\n", m_tile_index_read, !tile_read_is_in_bound()); + PPRINT_HEADER(); print("tiles_this_block: %d\n", m_tiles_this_block); + } +#endif + + if constexpr (DecompositionMode == DecompositionModeEnum::StreamK) + { + // the starting K tile index might not be aligned wth the G tile + auto tile_coord_init = get_tile_coord(0); + auto tile_K_index_init = get<2>(tile_coord_init); + m_smem_pipe_read_G_offset = tile_K_index_init % TileKsPerTileG{}; + +#if TS_DEBUG + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + PPRINT_HEADER(); print("tile_coord_init: "); print(tile_coord_init); print("\n"); + PPRINT_HEADER(); print("smem_pipe_read_G_offset: %d\n", m_smem_pipe_read_G_offset); + } +#endif + + } + + set_predicates(pred_A , coord_A , get_residue_M (m_tile_index_read)); + set_predicates(pred_Q , coord_Q , get_residue_P (m_tile_index_read)); + set_predicates(pred_Q2, coord_Q2, get_residue_P2(m_tile_index_read)); + set_predicates(pred_S , coord_S , get_residue_N (m_tile_index_read)); + +#if TS_DEBUG + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + PPRINT_HEADER(); print("A : "); print(get_tile_coord_A ()); print(" \t residue_M : %5d \n", get_residue_M (m_tile_index_read)); + PPRINT_HEADER(); print("Q : "); print(get_tile_coord_Q ()); print(" \t residue_P : %5d \n", get_residue_P (m_tile_index_read)); + PPRINT_HEADER(); print("Q2: "); print(get_tile_coord_Q2()); print(" \t residue_P2: %5d \n", get_residue_P2(m_tile_index_read)); + PPRINT_HEADER(); print("S : "); print(get_tile_coord_S ()); print(" \t residue_N : %5d \n", get_residue_N (m_tile_index_read)); + } +#endif + } + + template < + typename PrdEngineA , typename PrdLayoutA , + typename PrdEngineQ , typename PrdLayoutQ , + typename PrdEngineQ2, typename PrdLayoutQ2, + typename PrdEngineS , typename PrdLayoutS , + typename CrdEngineA , typename CrdLayoutA , + typename CrdEngineQ , typename CrdLayoutQ , + typename CrdEngineQ2, typename CrdLayoutQ2, + typename CrdEngineS , typename CrdLayoutS + > + CUTE_DEVICE + void + step_read( + cute::Tensor & pred_A, + cute::Tensor & pred_Q, + cute::Tensor & pred_Q2, + cute::Tensor & pred_S, + cute::Tensor const& coord_A, + cute::Tensor const& coord_Q, + cute::Tensor const& coord_Q2, + cute::Tensor const& coord_S) + { + + ++m_tile_index_read; + + // we don't need to reset predicates in Split-K + if constexpr (DecompositionMode == DecompositionModeEnum::StreamK) + { + auto tile_coord_old = get_tile_coord(m_tile_index_read - 1); + auto tile_coord_new = get_tile_coord(m_tile_index_read); + // the K index changes, we need to update the predicates + if (get<2>(tile_coord_old) != get<2>(tile_coord_new)) + { + set_predicates(pred_A , coord_A , get_residue_M (m_tile_index_read)); + set_predicates(pred_Q , coord_Q , get_residue_P (m_tile_index_read)); + set_predicates(pred_Q2, coord_Q2, get_residue_P2(m_tile_index_read)); + set_predicates(pred_S , coord_S , get_residue_N (m_tile_index_read)); + } + } + +#if TS_DEBUG + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + print("\n--------------- Step Read ---------------\n"); + auto tile_coord = get_tile_coord(m_tile_index_read); + PPRINT_HEADER(); print("tile_index_read: %d (ended = %d)\n", m_tile_index_read, !tile_read_is_in_bound()); + PPRINT_HEADER(); print("tile_coord: "); print(tile_coord); print("\n"); + PPRINT_HEADER(); print("A : "); print(get_tile_coord_A ()); print(" \t residue_M : %5d \n", get_residue_M (m_tile_index_read)); + PPRINT_HEADER(); print("Q : "); print(get_tile_coord_Q ()); print(" \t residue_P : %5d \n", get_residue_P (m_tile_index_read)); + PPRINT_HEADER(); print("Q2: "); print(get_tile_coord_Q2()); print(" \t residue_P2: %5d \n", get_residue_P2(m_tile_index_read)); + PPRINT_HEADER(); print("S : "); print(get_tile_coord_S ()); print(" \t residue_N : %5d \n", get_residue_N (m_tile_index_read)); + } +#endif + + } + + CUTE_DEVICE + void + step() + { + ++m_tile_index; + +#if TS_DEBUG + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + print("\n--------------- Step ---------------\n"); + auto tile_coord = get_tile_coord(m_tile_index); + PPRINT_HEADER(); print("tile_index: %d (ended = %d)\n", m_tile_index , !tile_is_in_bound()); + PPRINT_HEADER(); print("tile_coord: "); print(tile_coord); print("\n"); + } +#endif + + } + +#if TS_DEBUG + + template + CUTE_DEVICE + void + maybe_print_workspace_size(Tensor const& accum) const + { + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + using AccumulatorArrayT = cutlass::Array; + dim3 grid_size = grid(); + auto blocks = grid_size.x * grid_size.y * grid_size.z; + auto accum_size = sizeof(AccumulatorArrayT); + auto workspace_size_barriers = get_workspace_size_barriers(blocks); + auto workspace_size_partials = get_workspace_size_partials(blocks, Threads{}, accum_size); + auto workspace_size = workspace_size_barriers + workspace_size_partials; + + print("\n--------------- Workspace Size ---------------\n"); + PPRINT_HEADER(); print("accum: "); print(accum.layout()); print("\n"); + PPRINT_HEADER(); print("blocks: %d\n", blocks); + PPRINT_HEADER(); print("accum_size: %d\n", accum_size); + PPRINT_HEADER(); print("workspace_size_barriers: %d\n", workspace_size_barriers); + PPRINT_HEADER(); print("workspace_size_partials: %d\n", workspace_size_partials); + PPRINT_HEADER(); print("workspace_size: %d\n", workspace_size); + } + } + + template + CUTE_DEVICE + bool + is_output_tile_coord(int tile_index, Coord const& coord) const + { + auto tile_coord = get_tile_coord(tile_index); + return (get<0>(tile_coord) == get<0>(coord) && + get<1>(tile_coord) == get<1>(coord)); + } + +#endif + + CUTE_DEVICE + auto + M() const + { + return m_M; + } + + CUTE_DEVICE + auto + N() const + { + return m_N; + } + + CUTE_DEVICE + auto + K() const + { + return m_K; + } + + CUTE_DEVICE + auto + P() const + { + return m_P; + } + + CUTE_DEVICE + auto + P2() const + { + return m_P2; + } + + CUTE_DEVICE + auto + G() const + { + return m_G; + } + + CUTE_DEVICE + auto + tile_index() const + { + return m_tile_index; + } + + CUTE_DEVICE + auto + tile_index_read() const + { + return m_tile_index_read; + } + + CUTE_DEVICE + auto + smem_pipe_read_G_offset() const + { + return m_smem_pipe_read_G_offset; + } + + CUTE_HOST_DEVICE + dim3 + grid() const + { + if constexpr (DecompositionMode == DecompositionModeEnum::SplitK) + { + // tiles_N == tiles_P == tiles_P2 (3-bit case) + return dim3(m_tiles_N, m_tiles_M, m_slices); + } + else + { + return dim3(m_blocks); + } + } + + CUTE_HOST + dim3 + block() const + { + return dim3(Threads{}); + } + + CUTE_HOST + int + smem_size() const + { + return kSmemSize; + } + + CUTE_DEVICE + auto + workspace_size_barriers() const + { + int blocks; + if constexpr (DecompositionMode == DecompositionModeEnum::SplitK) + { + blocks = m_tiles_N * m_tiles_M * m_slices; + } + else + { + blocks = m_blocks; + } + + return get_workspace_size_barriers(blocks); + } + + CUTE_DEVICE + auto + get_tile_coord_A() const + { + auto tile_coord = get_tile_coord(m_tile_index_read); + auto tile_M_index = get<0>(tile_coord); + auto tile_K_index = get<2>(tile_coord); + return make_coord(_, _, _, tile_M_index, tile_K_index); + } + + CUTE_DEVICE + auto + get_tile_coord_Q() const + { + auto tile_coord = get_tile_coord(m_tile_index_read); + auto tile_P_index = get<1>(tile_coord); + auto tile_K_index = get<2>(tile_coord); + return make_coord(_, _, _, tile_P_index, tile_K_index); + } + + CUTE_DEVICE + auto + get_tile_coord_Q2() const + { + return get_tile_coord_Q(); + } + + CUTE_DEVICE + auto + get_tile_coord_S() const + { + auto tile_coord = get_tile_coord(m_tile_index_read); + auto tile_N_index = get<1>(tile_coord); + auto tile_K_index = get<2>(tile_coord); + auto tile_G_index = tile_K_index / TileKsPerTileG{}; + return make_coord(_, _, _, tile_N_index, tile_G_index); + } + + CUTE_DEVICE + bool + tile_read_is_in_bound() const + { + return (m_tile_index_read < m_tiles_this_block); + } + + CUTE_DEVICE + bool + tile_is_in_bound() const + { + return (m_tile_index < m_tiles_this_block); + } + + CUTE_DEVICE + bool + start_of_group() const + { + // the starting K tile index might not be aligned wth the G tile + if (m_tile_index_read == 0) + { + return true; + } + auto tile_coord = get_tile_coord(m_tile_index_read); + auto tile_K_index = get<2>(tile_coord); + return (tile_K_index % TileKsPerTileG{} == 0); + } + + CUTE_DEVICE + bool + needs_fixup() const + { + // We needs fixup if either + // 1. we are done with the last tile of the output tile + // 2. we are done with the last tile of all of CTA's tiles + return (finished_output_tile(m_tile_index) || is_last_tile(m_tile_index)); + } + + CUTE_DEVICE + bool + needs_epilogue() const + { + if constexpr (DecompositionMode == DecompositionModeEnum::StreamK && BACKWARDS == 1) + { + // the CTA that started the first tile of the output tile will do the epilogue + return (needs_fixup() && started_output_tile(m_tile_index)); + } + else + { + // the CTA that finished the last tile of the output tile will do the epilogue + return finished_output_tile(m_tile_index); + } + } + + CUTE_DEVICE + bool + needs_to_clear_accum() const + { + // we need to clear the accumulator if both + // 1. we are done with the output tile (perhaps partially), i.e., we did the fixup/epilogue + // 2. we are not done with all of CTA's tiles + return (needs_fixup() && (!is_last_tile(m_tile_index))); + } + + template + CUTE_DEVICE + void + maybe_fixup( + Tensor& accum, + int thread_index, + void* workspace_partials, + void* workspace_barriers) const + { + + int block_index; + int block_index_first; + auto block_started_output_tile = started_output_tile(m_tile_index); + auto block_finished_output_tile = finished_output_tile(m_tile_index); + + if constexpr (DecompositionMode == DecompositionModeEnum::SplitK) + { + if (m_slices == 1) + { + // Slice-K does not require fixup + return; + } + + // Split-K + block_index = blockIdx.x * gridDim.y * gridDim.z + blockIdx.y * gridDim.z + blockIdx.z; + block_index_first = blockIdx.x * gridDim.y * gridDim.z + blockIdx.y * gridDim.z; + +#if TS_DEBUG + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + print("\n--------------- Maybe Fixup (Split-K) ---------------\n"); + PPRINT_HEADER(); print("tile_index: %d\n", m_tile_index); + PPRINT_HEADER(); print("tile_index_read: %d\n", m_tile_index_read); + PPRINT_HEADER(); print("block_index: %d\n", block_index); + PPRINT_HEADER(); print("block_index_first: %d\n", block_index_first); + PPRINT_HEADER(); print("block_started_output_tile: %d\n", block_started_output_tile); + PPRINT_HEADER(); print("block_finished_output_tile: %d\n", block_finished_output_tile); + } +#endif + + } + else + { + // Stream-K + block_index = get_block_index_streamk(); + auto tiles_shape = make_shape(m_tiles_M, m_tiles_N, m_tiles_K); + auto tiles_layout = make_layout(tiles_shape, LayoutRight{}); + auto global_tile_index = get_global_tile_index_streamk(m_tile_index); + auto tile_coord = tiles_layout.get_hier_coord(global_tile_index); +#if BACKWARDS + // assuming all blocks are launched and run in lock-step, then the first block + // to finish its portion of the output tile is the one that finishes the last + // logical K-tile of the output tile. Similarly, the last block to finish its + // portion of the output tile is the one that finishes the first logical K-tile + auto tile_coord_first = make_coord(get<0>(tile_coord), get<1>(tile_coord), m_tiles_K - 1); +#else + // technically, even with BACKWARDS mode off, we still should define `tile_coord_first` + // as above. However, when not all blocks are launched, this might cause deadlocks due + // to circular dependencies. Hence, this is a workaround to avoid deadlocks. + auto tile_coord_first = make_coord(get<0>(tile_coord), get<1>(tile_coord), 0); +#endif + auto tiles_to_first = tiles_layout(tile_coord_first) + 1; + auto blocks_typical_first = ceil_div(cute::min(tiles_to_first, m_tiles_typical_streamk) , m_tiles_per_block); + auto blocks_special_first = ceil_div(cute::max(tiles_to_first, m_tiles_typical_streamk) - m_tiles_typical_streamk, m_tiles_per_block + 1); + // almost `floor_div`, except when `tiles_to_first` is a multiple of `m_tiles_per_block` + block_index_first = blocks_typical_first + blocks_special_first - 1; + +#if BACKWARDS + // in the BACKWARDS case, the "finishing" block is the one that started the output tile, + // who is likely the last block to finish its portion of the output tile. Similarly, the + // non "finishing" blocks are the ones that did not start the output tile, and includes + // the first block to finish its portion of the output tile. + cutlass::swap(block_started_output_tile, block_finished_output_tile); +#endif + + +#if TS_DEBUG + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + print("\n--------------- Maybe Fixup (Stream-K) ---------------\n"); + PPRINT_HEADER(); print("tile_index: %d\n", m_tile_index); + PPRINT_HEADER(); print("tile_index_read: %d\n", m_tile_index_read); + PPRINT_HEADER(); print("block_index: %d\n", block_index); + PPRINT_HEADER(); print("block_index_first: %d\n", block_index_first); + PPRINT_HEADER(); print("block_started_output_tile: %d\n", block_started_output_tile); + PPRINT_HEADER(); print("block_finished_output_tile: %d\n", block_finished_output_tile); + PPRINT_HEADER(); print("tile_coord: "); print(tile_coord); print("\n"); + PPRINT_HEADER(); print("tile_coord_first: "); print(tile_coord_first); print("\n"); + PPRINT_HEADER(); print("tiles_to_first: %d\n", tiles_to_first); + PPRINT_HEADER(); print("blocks_typical_first: %d\n", blocks_typical_first); + PPRINT_HEADER(); print("blocks_special_first: %d\n", blocks_special_first); + } +#endif + } + + + if (!block_finished_output_tile) + { + // Non "finishing" SK blocks must share their partial accumulator sums through global scratch workspace + FixupHelperT::initialize_or_accumulate( + accum, + thread_index, + block_index, + block_index_first, + workspace_partials, + workspace_barriers); + } + else + { + // DP blocks and "finishing" SK blocks must perform epilogue operations and write the output tile + if (!block_started_output_tile) + { + // A "finishing" SK block must first aggregate its accumulator partial sums with those shared by peer threadblocks + FixupHelperT::acquire( + accum, + thread_index, + block_index, + block_index_first, + workspace_partials, + workspace_barriers); + } + } + } + + CUTE_DEVICE + void + prepare_epilogue( + int & tile_M_index, + int & tile_N_index, + int & residue_M, + int & residue_N) const + { + auto tile_coord = get_tile_coord(m_tile_index); + tile_M_index = get<0>(tile_coord); + tile_N_index = get<1>(tile_coord); + residue_M = get_residue_M(m_tile_index); + residue_N = get_residue_N(m_tile_index); + +#if TS_DEBUG + + if(thread(TS_DEBUG_THR, TS_DEBUG_BLK)) + { + print("\n--------------- Epilogue ---------------\n"); + PPRINT_HEADER(); print("tile_coord: "); print(tile_coord); print("\n"); + PPRINT_HEADER(); print("tile_index: %d\n", m_tile_index); + PPRINT_HEADER(); print("tile_index_read: %d\n", m_tile_index_read); + PPRINT_HEADER(); print("residue_M: %d\n", residue_M); + PPRINT_HEADER(); print("residue_N: %d\n", residue_N); + } +#endif + + + } + +}; + +} // namespace config \ No newline at end of file