diff --git a/CMakeLists.txt b/CMakeLists.txt index 52a4560387..9c03d2dbfc 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -538,6 +538,34 @@ if(APHRODITE_GPU_LANG STREQUAL "CUDA" OR APHRODITE_GPU_LANG STREQUAL "HIP") SRCS "${APHRODITE_STABLE_EXT_SRC}" CUDA_ARCHS "${CUDA_ARCHS}") + # + # Swordfish kernels: w4a16 GEMM for the Blackwell sm100 family + # (datacenter sm100 + Thor sm110). Family (f) variants need CUDA >= 13.0; + # older toolchains fall back to arch-exact sm100. + # + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(SWORDFISH_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(SWORDFISH_ARCHS "10.0a" "${CUDA_ARCHS}") + endif() + if(SWORDFISH_ARCHS) + set(SWORDFISH_SRCS + "csrc/libtorch_stable/quantization/swordfish/swordfish_prepack.cu" + "csrc/libtorch_stable/quantization/swordfish/swordfish_mm.cu" + "csrc/libtorch_stable/quantization/swordfish/swordfish_moe.cu" + "csrc/libtorch_stable/quantization/swordfish/swordfish_dense_tier.cu" + "csrc/libtorch_stable/quantization/swordfish/swordfish_prefill.cu" + "csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_f16.cu") + set_gencode_flags_for_srcs( + SRCS "${SWORDFISH_SRCS}" + CUDA_ARCHS "${SWORDFISH_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${SWORDFISH_SRCS}") + message(STATUS "Building Swordfish kernels for archs: ${SWORDFISH_ARCHS}") + else() + message(STATUS "Not building Swordfish kernels as no compatible archs " + "found in CUDA target architectures") + endif() + if(COOPERATIVE_TOPK_ARCHS) list(APPEND APHRODITE_STABLE_EXT_SRC "csrc/libtorch_stable/cooperative_topk.cu") @@ -737,7 +765,10 @@ if(APHRODITE_GPU_LANG STREQUAL "CUDA" OR APHRODITE_GPU_LANG STREQUAL "HIP") endif() # Only build AllSpark kernels if we are building for at least some compatible archs. - cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}") + # 8.0+PTX JIT-forwards onto Blackwell like marlin; the kernel is plain + # mma.sync + cp.async and its large-M path is dequant + cuBLAS, which owns + # channelwise w8a16 prefill on sm100/sm110. + cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0+PTX" "${CUDA_ARCHS}") if (ALLSPARK_ARCHS) set(ALLSPARK_SRCS "csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu" diff --git a/aphrodite/_custom_ops.py b/aphrodite/_custom_ops.py index d3d4be1c5b..4074370b48 100644 --- a/aphrodite/_custom_ops.py +++ b/aphrodite/_custom_ops.py @@ -1550,6 +1550,209 @@ def machete_prepack_B_fake( return torch.empty_like(b_q_weight, memory_format=torch.contiguous_format) +# Swordfish (Blackwell sm100/sm110 w4a16) +def swordfish_prepack_B( + b_q_weight: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int = 4, + perm: torch.Tensor | None = None, +) -> torch.Tensor: + """Pack a GPTQ int weight (int32 [K*bits/32, N]) into Swordfish ABI v1 + (int32 [NB, KB, 512*bits/4] block-linear). perm applies the act_order + row sort during the repack.""" + return torch.ops._C.swordfish_prepack_B(b_q_weight, perm, size_k, size_n, num_bits) + + +if hasattr(torch.ops._C, "swordfish_prepack_B"): + + @register_fake("_C::swordfish_prepack_B") + def swordfish_prepack_B_fake( + b_q_weight: torch.Tensor, + perm: torch.Tensor | None, + size_k: int, + size_n: int, + num_bits: int, + ) -> torch.Tensor: + return torch.empty( + (size_n // 64, size_k // 64, 128 * num_bits), + dtype=torch.int32, + device=b_q_weight.device, + ) + + +def swordfish_mm( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_size: int, + size_k: int, + size_n: int, + group_zps: torch.Tensor | None = None, + num_bits: int = 4, + perm: torch.Tensor | None = None, +) -> torch.Tensor: + """w4a16/w8a16 GEMM: a [M, K] fp16/bf16 times a Swordfish ABI v1 packed + weight with per-group scales [groups, N]. group_zps holds prescaled + (8 - zp) * scale rows for zero-point checkpoints (AWQ/HQQ, 4-bit only). + perm is the act_order column sort; the op permutes the activations for + the fused paths and folds the sort into the dense tier's weight scatter.""" + return torch.ops._C.swordfish_mm(a, b_packed, group_scales, group_zps, perm, num_bits, group_size, size_k, size_n) + + +if hasattr(torch.ops._C, "swordfish_mm"): + + @register_fake("_C::swordfish_mm") + def swordfish_mm_fake( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_zps: torch.Tensor | None, + perm: torch.Tensor | None, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, + ) -> torch.Tensor: + return torch.empty((a.shape[0], size_n), dtype=a.dtype, device=a.device) + + +def swordfish_dequant_dense( + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_zps: torch.Tensor | None, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, + transpose: bool, +) -> torch.Tensor: + """Dequantize Swordfish-packed weights to dense fp16/bf16, [K, N] or + out-major [N, K], with an optional leading expert dimension.""" + return torch.ops._C.swordfish_dequant_dense( + b_packed, + group_scales, + group_zps, + num_bits, + group_size, + size_k, + size_n, + transpose, + ) + + +if hasattr(torch.ops._C, "swordfish_dequant_dense"): + + @register_fake("_C::swordfish_dequant_dense") + def swordfish_dequant_dense_fake( + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_zps: torch.Tensor | None, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, + transpose: bool, + ) -> torch.Tensor: + shape: tuple[int, ...] = (size_n, size_k) if transpose else (size_k, size_n) + if b_packed.dim() == 4: + shape = (b_packed.shape[0],) + shape + return torch.empty(shape, dtype=group_scales.dtype, device=b_packed.device) + + +def swordfish_moe_mm( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + topk_weights: torch.Tensor | None, + moe_block_size: int, + top_k: int, + mul_topk_weights: bool, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, +) -> torch.Tensor: + """Fused-MoE GEMM over per-expert Swordfish ABI v1 weights + (int32 [E, NB, KB, words]) with token blocks prepared by + moe_align_block_size(moe_block_size in {16, 64}). The 64-token blocks + run the CTA-wide weight-staging kernel for batched shapes. Returns + [M * top_k, size_n].""" + return torch.ops._C.swordfish_moe_mm( + a, + b_packed, + group_scales, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + num_bits, + group_size, + size_k, + size_n, + ) + + +if hasattr(torch.ops._C, "swordfish_moe_mm"): + + @register_fake("_C::swordfish_moe_mm") + def swordfish_moe_mm_fake( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + topk_weights: torch.Tensor | None, + moe_block_size: int, + top_k: int, + mul_topk_weights: bool, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, + ) -> torch.Tensor: + return torch.empty((a.shape[0] * top_k, size_n), dtype=a.dtype, device=a.device) + + +def swordfish_prefill_mm( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_size: int, + size_k: int, + size_n: int, + group_zps: torch.Tensor | None = None, + num_bits: int = 4, +) -> torch.Tensor: + """w4a16 prefill GEMM (sm100 tcgen05 mixed-input mainloop fork): a [M, K] + bf16 times a Swordfish ABI v1 packed weight (int32 [NB, KB, 512]) with + bf16 per-group scales [groups, N]. v1: group_size 128, K/N % 128 == 0.""" + return torch.ops._C.swordfish_prefill_mm(a, b_packed, group_scales, group_zps, num_bits, group_size, size_k, size_n) + + +if hasattr(torch.ops._C, "swordfish_prefill_mm"): + + @register_fake("_C::swordfish_prefill_mm") + def swordfish_prefill_mm_fake( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_zps: torch.Tensor | None, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, + ) -> torch.Tensor: + return torch.empty((a.shape[0], size_n), dtype=a.dtype, device=a.device) + + # CUTLASS W4A8 def cutlass_w4a8_mm( a: torch.Tensor, diff --git a/aphrodite/config/kernel.py b/aphrodite/config/kernel.py index ec4be93f45..8edf40a893 100644 --- a/aphrodite/config/kernel.py +++ b/aphrodite/config/kernel.py @@ -147,6 +147,7 @@ def with_default(cls, default: list[str], /, **kwargs: list[str]) -> "IrOpPriori "torch", "aiter", "machete", + "swordfish", "fbgemm", "conch", "exllama", diff --git a/aphrodite/model_executor/kernels/linear/__init__.py b/aphrodite/model_executor/kernels/linear/__init__.py index 23d04bdb07..26cfea421f 100644 --- a/aphrodite/model_executor/kernels/linear/__init__.py +++ b/aphrodite/model_executor/kernels/linear/__init__.py @@ -57,6 +57,9 @@ from aphrodite.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( RDNA3W4A16LinearKernel, ) +from aphrodite.model_executor.kernels.linear.mixed_precision.swordfish import ( + SwordfishLinearKernel, +) from aphrodite.model_executor.kernels.linear.mixed_precision.triton_w4a16 import ( TritonW4A16LinearKernel, ) @@ -275,6 +278,12 @@ def _get_linear_backend() -> str: "machete": { MacheteLinearKernel, }, + "swordfish": { + SwordfishLinearKernel, + }, + "allspark": { + AllSparkLinearKernel, + }, "fbgemm": { FbgemmNvFp4LinearKernel, }, @@ -391,6 +400,7 @@ def _filter_kernels_by_backend( _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = { PlatformEnum.CUDA: [ CutlassW4A8LinearKernel, + SwordfishLinearKernel, MacheteLinearKernel, AllSparkLinearKernel, MarlinLinearKernel, @@ -1037,6 +1047,7 @@ def register_linear_kernel( "ExllamaLinearKernel", "MacheteLinearKernel", "MarlinLinearKernel", + "SwordfishLinearKernel", "TritonW4A16LinearKernel", "XPUW4A8IntLinearKernel", "XPUwNa16LinearKernel", diff --git a/aphrodite/model_executor/kernels/linear/mixed_precision/swordfish.py b/aphrodite/model_executor/kernels/linear/mixed_precision/swordfish.py new file mode 100644 index 0000000000..7aa33bb3d7 --- /dev/null +++ b/aphrodite/model_executor/kernels/linear/mixed_precision/swordfish.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SwordfishLinearKernel: w4a16 GEMM for the Blackwell sm100 family +(datacenter sm100 + Thor sm110). GPTQ u4b8 and AWQ uint4+zp.""" + +import torch + +from aphrodite import _custom_ops as ops +from aphrodite.model_executor.layers.quantization.utils.marlin_utils import ( + marlin_sort_g_idx, +) +from aphrodite.model_executor.layers.quantization.utils.quant_utils import unpack_cols +from aphrodite.model_executor.layers.quantization.utils.swordfish_utils import ( + check_swordfish_supports_shape, + query_swordfish_supported_group_sizes, + query_swordfish_supported_quant_types, +) +from aphrodite.model_executor.parameter import ( + BaseAphroditeParameter, + permute_param_layout_, +) +from aphrodite.platforms import current_platform + +from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig + + +class SwordfishLinearKernel(MPLinearKernel): + @classmethod + def get_min_capability(cls) -> int: + return 100 + + @classmethod + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Swordfish only supported on CUDA" + + capability = current_platform.get_device_capability() + # sm100 family only, datacenter 10.x and Thor 11.x. Consumer + # Blackwell (12.x) is a different SM and is untested. + if capability is None or capability.major not in (10, 11): + return ( + False, + f"Swordfish requires the sm100 family (compute capability 10.x or 11.x), got {capability}", + ) + + if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]: + return ( + False, + "Act reordering with a partial K (row-parallel TP) not supported by Swordfish", + ) + + supported_types = query_swordfish_supported_quant_types(c.zero_points) + if c.weight_type not in supported_types: + return ( + False, + f"Quant type ({c.weight_type}) not supported by Swordfish v1, supported: {supported_types}", + ) + + if c.group_size not in query_swordfish_supported_group_sizes(c.act_type): + return ( + False, + f"Group size ({c.group_size}) / act type ({c.act_type}) not supported by Swordfish v1", + ) + + return check_swordfish_supports_shape(c.partition_weight_shape[0], c.partition_weight_shape[1]) + + # weight_packed has {input_dim 0, output_dim 1, packed_dim 0} and + # weight_scale has {input_dim 0, output_dim 1}. + def process_weights_after_loading(self, layer: torch.nn.Module): + c = self.config + size_k, size_n = c.partition_weight_shape + + # Act-order rows sort by group at prepack; scales then apply in the + # plain grouped order and only the activation columns need the sort + # permutation at run time. + if c.has_g_idx: + assert self.w_gidx_name is not None + g_idx, g_idx_sort_indices = marlin_sort_g_idx(getattr(layer, self.w_gidx_name)) + self._transform_param(layer, self.w_gidx_name, lambda _: g_idx) + layer.g_idx_sort_indices = g_idx_sort_indices + + def transform_w_q(x): + assert isinstance(x, BaseAphroditeParameter) + permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) + x.data = ops.swordfish_prepack_B( + x.data.contiguous(), + size_k, + size_n, + c.weight_type.size_bits, + perm=layer.g_idx_sort_indices if c.has_g_idx else None, + ) + return x + + def transform_w_s(x): + assert isinstance(x, BaseAphroditeParameter) + permute_param_layout_(x, input_dim=0, output_dim=1) + x.data = x.data.contiguous() + # Channelwise checkpoints replicate their single scale row to + # group 128, which buys the full grouped machinery (tcgen05 + # prefill, the dense tier) for kilobytes of duplicate scales. + if c.group_size == -1: + x.data = x.data.expand(size_k // 128, size_n).contiguous() + return x + + self._transform_param(layer, self.w_q_name, transform_w_q) + self._transform_param(layer, self.w_s_name, transform_w_s) + + if c.zero_points: + # The kernel dequantizes to (w - 8) * s and adds a per-group + # (8 - zp) * s, so the zp tensor becomes scale-shaped [groups, N] + # in the activation dtype. qzeros arrives in the standard packed + # layout [N / 8, groups]. + scales = getattr(layer, self.w_s_name) + num_groups = scales.shape[0] + + def transform_w_zp(x): + zp = unpack_cols( + x.data.t().contiguous(), + c.weight_type.size_bits, + num_groups, + size_n, + ) + x.data = (8.0 - zp.to(scales.dtype)) * scales.data + return x + + self._transform_param(layer, self.w_zp_name, transform_w_zp) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + c = self.config + w_q, w_s, w_zp, _ = self._get_weight_params(layer) + # Symmetric GPTQ checkpoints still materialize a qzeros param; only + # zero-point configs transformed it into the (8 - zp) * s tensor. + if not c.zero_points: + w_zp = None + + x_2d = x.reshape(-1, x.shape[-1]) + out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) + + # The decode/prefill crossover lives inside the C++ op. A Python + # branch would be baked in at torch.compile trace time. + output = ops.swordfish_mm( + x_2d, + w_q, + w_s, + c.group_size, + c.partition_weight_shape[0], + c.partition_weight_shape[1], + group_zps=w_zp, + num_bits=c.weight_type.size_bits, + perm=layer.g_idx_sort_indices if c.has_g_idx else None, + ) + + if bias is not None: + output.add_(bias) + + return output.reshape(out_shape) diff --git a/aphrodite/model_executor/layers/fused_moe/experts/swordfish_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/swordfish_moe.py new file mode 100644 index 0000000000..6b620ced65 --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/experts/swordfish_moe.py @@ -0,0 +1,421 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Swordfish fused MoE experts (Blackwell sm100/sm110 wNa16).""" + +import functools + +import torch + +import aphrodite._custom_ops as ops +import aphrodite.model_executor.layers.fused_moe.modular_kernel as mk +from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation +from aphrodite.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from aphrodite.model_executor.layers.fused_moe.moe_align_block_size import ( + moe_align_block_size, +) +from aphrodite.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kInt4Static, + kInt8Static, +) +from aphrodite.model_executor.layers.quantization.utils.swordfish_utils import ( + SWORDFISH_BLOCK_N, +) +from aphrodite.platforms import current_platform + +# Token-block size fed to moe_align_block_size and the kernel. 16-token +# blocks run the per-warp Stream-K kernel that wins small batches; 32-token +# blocks fuse two m16 tiles per staged weight chunk, amortizing the dequant +# once experts average enough tokens to fill them. +SWORDFISH_MOE_BLOCK_SIZE = 16 +SWORDFISH_MOE_BLOCK_SIZE_BATCHED = 32 +# Average tokens per expert above which the batched block size pays off. +SWORDFISH_MOE_BATCHED_THRESHOLD = 32 +# Average tokens per expert above which sorting the tokens once and running +# the tcgen05 prefill GEMM per expert segment beats the fused kernel. Pays +# one host sync for the segment bounds, acceptable in prefill. The serial +# per-expert launches underfill large parts, so the path only engages on +# small ones (Thor class); big parts stay on the fused 32-token blocks. +SWORDFISH_MOE_GROUPED_THRESHOLD = 128 +SWORDFISH_MOE_GROUPED_MAX_SMS = 40 +# Average tokens per expert above which the problem is compute-bound and +# dequantizing every expert once into a dense buffer, then running the +# stock bf16 triton fused-MoE kernels, beats the fused int paths. Small +# parts lack the dense rate to amortize the dequant (Thor loses at every M) +# so the tier needs a datacenter-class part. +SWORDFISH_MOE_DENSE_THRESHOLD = 192 +SWORDFISH_MOE_DENSE_MIN_SMS = 100 + + +@functools.cache +def _sm_count() -> int: + return torch.cuda.get_device_properties(0).multi_processor_count + + +class SwordfishExperts(mk.FusedMoEExpertsModular): + """Swordfish-based fused MoE expert implementation.""" + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int | None = None, + num_dispatchers: int | None = None, + ): + assert quant_config.use_int4_w4a16 or quant_config.use_int8_w8a16, "Supports only int4_w4a16 or int8_w8a16" + assert quant_config.w1_zp is None and quant_config.w2_zp is None, ( + "Swordfish MoE v1 supports only symmetric quantization (no zero points)" + ) + assert quant_config.w1_bias is None and quant_config.w2_bias is None, "Swordfish MoE v1 does not support bias" + self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit + self.gemm1_alpha = quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0 + self.gemm1_beta = quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0 + + super().__init__( + moe_config=moe_config, + quant_config=quant_config, + max_num_tokens=max_num_tokens, + num_dispatchers=num_dispatchers, + ) + + @staticmethod + def _supports_current_device() -> bool: + p = current_platform + if not (p.is_cuda() and p.has_device_capability(100)): + return False + # sm100 family only, datacenter 10.x and Thor 11.x. Consumer + # Blackwell (12.x) is a different SM and is untested. + capability = p.get_device_capability() + return capability is not None and capability.major in (10, 11) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return weight_key in [kInt4Static, kInt8Static] + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + # Gated activations only; moe_problem_size derives the intermediate + # size from the packed w1 assuming the 2N gate/up layout. Activation + # itself goes through the shared apply_moe_activation() callback. + return activation in [ + MoEActivation.SILU, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + MoEActivation.SWIGLUSTEP, + ] + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels or moe_parallel_config.use_fi_nvl_one_sided_kernels + ) + + @staticmethod + def _supports_shape(hidden_dim: int) -> bool: + return hidden_dim % 64 == 0 + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @property + def num_bits(self) -> int: + return 4 if self.quant_config.use_int4_w4a16 else 8 + + @property + def group_size(self) -> int: + block_shape = self.quant_config.block_shape + return block_shape[1] if block_shape is not None else -1 + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + # Swordfish packed weights are [E, N/64, K/64, words], with w1 holding + # the gate/up shards along N. + assert w1.dim() == 4 and w2.dim() == 4 + + E = w1.size(0) + K = a1.size(-1) + N = (w1.size(1) * SWORDFISH_BLOCK_N) // 2 + + assert a1.dim() == 2 + # Make sure we are using the correct a1 (pre-permute). + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" + M = a1.size(0) + + assert topk_ids.dim() == 2 + topk = topk_ids.size(1) + + return E, M, N, K, topk + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # swordfish_moe_mm allocates its own GEMM outputs, so the workspaces + # only back the final output buffer provisioned by the modular kernel. + workspace1 = (M * topk, K) + workspace2 = (M * topk, K) + output = (M, K) + return (workspace1, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert self.w1_scale is not None + assert self.w2_scale is not None + assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" + assert hidden_states.dtype in [torch.float16, torch.bfloat16] + assert topk_weights.dtype == torch.float32 + assert activation.is_gated + + E, M, N, K, topk = self.moe_problem_size(hidden_states, w1, w2, topk_ids) + + if global_num_experts == -1: + global_num_experts = E + if M * topk >= SWORDFISH_MOE_DENSE_THRESHOLD * E and _sm_count() >= SWORDFISH_MOE_DENSE_MIN_SMS: + self._apply_dense( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + activation, + M, + N, + K, + topk, + global_num_experts, + expert_map, + apply_router_weight_on_input, + ) + return + if ( + M * topk >= SWORDFISH_MOE_GROUPED_THRESHOLD * E + and _sm_count() <= SWORDFISH_MOE_GROUPED_MAX_SMS + and expert_map is None + and not apply_router_weight_on_input + and hidden_states.dtype == torch.bfloat16 + and self.num_bits == 4 + and K % 128 == 0 + and 2 * N % 128 == 0 + and N % 128 == 0 + ): + self._apply_grouped( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + activation, + M, + N, + K, + topk, + apply_router_weight_on_input, + ) + return + block_size = ( + SWORDFISH_MOE_BLOCK_SIZE_BATCHED + if M * topk >= SWORDFISH_MOE_BATCHED_THRESHOLD * E + else SWORDFISH_MOE_BLOCK_SIZE + ) + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, + block_size, + global_num_experts, + expert_map, + ignore_invalid_experts=True, + ) + + num_bits = self.num_bits + group_size = self.group_size + + intermediate_cache1 = ops.swordfish_moe_mm( + hidden_states, + w1, + self.w1_scale, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_weights if apply_router_weight_on_input else None, + block_size, + topk, + apply_router_weight_on_input, + num_bits, + group_size, + K, + 2 * N, + ) + + intermediate_cache2 = torch.empty( + (M * topk, N), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + self.activation( + activation, + intermediate_cache2, + intermediate_cache1.view(-1, 2 * N), + clamp_limit=self.gemm1_clamp_limit, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, + ) + + intermediate_cache3 = ops.swordfish_moe_mm( + intermediate_cache2, + w2, + self.w2_scale, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + None if apply_router_weight_on_input else topk_weights, + block_size, + 1, + not apply_router_weight_on_input, + num_bits, + group_size, + N, + K, + ) + + ops.moe_sum(intermediate_cache3.view(M, topk, K), output) + + def _apply_dense( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + ) -> None: + from aphrodite.model_executor.layers.fused_moe.fused_moe import ( + fused_experts, + ) + + group_size = self.group_size + num_bits = self.num_bits + w1_dense = ops.swordfish_dequant_dense(w1, self.w1_scale, None, num_bits, group_size, K, 2 * N, True) + w2_dense = ops.swordfish_dequant_dense(w2, self.w2_scale, None, num_bits, group_size, N, K, True) + out = fused_experts( + hidden_states, + w1_dense, + w2_dense, + topk_weights, + topk_ids, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=global_num_experts, + expert_map=expert_map, + ) + output.copy_(out) + + def _apply_grouped( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + M: int, + N: int, + K: int, + topk: int, + apply_router_weight_on_input: bool, + ) -> None: + E = w1.size(0) + group_size = self.group_size + flat = topk_ids.flatten().to(torch.int64) + order = torch.argsort(flat, stable=True) + bounds = torch.cumsum(torch.bincount(flat, minlength=E), 0).cpu().tolist() + a_sorted = hidden_states.index_select(0, order // topk).contiguous() + + assert self.w1_scale is not None and self.w2_scale is not None + cache1 = torch.empty((M * topk, 2 * N), dtype=hidden_states.dtype, device=hidden_states.device) + beg = 0 + for e, end in enumerate(bounds): + if end > beg: + cache1[beg:end] = ops.swordfish_prefill_mm( + a_sorted[beg:end], w1[e], self.w1_scale[e], group_size, K, 2 * N + ) + beg = end + cache2 = torch.empty((M * topk, N), dtype=hidden_states.dtype, device=hidden_states.device) + self.activation( + activation, + cache2, + cache1, + clamp_limit=self.gemm1_clamp_limit, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, + ) + out_sorted = torch.empty((M * topk, K), dtype=hidden_states.dtype, device=hidden_states.device) + beg = 0 + for e, end in enumerate(bounds): + if end > beg: + out_sorted[beg:end] = ops.swordfish_prefill_mm( + cache2[beg:end], w2[e], self.w2_scale[e], group_size, N, K + ) + beg = end + wt = topk_weights.flatten().index_select(0, order).unsqueeze(1) + out_sorted = out_sorted * wt.to(out_sorted.dtype) + out_flat = torch.zeros((M * topk, K), dtype=out_sorted.dtype, device=out_sorted.device) + out_flat.index_copy_(0, order, out_sorted) + ops.moe_sum(out_flat.view(M, topk, K), output) diff --git a/aphrodite/model_executor/layers/fused_moe/oracle/int_wna16.py b/aphrodite/model_executor/layers/fused_moe/oracle/int_wna16.py index a28d9555d8..cc57adc187 100644 --- a/aphrodite/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/aphrodite/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -24,6 +24,9 @@ MarlinExperts, MarlinExpertsBase, ) +from aphrodite.model_executor.layers.fused_moe.experts.swordfish_moe import ( + SwordfishExperts, +) from aphrodite.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( TrtLlmMxint4ExpertsMonolithic, ) @@ -45,6 +48,7 @@ class WNA16MoEBackend(Enum): + SWORDFISH = "SWORDFISH" MARLIN = "MARLIN" BATCHED_MARLIN = "BATCHED_MARLIN" HUMMING = "HUMMING" @@ -69,6 +73,8 @@ def backend_to_kernel_cls( HummingGroupedExperts, HummingIndexedExperts, ] + elif backend == WNA16MoEBackend.SWORDFISH: + return [SwordfishExperts] elif backend == WNA16MoEBackend.MARLIN: return [MarlinExperts] elif backend == WNA16MoEBackend.BATCHED_MARLIN: @@ -102,6 +108,7 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: _AVAILABLE_BACKENDS = [ WNA16MoEBackend.FLASHINFER_TRTLLM, + WNA16MoEBackend.SWORDFISH, WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, WNA16MoEBackend.HUMMING, @@ -112,6 +119,7 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: """Map user's MoEBackend to WNA16MoEBackend.""" mapping = { + "swordfish": WNA16MoEBackend.SWORDFISH, "marlin": WNA16MoEBackend.MARLIN, "humming": WNA16MoEBackend.HUMMING, "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM, @@ -253,11 +261,12 @@ def make_wna16_moe_kernel( ) # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts, - # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4, and the Humming - # grouped/indexed experts. + # BatchedMarlinExperts, SwordfishExperts, XPUExpertsWNA16, CPUExpertsInt4, + # and the Humming grouped/indexed experts. allowed_experts: tuple[type[mk.FusedMoEExperts], ...] = ( MarlinExperts, BatchedMarlinExperts, + SwordfishExperts, TrtLlmMxint4ExpertsMonolithic, XPUExpertsWNA16, CPUExpertsInt4, @@ -623,6 +632,96 @@ def _process_weights_marlin( ) +def _process_weights_swordfish( + input_dtype: torch.dtype | None, + num_bits: int, + pack_factor: int, + group_size: int, + actorder: str | None, + is_sym: bool, + w13_qweight: torch.Tensor, + w2_qweight: torch.Tensor, + w13_scales: torch.Tensor, + w2_scales: torch.Tensor, + w13_qzeros: torch.Tensor | None = None, + w2_qzeros: torch.Tensor | None = None, + w13_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, +) -> tuple[ + torch.Tensor, # w13_qweight + torch.Tensor, # w2_qweight + torch.Tensor, # w13_scales + torch.Tensor, # w2_scales + torch.Tensor, # w13_g_idx + torch.Tensor, # w2_g_idx + torch.Tensor, # w13_g_idx_sort_indices + torch.Tensor, # w2_g_idx_sort_indices + torch.Tensor | None, # w13_qzeros + torch.Tensor | None, # w2_qzeros + torch.Tensor | None, # w13_input_global_scale + torch.Tensor | None, # w2_input_global_scale + torch.Tensor | None, # w13_bias + torch.Tensor | None, # w2_bias +]: + """Swordfish weight post-processing shared by the SWORDFISH backend. + + Prepacks the GPTQ int32 [K / pack_factor, N] weights per expert into the + Swordfish block-linear layout. Scales stay in the plain checkpoint layout + [E, K / group_size, N], only made contiguous. + """ + # These limits are not visible to the oracle kernel predicates (the + # weight QuantKey does not carry group_size or act_order), so they are + # validated here with a pointer to the marlin fallback. + if not is_sym or w13_qzeros is not None or w2_qzeros is not None: + raise ValueError("Swordfish MoE v1 supports only symmetric quantization. Set moe_backend='marlin'.") + if actorder == "group": + raise ValueError("Swordfish MoE v1 does not support act_order. Set moe_backend='marlin'.") + if group_size not in (-1, 64, 128): + raise ValueError(f"Swordfish MoE v1 does not support group_size={group_size}. Set moe_backend='marlin'.") + if w13_bias is not None or w2_bias is not None: + raise ValueError("Swordfish MoE v1 does not support bias. Set moe_backend='marlin'.") + assert input_dtype is None or input_dtype.itemsize == 2, "Swordfish MoE runs on 16-bit activations" + + num_experts = w13_qweight.shape[0] + device = w13_qweight.device + + def prepack(qweight: torch.Tensor) -> torch.Tensor: + size_k = qweight.shape[1] * pack_factor + size_n = qweight.shape[2] + return torch.stack( + [ops.swordfish_prepack_B(qweight[e].contiguous(), size_k, size_n, num_bits) for e in range(num_experts)], + dim=0, + ).contiguous() + + swordfish_w13_qweight = prepack(w13_qweight) + swordfish_w2_qweight = prepack(w2_qweight) + swordfish_w13_scales = w13_scales.contiguous() + swordfish_w2_scales = w2_scales.contiguous() + + def empty_g_idx() -> torch.nn.Parameter: + return torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + + return ( + swordfish_w13_qweight, + swordfish_w2_qweight, + swordfish_w13_scales, + swordfish_w2_scales, + empty_g_idx(), + empty_g_idx(), + empty_g_idx(), + empty_g_idx(), + None, # w13_qzeros + None, # w2_qzeros + None, # w13_input_global_scale + None, # w2_input_global_scale + None, # w13_bias + None, # w2_bias + ) + + def _process_awq_weights_marlin( layer: torch.nn.Module, weight_bits: int, @@ -1025,6 +1124,45 @@ def convert_to_wna16_moe_kernel_format( convert_to_humming_moe_kernel_format(layer, quant_config=_humming_wna16_weight_schema(quant_config)) return None + if backend == WNA16MoEBackend.SWORDFISH: + from aphrodite.model_executor.layers.quantization.auto_gptq import ( + AutoGPTQConfig, + ) + + if isinstance(quant_config, AutoGPTQConfig): + num_bits = quant_config.quant_type.size_bits + pack_factor = quant_config.pack_factor + group_size = quant_config.group_size + actorder = "group" if quant_config.desc_act else None + is_sym = quant_config.is_sym + elif isinstance(quant_config, QuantizationArgs): + num_bits = quant_config.num_bits + pack_factor = 32 // quant_config.num_bits + group_size = quant_config.group_size + actorder = quant_config.actorder + is_sym = quant_config.symmetric + else: + raise TypeError( + "Swordfish WNA16 MoE backend requires AutoGPTQConfig or " + f"QuantizationArgs, got {type(quant_config).__name__}." + ) + return _process_weights_swordfish( + input_dtype, + num_bits, + pack_factor, + group_size, + actorder, + is_sym, + w13, + w2, + w13_scale, + w2_scale, + w13_qzeros, + w2_qzeros, + w13_bias, + w2_bias, + ) + if backend in ( WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, diff --git a/aphrodite/model_executor/layers/fused_moe/routed_experts.py b/aphrodite/model_executor/layers/fused_moe/routed_experts.py index 4a88634d56..98005c7765 100644 --- a/aphrodite/model_executor/layers/fused_moe/routed_experts.py +++ b/aphrodite/model_executor/layers/fused_moe/routed_experts.py @@ -844,7 +844,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterable[ matched = True weight_name = qual_name.replace(weight_name, param_name) param_name = weight_name.removeprefix(f"{self.layer_name}.") - param = getattr(self, param_name) + param = getattr(self, param_name, None) + if param is None: + # GPTQ exports write placeholder bias tensors for + # bias-free models; the quant method registers no such + # param, so the tensor has nowhere to land. + continue if is_fused: # w1 and w3 share one fused tensor; use a local copy so the # transpose below doesn't mutate loaded_weight across diff --git a/aphrodite/model_executor/layers/quantization/utils/allspark_utils.py b/aphrodite/model_executor/layers/quantization/utils/allspark_utils.py index 81b49e8414..b7e350e5b7 100644 --- a/aphrodite/model_executor/layers/quantization/utils/allspark_utils.py +++ b/aphrodite/model_executor/layers/quantization/utils/allspark_utils.py @@ -22,8 +22,8 @@ def check_allspark_supported_dtype_shape( capability_tuple = current_platform.get_device_capability() device_capability = -1 if capability_tuple is None else capability_tuple.to_int() - # For Ampere GPU - if device_capability >= 80 and device_capability < 90: + # The Ampere kernels run unchanged on later architectures. + if device_capability >= 80: if group_size != -1: return ( False, diff --git a/aphrodite/model_executor/layers/quantization/utils/swordfish_utils.py b/aphrodite/model_executor/layers/quantization/utils/swordfish_utils.py new file mode 100644 index 0000000000..fdfdceae0b --- /dev/null +++ b/aphrodite/model_executor/layers/quantization/utils/swordfish_utils.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Swordfish (Blackwell sm100/sm110 w4a16) support queries. + +ABI v1 (frozen): int4 weights x fp16/bf16 activations, group scales +{-1, 64, 128}, no act_order. Zero-point checkpoints (AWQ) fold +(8 - zp) * scale into a second scale-shaped tensor at load time. +""" + +import torch + +from aphrodite.scalar_type import ScalarType, scalar_types + +SWORDFISH_BLOCK_N = 64 +SWORDFISH_BLOCK_K = 64 + + +def query_swordfish_supported_quant_types(zero_points: bool) -> list[ScalarType]: + if zero_points: + return [scalar_types.uint4] + return [scalar_types.uint4b8, scalar_types.uint8b128] + + +def query_swordfish_supported_group_sizes(act_type: torch.dtype) -> list[int]: + if act_type not in (torch.float16, torch.bfloat16): + return [] + # Group 32 and channelwise run on the decode kernels at every M; the + # tcgen05 prefill covers groups 64 and 128. + return [-1, 32, 64, 128] + + +def check_swordfish_supports_shape(in_features: int, out_features: int) -> tuple[bool, str | None]: + if in_features % SWORDFISH_BLOCK_K != 0: + return ( + False, + f"in_features ({in_features}) not divisible by {SWORDFISH_BLOCK_K}", + ) + if out_features % SWORDFISH_BLOCK_N != 0: + return ( + False, + f"out_features ({out_features}) not divisible by {SWORDFISH_BLOCK_N}", + ) + return True, None diff --git a/aphrodite/model_executor/layers/quantization/utils/swordfish_utils_test.py b/aphrodite/model_executor/layers/quantization/utils/swordfish_utils_test.py new file mode 100644 index 0000000000..25312427d2 --- /dev/null +++ b/aphrodite/model_executor/layers/quantization/utils/swordfish_utils_test.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pure-torch reference implementation of the Swordfish packed-weight ABI v1. + +Test oracle for `ops.swordfish_prepack_B` (bit-exact contract) and for the +mm kernels' dequant reference. + +ABI v1 is Marlin's in-tile permutation (16x64 tiles, pack_idx nibble +interleave, reused verbatim from marlin_utils_test) re-tiled into +(NB, KB, 512) int32 block-linear order: + block (nb, kb) <- marlin rows [4*kb, 4*kb+4) x int32 cols [128*nb, 128*(nb+1)) +""" + +import torch + +from aphrodite.model_executor.layers.quantization.utils.marlin_utils_test import ( + get_weight_perm, + marlin_weights, +) +from aphrodite.scalar_type import ScalarType, scalar_types + +SWORDFISH_BLOCK_N = 64 +SWORDFISH_BLOCK_K = 64 +SWORDFISH_BLOCK_INT32 = 512 +SWORDFISH_SUPPORTED_QUANT_TYPES = [scalar_types.uint4b8, scalar_types.uint8b128] +SWORDFISH_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] + + +def swordfish_shape_ok(size_k: int, size_n: int) -> bool: + return size_k > 0 and size_n > 0 and size_k % SWORDFISH_BLOCK_K == 0 and size_n % SWORDFISH_BLOCK_N == 0 + + +def swordfish_pack_weights_ref(q_w: torch.Tensor, size_k: int, size_n: int, num_bits: int = 4) -> torch.Tensor: + """Pack unpacked int codes q_w [K, N] into the Swordfish ABI v1 tensor: + int32 [NB, KB, 512] at 4-bit or [NB, KB, 1024] at 8-bit.""" + assert swordfish_shape_ok(size_k, size_n), (size_k, size_n) + assert q_w.shape == (size_k, size_n) + + # Stage 1, Marlin permutation and pack, int32 [K/16, N*32/(32/bits)/16]. + tile_int32 = 128 * num_bits // 4 + perm = get_weight_perm(num_bits=num_bits) + marlin_flat = marlin_weights(q_w, size_k, size_n, num_bits=num_bits, perm=perm) + assert marlin_flat.shape == (size_k // 16, size_n * tile_int32 // 64) + + # Stage 2, block re-tile to (NB, KB, words). + kb = size_k // SWORDFISH_BLOCK_K + nb = size_n // SWORDFISH_BLOCK_N + # rows split as (KB, 4) k16-slices, cols as (NB, tile) int32 runs + x = marlin_flat.reshape(kb, 4, nb, tile_int32) + x = x.permute(2, 0, 1, 3) # (NB, KB, 4, tile) + return x.reshape(nb, kb, 4 * tile_int32).contiguous() + + +def swordfish_quantize( + w: torch.Tensor, + quant_type: ScalarType, + group_size: int, +): + """Quantize fp weight [K, N] GPTQ-style (no act_order) and pack to the + Swordfish ABI. Returns (w_ref, packed, scales).""" + from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + gptq_quantize_weights, + ) + + assert quant_type in SWORDFISH_SUPPORTED_QUANT_TYPES + size_k, size_n = w.shape + + w_ref, q_w, s, _, _ = gptq_quantize_weights(w, quant_type, group_size, act_order=False) + packed = swordfish_pack_weights_ref(q_w, size_k, size_n, quant_type.size_bits) + return w_ref, packed, s + + +def swordfish_quantize_act_order( + w: torch.Tensor, + quant_type: ScalarType, + group_size: int, +): + """Quantize fp weight [K, N] GPTQ-style WITH act_order and pack the + row-sorted weight. Returns (w_ref, packed, scales, sort_indices); the + caller permutes activation columns by sort_indices before the GEMM.""" + from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + gptq_quantize_weights, + sort_weights, + ) + + assert quant_type in SWORDFISH_SUPPORTED_QUANT_TYPES + size_k, size_n = w.shape + + w_ref, q_w, s, g_idx, _ = gptq_quantize_weights(w, quant_type, group_size, act_order=True) + q_w, g_idx, sort_indices = sort_weights(q_w, g_idx) + packed = swordfish_pack_weights_ref(q_w, size_k, size_n, quant_type.size_bits) + return w_ref, packed, s, sort_indices.to(torch.int) + + +def swordfish_quantize_awq( + w: torch.Tensor, + group_size: int, +): + """Quantize fp weight [K, N] AWQ-style (uint4 + group zero points) and + pack to the Swordfish ABI. Returns (w_ref, packed, scales, zps_neg) + where zps_neg holds the prescaled (8 - zp) * scale rows the kernel + consumes.""" + from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + quantize_weights, + ) + + size_k, size_n = w.shape + # The kernel applies the zero point after scaling, (w - 8) * s plus + # (8 - zp) * s, so the matching reference keeps tolerances tight. + w_ref, q_w, s, zp = quantize_weights( + w, + scalar_types.uint4, + group_size, + zero_points=True, + ref_zero_points_after_scales=True, + ) + packed = swordfish_pack_weights_ref(q_w, size_k, size_n) + zps_neg = ((8.0 - zp.to(s.dtype)) * s).contiguous() + return w_ref, packed, s, zps_neg diff --git a/aphrodite/utils/humming.py b/aphrodite/utils/humming.py index f8b0805558..10d23c3b1d 100644 --- a/aphrodite/utils/humming.py +++ b/aphrodite/utils/humming.py @@ -38,10 +38,36 @@ } +_sm110_compat_applied = False + + +def _apply_sm110_compat() -> None: + """Humming ships no sm110 support: its heuristics map has no 11.0 entry + and Jetson does not expose the NVML clock/bus-width queries its roofline + estimates rely on. The sm100 heuristics work unchanged on sm110, so map + them in and pin the two roofline numbers to Thor's.""" + global _sm110_compat_applied + if _sm110_compat_applied: + return + _sm110_compat_applied = True + import torch + + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (11, 0): + return + from humming.tune import heuristics_map + from humming.tune.sm100 import Sm100Heuristics + from humming.utils import device as humming_device + + heuristics_map.setdefault(110, Sm100Heuristics) + humming_device.calculate_gpu_bandwidth = lambda gpu_index=0: 273.0 + humming_device.estimate_tensorcore_max_tops = lambda gpu_index=0: 250 + + def __getattr__(name: str) -> Any: spec = _EXPORTS.get(name) if spec is None: raise AttributeError(f"module 'aphrodite.utils.humming' has no attribute {name!r}") + _apply_sm110_compat() if ":" in spec: mod_path, attr = spec.split(":", 1) obj = getattr(importlib.import_module(mod_path), attr) diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_abi.cuh b/csrc/libtorch_stable/quantization/swordfish/swordfish_abi.cuh new file mode 100644 index 0000000000..398593fe0e --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_abi.cuh @@ -0,0 +1,44 @@ +// Swordfish ABI v1 address math, shared by the prepack kernel and both +// mainloops. Single source of truth for the (NB, KB, tile) block-linear +// layout; the in-tile permutation is Marlin's. +#pragma once + +#include "swordfish_types.cuh" + +namespace swordfish { + +// Byte offset of packed block (nb, kb) in a (NB, KB, kBlockBytes) tensor. +__host__ __device__ inline constexpr int64_t block_byte_offset(int64_t nb, + int64_t kb, + int64_t num_kb) { + return (nb * num_kb + kb) * kBlockBytes; +} + +// Byte offset of the k16-slice sub-tile `t` (0..3) within a block. +__host__ __device__ inline constexpr int64_t subtile_byte_offset(int t) { + return int64_t(t) * kSubTileBytes; +} + +// Mapping from the Marlin flat repack layout to Swordfish blocks. +// +// gptq_marlin_repack emits int32[K/16][N*16/8] = int32[K/16][N*2]; row r is +// the k16-slice starting at k = 16*r; within a row, each Marlin 16x64 n-tile +// occupies 128 consecutive int32 (64 cols * 16 rows / 8 nibbles-per-int32). +// Swordfish block (nb, kb) gathers rows {4*kb .. 4*kb+3}, int32 columns +// [128*nb, 128*(nb+1)), i.e. int32 index +// marlin_idx(row, col) = row * (N*2) + col +// swordfish word w of block (nb,kb): t = w / 128, c = w % 128 +// -> marlin row = 4*kb + t, marlin col = 128*nb + c +// kTileInt32 = int32 words per marlin 16x64 tile: 128 at 4-bit, 256 at +// 8-bit (where the flat layout is int32[K/16][N*4]). +template +__host__ __device__ inline constexpr int64_t marlin_word_index(int64_t nb, + int64_t kb, + int w, + int64_t n) { + const int t = w / kTileInt32; + const int c = w % kTileInt32; + return (4 * kb + t) * (n * kTileInt32 / 64) + kTileInt32 * nb + c; +} + +} // namespace swordfish diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_decode.cuh b/csrc/libtorch_stable/quantization/swordfish/swordfish_decode.cuh new file mode 100644 index 0000000000..6ee8fc67c3 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_decode.cuh @@ -0,0 +1,873 @@ +// Swordfish decode GEMM. w4a16 mma.sync mainloop over the ABI v1 +// block-linear layout. +// +// One CTA (4 warps) covers M_TILES m16 tiles of one n64 block-column against +// a single weight stream. Warps split K into contiguous pair ranges (pair = +// 2 consecutive k16 slices = 1024 B of packed weights), which keeps each +// warp's group scales register-resident. Weights and the tiles' activation +// rows ride one cp.async pipeline, each lane copying exactly the bytes it +// later consumes into its own smem slot, so bytes in flight cost no +// registers and no warp synchronization. Each staged word is dequantized +// once and fans out to M_TILES mmas. +// +// The in-tile read contract is Marlin's, proven bit-exact by the prepack +// tests. Lane T consumes I4[T] at word 4T of each 512 B sub-tile; word j +// dequants into the fragments of n8-tiles 2j and 2j+1, where lane T owns +// column T/4 and rows 2*(T%4) + {0,1,8,9} of the k16xn8 fragment. +#pragma once + +#include +#include + +#include +#include + +#include "libtorch_stable/quantization/marlin/marlin.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/dequant.h" +#include "libtorch_stable/quantization/marlin/marlin_mma.h" + +#include "swordfish_abi.cuh" + +namespace swordfish { + +inline constexpr int kDecodeWarps = 4; +inline constexpr int kDecodeThreads = kDecodeWarps * 32; +// One k16-slice pair = two consecutive 512 B sub-tiles. +inline constexpr int kPairInt32 = 2 * (kSubTileBytes / 4); +// cp.async pipeline depth in pairs (stages in flight per warp). +inline constexpr int kStages = 5; + +// One ldmatrix.x4 loads a full m16k16 activation fragment from smem in +// tensor-core register order. +template +__device__ __forceinline__ void ldsm4( + typename marlin::MarlinScalarType::FragA& frag_a, + const void* smem_ptr) { + uint32_t* a = reinterpret_cast(&frag_a); + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(smem)); +} + +// cp.async.cg with an evict_first L2 hint. Packed weights are read once per +// launch and must not evict the re-read activation rows and scales. +__device__ __forceinline__ void cp_async4_evict_first(void* smem_ptr, + const void* glob_ptr, + uint64_t pol) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "cp.async.cg.shared.global.L2::cache_hint [%0], [%1], 16, %2;\n" ::"r"( + smem), + "l"(glob_ptr), "l"(pol)); +#else + marlin::cp_async4(smem_ptr, glob_ptr); +#endif +} + +__device__ __forceinline__ uint64_t l2_evict_first_policy() { + uint64_t pol = 0; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" + : "=l"(pol)); +#endif + return pol; +} + +// Packed-pair global atomic add. The atomicAdd(half2*) intrinsic compiles to +// a CAS loop on this toolchain, so emit red.global directly. +template +__device__ __forceinline__ void red_add2(scalar_t2* p, scalar_t2 v) { + if constexpr (sizeof(scalar_t2) == 4) { + if constexpr (std::is_same_v) { + asm volatile("red.global.add.noftz.bf16x2 [%0], %1;" ::"l"(p), + "r"(*reinterpret_cast(&v)) + : "memory"); + } else { + asm volatile("red.global.add.noftz.f16x2 [%0], %1;" ::"l"(p), + "r"(*reinterpret_cast(&v)) + : "memory"); + } + } +} + +// Grid-stride C zeroing for the atomic paths. A cudaMemsetAsync node runs +// on the copy engine, and in a captured decode graph every compute-to-copy +// transition costs an engine-switch gap -- at one memset per GEMM that idle +// time measured ~90 us per decoded token on B200. A trivial kernel stays on +// the compute engine. +template +__global__ void swordfish_zero_c_kernel(int4* __restrict__ c, int64_t vecs) { + const int64_t i = blockIdx.x * int64_t(blockDim.x) + threadIdx.x; + const int4 z = {0, 0, 0, 0}; + if (i < vecs) c[i] = z; +} + +template +inline void launch_zero_c(void* c, int64_t m, int64_t n, cudaStream_t stream) { + const int64_t vecs = m * n * int64_t(sizeof(scalar_t)) / sizeof(int4); + const int threads = 256; + const int blocks = int((vecs + threads - 1) / threads); + swordfish_zero_c_kernel + <<>>(reinterpret_cast(c), vecs); +} + +// Fused act_order prep for the decode window: writes the column-sorted +// activation copy and zeroes the output in one launch, replacing the +// separate permute_cols and zero nodes (each in-graph node and its +// engine/launch gap costs about a microsecond per GEMM at bs=1). Thread i +// of the permute range loads its perm entry coalesced and gathers one +// activation half; the tail range zeroes C in int4 stores. +template +__global__ void swordfish_prep_kernel(const scalar_t* __restrict__ a, + const int32_t* __restrict__ perm, + scalar_t* __restrict__ a_perm, + int4* __restrict__ c, int64_t m, + int64_t k, int64_t c_vecs) { + const int64_t idx = blockIdx.x * int64_t(blockDim.x) + threadIdx.x; + const int64_t perm_work = m * k; + if (idx < perm_work) { + const int64_t r = idx / k; + const int64_t i = idx - r * k; + a_perm[r * k + i] = a[r * k + perm[i]]; + } else if (idx - perm_work < c_vecs) { + const int4 z = {0, 0, 0, 0}; + c[idx - perm_work] = z; + } +} + +template +inline void launch_prep(const void* a, const int32_t* perm, void* a_perm, + void* c, int64_t m, int64_t k, int64_t n, + cudaStream_t stream) { + const int64_t c_vecs = m * n * int64_t(sizeof(scalar_t)) / sizeof(int4); + const int64_t work = m * k + c_vecs; + const int threads = 256; + const int blocks = int((work + threads - 1) / threads); + swordfish_prep_kernel<<>>( + reinterpret_cast(a), perm, + reinterpret_cast(a_perm), reinterpret_cast(c), m, k, + c_vecs); +} + +// ATOMIC_EPI replaces the cross-warp smem reduction with red.global adds +// into a zeroed C, freeing 16 KB smem per CTA and both __syncthreads, and +// making cross-CTA split-K free. Summation order is nondeterministic. The +// launcher uses it for the decode window (M <= 47). +template +__global__ void swordfish_decode_kernel( + const typename marlin::MarlinScalarType::scalar_t* __restrict__ A, + const int32_t* __restrict__ B, + const typename marlin::MarlinScalarType::scalar_t* __restrict__ S, + const typename marlin::MarlinScalarType::scalar_t* __restrict__ Z, + typename marlin::MarlinScalarType::scalar_t* __restrict__ C, int M, + int K, int N, int group_size) { + static_assert(M_TILES >= 1 && M_TILES <= 3, "supported m-tile fusion: 1-3"); + static_assert(M_TILES == 1 || ATOMIC_EPI, + "multi-m-tile is an atomic-epilogue configuration"); + // Pipeline depth scales inversely with M_TILES. Compute per staged pair + // grows with the tile count, so fewer stages hide the same copy latency, + // and the astage footprint stays inside the 48 KB static smem budget. + // 8-bit units are half the activation bytes, so the fused tiers run + // deeper pipelines in the same smem. + constexpr int kStagesT = + M_TILES == 1 + ? kStages + : (M_TILES == 2 ? (W8 ? 5 : 4) : (M_TILES == 3 ? 3 : (W8 ? 4 : 2))); + // A staging unit stays 1 KB of packed weights at both widths: a k32 slice + // pair at 4-bit, a single k16 slice at 8-bit (deeper W8 pipelines fit the + // freed smem but measured flat). + constexpr int kUnitK = W8 ? 16 : 32; + static_assert(!(W8 && HAS_ZP), "8-bit weights carry no zero points"); + using Dtype = marlin::MarlinScalarType; + using scalar_t = typename Dtype::scalar_t; + using scalar_t2 = typename Dtype::scalar_t2; + using FragA = typename Dtype::FragA; + using FragB = typename Dtype::FragB; + using FragC = typename Dtype::FragC; + + const int lane = threadIdx.x % 32; + const int warp = threadIdx.x / 32; + const int group_id = lane / 4; + const int tig = lane % 4; // thread-in-group + + const int nb = blockIdx.y; // n64 block-column + const int m_base = blockIdx.x * (16 * M_TILES); // first m16 tile + const int col_base = nb * kBlockN; + const int num_kb = K / kBlockK; + + // Output rows owned by this lane's C fragments, per m16 tile. + int row0[M_TILES]; + bool r0ok[M_TILES], r1ok[M_TILES]; +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + row0[t] = m_base + 16 * t + group_id; + r0ok[t] = row0[t] < M; + r1ok[t] = row0[t] + 8 < M; + } + + // acc[t][j][b]: tile t's n8-tile at columns col_base + 16*j + 8*b + [0, 8). + FragC acc[M_TILES][4][2]; +#pragma unroll + for (int t = 0; t < M_TILES; t++) +#pragma unroll + for (int j = 0; j < 4; j++) +#pragma unroll + for (int b = 0; b < 2; b++) +#pragma unroll + for (int i = 0; i < 4; i++) acc[t][j][b][i] = 0.0f; + + const scalar_t2 zero2 = Dtype::num2num2(Dtype::float2num(0.0f)); + + // At split-K 1 this CTA owns its C tile exclusively and zeroes it here, + // sparing the launcher a flat ~2 us cudaMemsetAsync. At split > 1 tiles + // are shared and the launcher memsets instead. + if constexpr (ATOMIC_EPI) { + if (gridDim.z == 1) { + const int rows = min(16 * M_TILES, M - m_base); + auto* c2 = reinterpret_cast(C); + for (int i = threadIdx.x; i < rows * 32; i += kDecodeThreads) { + c2[int64_t(m_base + (i >> 5)) * (N >> 1) + (col_base >> 1) + (i & 31)] = + zero2; + } + __syncthreads(); // tile zeroed before any warp's atomic epilogue + } + } + + // Global just-in-time A fragment load, used only by the deterministic + // path. Register order per mma contract, reg0/reg1 at k0, reg2/reg3 at + // k0+8. + auto load_a_global = [&](int k0, FragA& fa) { + const int ca = k0 + 2 * tig; + const int cb = ca + 8; + const auto* a_row0 = reinterpret_cast( + A + int64_t(r0ok[0] ? row0[0] : 0) * K); + const auto* a_row1 = reinterpret_cast( + A + int64_t(r1ok[0] ? row0[0] + 8 : 0) * K); + fa[0] = r0ok[0] ? a_row0[ca / 2] : zero2; + fa[1] = r1ok[0] ? a_row1[ca / 2] : zero2; + fa[2] = r0ok[0] ? a_row0[cb / 2] : zero2; + fa[3] = r1ok[0] ? a_row1[cb / 2] : zero2; + }; + // ldmatrix per-lane addressing. Lanes 0-15 cover rows 0-15 cols 0-7 and + // lanes 16-31 cols 8-15, yielding FragA's register order. Rows >= M hold + // unstaged garbage whose products only reach acc components the epilogue + // guards never store. + const int ldsm_row = lane & 15; + const int ldsm_col = (lane >> 4) << 3; + auto load_a = [&](const scalar_t(*sa)[kUnitK], int ks, FragA& fa) { + ldsm4(fa, &sa[ldsm_row][ldsm_col + ks]); + }; + + // Group scales, hoisted. The contiguous k-split makes each warp see each + // group exactly once. One 4-byte load per lane covers the 64-wide row and + // shfl broadcasts the pairs. Fetch and expand are split so the next + // group's row prefetches a full group early, hiding its latency under + // compute (on-demand loads were the dominant stall on K-heavy shapes). + // HAS_ZP runs the same machinery over Z, whose rows hold prescaled + // (8 - zp) * scale, turning the dequant scaling into an fma. + scalar_t2 s_reg[4][2]; // [j][b] broadcast pairs for this lane's column + scalar_t2 z_reg[4][2]; + auto fetch_row = [&](const scalar_t* base, int g) -> scalar_t2 { + return reinterpret_cast(base + int64_t(g) * N + + col_base)[lane]; + }; + auto expand_row = [&](scalar_t2 mine, scalar_t2(&dst)[4][2]) { + const uint32_t sel = group_id & 1; // half index within the shfl'd pair +#pragma unroll + for (int j = 0; j < 4; j++) { +#pragma unroll + for (int b = 0; b < 2; b++) { + // half index i = 16j + 8b + group_id -> lane i/2, element i%2 + const scalar_t2 v = + __shfl_sync(0xffffffffu, mine, 8 * j + 4 * b + (group_id >> 1)); + const uint32_t bits = reinterpret_cast(v); + const uint16_t h16 = sel ? uint16_t(bits >> 16) : uint16_t(bits); + dst[j][b] = Dtype::num2num2(reinterpret_cast(h16)); + } + } + }; + + // One k16 slice. Dequant the lane's word once, scale once, then fan the + // mma out across the CTA's m16 tiles. + auto process_slice = [&](const FragA(&fa)[M_TILES], const marlin::I4& bqa, + const marlin::I4& bqb) { +#pragma unroll + for (int j = 0; j < 4; j++) { + FragB frag_b0, frag_b1; + if constexpr (W8) { + // One int32 per n8 tile; n16 groups 2 and 3 sit in the second I4. + const marlin::I4& src = j < 2 ? bqa : bqb; + marlin::dequant( + src.elems[(2 * j) & 3], reinterpret_cast(&frag_b0)); + marlin::dequant( + src.elems[(2 * j + 1) & 3], reinterpret_cast(&frag_b1)); + } else { + const int b_quant_0 = bqa.elems[j]; + const int b_quant_1 = b_quant_0 >> 8; + marlin::dequant( + b_quant_0, reinterpret_cast(&frag_b0)); + marlin::dequant( + b_quant_1, reinterpret_cast(&frag_b1)); + } + + if constexpr (HAS_ZP) { + frag_b0[0] = __hfma2(frag_b0[0], s_reg[j][0], z_reg[j][0]); + frag_b0[1] = __hfma2(frag_b0[1], s_reg[j][0], z_reg[j][0]); + frag_b1[0] = __hfma2(frag_b1[0], s_reg[j][1], z_reg[j][1]); + frag_b1[1] = __hfma2(frag_b1[1], s_reg[j][1], z_reg[j][1]); + } else { + frag_b0[0] = __hmul2(frag_b0[0], s_reg[j][0]); + frag_b0[1] = __hmul2(frag_b0[1], s_reg[j][0]); + frag_b1[0] = __hmul2(frag_b1[0], s_reg[j][1]); + frag_b1[1] = __hmul2(frag_b1[1], s_reg[j][1]); + } + +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + marlin::mma(fa[t], frag_b0, + acc[t][j][0]); + marlin::mma(fa[t], frag_b1, + acc[t][j][1]); + } + } + }; + + // One pair = two consecutive k16 slices from one staged buffer. + auto process_pair = [&](int p, const scalar_t(*sa)[16][kUnitK], + const int4* buf) { + FragA fa0[M_TILES], fa1[M_TILES]; + if constexpr (ATOMIC_EPI) { +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + load_a(sa[t], 0, fa0[t]); + if constexpr (!W8) load_a(sa[t], 16, fa1[t]); + } + } else { + load_a_global(kUnitK * p, fa0[0]); + if constexpr (!W8) load_a_global(kUnitK * p + 16, fa1[0]); + } + if constexpr (W8) { + const marlin::I4 bq0 = + *reinterpret_cast(&buf[2 * lane]); + const marlin::I4 bq1 = + *reinterpret_cast(&buf[2 * lane + 1]); + process_slice(fa0, bq0, bq1); + } else { + const marlin::I4 bq0 = *reinterpret_cast(&buf[lane]); + const marlin::I4 bq1 = + *reinterpret_cast(&buf[32 + lane]); + process_slice(fa0, bq0, bq0); + process_slice(fa1, bq1, bq1); + } + }; + + // Weight staging, one self-slot buffer per warp and stage. + __shared__ int4 bstage[kDecodeWarps][kStagesT][2 * 32]; + // Activation slices for the in-flight pairs, 1 KB per tile and stage, + // copied in the same commit group as the weights so one wait covers both. + __shared__ scalar_t + astage[kDecodeWarps][ATOMIC_EPI ? kStagesT : 1][M_TILES][16][kUnitK]; + const int32_t* b_col = + B + int64_t(nb) * num_kb * (W8 ? kBlockInt32_8 : kBlockInt32); + const int num_pairs = K / kUnitK; + // Cross-CTA split-K. blockIdx.z slices the pair space and the atomic + // epilogue merges slices in the zeroed C. gridDim.z > 1 only on the + // ATOMIC_EPI path. + const int slice_pairs = (num_pairs + gridDim.z - 1) / gridDim.z; + const int s_beg = blockIdx.z * slice_pairs; + const int s_end = min(s_beg + slice_pairs, num_pairs); + const int pairs_per_warp = (s_end - s_beg + kDecodeWarps - 1) / kDecodeWarps; + const int p_beg = s_beg + warp * pairs_per_warp; + const int p_end = min(p_beg + pairs_per_warp, s_end); + + // Scale-group bookkeeping: `left` = pairs left in the current group + // (host guarantees group_size % 32 == 0, so pairs never straddle groups). + const int ppg = group_size > 0 ? group_size / kUnitK : INT_MAX; + int g = 0; + int left = INT_MAX; + + // One commit group per stage. Fences are unconditional so the wait + // accounting stays uniform through the tail (empty groups are legal). + const int g_last = + group_size > 0 && p_beg < p_end ? (kUnitK * (p_end - 1)) / group_size : 0; + scalar_t2 s_next = zero2; // prefetched next-group row (lane's slice) + scalar_t2 z_next = zero2; + if (p_beg < p_end) { + if (group_size > 0) { + g = p_beg / ppg; + left = ppg - p_beg % ppg; + } + expand_row(fetch_row(S, g), s_reg); + if constexpr (HAS_ZP) expand_row(fetch_row(Z, g), z_reg); + if (g + 1 <= g_last) { + s_next = fetch_row(S, g + 1); + if constexpr (HAS_ZP) z_next = fetch_row(Z, g + 1); + } + } + + // Incremental int32 issue cursors. A modulo-indexed ring compiles to + // magic-number division and per-pair int64 address math costs multiplies, + // both measurable in this loop. Offsets fit int32 for every shape the + // host validates. + const uint64_t bpol = l2_evict_first_policy(); + const int32_t* ipair_ptr = b_col + p_beg * kPairInt32 + (W8 ? 8 : 4) * lane; + const int ia_row = lane >> 1; + const int ia_c0 = (kUnitK / 2) * (lane & 1); + // Per-tile activation cursors. Only valid rows are staged, since clamping + // out-of-range rows to row 0 multiplies activation traffic at small M. + bool ia_okt[M_TILES]; + const scalar_t* ia_ptr[M_TILES]; +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + const int r = m_base + 16 * t + ia_row; + ia_okt[t] = r < M; + ia_ptr[t] = A + (ia_okt[t] ? r : 0) * K + kUnitK * p_beg + ia_c0; + } + int ipend = p_end - p_beg; // pairs left to issue + + auto issue_pair = [&](int slot) { + if (ipend > 0) { + ipend--; + if constexpr (W8) { + cp_async4_evict_first(&bstage[warp][slot][2 * lane], ipair_ptr, bpol); + cp_async4_evict_first(&bstage[warp][slot][2 * lane + 1], ipair_ptr + 4, + bpol); + } else { + cp_async4_evict_first(&bstage[warp][slot][lane], ipair_ptr, bpol); + cp_async4_evict_first(&bstage[warp][slot][32 + lane], + ipair_ptr + kPairInt32 / 2, bpol); + } + ipair_ptr += kPairInt32; + if constexpr (ATOMIC_EPI) { +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + if (ia_okt[t]) { + marlin::cp_async4(&astage[warp][slot][t][ia_row][ia_c0], ia_ptr[t]); + if constexpr (!W8) { + marlin::cp_async4(&astage[warp][slot][t][ia_row][ia_c0 + 8], + ia_ptr[t] + 8); + } + } + ia_ptr[t] += kUnitK; + } + } + } + marlin::cp_async_fence(); + }; + +#pragma unroll + for (int s = 0; s < kStagesT - 1; s++) issue_pair(s); + + int slot = 0; + int islot = kStagesT - 1; + for (int p = p_beg; p < p_end; p++) { + issue_pair(islot); + if (++islot == kStagesT) islot = 0; + marlin::cp_async_wait(); // oldest stage (slot) complete + if (left == 0) { + ++g; + expand_row(s_next, s_reg); // value already in flight since boundary + if constexpr (HAS_ZP) expand_row(z_next, z_reg); + if (g + 1 <= g_last) { + s_next = fetch_row(S, g + 1); + if constexpr (HAS_ZP) z_next = fetch_row(Z, g + 1); + } + left = ppg; + } + process_pair(p, astage[warp][ATOMIC_EPI ? slot : 0], bstage[warp][slot]); + left--; + if (++slot == kStagesT) slot = 0; + } + + if constexpr (ATOMIC_EPI) { + // Direct atomic epilogue: every warp adds its partial fragments to C. +#pragma unroll + for (int t = 0; t < M_TILES; t++) { +#pragma unroll + for (int j = 0; j < 4; j++) { +#pragma unroll + for (int b = 0; b < 2; b++) { + const int col = col_base + 8 * (2 * j + b) + 2 * tig; + const float4 v = *reinterpret_cast(&acc[t][j][b]); + if (r0ok[t]) + red_add2( + reinterpret_cast(C + int64_t(row0[t]) * N + col), + Dtype::nums2num2(Dtype::float2num(v.x), Dtype::float2num(v.y))); + if (r1ok[t]) + red_add2( + reinterpret_cast(C + int64_t(row0[t] + 8) * N + + col), + Dtype::nums2num2(Dtype::float2num(v.z), Dtype::float2num(v.w))); + } + } + } + return; + } + + // Deterministic epilogue. Partials meet in smem and warp w reduces and + // writes n8-tiles {2w, 2w+1}. Tile index is 2*j + b. + __shared__ float4 red[kDecodeWarps][8][32]; +#pragma unroll + for (int j = 0; j < 4; j++) +#pragma unroll + for (int b = 0; b < 2; b++) + red[warp][2 * j + b][lane] = *reinterpret_cast(&acc[0][j][b]); + __syncthreads(); + +#pragma unroll + for (int i = 0; i < 2; i++) { + const int tile = 2 * warp + i; + float4 sum = red[0][tile][lane]; +#pragma unroll + for (int w = 1; w < kDecodeWarps; w++) { + const float4 v = red[w][tile][lane]; + sum.x += v.x; + sum.y += v.y; + sum.z += v.z; + sum.w += v.w; + } + const int col = col_base + 8 * tile + 2 * tig; + if (r0ok[0]) { + *reinterpret_cast(C + int64_t(row0[0]) * N + col) = + Dtype::nums2num2(Dtype::float2num(sum.x), Dtype::float2num(sum.y)); + } + if (r1ok[0]) { + *reinterpret_cast(C + int64_t(row0[0] + 8) * N + col) = + Dtype::nums2num2(Dtype::float2num(sum.z), Dtype::float2num(sum.w)); + } + } +} + +// Stream-K decode for the fused window (M 17 to 96). Persistent CTAs; one +// unit of flat work is one k16-slice pair of one (m-group, n64-column) +// tile, and each warp owns a contiguous range of units. The m-group index +// varies fastest after the pair so consecutive segments in a warp's range +// revisit the same weight column while it is L2-resident. Accumulators +// flush through red.global at segment boundaries into a launcher-zeroed C, +// so a tile's K slices may come from any number of warps with no locks and +// no fixup pass. This removes split-K heuristics and wave quantization for +// the window entirely. +// CTA_QUAD maps claims at CTA granularity over n256 column quads: the four +// warps take the four adjacent n64 columns of one quad over a common +// (m-group, k-range) claim. Per-warp staging, pipeline, and epilogue are +// untouched -- the quad buys 4x-amortized claim bookkeeping and adjacent +// column addressing (the four B streams walk consecutive column blocks and +// all four warps touch the same A slices while they are L2-resident). +template +__global__ void swordfish_decode_streamk_kernel( + const typename marlin::MarlinScalarType::scalar_t* __restrict__ A, + const int32_t* __restrict__ B, + const typename marlin::MarlinScalarType::scalar_t* __restrict__ S, + const typename marlin::MarlinScalarType::scalar_t* __restrict__ Z, + typename marlin::MarlinScalarType::scalar_t* __restrict__ C, int M, + int K, int N, int group_size, int m_groups) { + // 8-bit units are half the activation bytes, so the fused tiers run + // deeper pipelines in the same smem. + constexpr int kStagesT = + M_TILES == 1 + ? kStages + : (M_TILES == 2 ? (W8 ? 5 : 4) : (M_TILES == 3 ? 3 : (W8 ? 4 : 2))); + // A staging unit stays 1 KB of packed weights at both widths: a k32 slice + // pair at 4-bit, a single k16 slice at 8-bit (deeper W8 pipelines fit the + // freed smem but measured flat). + constexpr int kUnitK = W8 ? 16 : 32; + static_assert(!(W8 && HAS_ZP), "8-bit weights carry no zero points"); + using Dtype = marlin::MarlinScalarType; + using scalar_t = typename Dtype::scalar_t; + using scalar_t2 = typename Dtype::scalar_t2; + using FragA = typename Dtype::FragA; + using FragB = typename Dtype::FragB; + using FragC = typename Dtype::FragC; + + const int lane = threadIdx.x % 32; + const int warp = threadIdx.x / 32; + const int group_id = lane / 4; + const int tig = lane % 4; + const int ldsm_row = lane & 15; + const int ldsm_col = (lane >> 4) << 3; + const int ia_row = lane >> 1; + const int ia_c0 = (kUnitK / 2) * (lane & 1); + + const int num_pairs = K / kUnitK; + const int nb_cnt = N / (CTA_QUAD ? kBlockN * kDecodeWarps : kBlockN); + const int num_kb = K / kBlockK; + const int ppg = group_size > 0 ? group_size / kUnitK : INT_MAX; + const scalar_t2 zero2 = Dtype::num2num2(Dtype::float2num(0.0f)); + const uint64_t bpol = l2_evict_first_policy(); + + const int64_t total = int64_t(m_groups) * nb_cnt * num_pairs; + const int total_claimers = CTA_QUAD ? gridDim.x : gridDim.x * kDecodeWarps; + const int64_t per = (total + total_claimers - 1) / total_claimers; + int64_t w = + int64_t(CTA_QUAD ? blockIdx.x : blockIdx.x * kDecodeWarps + warp) * per; + const int64_t w_end = w + per < total ? w + per : total; + + __shared__ int4 bstage[kDecodeWarps][kStagesT][2 * 32]; + __shared__ scalar_t astage[kDecodeWarps][kStagesT][M_TILES][16][kUnitK]; + + FragC acc[M_TILES][4][2]; + scalar_t2 s_reg[4][2]; + scalar_t2 z_reg[4][2]; + + while (w < w_end) { + const int64_t cg = w / num_pairs; + const int p_beg = int(w - cg * num_pairs); + const int col_claim = int(cg / m_groups); + const int col = CTA_QUAD ? col_claim * kDecodeWarps + warp : col_claim; + const int g_idx = int(cg - int64_t(col_claim) * m_groups); + const int p_end = + int(int64_t(num_pairs) < p_beg + (w_end - w) ? int64_t(num_pairs) + : p_beg + (w_end - w)); + w += p_end - p_beg; + + const int m_base = g_idx * (16 * M_TILES); + const int col_base = col * kBlockN; + + int row0[M_TILES]; + bool r0ok[M_TILES], r1ok[M_TILES]; +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + row0[t] = m_base + 16 * t + group_id; + r0ok[t] = row0[t] < M; + r1ok[t] = row0[t] + 8 < M; + } +#pragma unroll + for (int t = 0; t < M_TILES; t++) +#pragma unroll + for (int j = 0; j < 4; j++) +#pragma unroll + for (int b = 0; b < 2; b++) +#pragma unroll + for (int i = 0; i < 4; i++) acc[t][j][b][i] = 0.0f; + + auto fetch_row = [&](const scalar_t* base, int g) -> scalar_t2 { + return reinterpret_cast(base + int64_t(g) * N + + col_base)[lane]; + }; + auto expand_row = [&](scalar_t2 mine, scalar_t2(&dst)[4][2]) { + const uint32_t sel = group_id & 1; +#pragma unroll + for (int j = 0; j < 4; j++) { +#pragma unroll + for (int b = 0; b < 2; b++) { + const scalar_t2 v = + __shfl_sync(0xffffffffu, mine, 8 * j + 4 * b + (group_id >> 1)); + const uint32_t bits = reinterpret_cast(v); + const uint16_t h16 = sel ? uint16_t(bits >> 16) : uint16_t(bits); + dst[j][b] = Dtype::num2num2(reinterpret_cast(h16)); + } + } + }; + auto load_a = [&](const scalar_t(*sa)[kUnitK], int ks, FragA& fa) { + ldsm4(fa, &sa[ldsm_row][ldsm_col + ks]); + }; + auto process_slice = [&](const FragA(&fa)[M_TILES], const marlin::I4& bqa, + const marlin::I4& bqb) { +#pragma unroll + for (int j = 0; j < 4; j++) { + FragB frag_b0, frag_b1; + if constexpr (W8) { + const marlin::I4& src = j < 2 ? bqa : bqb; + marlin::dequant( + src.elems[(2 * j) & 3], reinterpret_cast(&frag_b0)); + marlin::dequant( + src.elems[(2 * j + 1) & 3], + reinterpret_cast(&frag_b1)); + } else { + const int b_quant_0 = bqa.elems[j]; + const int b_quant_1 = b_quant_0 >> 8; + marlin::dequant( + b_quant_0, reinterpret_cast(&frag_b0)); + marlin::dequant( + b_quant_1, reinterpret_cast(&frag_b1)); + } + if constexpr (HAS_ZP) { + frag_b0[0] = __hfma2(frag_b0[0], s_reg[j][0], z_reg[j][0]); + frag_b0[1] = __hfma2(frag_b0[1], s_reg[j][0], z_reg[j][0]); + frag_b1[0] = __hfma2(frag_b1[0], s_reg[j][1], z_reg[j][1]); + frag_b1[1] = __hfma2(frag_b1[1], s_reg[j][1], z_reg[j][1]); + } else { + frag_b0[0] = __hmul2(frag_b0[0], s_reg[j][0]); + frag_b0[1] = __hmul2(frag_b0[1], s_reg[j][0]); + frag_b1[0] = __hmul2(frag_b1[0], s_reg[j][1]); + frag_b1[1] = __hmul2(frag_b1[1], s_reg[j][1]); + } +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + marlin::mma(fa[t], frag_b0, acc[t][j][0]); + marlin::mma(fa[t], frag_b1, acc[t][j][1]); + } + } + }; + auto process_pair = [&](int aslot, const int4* buf) { + FragA fa0[M_TILES], fa1[M_TILES]; +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + load_a(astage[warp][aslot][t], 0, fa0[t]); + if constexpr (!W8) load_a(astage[warp][aslot][t], 16, fa1[t]); + } + if constexpr (W8) { + const marlin::I4 bq0 = + *reinterpret_cast(&buf[2 * lane]); + const marlin::I4 bq1 = + *reinterpret_cast(&buf[2 * lane + 1]); + process_slice(fa0, bq0, bq1); + } else { + const marlin::I4 bq0 = *reinterpret_cast(&buf[lane]); + const marlin::I4 bq1 = + *reinterpret_cast(&buf[32 + lane]); + process_slice(fa0, bq0, bq0); + process_slice(fa1, bq1, bq1); + } + }; + + const int32_t* ipair_ptr = + B + int64_t(col) * num_kb * (W8 ? kBlockInt32_8 : kBlockInt32) + + p_beg * kPairInt32 + (W8 ? 8 : 4) * lane; + bool ia_okt[M_TILES]; + const scalar_t* ia_ptr[M_TILES]; +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + const int r = m_base + 16 * t + ia_row; + ia_okt[t] = r < M; + ia_ptr[t] = A + (ia_okt[t] ? r : 0) * K + kUnitK * p_beg + ia_c0; + } + int ipend = p_end - p_beg; + + auto issue_pair = [&](int slot) { + if (ipend > 0) { + ipend--; + if constexpr (W8) { + cp_async4_evict_first(&bstage[warp][slot][2 * lane], ipair_ptr, bpol); + cp_async4_evict_first(&bstage[warp][slot][2 * lane + 1], + ipair_ptr + 4, bpol); + } else { + cp_async4_evict_first(&bstage[warp][slot][lane], ipair_ptr, bpol); + cp_async4_evict_first(&bstage[warp][slot][32 + lane], + ipair_ptr + kPairInt32 / 2, bpol); + } + ipair_ptr += kPairInt32; +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + if (ia_okt[t]) { + marlin::cp_async4(&astage[warp][slot][t][ia_row][ia_c0], ia_ptr[t]); + if constexpr (!W8) { + marlin::cp_async4(&astage[warp][slot][t][ia_row][ia_c0 + 8], + ia_ptr[t] + 8); + } + } + ia_ptr[t] += kUnitK; + } + } + marlin::cp_async_fence(); + }; + + int g = 0; + int left = INT_MAX; + const int g_last = group_size > 0 && p_beg < p_end + ? (kUnitK * (p_end - 1)) / group_size + : 0; + scalar_t2 s_next = zero2; + scalar_t2 z_next = zero2; + if (p_beg < p_end) { + if (group_size > 0) { + g = p_beg / ppg; + left = ppg - p_beg % ppg; + } + expand_row(fetch_row(S, g), s_reg); + if constexpr (HAS_ZP) expand_row(fetch_row(Z, g), z_reg); + if (g + 1 <= g_last) { + s_next = fetch_row(S, g + 1); + if constexpr (HAS_ZP) z_next = fetch_row(Z, g + 1); + } + } + + // The two-unit W8 step issues two copies per iteration, so its + // prologue holds one slot fewer or the second issue of the first + // iteration would overwrite the oldest unprocessed stage. + constexpr int kPrologue = W8 ? kStagesT - 2 : kStagesT - 1; +#pragma unroll + for (int st = 0; st < kPrologue; st++) issue_pair(st); + + // At two stages the historical wait is wait<0>, which + // drains the copy issued THIS iteration and serializes the pipeline; + // one group must stay in flight. + constexpr int kWaitN = kStagesT == 2 ? 1 : kStagesT - 2; + int slot = 0; + int islot = kPrologue; + auto step = [&](int p) { + if (left == 0) { + ++g; + expand_row(s_next, s_reg); + if constexpr (HAS_ZP) expand_row(z_next, z_reg); + if (g + 1 <= g_last) { + s_next = fetch_row(S, g + 1); + if constexpr (HAS_ZP) z_next = fetch_row(Z, g + 1); + } + left = ppg; + } + process_pair(slot, bstage[warp][slot]); + left--; + if (++slot == kStagesT) slot = 0; + }; + if constexpr (W8) { + // A single k16 unit leaves one thin dequant chain exposed to the + // staging-buffer load latency (short-scoreboard dominates profiles); + // stepping two units per wait restores k32 of independent work. + // Both stepped slots must be complete, so the two-step wait leaves + // kStagesT - 2 groups in flight (zero at two stages). + constexpr int kWaitN2 = kStagesT - 2; + int p = p_beg; + for (; p + 1 < p_end; p += 2) { + issue_pair(islot); + if (++islot == kStagesT) islot = 0; + issue_pair(islot); + if (++islot == kStagesT) islot = 0; + marlin::cp_async_wait(); + step(p); + step(p + 1); + } + for (; p < p_end; p++) { + issue_pair(islot); + if (++islot == kStagesT) islot = 0; + marlin::cp_async_wait(); + step(p); + } + } else { + for (int p = p_beg; p < p_end; p++) { + issue_pair(islot); + if (++islot == kStagesT) islot = 0; + marlin::cp_async_wait(); + step(p); + } + } + + // Segment flush into the launcher-zeroed C. +#pragma unroll + for (int t = 0; t < M_TILES; t++) { +#pragma unroll + for (int j = 0; j < 4; j++) { +#pragma unroll + for (int b = 0; b < 2; b++) { + const int cc = col_base + 8 * (2 * j + b) + 2 * tig; + const float4 v = *reinterpret_cast(&acc[t][j][b]); + if (r0ok[t]) + red_add2( + reinterpret_cast(C + int64_t(row0[t]) * N + cc), + Dtype::nums2num2(Dtype::float2num(v.x), Dtype::float2num(v.y))); + if (r1ok[t]) + red_add2( + reinterpret_cast(C + int64_t(row0[t] + 8) * N + cc), + Dtype::nums2num2(Dtype::float2num(v.z), Dtype::float2num(v.w))); + } + } + } + } +} + +} // namespace swordfish diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_dense_tier.cu b/csrc/libtorch_stable/quantization/swordfish/swordfish_dense_tier.cu new file mode 100644 index 0000000000..f9d7fa9b28 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_dense_tier.cu @@ -0,0 +1,309 @@ +// Dense-GEMM tier for very large M. Past a few thousand rows the problem is +// compute-bound and Blackwell's dense fp16/bf16 rate through cuBLAS beats +// any fused mixed-input mainloop, so the weight dequantizes once into a +// transient dense buffer and cuBLAS takes the GEMM. The dequant cost is a +// few weight-reads and amortizes to nothing at this scale. + +#include + +#include + +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" +#include "swordfish_decode.cuh" + +#define SWORDFISH_CHECK_CUBLAS(cmd) \ + do { \ + cublasStatus_t e = cmd; \ + STD_TORCH_CHECK(e == CUBLAS_STATUS_SUCCESS, \ + "swordfish dense tier cuBLAS error ", int(e)); \ + } while (0) + +namespace swordfish { + +namespace { + +// One warp per marlin 16x64 sub-tile, staged through smem so the dense +// stores coalesce. The fragment contract matches the decode kernels: word j +// of lane T covers columns {16j + T/4, 16j + 8 + T/4} at sub-tile rows +// 2(T%4) + {0, 1, 8, 9}. +template +__global__ void swordfish_dequant_dense_kernel( + const int32_t* __restrict__ B, + const typename marlin::MarlinScalarType::scalar_t* __restrict__ S, + const typename marlin::MarlinScalarType::scalar_t* __restrict__ Z, + const int32_t* __restrict__ perm, + typename marlin::MarlinScalarType::scalar_t* __restrict__ W, int K, + int N, int group_size) { + using Dtype = marlin::MarlinScalarType; + using scalar_t = typename Dtype::scalar_t; + using scalar_t2 = typename Dtype::scalar_t2; + using FragB = typename Dtype::FragB; + + const int lane = threadIdx.x % 32; + const int warp = threadIdx.x / 32; // sub-tile within the block + const int c = lane >> 2; + const int t = lane & 3; + const int nb = blockIdx.x; + const int kb = blockIdx.y; + const int num_kb = K / kBlockK; + const int num_nb = N / kBlockN; + constexpr int64_t kBlockW = W8 ? kBlockInt32_8 : kBlockInt32; + + // blockIdx.z indexes the expert for stacked MoE weights. + const int64_t expert = blockIdx.z; + B += expert * num_nb * num_kb * kBlockW; + const int num_groups = group_size > 0 ? K / group_size : 1; + S += expert * int64_t(num_groups) * N; + if constexpr (HAS_ZP) Z += expert * int64_t(num_groups) * N; + W += expert * int64_t(K) * N; + + // Rows padded 64 -> 72: a 128 B row stride puts every same-column store + // across rows on one shared-memory bank; 144 B keeps the 16-byte + // alignment the vector reads need while spreading rows over the banks. + __shared__ scalar_t tile[4][16][72]; + + const int4* buf = reinterpret_cast( + B + (int64_t(nb) * num_kb + kb) * kBlockW + + warp * (W8 ? 2 * kPairInt32 / 2 : kPairInt32 / 2)); + + // One scale group covers the whole 16-row sub-tile for every supported + // group size. + const int g = group_size > 0 ? (kb * kBlockK + 16 * warp) / group_size : 0; + const scalar_t* srow = S + int64_t(g) * N + nb * kBlockN; + const scalar_t* zrow = HAS_ZP ? Z + int64_t(g) * N + nb * kBlockN : nullptr; + + const marlin::I4 bq0 = + *reinterpret_cast(&buf[W8 ? 2 * lane : lane]); + marlin::I4 bq1; + if constexpr (W8) { + bq1 = *reinterpret_cast(&buf[2 * lane + 1]); + } + + auto put = [&](int row, int col, scalar_t2 v2, bool hi) { + tile[warp][row][col] = reinterpret_cast(&v2)[hi ? 1 : 0]; + }; + +#pragma unroll + for (int j = 0; j < 4; j++) { + FragB frag_b0, frag_b1; + if constexpr (W8) { + const marlin::I4& src = j < 2 ? bq0 : bq1; + marlin::dequant( + src.elems[(2 * j) & 3], reinterpret_cast(&frag_b0)); + marlin::dequant( + src.elems[(2 * j + 1) & 3], reinterpret_cast(&frag_b1)); + } else { + const int q = bq0.elems[j]; + marlin::dequant( + q, reinterpret_cast(&frag_b0)); + marlin::dequant( + q >> 8, reinterpret_cast(&frag_b1)); + } + const int n0 = 16 * j + c; + const int n1 = n0 + 8; + const scalar_t2 s0 = Dtype::num2num2(srow[n0]); + const scalar_t2 s1 = Dtype::num2num2(srow[n1]); + if constexpr (HAS_ZP) { + const scalar_t2 z0 = Dtype::num2num2(zrow[n0]); + const scalar_t2 z1 = Dtype::num2num2(zrow[n1]); + frag_b0[0] = __hfma2(frag_b0[0], s0, z0); + frag_b0[1] = __hfma2(frag_b0[1], s0, z0); + frag_b1[0] = __hfma2(frag_b1[0], s1, z1); + frag_b1[1] = __hfma2(frag_b1[1], s1, z1); + } else { + frag_b0[0] = __hmul2(frag_b0[0], s0); + frag_b0[1] = __hmul2(frag_b0[1], s0); + frag_b1[0] = __hmul2(frag_b1[0], s1); + frag_b1[1] = __hmul2(frag_b1[1], s1); + } + // frag[0] holds rows {2t, 2t+1}, frag[1] rows {2t+8, 2t+9}. + put(2 * t, n0, frag_b0[0], false); + put(2 * t + 1, n0, frag_b0[0], true); + put(2 * t + 8, n0, frag_b0[1], false); + put(2 * t + 9, n0, frag_b0[1], true); + put(2 * t, n1, frag_b1[0], false); + put(2 * t + 1, n1, frag_b1[0], true); + put(2 * t + 8, n1, frag_b1[1], false); + put(2 * t + 9, n1, frag_b1[1], true); + } + __syncwarp(); + + const int k_base = kb * kBlockK + 16 * warp; + if constexpr (TRANSPOSED) { + // [N, K] output for consumers that want out-major weights (the triton + // fused-MoE kernels). 16 k-elements per n form two 16-byte runs. + for (int nn = lane; nn < 64; nn += 32) { + scalar_t run[16]; +#pragma unroll + for (int r = 0; r < 16; r++) run[r] = tile[warp][r][nn]; + int4* dst = + reinterpret_cast(&W[int64_t(nb * kBlockN + nn) * K + k_base]); + dst[0] = *reinterpret_cast(&run[0]); + dst[1] = *reinterpret_cast(&run[8]); + } + } else { + // Coalesced stores: 16-byte runs per lane over the sub-tile's rows. + // Under act_order the packed rows are group-sorted, so row r scatters + // to its original position and the activations stay unpermuted. + const int c0 = 8 * (lane % 8); +#pragma unroll + for (int r = lane / 8; r < 16; r += 4) { + const int64_t k_dst = perm != nullptr ? perm[k_base + r] : k_base + r; + *reinterpret_cast(&W[k_dst * N + nb * kBlockN + c0]) = + *reinterpret_cast(&tile[warp][r][c0]); + } + } +} + +} // namespace + +void swordfish_dense_tier_mm(const void* a, const int32_t* b, const void* s, + const void* z, const int32_t* perm, void* c, + void* w_dense, bool is_half, bool w8, bool has_zp, + int m, int k, int n, int group_size, + cudaStream_t stream) { + dim3 grid(n / kBlockN, k / kBlockK, 1); + const auto run = [&](auto tid, auto w8c, auto zpc) { + constexpr aphrodite::ScalarTypeId kTid = decltype(tid)::value; + using scalar_t = typename marlin::MarlinScalarType::scalar_t; + swordfish_dequant_dense_kernel + <<>>(b, reinterpret_cast(s), + reinterpret_cast(z), perm, + reinterpret_cast(w_dense), k, n, + group_size); + }; + using kF16 = + std::integral_constant; + using kBF16 = std::integral_constant; + using kT = std::true_type; + using kF = std::false_type; + if (is_half) { + if (has_zp) { + run(kF16{}, kF{}, kT{}); + } else if (w8) { + run(kF16{}, kT{}, kF{}); + } else { + run(kF16{}, kF{}, kF{}); + } + } else { + if (has_zp) { + run(kBF16{}, kF{}, kT{}); + } else if (w8) { + run(kBF16{}, kT{}, kF{}); + } else { + run(kBF16{}, kF{}, kF{}); + } + } + + // C = A W with row-major tensors through column-major cuBLAS as + // C^T = W^T A^T. + cublasHandle_t handle = get_current_cuda_blas_handle(); + SWORDFISH_CHECK_CUBLAS(cublasSetStream(handle, stream)); + const float alpha = 1.0f, beta = 0.0f; + const cudaDataType_t ct = is_half ? CUDA_R_16F : CUDA_R_16BF; + SWORDFISH_CHECK_CUBLAS(cublasGemmEx( + handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, w_dense, ct, n, a, ct, + k, &beta, c, ct, n, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +} + +// Standalone dequant for consumers of dense weights, notably the MoE dense +// tier which feeds the triton fused-MoE kernels ([E, out, in] transposed +// layout). Accepts single [NB, KB, words] or stacked [E, NB, KB, words] +// packed tensors. +torch::stable::Tensor swordfish_dequant_dense( + torch::stable::Tensor& b_packed, torch::stable::Tensor& group_scales, + std::optional const& group_zps, int64_t num_bits, + int64_t group_size, int64_t size_k, int64_t size_n, bool transpose) { + STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, + "swordfish supports 4-bit and 8-bit weights"); + const bool w8 = num_bits == 8; + const bool stacked = b_packed.dim() == 4; + const int64_t experts = stacked ? b_packed.size(0) : 1; + const auto a_st = group_scales.scalar_type(); + STD_TORCH_CHECK(a_st == torch::headeronly::ScalarType::Half || + a_st == torch::headeronly::ScalarType::BFloat16, + "group_scales must be fp16 or bf16"); + const bool has_zp = group_zps.has_value(); + + const int32_t device_index = b_packed.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); + + torch::stable::Tensor out = + stacked + ? torch::stable::empty( + transpose + ? std::initializer_list{experts, size_n, size_k} + : std::initializer_list{experts, size_k, size_n}, + a_st, std::nullopt, b_packed.device()) + : torch::stable::empty( + transpose ? std::initializer_list{size_n, size_k} + : std::initializer_list{size_k, size_n}, + a_st, std::nullopt, b_packed.device()); + + dim3 grid(size_n / kBlockN, size_k / kBlockK, experts); + const auto* b = reinterpret_cast(b_packed.const_data_ptr()); + const void* sp = group_scales.const_data_ptr(); + const void* zp = has_zp ? group_zps->const_data_ptr() : nullptr; + void* wp = out.mutable_data_ptr(); + const int k = int(size_k), n = int(size_n), gs = int(group_size); + + const auto run = [&](auto tid, auto w8c, auto zpc, auto trc) { + constexpr aphrodite::ScalarTypeId kTid = decltype(tid)::value; + using scalar_t = typename marlin::MarlinScalarType::scalar_t; + swordfish_dequant_dense_kernel + <<>>(b, reinterpret_cast(sp), + reinterpret_cast(zp), + nullptr, reinterpret_cast(wp), k, + n, gs); + }; + using kF16 = + std::integral_constant; + using kBF16 = std::integral_constant; + using kT = std::true_type; + using kF = std::false_type; + const auto dispatch = [&](auto tid) { + if (transpose) { + if (has_zp) { + run(tid, kF{}, kT{}, kT{}); + } else if (w8) { + run(tid, kT{}, kF{}, kT{}); + } else { + run(tid, kF{}, kF{}, kT{}); + } + } else { + if (has_zp) { + run(tid, kF{}, kT{}, kF{}); + } else if (w8) { + run(tid, kT{}, kF{}, kF{}); + } else { + run(tid, kF{}, kF{}, kF{}); + } + } + }; + if (a_st == torch::headeronly::ScalarType::Half) { + dispatch(kF16{}); + } else { + dispatch(kBF16{}); + } + return out; +} + +} // namespace swordfish + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("swordfish_dequant_dense", + TORCH_BOX(&swordfish::swordfish_dequant_dense)); +} diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_mm.cu b/csrc/libtorch_stable/quantization/swordfish/swordfish_mm.cu new file mode 100644 index 0000000000..9e29a39f75 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_mm.cu @@ -0,0 +1,494 @@ +// w4a16 decode GEMM over the Swordfish ABI v1 packed weight +//. The kernel lives in swordfish_decode.cuh. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/ops.h" +#include "libtorch_stable/torch_utils.h" +#include "swordfish_decode.cuh" + +namespace swordfish { + +// Defined in swordfish_dense_tier.cu (same extension). +void swordfish_dense_tier_mm(const void* a, const int32_t* b, const void* s, + const void* z, const int32_t* perm, void* c, + void* w_dense, bool is_half, bool w8, bool has_zp, + int m, int k, int n, int group_size, + cudaStream_t stream); + +// Defined in swordfish_prefill.cu (same extension). +torch::stable::Tensor swordfish_prefill_mm( + torch::stable::Tensor& a, torch::stable::Tensor& b_packed, + torch::stable::Tensor& group_scales, + std::optional const& group_zps, int64_t num_bits, + int64_t group_size, int64_t size_k, int64_t size_n); + +namespace { + +// Decode/prefill crossover. It must live in C++ because a Python-side +// branch is traced by torch.compile at one representative M and baked into +// the compiled graph. Here the true runtime M decides on every call and on +// every captured CUDA graph. +inline bool use_prefill(int64_t m, torch::headeronly::ScalarType a_st, bool w8, + int64_t group_size, int64_t k, int64_t n) { + if (a_st != torch::headeronly::ScalarType::BFloat16 && + a_st != torch::headeronly::ScalarType::Half) { + return false; + } + if (w8 ? group_size != 128 + : (group_size != 32 && group_size != 64 && group_size != 128)) { + return false; + } + if (k % 128 != 0 || n % 128 != 0) return false; + // The prefill grid launches about n/128 CTAs per M tile. When that fills + // the machine the tcgen05 path wins from M 48 up; when it underfills + // (many SMs, narrow N), the Stream-K decode window carries [17, 96). + static int sms = 0; + if (sms == 0) { + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0); + if (sms <= 0) sms = 1; + } + const bool prefill_fills = n / 128 >= sms; + // The four-tile decode window carries [48, 56) even at wide N; the + // tcgen05 wave only pulls ahead of it from 56 rows up. K-heavy narrow-N + // shapes keep Stream-K through [96, 128) on many-SM parts, where a single + // underfilled tcgen05 wave loses to it. + if (m >= 96 && m < 128 && k >= 2 * n && n / 128 < sms / 4) return false; + return m >= 96 || (m >= 56 && prefill_fills); +} + +// Above this M the problem is compute-bound and dequant-once + dense +// cuBLAS outruns the fused mixed-input mainloops, so the tier takes over. +// 8-bit crosses earlier than 4-bit because the mixed-input pipeline moves +// twice the weight bytes through the transform. On few-SM parts the dense +// rate only clears the 4-bit weight stream's bandwidth win at 8 bits, and +// K-heavy shapes pay proportionally more dequant per flop, so they cross +// later still. +inline int dense_tier_min_m(int sms, bool w8, int64_t size_k, int64_t size_n, + bool has_perm) { + static const int v = [] { + const char* e = std::getenv("APHRODITE_SWORDFISH_DENSE_M"); + return e != nullptr && e[0] != '\0' ? std::atoi(e) : 0; + }(); + if (v != 0) return v; + if (sms < 100) { + if (!w8) return INT_MAX; + return size_k >= 2 * size_n ? 8192 : 2048; + } + // act_order crosses much earlier outside wide N: the fused paths pay the + // activation sort and per-group scale gathers that the dense tier folds + // into its weight scatter for free (K-heavy shapes halve at M >= 512). + // Wide N stays fused, where the K*N dequant write dominates instead. + if (has_perm && size_n <= 2 * size_k) return 512; + return w8 ? 1024 : 4096; +} + +// APHRODITE_SWORDFISH_DETERMINISTIC forces the run-stable decode paths: +// the smem-reduction epilogue replaces every atomic-window kernel, at the +// decode window's cost. The tcgen05 prefill is deterministic either way. +// The fused-MoE kernels keep their atomic merge regardless. +inline bool force_deterministic() { + static const bool v = [] { + const char* e = std::getenv("APHRODITE_SWORDFISH_DETERMINISTIC"); + return e != nullptr && e[0] != '\0' && std::strcmp(e, "0") != 0; + }(); + return v; +} + +inline int cached_sm_count() { + static int sms = 0; + if (sms == 0) { + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0); + if (sms <= 0) sms = 1; + } + return sms; +} + +template +void launch_decode_streamk_t(const void* a, const int32_t* b, const void* s, + const void* z, void* c, int m, int k, int n, + int group_size, cudaStream_t stream, + bool c_zeroed = false) { + using scalar_t = typename marlin::MarlinScalarType::scalar_t; + constexpr int kStagesT = + T == 1 ? kStages : (T == 2 ? (W8 ? 5 : 4) : (T == 3 ? 3 : (W8 ? 4 : 2))); + constexpr int kUnitK = W8 ? 16 : 32; + static int ctas_per_sm = 0; + if (ctas_per_sm == 0) { + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &ctas_per_sm, + swordfish_decode_streamk_kernel, + kDecodeThreads, 0); + if (ctas_per_sm <= 0) ctas_per_sm = 2; + } + int sms = 0; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0); + const int m_tiles = (m + 15) / 16; + const int m_groups = (m_tiles + T - 1) / T; + const int nb = n / (CTA_QUAD ? kBlockN * kDecodeWarps : kBlockN); + const int num_pairs = k / kUnitK; + const int64_t total = int64_t(m_groups) * nb * num_pairs; + int ctas = ctas_per_sm * sms; + if (CTA_QUAD) { + // Quad claims are CTA-granular: keep them tens of units long so the + // claim, flush, and atomic costs amortize, and stop at the number of + // claims rather than idling CTAs. + const int64_t max_ctas = total / 32 > 0 ? total / 32 : 1; + if (ctas > max_ctas) ctas = int(max_ctas); + } else { + // Cap the grid so every warp gets at least a pipeline's worth of pairs. + const int64_t max_warps = + total / (2 * kStagesT) > 0 ? total / (2 * kStagesT) : 1; + if (int64_t(ctas) * kDecodeWarps > max_warps) { + ctas = int((max_warps + kDecodeWarps - 1) / kDecodeWarps); + } + } + if (ctas < 1) ctas = 1; + if (!c_zeroed) launch_zero_c(c, m, n, stream); + swordfish_decode_streamk_kernel + <<>>( + reinterpret_cast(a), b, + reinterpret_cast(s), + reinterpret_cast(z), reinterpret_cast(c), + m, k, n, group_size, m_groups); +} + +template +void launch_decode_atomic_t(const void* a, const int32_t* b, const void* s, + const void* z, void* c, int m, int k, int n, + int group_size, cudaStream_t stream, + bool c_zeroed = false) { + using scalar_t = typename marlin::MarlinScalarType::scalar_t; + constexpr int kStagesT = T == 1 ? kStages : (T == 2 ? 4 : 3); + static int ctas_per_sm = 0; // per (type, T) instantiation + if (ctas_per_sm == 0) { + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &ctas_per_sm, swordfish_decode_kernel, + kDecodeThreads, 0); + if (ctas_per_sm <= 0) ctas_per_sm = 4; + } + int sms = 0; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0); + const int m_ctas = (m + 16 * T - 1) / (16 * T); + const int nb = n / kBlockN; + const int num_pairs = k / (W8 ? 16 : 32); + const int target = ctas_per_sm * sms; + // Floor division. Splitting K pays only when columns fall well short of + // the machine, since slicing shortens the per-warp pipeline and raises + // atomic contention. + int split = target / (nb * m_ctas); + // On few-SM parts, columns alone reaching the machine is enough. K-slices + // beyond that shorten the latency-bound per-warp chains and buy the + // launcher memset, measured 8-15 percent slower at M=1 narrow N on 20 + // SMs. Many-SM parts want the occupancy target: the same sweep run the + // other way is 2-3x slower at split 1 on 148 SMs. + if (sms < 100 && nb * m_ctas >= sms) split = 1; + // Keep at least 2*kStagesT pairs per warp so the pipeline reaches steady + // state within a slice. + const int max_split = std::max(1, num_pairs / (kDecodeWarps * 2 * kStagesT)); + split = std::min(std::max(split, 1), max_split); + dim3 sgrid(m_ctas, nb, split); + // At split 1 each CTA zeroes its exclusive C tile in-kernel. At split > 1 + // tiles are shared and the memset is required. + if (split > 1 && !c_zeroed) { + launch_zero_c(c, m, n, stream); + } + swordfish_decode_kernel + <<>>( + reinterpret_cast(a), b, + reinterpret_cast(s), + reinterpret_cast(z), reinterpret_cast(c), + m, k, n, group_size); +} + +template +void launch_decode_atomic(int T, const void* a, const int32_t* b, const void* s, + const void* z, void* c, int m, int k, int n, + int group_size, cudaStream_t stream, + bool c_zeroed = false) { + if (T == 1) { + launch_decode_atomic_t( + a, b, s, z, c, m, k, n, group_size, stream, c_zeroed); + } else if (T == 2) { + launch_decode_atomic_t( + a, b, s, z, c, m, k, n, group_size, stream, c_zeroed); + } else { + launch_decode_atomic_t( + a, b, s, z, c, m, k, n, group_size, stream, c_zeroed); + } +} + +template +void launch_decode(const void* a, const int32_t* b, const void* s, + const void* z, void* c, int m, int k, int n, int group_size, + cudaStream_t stream, bool c_zeroed = false) { + using scalar_t = typename marlin::MarlinScalarType::scalar_t; + if (force_deterministic()) { + dim3 grid((m + 15) / 16, n / kBlockN); + swordfish_decode_kernel + <<>>( + reinterpret_cast(a), b, + reinterpret_cast(s), + reinterpret_cast(z), + reinterpret_cast(c), m, k, n, group_size); + } else if (m <= 16) { + // Tuned single-tile path with in-kernel C zeroing and heuristic split-K. + launch_decode_atomic(1, a, b, s, z, c, m, k, n, + group_size, stream, c_zeroed); + } else if (m <= 127) { + // Window dispatch. When columns alone fill the machine the fused atomic + // grid (in-kernel zeroing, no memset) is already balanced; otherwise + // Stream-K hands each warp a contiguous flat range of (fused tile, + // column, pair) work and the atomic epilogue merges segments, removing + // split-K heuristics and wave quantization. + const bool wide_n = n / kBlockN >= 4 * cached_sm_count(); + // On many-SM parts the whole [17, 48] band belongs to the fused atomic + // grid at any width: its CTA shares one weight stream across four + // warps, where Stream-K's warp-private B and A staging pays for itself + // only when the machine is small enough for warps to get long claims. + // Few-SM parts keep Stream-K outside wide N (measured 0.73-0.96 of + // marlin the other way around on 20 SMs). + const bool band_atomic = cached_sm_count() >= 100 && m <= 48; + if (band_atomic || (wide_n && m <= 47)) { + launch_decode_atomic(m <= 32 ? 2 : 3, a, b, s, z, c, + m, k, n, group_size, stream, + c_zeroed); + } else if (m <= 48 && n % (4 * kBlockN) == 0) { + // Few-SM band: CTA-granular claims over n256 column quads. + if (m <= 32) { + launch_decode_streamk_t( + a, b, s, z, c, m, k, n, group_size, stream, c_zeroed); + } else { + launch_decode_streamk_t( + a, b, s, z, c, m, k, n, group_size, stream, c_zeroed); + } + } else if (m <= 32) { + launch_decode_streamk_t( + a, b, s, z, c, m, k, n, group_size, stream, c_zeroed); + } else if (m <= 48) { + launch_decode_streamk_t( + a, b, s, z, c, m, k, n, group_size, stream, c_zeroed); + } else if (m <= 64) { + // Four-tile fusion amortizes the dequant across the whole band; the + // m-shared CTA it replaces dequantized the same weights once per warp + // and issued 2.7x the instructions for it. + launch_decode_streamk_t( + a, b, s, z, c, m, k, n, group_size, stream, c_zeroed); + } else { + launch_decode_streamk_t( + a, b, s, z, c, m, k, n, group_size, stream, c_zeroed); + } + } else { + dim3 grid((m + 15) / 16, n / kBlockN); + swordfish_decode_kernel + <<>>( + reinterpret_cast(a), b, + reinterpret_cast(s), + reinterpret_cast(z), + reinterpret_cast(c), m, k, n, group_size); + } +} + +} // namespace + +torch::stable::Tensor swordfish_mm( + torch::stable::Tensor& a, torch::stable::Tensor& b_packed, + torch::stable::Tensor& group_scales, + std::optional const& group_zps, + std::optional const& perm, int64_t num_bits, + int64_t group_size, int64_t size_k, int64_t size_n) { + STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, + "swordfish supports 4-bit and 8-bit weights"); + const bool w8 = num_bits == 8; + STD_TORCH_CHECK(shape_ok(size_k, size_n), "swordfish ABI v1 requires K % ", + kBlockK, " == 0 and N % ", kBlockN, " == 0; got K=", size_k, + " N=", size_n); + STD_TORCH_CHECK(a.dim() == 2 && a.size(1) == size_k, + "a must be [M, K] with K=", size_k); + STD_TORCH_CHECK(a.stride(1) == 1 && a.stride(0) == size_k, + "a must be contiguous"); + const auto a_st = a.scalar_type(); + STD_TORCH_CHECK(a_st == torch::headeronly::ScalarType::Half || + a_st == torch::headeronly::ScalarType::BFloat16, + "a must be fp16 or bf16"); + STD_TORCH_CHECK(group_scales.scalar_type() == a_st, + "group_scales dtype must match a"); + + const int64_t nb = num_blocks_n(size_n); + const int64_t kb = num_blocks_k(size_k); + const int64_t words = w8 ? kBlockInt32_8 : kBlockInt32; + STD_TORCH_CHECK( + b_packed.scalar_type() == torch::headeronly::ScalarType::Int && + b_packed.dim() == 3 && b_packed.size(0) == nb && + b_packed.size(1) == kb && b_packed.size(2) == words, + "b_packed must be int32 [", nb, ", ", kb, ", ", words, "]"); + + int64_t num_groups = 1; + // Channelwise checkpoints arrive as group -1 with the single scale row + // replicated to group 128 by the weight loader. The grouped tiers + // (tcgen05 prefill, dense dequant) consume the replicated rows as g128; + // the decode kernels take the native -1 path, which reads row 0 once and + // skips the per-group fetch/expand bookkeeping the duplicate rows would + // otherwise re-run at every 128-row boundary. + int64_t tier_group = group_size; + if (group_size == -1) { + if (group_scales.size(0) == size_k / 128) { + num_groups = size_k / 128; + tier_group = 128; + } + } else { + // The decode mainloop consumes k16-slice PAIRS (k32) with scales hoisted + // per pair, so a scale group must cover whole pairs. + STD_TORCH_CHECK(group_size > 0 && size_k % group_size == 0 && + group_size % (2 * kMarlinTileK) == 0, + "group_size must be -1 or a multiple of ", 2 * kMarlinTileK, + " dividing K; got ", group_size); + num_groups = size_k / group_size; + } + STD_TORCH_CHECK(group_scales.dim() == 2 && + group_scales.size(0) == num_groups && + group_scales.size(1) == size_n, + "group_scales must be [", num_groups, ", ", size_n, "]"); + const bool has_zp = group_zps.has_value(); + if (has_zp) { + STD_TORCH_CHECK(!w8, "zero points are a 4-bit (AWQ/HQQ) feature"); + STD_TORCH_CHECK( + group_zps->scalar_type() == a_st && group_zps->dim() == 2 && + group_zps->size(0) == num_groups && group_zps->size(1) == size_n, + "group_zps must be [", num_groups, ", ", size_n, "] with a's dtype"); + } + + const int64_t size_m = a.size(0); + + const bool has_perm = perm.has_value() && perm->numel() > 0; + if (has_perm) { + STD_TORCH_CHECK(perm->scalar_type() == torch::headeronly::ScalarType::Int && + perm->numel() == size_k, + "perm must be int32 [", size_k, "]"); + } + + if (size_m >= + dense_tier_min_m(cached_sm_count(), w8, size_k, size_n, has_perm) && + !force_deterministic()) { + const int32_t device_index = a.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); + torch::stable::Tensor c = + torch::stable::empty({size_m, size_n}, a_st, std::nullopt, a.device()); + torch::stable::Tensor w_dense = + torch::stable::empty({size_k, size_n}, a_st, std::nullopt, a.device()); + // Act_order folds into the weight scatter here, so the activations are + // consumed unpermuted. + swordfish_dense_tier_mm( + a.const_data_ptr(), + reinterpret_cast(b_packed.const_data_ptr()), + group_scales.const_data_ptr(), + has_zp ? group_zps->const_data_ptr() : nullptr, + has_perm ? reinterpret_cast(perm->const_data_ptr()) + : nullptr, + c.mutable_data_ptr(), w_dense.mutable_data_ptr(), + a_st == torch::headeronly::ScalarType::Half, w8, has_zp, int(size_m), + int(size_k), int(size_n), int(tier_group), stream); + return c; + } + + // The fused paths consume group-sorted K. In the decode window the sort + // and the output zeroing fuse into one prep launch (a separate + // permute_cols node plus a zero node cost several microseconds of launch + // and engine gaps per GEMM at bs=1); prefill and the tall tail keep the + // plain sorted copy. + const bool will_prefill = + use_prefill(size_m, a_st, w8, tier_group, size_k, size_n); + const bool prep_perm = + has_perm && !will_prefill && size_m <= 127 && !force_deterministic(); + torch::stable::Tensor a_used = + has_perm && !prep_perm ? permute_cols(a, *perm) : a; + + if (will_prefill) { + return swordfish_prefill_mm(a_used, b_packed, group_scales, group_zps, + num_bits, tier_group, size_k, size_n); + } + + const int32_t device_index = a.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); + + torch::stable::Tensor c = + torch::stable::empty({size_m, size_n}, a_st, std::nullopt, a.device()); + if (size_m == 0) return c; + + torch::stable::Tensor a_perm = a_used; + if (prep_perm) { + a_perm = + torch::stable::empty({size_m, size_k}, a_st, std::nullopt, a.device()); + if (a_st == torch::headeronly::ScalarType::Half) { + launch_prep( + a_used.const_data_ptr(), + reinterpret_cast(perm->const_data_ptr()), + a_perm.mutable_data_ptr(), c.mutable_data_ptr(), size_m, size_k, + size_n, stream); + } else { + launch_prep( + a_used.const_data_ptr(), + reinterpret_cast(perm->const_data_ptr()), + a_perm.mutable_data_ptr(), c.mutable_data_ptr(), size_m, size_k, + size_n, stream); + } + } + + const auto* b_ptr = + reinterpret_cast(b_packed.const_data_ptr()); + const void* a_ptr = a_perm.const_data_ptr(); + const void* s_ptr = group_scales.const_data_ptr(); + const void* z_ptr = has_zp ? group_zps->const_data_ptr() : nullptr; + void* c_ptr = c.mutable_data_ptr(); + + if (a_st == torch::headeronly::ScalarType::Half) { + if (has_zp) { + launch_decode( + a_ptr, b_ptr, s_ptr, z_ptr, c_ptr, size_m, size_k, size_n, group_size, + stream, prep_perm); + } else if (w8) { + launch_decode( + a_ptr, b_ptr, s_ptr, z_ptr, c_ptr, size_m, size_k, size_n, group_size, + stream, prep_perm); + } else { + launch_decode( + a_ptr, b_ptr, s_ptr, z_ptr, c_ptr, size_m, size_k, size_n, group_size, + stream, prep_perm); + } + } else { + if (has_zp) { + launch_decode( + a_ptr, b_ptr, s_ptr, z_ptr, c_ptr, size_m, size_k, size_n, group_size, + stream, prep_perm); + } else if (w8) { + launch_decode( + a_ptr, b_ptr, s_ptr, z_ptr, c_ptr, size_m, size_k, size_n, group_size, + stream, prep_perm); + } else { + launch_decode( + a_ptr, b_ptr, s_ptr, z_ptr, c_ptr, size_m, size_k, size_n, group_size, + stream, prep_perm); + } + } + + return c; +} + +} // namespace swordfish + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("swordfish_mm", TORCH_BOX(&swordfish::swordfish_mm)); +} diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_moe.cu b/csrc/libtorch_stable/quantization/swordfish/swordfish_moe.cu new file mode 100644 index 0000000000..931af66928 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_moe.cu @@ -0,0 +1,480 @@ +// Fused-MoE decode GEMM over per-expert Swordfish ABI v1 packed weights. +// One persistent Stream-K kernel covers every (expert block, column, unit) +// of the token-sorted problem that moe_align_block_size(block_size=16) +// prepares: each 16-token block is one m16 tile of exactly one expert, so +// the dense Stream-K decode mainloop applies with the activation rows and +// output rows indirected through sorted_token_ids and the weight and scale +// bases indexed by expert_ids. Work totals depend on the device-side padded +// token count, so per-warp ranges derive in-kernel from +// num_tokens_post_padded rather than on the host. + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" +#include "swordfish_decode.cuh" + +namespace swordfish { + +namespace { + +template +__global__ void swordfish_decode_moe_kernel( + const typename marlin::MarlinScalarType::scalar_t* __restrict__ A, + const int32_t* __restrict__ B, + const typename marlin::MarlinScalarType::scalar_t* __restrict__ S, + typename marlin::MarlinScalarType::scalar_t* __restrict__ C, + const int32_t* __restrict__ sorted_token_ids, + const int32_t* __restrict__ expert_ids, + const int32_t* __restrict__ num_tokens_post_padded, + const float* __restrict__ topk_weights, int top_k, bool mul_topk_weights, + int total_tokens, int K, int N, int group_size) { + static_assert(M_TILES == 1 || M_TILES == 2, "moe m-tile fusion is 1 or 2"); + constexpr int kStagesT = M_TILES == 1 ? kStages : 4; + constexpr int kUnitK = W8 ? 16 : 32; + using Dtype = marlin::MarlinScalarType; + using scalar_t = typename Dtype::scalar_t; + using scalar_t2 = typename Dtype::scalar_t2; + using FragA = typename Dtype::FragA; + using FragB = typename Dtype::FragB; + using FragC = typename Dtype::FragC; + + const int lane = threadIdx.x % 32; + const int warp = threadIdx.x / 32; + const int group_id = lane / 4; + const int tig = lane % 4; + const int ldsm_row = lane & 15; + const int ldsm_col = (lane >> 4) << 3; + const int ia_row = lane >> 1; + const int ia_c0 = (kUnitK / 2) * (lane & 1); + + const int num_units = K / kUnitK; + const int nb_cnt = N / kBlockN; + const int num_kb = K / kBlockK; + const int ppg = group_size > 0 ? group_size / kUnitK : INT_MAX; + const scalar_t2 zero2 = Dtype::num2num2(Dtype::float2num(0.0f)); + const uint64_t bpol = l2_evict_first_policy(); + constexpr int64_t kBlockW = W8 ? kBlockInt32_8 : kBlockInt32; + + const int num_blocks = *num_tokens_post_padded / (16 * M_TILES); + const int64_t total = int64_t(num_blocks) * nb_cnt * num_units; + const int total_warps = gridDim.x * kDecodeWarps; + const int64_t per = (total + total_warps - 1) / total_warps; + int64_t w = int64_t(blockIdx.x * kDecodeWarps + warp) * per; + const int64_t w_end = w + per < total ? w + per : total; + + __shared__ int4 bstage[kDecodeWarps][kStagesT][2 * 32]; + __shared__ scalar_t astage[kDecodeWarps][kStagesT][M_TILES][16][kUnitK]; + + FragC acc[M_TILES][4][2]; + scalar_t2 s_reg[4][2]; + + while (w < w_end) { + const int64_t cg = w / num_units; + const int p_beg = int(w - cg * num_units); + const int col = int(cg / num_blocks); + const int blk = int(cg - int64_t(col) * num_blocks); + const int p_end = + int(int64_t(num_units) < p_beg + (w_end - w) ? int64_t(num_units) + : p_beg + (w_end - w)); + w += p_end - p_beg; + + const int eid = expert_ids[blk]; + if (eid < 0) continue; // expert-parallel block owned by another rank + + const int m_base = 16 * M_TILES * blk; + const int col_base = col * kBlockN; + + // Output rows for this lane's C fragments, through the token sort. + int id0[M_TILES], id1[M_TILES]; + bool r0ok[M_TILES], r1ok[M_TILES]; +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + id0[t] = sorted_token_ids[m_base + 16 * t + group_id]; + id1[t] = sorted_token_ids[m_base + 16 * t + group_id + 8]; + r0ok[t] = id0[t] < total_tokens; + r1ok[t] = id1[t] < total_tokens; + } + +#pragma unroll + for (int t = 0; t < M_TILES; t++) +#pragma unroll + for (int j = 0; j < 4; j++) +#pragma unroll + for (int b = 0; b < 2; b++) +#pragma unroll + for (int i = 0; i < 4; i++) acc[t][j][b][i] = 0.0f; + + const scalar_t* S_e = + S + int64_t(eid) * (group_size > 0 ? K / group_size : 1) * N; + auto fetch_row = [&](int g) -> scalar_t2 { + return reinterpret_cast(S_e + int64_t(g) * N + + col_base)[lane]; + }; + auto expand_row = [&](scalar_t2 mine) { + const uint32_t sel = group_id & 1; +#pragma unroll + for (int j = 0; j < 4; j++) { +#pragma unroll + for (int b = 0; b < 2; b++) { + const scalar_t2 v = + __shfl_sync(0xffffffffu, mine, 8 * j + 4 * b + (group_id >> 1)); + const uint32_t bits = reinterpret_cast(v); + const uint16_t h16 = sel ? uint16_t(bits >> 16) : uint16_t(bits); + s_reg[j][b] = Dtype::num2num2(reinterpret_cast(h16)); + } + } + }; + auto process_slice = [&](const FragA(&fa)[M_TILES], const marlin::I4& bqa, + const marlin::I4& bqb) { +#pragma unroll + for (int j = 0; j < 4; j++) { + FragB frag_b0, frag_b1; + if constexpr (W8) { + const marlin::I4& src = j < 2 ? bqa : bqb; + marlin::dequant( + src.elems[(2 * j) & 3], reinterpret_cast(&frag_b0)); + marlin::dequant( + src.elems[(2 * j + 1) & 3], + reinterpret_cast(&frag_b1)); + } else { + const int b_quant_0 = bqa.elems[j]; + const int b_quant_1 = b_quant_0 >> 8; + marlin::dequant( + b_quant_0, reinterpret_cast(&frag_b0)); + marlin::dequant( + b_quant_1, reinterpret_cast(&frag_b1)); + } + frag_b0[0] = __hmul2(frag_b0[0], s_reg[j][0]); + frag_b0[1] = __hmul2(frag_b0[1], s_reg[j][0]); + frag_b1[0] = __hmul2(frag_b1[0], s_reg[j][1]); + frag_b1[1] = __hmul2(frag_b1[1], s_reg[j][1]); +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + marlin::mma(fa[t], frag_b0, acc[t][j][0]); + marlin::mma(fa[t], frag_b1, acc[t][j][1]); + } + } + }; + auto process_pair = [&](int aslot, const int4* buf) { + FragA fa0[M_TILES], fa1[M_TILES]; +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + ldsm4(fa0[t], &astage[warp][aslot][t][ldsm_row][ldsm_col]); + if constexpr (!W8) { + ldsm4(fa1[t], + &astage[warp][aslot][t][ldsm_row][ldsm_col + 16]); + } + } + if constexpr (W8) { + const marlin::I4 bq0 = + *reinterpret_cast(&buf[2 * lane]); + const marlin::I4 bq1 = + *reinterpret_cast(&buf[2 * lane + 1]); + process_slice(fa0, bq0, bq1); + } else { + const marlin::I4 bq0 = *reinterpret_cast(&buf[lane]); + const marlin::I4 bq1 = + *reinterpret_cast(&buf[32 + lane]); + process_slice(fa0, bq0, bq0); + process_slice(fa1, bq1, bq1); + } + }; + + const int32_t* ipair_ptr = + B + (int64_t(eid) * nb_cnt + col) * num_kb * kBlockW + + p_beg * kPairInt32 + (W8 ? 8 : 4) * lane; + // Activation rows through the sort; padding rows stage row 0 disabled. + bool a_okt[M_TILES]; + const scalar_t* ia_ptr[M_TILES]; +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + const int a_id = sorted_token_ids[m_base + 16 * t + ia_row]; + a_okt[t] = a_id < total_tokens; + ia_ptr[t] = A + (a_okt[t] ? int64_t(a_id / top_k) : 0) * K + + kUnitK * p_beg + ia_c0; + } + int ipend = p_end - p_beg; + + auto issue_pair = [&](int slot) { + if (ipend > 0) { + ipend--; + if constexpr (W8) { + cp_async4_evict_first(&bstage[warp][slot][2 * lane], ipair_ptr, bpol); + cp_async4_evict_first(&bstage[warp][slot][2 * lane + 1], + ipair_ptr + 4, bpol); + } else { + cp_async4_evict_first(&bstage[warp][slot][lane], ipair_ptr, bpol); + cp_async4_evict_first(&bstage[warp][slot][32 + lane], + ipair_ptr + kPairInt32 / 2, bpol); + } + ipair_ptr += kPairInt32; +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + if (a_okt[t]) { + marlin::cp_async4(&astage[warp][slot][t][ia_row][ia_c0], ia_ptr[t]); + if constexpr (!W8) { + marlin::cp_async4(&astage[warp][slot][t][ia_row][ia_c0 + 8], + ia_ptr[t] + 8); + } + } + ia_ptr[t] += kUnitK; + } + } + marlin::cp_async_fence(); + }; + + int g = 0; + int left = INT_MAX; + const int g_last = group_size > 0 && p_beg < p_end + ? (kUnitK * (p_end - 1)) / group_size + : 0; + scalar_t2 s_next = zero2; + if (p_beg < p_end) { + if (group_size > 0) { + g = p_beg / ppg; + left = ppg - p_beg % ppg; + } + expand_row(fetch_row(g)); + if (g + 1 <= g_last) s_next = fetch_row(g + 1); + } + +#pragma unroll + for (int st = 0; st < kStagesT - 1; st++) issue_pair(st); + + int slot = 0; + int islot = kStagesT - 1; + for (int p = p_beg; p < p_end; p++) { + issue_pair(islot); + if (++islot == kStagesT) islot = 0; + marlin::cp_async_wait(); + if (left == 0) { + ++g; + expand_row(s_next); + if (g + 1 <= g_last) s_next = fetch_row(g + 1); + left = ppg; + } + process_pair(slot, bstage[warp][slot]); + left--; + if (++slot == kStagesT) slot = 0; + } + + // Segment flush into the launcher-zeroed C, rows indirected through the + // token sort, with the router weight folded in when requested. +#pragma unroll + for (int t = 0; t < M_TILES; t++) { + const float wt0 = + mul_topk_weights && r0ok[t] ? topk_weights[id0[t]] : 1.0f; + const float wt1 = + mul_topk_weights && r1ok[t] ? topk_weights[id1[t]] : 1.0f; +#pragma unroll + for (int j = 0; j < 4; j++) { +#pragma unroll + for (int b = 0; b < 2; b++) { + const int cc = col_base + 8 * (2 * j + b) + 2 * tig; + const float4 v = *reinterpret_cast(&acc[t][j][b]); + if (r0ok[t]) + red_add2(reinterpret_cast(C + int64_t(id0[t]) * N + cc), + Dtype::nums2num2(Dtype::float2num(v.x * wt0), + Dtype::float2num(v.y * wt0))); + if (r1ok[t]) + red_add2(reinterpret_cast(C + int64_t(id1[t]) * N + cc), + Dtype::nums2num2(Dtype::float2num(v.z * wt1), + Dtype::float2num(v.w * wt1))); + } + } + } + } +} + +template +void launch_moe(const void* a, const int32_t* b, const void* s, void* c, + const int32_t* sorted_ids, const int32_t* expert_ids, + const int32_t* num_post_padded, const float* topk_w, int top_k, + bool mul_topk, int total_tokens, int max_blocks, int block_size, + int k, int n, int group_size, cudaStream_t stream) { + using scalar_t = typename marlin::MarlinScalarType::scalar_t; + int sms = 0; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0); + constexpr int kUnitK = W8 ? 16 : 32; + launch_zero_c(c, total_tokens, n, stream); + // Host-side upper bound on work keeps tiny batches from launching idle + // CTAs; the true total is device-side. + const int64_t max_total = int64_t(max_blocks) * (n / kBlockN) * (k / kUnitK); + const int64_t max_warps = + max_total / (2 * kStages) > 0 ? max_total / (2 * kStages) : 1; + const auto grid = [&](int ctas_per_sm) { + int ctas = ctas_per_sm * sms; + if (int64_t(ctas) * kDecodeWarps > max_warps) { + ctas = int((max_warps + kDecodeWarps - 1) / kDecodeWarps); + } + return ctas < 1 ? 1 : ctas; + }; + if (block_size == 32) { + // 32-token blocks fuse two m16 tiles per weight fetch, amortizing the + // dequant for batched shapes. + static int ctas_per_sm2 = 0; + if (ctas_per_sm2 == 0) { + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &ctas_per_sm2, swordfish_decode_moe_kernel, + kDecodeThreads, 0); + if (ctas_per_sm2 <= 0) ctas_per_sm2 = 2; + } + swordfish_decode_moe_kernel + <<>>( + reinterpret_cast(a), b, + reinterpret_cast(s), + reinterpret_cast(c), sorted_ids, expert_ids, + num_post_padded, topk_w, top_k, mul_topk, total_tokens, k, n, + group_size); + return; + } + static int ctas_per_sm = 0; + if (ctas_per_sm == 0) { + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &ctas_per_sm, swordfish_decode_moe_kernel, kDecodeThreads, + 0); + if (ctas_per_sm <= 0) ctas_per_sm = 2; + } + swordfish_decode_moe_kernel + <<>>( + reinterpret_cast(a), b, + reinterpret_cast(s), reinterpret_cast(c), + sorted_ids, expert_ids, num_post_padded, topk_w, top_k, mul_topk, + total_tokens, k, n, group_size); +} + +} // namespace + +torch::stable::Tensor swordfish_moe_mm( + torch::stable::Tensor& a, torch::stable::Tensor& b_packed, + torch::stable::Tensor& group_scales, + torch::stable::Tensor& sorted_token_ids, torch::stable::Tensor& expert_ids, + torch::stable::Tensor& num_tokens_post_padded, + std::optional const& topk_weights, + int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, + int64_t num_bits, int64_t group_size, int64_t size_k, int64_t size_n) { + STD_TORCH_CHECK(moe_block_size == 16 || moe_block_size == 32, + "swordfish_moe_mm supports block sizes 16 and 32"); + STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, + "swordfish supports 4-bit and 8-bit weights"); + const bool w8 = num_bits == 8; + STD_TORCH_CHECK(shape_ok(size_k, size_n), "swordfish ABI v1 requires K % ", + kBlockK, " == 0 and N % ", kBlockN, " == 0; got K=", size_k, + " N=", size_n); + STD_TORCH_CHECK(a.dim() == 2 && a.size(1) == size_k && a.stride(1) == 1 && + a.stride(0) == size_k, + "a must be contiguous [M, K]"); + const auto a_st = a.scalar_type(); + STD_TORCH_CHECK(a_st == torch::headeronly::ScalarType::Half || + a_st == torch::headeronly::ScalarType::BFloat16, + "a must be fp16 or bf16"); + STD_TORCH_CHECK(group_scales.scalar_type() == a_st, + "group_scales dtype must match a"); + + const int64_t nb = num_blocks_n(size_n); + const int64_t kb = num_blocks_k(size_k); + const int64_t words = w8 ? kBlockInt32_8 : kBlockInt32; + STD_TORCH_CHECK( + b_packed.scalar_type() == torch::headeronly::ScalarType::Int && + b_packed.dim() == 4 && b_packed.size(1) == nb && + b_packed.size(2) == kb && b_packed.size(3) == words, + "b_packed must be int32 [E, ", nb, ", ", kb, ", ", words, "]"); + const int64_t num_experts = b_packed.size(0); + + int64_t num_groups = 1; + if (group_size != -1) { + STD_TORCH_CHECK(group_size > 0 && size_k % group_size == 0 && + group_size % (2 * kMarlinTileK) == 0, + "group_size must be -1 or a multiple of ", 2 * kMarlinTileK, + " dividing K; got ", group_size); + num_groups = size_k / group_size; + } + STD_TORCH_CHECK( + group_scales.dim() == 3 && group_scales.size(0) == num_experts && + group_scales.size(1) == num_groups && group_scales.size(2) == size_n, + "group_scales must be [", num_experts, ", ", num_groups, ", ", size_n, + "]"); + STD_TORCH_CHECK( + sorted_token_ids.scalar_type() == torch::headeronly::ScalarType::Int && + expert_ids.scalar_type() == torch::headeronly::ScalarType::Int && + num_tokens_post_padded.scalar_type() == + torch::headeronly::ScalarType::Int, + "alignment buffers must be int32"); + if (mul_topk_weights) { + STD_TORCH_CHECK( + topk_weights.has_value() && + topk_weights->scalar_type() == torch::headeronly::ScalarType::Float, + "topk_weights must be fp32 when mul_topk_weights"); + } + + const int64_t total_tokens = a.size(0) * top_k; + + const int32_t device_index = a.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); + + torch::stable::Tensor c = torch::stable::empty({total_tokens, size_n}, a_st, + std::nullopt, a.device()); + if (total_tokens == 0) return c; + + const auto* b_ptr = + reinterpret_cast(b_packed.const_data_ptr()); + const auto* sid_ptr = + reinterpret_cast(sorted_token_ids.const_data_ptr()); + const auto* eid_ptr = + reinterpret_cast(expert_ids.const_data_ptr()); + const auto* npp_ptr = + reinterpret_cast(num_tokens_post_padded.const_data_ptr()); + const float* tw_ptr = + mul_topk_weights + ? reinterpret_cast(topk_weights->const_data_ptr()) + : nullptr; + const int max_blocks = + int((sorted_token_ids.numel() + moe_block_size - 1) / moe_block_size); + + if (a_st == torch::headeronly::ScalarType::Half) { + if (w8) { + launch_moe( + a.const_data_ptr(), b_ptr, group_scales.const_data_ptr(), + c.mutable_data_ptr(), sid_ptr, eid_ptr, npp_ptr, tw_ptr, int(top_k), + mul_topk_weights, int(total_tokens), max_blocks, int(moe_block_size), + int(size_k), int(size_n), int(group_size), stream); + } else { + launch_moe( + a.const_data_ptr(), b_ptr, group_scales.const_data_ptr(), + c.mutable_data_ptr(), sid_ptr, eid_ptr, npp_ptr, tw_ptr, int(top_k), + mul_topk_weights, int(total_tokens), max_blocks, int(moe_block_size), + int(size_k), int(size_n), int(group_size), stream); + } + } else { + if (w8) { + launch_moe( + a.const_data_ptr(), b_ptr, group_scales.const_data_ptr(), + c.mutable_data_ptr(), sid_ptr, eid_ptr, npp_ptr, tw_ptr, int(top_k), + mul_topk_weights, int(total_tokens), max_blocks, int(moe_block_size), + int(size_k), int(size_n), int(group_size), stream); + } else { + launch_moe( + a.const_data_ptr(), b_ptr, group_scales.const_data_ptr(), + c.mutable_data_ptr(), sid_ptr, eid_ptr, npp_ptr, tw_ptr, int(top_k), + mul_topk_weights, int(total_tokens), max_blocks, int(moe_block_size), + int(size_k), int(size_n), int(group_size), stream); + } + } + + return c; +} + +} // namespace swordfish + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("swordfish_moe_mm", TORCH_BOX(&swordfish::swordfish_moe_mm)); +} diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill.cu b/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill.cu new file mode 100644 index 0000000000..122b3a1524 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill.cu @@ -0,0 +1,108 @@ +// w4a16/w8a16 prefill GEMM over the Swordfish ABI v1 packed weight, using +// the forked sm100 tcgen05 mixed-input collective. This TU instantiates the +// bf16-activation configurations and hosts the op entry; the fp16 set +// compiles in swordfish_prefill_f16.cu. + +#include "swordfish_prefill_impl.cuh" + +namespace swordfish { + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) +// Defined in swordfish_prefill_f16.cu. +namespace prefill { +extern template void run_prefill_all( + torch::stable::Tensor&, torch::stable::Tensor&, torch::stable::Tensor&, + const void*, bool, bool, int, torch::stable::Tensor&, int, int, int, + cudaStream_t); +} +#endif + +torch::stable::Tensor swordfish_prefill_mm( + torch::stable::Tensor& a, torch::stable::Tensor& b_packed, + torch::stable::Tensor& group_scales, + std::optional const& group_zps, int64_t num_bits, + int64_t group_size, int64_t size_k, int64_t size_n) { +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + STD_TORCH_CHECK( + shape_ok(size_k, size_n) && size_k % 128 == 0 && size_n % 128 == 0, + "swordfish prefill v1 requires K % 128 == 0 and " + "N % 128 == 0; got K=", + size_k, " N=", size_n); + STD_TORCH_CHECK(a.dim() == 2 && a.size(1) == size_k, + "a must be [M, K] with K=", size_k); + STD_TORCH_CHECK(a.stride(1) == 1 && a.stride(0) == size_k, + "a must be contiguous"); + const auto a_st = a.scalar_type(); + STD_TORCH_CHECK(a_st == torch::headeronly::ScalarType::BFloat16 || + a_st == torch::headeronly::ScalarType::Half, + "swordfish prefill requires fp16 or bf16 activations"); + STD_TORCH_CHECK(group_scales.scalar_type() == a_st, + "group_scales dtype must match a"); + STD_TORCH_CHECK(group_size == 32 || group_size == 64 || group_size == 128, + "swordfish prefill supports group sizes 32, 64 and 128; " + "got ", + group_size); + + STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, + "swordfish supports 4-bit and 8-bit weights"); + const bool w8 = num_bits == 8; + const int64_t nb = num_blocks_n(size_n); + const int64_t kb = num_blocks_k(size_k); + const int64_t words = w8 ? kBlockInt32_8 : kBlockInt32; + STD_TORCH_CHECK( + b_packed.scalar_type() == torch::headeronly::ScalarType::Int && + b_packed.dim() == 3 && b_packed.size(0) == nb && + b_packed.size(1) == kb && b_packed.size(2) == words, + "b_packed must be int32 [", nb, ", ", kb, ", ", words, "]"); + + const int64_t num_groups = size_k / group_size; + STD_TORCH_CHECK(group_scales.dim() == 2 && + group_scales.size(0) == num_groups && + group_scales.size(1) == size_n, + "group_scales must be [", num_groups, ", ", size_n, "]"); + const bool has_zp = group_zps.has_value(); + if (has_zp) { + STD_TORCH_CHECK( + group_zps->scalar_type() == a_st && group_zps->dim() == 2 && + group_zps->size(0) == num_groups && group_zps->size(1) == size_n, + "group_zps must be [", num_groups, ", ", size_n, "] with a's dtype"); + } + + const int64_t size_m = a.size(0); + + const int32_t device_index = a.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); + + torch::stable::Tensor c = torch::stable::empty( + {size_m, size_n}, a.scalar_type(), std::nullopt, a.device()); + if (size_m == 0) return c; + + const int M = int(size_m), N = int(size_n), K = int(size_k); + + STD_TORCH_CHECK(!(has_zp && w8), "zero points are a 4-bit feature"); + STD_TORCH_CHECK(!(w8 && group_size != 128), + "8-bit prefill supports group_size 128 only"); + const void* zp_ptr = has_zp ? group_zps->const_data_ptr() : nullptr; + if (a_st == torch::headeronly::ScalarType::Half) { + prefill::run_prefill_all(a, b_packed, group_scales, zp_ptr, + has_zp, w8, int(group_size), c, M, + N, K, stream); + } else { + prefill::run_prefill_all( + a, b_packed, group_scales, zp_ptr, has_zp, w8, int(group_size), c, M, N, + K, stream); + } + return c; +#else + STD_TORCH_CHECK(false, + "swordfish_prefill_mm requires a CUDA >= 12.8 build with " + "sm100-family support"); +#endif +} + +} // namespace swordfish + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("swordfish_prefill_mm", TORCH_BOX(&swordfish::swordfish_prefill_mm)); +} diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_f16.cu b/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_f16.cu new file mode 100644 index 0000000000..c022897375 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_f16.cu @@ -0,0 +1,16 @@ +// fp16-activation instantiations of the Swordfish prefill configurations. + +#include "swordfish_prefill_impl.cuh" + +namespace swordfish { +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) +namespace prefill { +template void run_prefill_all(torch::stable::Tensor&, + torch::stable::Tensor&, + torch::stable::Tensor&, + const void*, bool, bool, int, + torch::stable::Tensor&, int, int, + int, cudaStream_t); +} // namespace prefill +#endif +} // namespace swordfish diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_impl.cuh b/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_impl.cuh new file mode 100644 index 0000000000..82b1f255b4 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_impl.cuh @@ -0,0 +1,314 @@ +// Prefill GEMM configuration and launch for the Swordfish packed ABI, shared +// by the per-dtype instantiation TUs (swordfish_prefill.cu for bf16 and +// swordfish_prefill_f16.cu for fp16). +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" +#include "swordfish_types.cuh" + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/util/packed_stride.hpp" + +#include "swordfish_prefill_mainloop.cuh" +#include "swordfish_prefill_kernel.cuh" + +namespace swordfish { + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +namespace prefill { + +using namespace cute; + +using ElementAccumulator = float; +using ArchTag = cutlass::arch::Sm100; +using OperatorClass = cutlass::arch::OpClassTensorOp; + +using ClusterShape = Shape<_2, _1, _1>; +using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; +using LayoutC = cutlass::layout::RowMajor; +constexpr int AlignmentC = 8; // 128b / 16b elems, fp16 and bf16 alike + +using StrideA = + cutlass::gemm::TagToStrideA_t; // packed slot + // (unused) +using StrideB = + cutlass::gemm::TagToStrideB_t::type>; // activations + +static constexpr cute::UMMA::Major UmmaMajorA = cute::UMMA::Major::K; +static constexpr cute::UMMA::Major UmmaMajorB = cute::UMMA::Major::K; + +// 2-SM (cta_group::2) MMA config, parameterized on the instruction N width. +// Tile-M (the weight N dimension) spans the two CTAs of an SM pair, 128 +// columns each. N=256 wins compute-bound shapes (per-instruction issue +// overhead caps a 256x128 UMMA at about two thirds of the 256x256 rate); +// N=128 wins K-heavy shapes, where the wide tile starves the K pipeline. +template +struct PrefillCfg { + using MmaType = TAct; + using MmaTileShape = Shape<_256, Int, _128>; + // A third element in the A tuple enables the collective's zero-point row. + using ElementPairA = + cute::conditional_t, + cute::tuple>; + using ScaleConfig = cutlass::detail::Sm100MixedInputBlockwiseScaleConfig< + /*GranN=*/1, kGran>; + using LayoutScale = decltype(ScaleConfig::deduce_layout_scale()); + using StridePairA = decltype(cute::make_tuple(StrideA{}, LayoutScale{})); + + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, + ElementAccumulator, MmaType, + typename cutlass::layout::LayoutTranspose::type, AlignmentC, + MmaType, typename cutlass::layout::LayoutTranspose::type, + AlignmentC, EpilogueSchedule>::CollectiveOp; + + // The CUTLASS convenience builder has no branch for 2-SM atoms with + // smem-sourced A, so the SS atom is constructed directly. M is the full + // cluster tile-M spanning both CTAs. + using TiledMma = decltype(cute::make_tiled_mma( + cute::SM100_MMA_F16BF16_2x1SM_SS{})); + + // partition_shape_A/B take CTA-local shapes for a 2-SM atom. + using CtaShape = decltype(shape_div(MmaTileShape{}, ClusterShape{})); + using MmaShapeA_MK = decltype(partition_shape_A( + TiledMma{}, + make_shape(cute::size<0>(CtaShape{}), cute::size<2>(CtaShape{})))); + using MmaShapeB_NK = decltype(partition_shape_B( + TiledMma{}, + make_shape(cute::size<1>(CtaShape{}), cute::size<2>(CtaShape{})))); + using BlockTileA_M = decltype(cute::size<0, 0>(MmaShapeA_MK{}) * + cute::size<1>(MmaShapeA_MK{})); + using BlockTileA_K = decltype(cute::size<0, 1>(MmaShapeA_MK{}) * + cute::size<2>(MmaShapeA_MK{})); + using BlockTileB_N = decltype(cute::size<0, 0>(MmaShapeB_NK{}) * + cute::size<1>(MmaShapeB_NK{})); + using BlockTileB_K = decltype(cute::size<0, 1>(MmaShapeB_NK{}) * + cute::size<2>(MmaShapeB_NK{})); + + using SmemLayoutAtomACompute = + decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + UmmaMajorA, MmaType, BlockTileA_M, BlockTileA_K>()); + using SmemLayoutAtomPairA = + cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType< + SmemLayoutAtomACompute, SmemLayoutAtomACompute>; + using CopyAtomPairA = + cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType< + Copy_Atom, uint8_t>, + Copy_Atom, MmaType>>; + + using SmemLayoutAtomB = + decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + UmmaMajorB, MmaType, BlockTileB_N, BlockTileB_K>()); + using SmemLayoutAtomPairB = + cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType< + SmemLayoutAtomB, SmemLayoutAtomB>; + using CopyAtomPairB = + cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType< + Copy_Atom, MmaType>, + Copy_Atom, MmaType>>; + + // Pipeline stage counts, derived from the smem budget. + static constexpr int kSmemCapacity = + cutlass::gemm::collective::detail::sm100_smem_capacity_bytes; + static constexpr int kKernelCarveout = 2048; + static constexpr int kEpilogueBytes = + int(sizeof(typename CollectiveEpilogue::SharedStorage)); + // TMEM is 512 columns and an accumulator stage needs kTileN of them. + static constexpr int kAccumStages = 512 / kTileN; + // B stage is the CTA's half of the N-split activation tile, kTileN/2 rows. + static constexpr int kInputStageBytes = (kWBits == 8 ? 16384 : 8192) + + kTileN * 128 + 256 + 64 + + (kHasZp ? 256 : 0); + static constexpr int kComputeStageBytes = 32768 + 32; + static constexpr int kT2MStages = 2; + static constexpr int kAvail = kSmemCapacity - kKernelCarveout - + kEpilogueBytes - kAccumStages * 32 - + kT2MStages * kComputeStageBytes; + static constexpr int kL2TStages = + kAvail / kInputStageBytes < 4 ? kAvail / kInputStageBytes : 4; + static_assert(kL2TStages >= 2, "not enough SMEM for two input stages"); + + using CollectiveMainloop = + cutlass::gemm::collective::SwordfishMainloopSm100MixedInput< + kL2TStages, kT2MStages, /*SchedulerStages=*/3, kAccumStages, + ClusterShape, MmaTileShape, ElementPairA, StridePairA, MmaType, + StrideB, TiledMma, cute::SM90_TMA_LOAD, SmemLayoutAtomPairA, + CopyAtomPairA, cute::identity, + // Activations must use the cta_group::2 TMA. The 2x1SM atom N-splits + // B across the CTA pair and the leader's MMA reads the peer's half, + // so the copy must deliver both halves with one arrival on the + // leader's barrier. A per-CTA TMA arrives at local barriers and + // races. + cute::SM100_TMA_2SM_LOAD, SmemLayoutAtomPairB, CopyAtomPairB, + cute::identity, kWBits>; + + // Forked kernel layer whose warp layout derives from the collective's + // NumTransformationThreads. + using GemmKernel = cutlass::gemm::kernel::SwordfishPrefillKernel< + Shape, CollectiveMainloop, CollectiveEpilogue, void>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +}; + +template +void run(torch::stable::Tensor& a, torch::stable::Tensor& b_packed, + torch::stable::Tensor& group_scales, const void* zp_ptr, + torch::stable::Tensor& c, int M, int N, int K, cudaStream_t stream) { + using Gemm = typename Cfg::Gemm; + using GemmKernel = typename Cfg::GemmKernel; + using MmaType = typename Cfg::MmaType; + using LayoutScale = typename Cfg::LayoutScale; + using ScaleConfig = typename Cfg::ScaleConfig; + + // L2-aware M chunking. The activation stream loses all L2 reuse once a + // chunk outgrows Thor's 32 MB L2, a two-thirds throughput loss, so the + // per-launch activation footprint is capped near 12 MB. + constexpr int64_t kAChunkBytes = int64_t(12) << 20; + int m_chunk = int(kAChunkBytes / (int64_t(K) * 2)); + m_chunk = std::max(256, (m_chunk / 128) * 128); + + LayoutScale layout_S = + ScaleConfig::tile_atom_to_shape_scale(cute::make_shape(N, K, 1)); + auto* c_base = reinterpret_cast(c.mutable_data_ptr()); + const auto* a_base = reinterpret_cast(a.const_data_ptr()); + + torch::stable::Tensor workspace; + size_t ws_alloc = 0; + + for (int m0 = 0; m0 < M; m0 += m_chunk) { + const int mc = std::min(m_chunk, M - m0); + + auto stride_act = + cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(mc, K, 1)); + auto stride_c = cutlass::make_cute_packed_stride( + typename GemmKernel::StrideC{}, cute::make_shape(N, mc, 1)); + auto stride_d = cutlass::make_cute_packed_stride( + typename GemmKernel::StrideD{}, cute::make_shape(N, mc, 1)); + + MmaType* c_ptr = c_base + int64_t(m0) * N; + + // Swapped problem shape (N, mc, K). The packed weight rides the A slot. + typename Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + {N, mc, K, 1}, + {reinterpret_cast(b_packed.const_data_ptr()), StrideA{}, + a_base + int64_t(m0) * K, stride_act, + reinterpret_cast(group_scales.const_data_ptr()), + layout_S, reinterpret_cast(zp_ptr)}, + {{1.0f, 0.0f}, c_ptr, stride_c, c_ptr, stride_d}}; + + Gemm gemm; + STD_TORCH_CHECK(gemm.can_implement(arguments) == cutlass::Status::kSuccess, + "swordfish_prefill_mm: unsupported problem"); + + const size_t ws_bytes = Gemm::get_workspace_size(arguments); + if (ws_bytes > ws_alloc) { + workspace = torch::stable::empty({int64_t(ws_bytes)}, + torch::headeronly::ScalarType::Byte, + std::nullopt, a.device()); + ws_alloc = ws_bytes; + } + + STD_TORCH_CHECK( + gemm.initialize(arguments, + ws_alloc ? workspace.mutable_data_ptr() : nullptr, + stream) == cutlass::Status::kSuccess, + "swordfish_prefill_mm: initialize failed"); + STD_TORCH_CHECK(gemm.run(stream) == cutlass::Status::kSuccess, + "swordfish_prefill_mm: launch failed"); + } +} + +// All prefill configurations for one activation dtype, one TU each so the +// fp16 and bf16 sets compile in parallel. +template +void run_prefill_all(torch::stable::Tensor& a, torch::stable::Tensor& b_packed, + torch::stable::Tensor& group_scales, const void* zp_ptr, + bool has_zp, bool w8, int gran, torch::stable::Tensor& c, + int M, int N, int K, cudaStream_t stream) { + // Tile-N dispatch. The 256-wide tile wins compute-bound shapes; K-heavy + // shapes starve its K pipeline and prefer 256x128 (measured on both + // Thor and B200 at K=14336). The 256-wide tile's input stages do not fit + // SMEM at 8 bits, and the doubled weight stream pressures the K pipeline + // the way K-heavy shapes do, so 8-bit always runs 256x128. + const bool narrow = w8 || (K >= 2 * N && K >= 8192); + if (w8) { + run>(a, b_packed, group_scales, nullptr, c, + M, N, K, stream); + } else if (gran == 32) { + if (has_zp) { + if (narrow) { + run>(a, b_packed, group_scales, + zp_ptr, c, M, N, K, stream); + } else { + run>(a, b_packed, group_scales, + zp_ptr, c, M, N, K, stream); + } + } else if (narrow) { + run>(a, b_packed, group_scales, + nullptr, c, M, N, K, stream); + } else { + run>(a, b_packed, group_scales, + nullptr, c, M, N, K, stream); + } + } else if (gran == 64) { + if (has_zp) { + if (narrow) { + run>(a, b_packed, group_scales, + zp_ptr, c, M, N, K, stream); + } else { + run>(a, b_packed, group_scales, + zp_ptr, c, M, N, K, stream); + } + } else if (narrow) { + run>(a, b_packed, group_scales, + nullptr, c, M, N, K, stream); + } else { + run>(a, b_packed, group_scales, + nullptr, c, M, N, K, stream); + } + } else { + if (has_zp) { + if (narrow) { + run>(a, b_packed, group_scales, zp_ptr, + c, M, N, K, stream); + } else { + run>(a, b_packed, group_scales, zp_ptr, + c, M, N, K, stream); + } + } else if (narrow) { + run>(a, b_packed, group_scales, nullptr, + c, M, N, K, stream); + } else { + run>(a, b_packed, group_scales, nullptr, + c, M, N, K, stream); + } + } +} + +} // namespace prefill + +#endif // CUTLASS_ARCH_MMA_SM100_SUPPORTED + +} // namespace swordfish diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_kernel.cuh b/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_kernel.cuh new file mode 100644 index 0000000000..641ac89182 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_kernel.cuh @@ -0,0 +1,1169 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + * this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + * POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.hpp" +#include "cutlass/detail/cluster.hpp" +#include "cutlass/arch/grid_dependency_control.h" +#include "cutlass/fast_math.h" +#include "cute/arch/cluster_sm90.hpp" +#include "cutlass/arch/arch.h" +#include "cutlass/arch/reg_reconfig.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/kernel/sm100_tile_scheduler.hpp" +#include "cutlass/pipeline/pipeline.hpp" + +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" +/////////////////////////////////////////////////////////////////////////////// +// SWORDFISH FORK of CUTLASS 4.4.2's +// include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_mixed_input_transform.hpp +// (BSD-3-Clause, Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES). +// One change: the warp-category boundaries derive from the collective's +// NumTransformationThreads instead of hardcoding 4 transform warps, so the +// Transform (dequant producer) stage can be widened for tuning +// experiments. +// Instantiated DIRECTLY (standalone class, no GemmUniversal dispatch). +/////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::kernel { + +/////////////////////////////////////////////////////////////////////////////// + +template +class SwordfishPrefillKernel { + public: + // + // Type Aliases + // + using ProblemShape = ProblemShape_; + static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4, + "ProblemShape{} should be or "); + static constexpr bool IsGdcEnabled = cutlass::arch::IsGdcGloballyEnabled; + + // Mainloop derived types + using CollectiveMainloop = CollectiveMainloop_; + using TileShape = typename CollectiveMainloop::TileShape; + + // Get Blk and Scheduling tile shapes + using CtaShape_MNK = typename CollectiveMainloop::CtaShape_MNK; + using AtomThrShapeMNK = typename CollectiveMainloop::AtomThrShapeMNK; + + using TiledMma = typename CollectiveMainloop::TiledMma; + using ArchTag = typename CollectiveMainloop::ArchTag; + using ElementA = typename CollectiveMainloop::ElementA; + using StrideA = typename CollectiveMainloop::StrideA; + using ElementB = typename CollectiveMainloop::ElementB; + using StrideB = typename CollectiveMainloop::StrideB; + using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy; + using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator; + using ClusterShape = typename DispatchPolicy::ClusterShape; + using MainloopArguments = typename CollectiveMainloop::Arguments; + using MainloopParams = typename CollectiveMainloop::Params; + static constexpr bool IsComplex = + DispatchPolicy::InputTransformType == + cutlass::gemm::detail::KernelInputTransformType::InterleavedComplexTF32; + static_assert(ArchTag::kMinComputeCapability >= 100); + + // Epilogue derived types + using CollectiveEpilogue = CollectiveEpilogue_; + using ElementC = typename CollectiveEpilogue::ElementC; + using StrideC = typename CollectiveEpilogue::StrideC; + using ElementD = typename CollectiveEpilogue::ElementD; + using StrideD = typename CollectiveEpilogue::StrideD; + using EpilogueArguments = typename CollectiveEpilogue::Arguments; + using EpilogueParams = typename CollectiveEpilogue::Params; + + // CLC pipeline depth + // determines how many waves (stages-1) a warp can race ahead + static constexpr uint32_t SchedulerPipelineStageCount = + DispatchPolicy::Schedule::SchedulerPipelineStageCount; + // TileID scheduler + using TileSchedulerTag = TileScheduler_; + using TileScheduler = typename detail::TileSchedulerSelector< + TileScheduler_, ArchTag, CtaShape_MNK, ClusterShape, + SchedulerPipelineStageCount>::Scheduler; + using TileSchedulerArguments = typename TileScheduler::Arguments; + using TileSchedulerParams = typename TileScheduler::Params; + + static constexpr bool IsDynamicCluster = not cute::is_static_v; + + // Warp specialization thread count per threadblock + static constexpr uint32_t NumSchedThreads = NumThreadsPerWarp; // 1 warp + static constexpr uint32_t NumMMAThreads = NumThreadsPerWarp; // 1 warp + static constexpr uint32_t NumMainloopLoadThreads = + NumThreadsPerWarp; // 1 warp + static constexpr uint32_t NumEpilogueLoadThreads = + NumThreadsPerWarp; // 1 warp + static constexpr uint32_t NumEpilogueThreads = + CollectiveMainloop::NumAccumThreads; // 4 warps + static constexpr uint32_t NumEpilogueWarps = + NumEpilogueThreads / NumThreadsPerWarp; + static constexpr uint32_t NumTransformationThreads = + CollectiveMainloop::NumTransformationThreads; // 4 warps + static constexpr uint32_t NumMainloopLoadBThreads = + NumThreadsPerWarp; // 1 warp + + static constexpr uint32_t MaxThreadsPerBlock = + NumSchedThreads + NumMainloopLoadThreads + NumMMAThreads + + NumEpilogueLoadThreads + NumEpilogueThreads + NumTransformationThreads + + NumMainloopLoadBThreads; + static constexpr uint32_t MinBlocksPerMultiprocessor = 1; + + static constexpr uint32_t AccumulatorPipelineStageCount = + DispatchPolicy::Schedule::AccumulatorPipelineStageCount; + static constexpr cutlass::gemm::detail::KernelInputTransformType + InputTransformType = DispatchPolicy::InputTransformType; + static constexpr uint32_t NumFixupBarriers = 1; + static constexpr uint32_t CLCResponseSize = + sizeof(typename TileScheduler::CLCResponse); + + static constexpr bool IsSchedDynamicPersistent = + TileScheduler::IsDynamicPersistent; + + // Pipeline and pipeline state types + using Load2TransformPipeline = + typename CollectiveMainloop::Load2TransformPipeline; + using Load2TransformPipelineState = + typename CollectiveMainloop::Load2TransformPipelineState; + + using Load2MmaPipeline = typename CollectiveMainloop::Load2MmaPipeline; + using Load2MmaPipelineState = + typename CollectiveMainloop::Load2MmaPipelineState; + + using Transform2MmaPipeline = + typename CollectiveMainloop::Transform2MmaPipeline; + using Transform2MmaPipelineState = + typename CollectiveMainloop::Transform2MmaPipelineState; + + using Mma2AccumPipeline = typename CollectiveMainloop::Mma2AccumPipeline; + using Mma2AccumPipelineState = + typename CollectiveMainloop::Mma2AccumPipelineState; + + using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline; + using EpiLoadPipelineState = typename CollectiveEpilogue::LoadPipelineState; + + using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline; + using EpiStorePipelineState = typename CollectiveEpilogue::StorePipelineState; + + using LoadOrderBarrier = cutlass::OrderedSequenceBarrier<1, 2>; + + using CLCPipeline = + cutlass::PipelineCLCFetchAsync; + using CLCPipelineState = cutlass::PipelineState; + + using CLCThrottlePipeline = + cutlass::PipelineAsync; + using CLCThrottlePipelineState = typename CLCThrottlePipeline::PipelineState; + + using TmemAllocator = + cute::conditional_t( + typename TiledMma::ThrLayoutVMNK{})) == 1, + cute::TMEM::Allocator1Sm, cute::TMEM::Allocator2Sm>; + + // Kernel level shared memory storage + struct SharedStorage { + struct PipelineStorage : cute::aligned_struct<16, _1> { + using MainloopPipelineStorage = + typename CollectiveMainloop::PipelineStorage; + using EpiLoadPipelineStorage = + typename CollectiveEpilogue::PipelineStorage; + using LoadOrderBarrierStorage = typename LoadOrderBarrier::SharedStorage; + using CLCPipelineStorage = typename CLCPipeline::SharedStorage; + using CLCThrottlePipelineStorage = + typename CLCThrottlePipeline::SharedStorage; + + alignas(16) MainloopPipelineStorage mainloop; + alignas(16) EpiLoadPipelineStorage epi_load; + alignas(16) LoadOrderBarrierStorage load_order; + alignas(16) CLCPipelineStorage clc; + alignas(16) CLCThrottlePipelineStorage clc_throttle; + alignas(16) arch::ClusterBarrier tmem_dealloc; + alignas(16) arch::ClusterBarrier epilogue_throttle; + } pipelines; + + alignas(16) typename TileScheduler::CLCResponse + clc_response[SchedulerPipelineStageCount]; + uint32_t tmem_base_ptr; + + struct TensorStorage : cute::aligned_struct<128, _1> { + using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage; + using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage; + + EpilogueTensorStorage epilogue; + MainloopTensorStorage mainloop; + } tensors; + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + + // Device side arguments + struct Arguments { + GemmUniversalMode mode{}; + ProblemShape problem_shape{}; + MainloopArguments mainloop{}; + EpilogueArguments epilogue{}; + KernelHardwareInfo hw_info{}; + TileSchedulerArguments scheduler{}; + }; + + // Kernel entry point API + struct Params { + GemmUniversalMode mode{}; + ProblemShape problem_shape{}; + MainloopParams mainloop{}; + EpilogueParams epilogue{}; + TileSchedulerParams scheduler{}; + KernelHardwareInfo hw_info{}; + }; + + enum class WarpCategory : int32_t { + MMA = 0, + Sched = 1, + MainloopLoad = 2, + EpilogueLoad = 3, + Epilogue = 4, + // Transformation starts at 256 thread alignment + Transformation = 8, + // SWORDFISH: derived from the collective (stock hardcoded 12 = 4 warps) + MainloopLoadB = 8 + int32_t(CollectiveMainloop::NumTransformationThreads / + NumThreadsPerWarp), + }; + + struct IsParticipant { + uint32_t mma = false; + uint32_t sched = false; + uint32_t main_load = false; + uint32_t main_loadA = false; + uint32_t main_loadB = false; + uint32_t epi_load = false; + uint32_t epilogue = false; + uint32_t transformation = false; + }; + + // + // Methods + // + + // Convert to underlying arguments. In this case, a simple copy for the + // aliased type. + static Params to_underlying_arguments(Arguments const& args, + void* workspace) { + static constexpr uint32_t NumEpilogueSubTiles = 1; + auto problem_shape = args.problem_shape; + if constexpr (detail::Has_SwapAB_v) { + // swap M/N + get<0>(problem_shape) = get<1>(args.problem_shape); + get<1>(problem_shape) = get<0>(args.problem_shape); + } + auto problem_shape_MNKL = append<4>(problem_shape, 1); + + // Get SM count if needed, otherwise use user supplied SM count + int sm_count = args.hw_info.sm_count; + if (sm_count <= 0) { + CUTLASS_TRACE_HOST( + " WARNING: Arguments do not include a valid SM count.\n" + " For optimal performance, populate the arguments " + "KernelHardwareInfo struct with the SM count."); + sm_count = KernelHardwareInfo::query_device_multiprocessor_count( + args.hw_info.device_id); + } + + CUTLASS_TRACE_HOST( + "to_underlying_arguments(): Setting persistent grid SM count to " + << sm_count); + // Calculate workspace pointers + uint8_t* workspace_ptr = reinterpret_cast(workspace); + size_t workspace_offset = 0; + + // Epilogue + void* epilogue_workspace = workspace_ptr + workspace_offset; + workspace_offset += CollectiveEpilogue::get_workspace_size( + args.problem_shape, args.epilogue); + workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); + + void* mainloop_workspace = nullptr; + + // Tile scheduler + void* scheduler_workspace = workspace_ptr + workspace_offset; + workspace_offset += + TileScheduler::template get_workspace_size( + args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers, + NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs); + workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); + + return { + args.mode, + args.problem_shape, + CollectiveMainloop::to_underlying_arguments( + args.problem_shape, args.mainloop, mainloop_workspace, + args.hw_info), + CollectiveEpilogue::to_underlying_arguments( + args.problem_shape, args.epilogue, epilogue_workspace), + TileScheduler::to_underlying_arguments( + problem_shape_MNKL, TileShape{}, AtomThrShapeMNK{}, ClusterShape{}, + args.hw_info, args.scheduler, scheduler_workspace), + args.hw_info}; + } + + static bool can_implement(Arguments const& args) { + bool implementable = + (args.mode == GemmUniversalMode::kGemm) or + (args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4); + if (!implementable) { + CUTLASS_TRACE_HOST( + " CAN IMPLEMENT: Arguments or Problem Shape don't meet the " + "requirements.\n"); + return implementable; + } + implementable &= + CollectiveMainloop::can_implement(args.problem_shape, args.mainloop); + implementable &= + CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue); + implementable &= TileScheduler::can_implement(args.scheduler); + + if constexpr (IsDynamicCluster) { + static constexpr int MaxClusterSize = 16; + implementable &= size(args.hw_info.cluster_shape) <= MaxClusterSize; + implementable &= + size(args.hw_info.cluster_shape_fallback) <= MaxClusterSize; + implementable &= + cutlass::detail::preferred_cluster_can_implement( + args.hw_info.cluster_shape, args.hw_info.cluster_shape_fallback); + } + + return implementable; + } + + static size_t get_workspace_size(Arguments const& args) { + static constexpr uint32_t NumEpilogueSubTiles = 1; + size_t workspace_size = 0; + + // Epilogue + workspace_size += CollectiveEpilogue::get_workspace_size(args.problem_shape, + args.epilogue); + workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment); + + // Tile scheduler + workspace_size += + TileScheduler::template get_workspace_size( + args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers, + NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs); + workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment); + + return workspace_size; + } + + static cutlass::Status initialize_workspace( + Arguments const& args, void* workspace = nullptr, + cudaStream_t stream = nullptr, CudaHostAdapter* cuda_adapter = nullptr) { + Status status = Status::kSuccess; + uint8_t* workspace_ptr = reinterpret_cast(workspace); + size_t workspace_offset = 0; + static constexpr uint32_t NumEpilogueSubTiles = 1; + + // Epilogue + status = CollectiveEpilogue::initialize_workspace( + args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, + stream, cuda_adapter); + workspace_offset += CollectiveEpilogue::get_workspace_size( + args.problem_shape, args.epilogue); + workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); + if (status != Status::kSuccess) { + return status; + } + + // Tile scheduler + status = TileScheduler::template initialize_workspace( + args.scheduler, workspace_ptr + workspace_offset, stream, + args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, + CollectiveEpilogue::NumAccumulatorMtxs, cuda_adapter); + workspace_offset += + TileScheduler::template get_workspace_size( + args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers, + NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs); + workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); + if (status != Status::kSuccess) { + return status; + } + + return status; + } + + // Computes the kernel launch grid shape based on runtime parameters + static dim3 get_grid_shape(Params const& params) { + auto cluster_shape = cutlass::detail::select_cluster_shape( + ClusterShape{}, params.hw_info.cluster_shape); + auto blk_shape = CtaShape_MNK{}; + auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{}); + return TileScheduler::get_grid_shape(params.scheduler, problem_shape_MNKL, + TileShape{}, AtomThrShapeMNK{}, + cluster_shape, params.hw_info); + } + + static dim3 get_block_shape() { return dim3(MaxThreadsPerBlock, 1, 1); } + + CUTLASS_DEVICE + void operator()(Params const& params, char* smem_buf) { + using namespace cute; + using X = Underscore; + + static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, + "SMEM usage exceeded capacity."); + // Separate out problem shape for convenience + // Optionally append 1s until problem shape is rank-4 in case its is only + // rank-3 (MNK) + auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto K = get<2>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + // Account for multiple epilogue and transformation warps + int warp_idx = canonical_warp_idx_sync(); + WarpCategory warp_category = + warp_idx < static_cast(WarpCategory::Epilogue) + ? WarpCategory(warp_idx) + : warp_idx < static_cast(WarpCategory::Transformation) + ? WarpCategory::Epilogue + : warp_idx < static_cast(WarpCategory::MainloopLoadB) + ? WarpCategory::Transformation + : WarpCategory::MainloopLoadB; + + int thread_idx = int(threadIdx.x); + int thread_idx_in_warp = thread_idx % 32; + uint32_t lane_predicate = cute::elect_one_sync(); + int cta_rank_in_cluster = cute::block_rank_in_cluster(); + auto cluster_shape = cutlass::detail::select_cluster_shape( + ClusterShape{}, cute::cluster_shape()); + int cluster_size = size(cluster_shape); + bool is_first_cta_in_cluster = (cta_rank_in_cluster == 0); + bool is_mma_leader_cta = (cta_rank_in_cluster % size<0>(TiledMma{}) == 0); + // Even if this variable is unused, shape_div still performs useful + // compile-time checks. + [[maybe_unused]] auto mma_leader_ctas = + size(shape_div(cluster_shape, AtomThrShapeMNK{})); + constexpr bool has_mma_peer_cta = size(AtomThrShapeMNK{}) == 2; + uint32_t mma_peer_cta_rank = + has_mma_peer_cta ? cta_rank_in_cluster ^ 1 : cta_rank_in_cluster; + + // Issue Tma Descriptor Prefetch from a single thread + if ((warp_category == WarpCategory::Sched) && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + } + if ((warp_category == WarpCategory::EpilogueLoad) && lane_predicate) { + CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue); + } + + // Kernel level shared memory storage + SharedStorage& shared_storage = *reinterpret_cast(smem_buf); + + CollectiveMainloop collective_mainloop(params.mainloop, cluster_shape, + cta_rank_in_cluster); + CollectiveEpilogue collective_epilogue{params.epilogue, + shared_storage.tensors.epilogue}; + + bool is_epi_load_needed = collective_epilogue.is_producer_load_needed(); + IsParticipant is_participant = { + (warp_category == WarpCategory::MMA), // mma + (warp_category == WarpCategory::Sched) && + (is_first_cta_in_cluster), // sched + (warp_category == WarpCategory::MainloopLoad || + warp_category == WarpCategory::MainloopLoadB), // main_load + (warp_category == WarpCategory::MainloopLoad), // main_loadA + (warp_category == WarpCategory::MainloopLoadB), // main_loadB + (warp_category == WarpCategory::EpilogueLoad) && + is_epi_load_needed, // epi_load + (warp_category == WarpCategory::Epilogue), // epilogue + (warp_category == WarpCategory::Transformation) // transformation + }; + + // MainloopLoad <--> Transformation Pipeline + typename Load2TransformPipeline::Params load2transform_pipeline_params; + if (warp_category == WarpCategory::MainloopLoad) { + load2transform_pipeline_params.role = + Load2TransformPipeline::ThreadCategory::Producer; + } else if (warp_category == WarpCategory::Transformation) { + load2transform_pipeline_params.role = + Load2TransformPipeline::ThreadCategory::Consumer; + } + load2transform_pipeline_params.is_leader = (thread_idx_in_warp == 0); + load2transform_pipeline_params.num_consumers = NumTransformationThreads; + load2transform_pipeline_params.transaction_bytes = + CollectiveMainloop::TmaTransactionBytes_A; + load2transform_pipeline_params.initializing_warp = 0; + Load2TransformPipeline load2transform_pipeline( + shared_storage.pipelines.mainloop.load2transform_pipeline, + load2transform_pipeline_params, cluster_shape, McastDirection::kRow, + cute::true_type{}, // Perform barrier init + cute::false_type{} // Delay mask calculation + ); + + Load2TransformPipelineState load2transform_pipeline_consumer_state; + Load2TransformPipelineState load2transform_pipeline_producer_state = + cutlass::make_producer_start_state(); + + // MainloopLoad <--> MMA Pipeline + typename Load2MmaPipeline::Params load2mma_pipeline_params; + if (warp_category == WarpCategory::MainloopLoadB) { + load2mma_pipeline_params.role = + Load2MmaPipeline::ThreadCategory::Producer; + } else if (warp_category == WarpCategory::MMA) { + load2mma_pipeline_params.role = + Load2MmaPipeline::ThreadCategory::Consumer; + } + load2mma_pipeline_params.is_leader = + lane_predicate && is_mma_leader_cta && is_participant.main_loadB; + load2mma_pipeline_params.num_consumers = NumMMAThreads; + load2mma_pipeline_params.transaction_bytes = + CollectiveMainloop::TmaTransactionBytes_B; + load2mma_pipeline_params.initializing_warp = 8; + Load2MmaPipeline load2mma_pipeline( + shared_storage.pipelines.mainloop.load2mma_pipeline, + load2mma_pipeline_params, cluster_shape, McastDirection::kCol, + cute::true_type{}, // Perform barrier init + cute::false_type{} // Delay mask calculation + ); + + Load2MmaPipelineState load2mma_pipeline_consumer_state; + Load2MmaPipelineState load2mma_pipeline_producer_state = + cutlass::make_producer_start_state(); + + // Transformation <--> MMA pipeline + typename Transform2MmaPipeline::Params transform2mma_pipeline_params; + if (warp_category == WarpCategory::Transformation) { + transform2mma_pipeline_params.role = + Transform2MmaPipeline::ThreadCategory::Producer; + } else if (warp_category == WarpCategory::MMA) { + transform2mma_pipeline_params.role = + Transform2MmaPipeline::ThreadCategory::Consumer; + } + transform2mma_pipeline_params.consumer_arv_count = 1; + transform2mma_pipeline_params.producer_arv_count = + size(AtomThrShapeMNK{}) * NumTransformationThreads; + transform2mma_pipeline_params.initializing_warp = 2; + Transform2MmaPipeline transform2mma_pipeline( + shared_storage.pipelines.mainloop.transform2mma_pipeline, + transform2mma_pipeline_params, cluster_shape, + cute::true_type{}, // Perform barrier init + cute::false_type{} // Delay mask calculation + ); + + Transform2MmaPipelineState transform2mma_pipeline_consumer_state; + Transform2MmaPipelineState transform2mma_pipeline_producer_state = + cutlass::make_producer_start_state(); + + // MMA <--> Accumulator pipeline + typename Mma2AccumPipeline::Params mma2accum_pipeline_params; + if (warp_category == WarpCategory::MMA) { + mma2accum_pipeline_params.role = + Mma2AccumPipeline::ThreadCategory::Producer; + } else if (warp_category == WarpCategory::Epilogue) { + mma2accum_pipeline_params.role = + Mma2AccumPipeline::ThreadCategory::Consumer; + } + mma2accum_pipeline_params.producer_arv_count = 1; + mma2accum_pipeline_params.consumer_arv_count = + size(AtomThrShapeMNK{}) * NumEpilogueThreads; + mma2accum_pipeline_params.initializing_warp = 6; + Mma2AccumPipeline mma2accum_pipeline( + shared_storage.pipelines.mainloop.mma2accum_pipeline, + mma2accum_pipeline_params, cluster_shape, + cute::true_type{}, // Perform barrier init + cute::false_type{} // Delay mask calculation + ); + + Mma2AccumPipelineState mma2accum_pipeline_consumer_state; + Mma2AccumPipelineState mma2accum_pipeline_producer_state = + cutlass::make_producer_start_state(); + + // Epilogue Load pipeline + typename EpiLoadPipeline::Params epi_load_pipeline_params; + if (WarpCategory::EpilogueLoad == warp_category) { + epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer; + } + if (WarpCategory::Epilogue == warp_category) { + epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer; + } + epi_load_pipeline_params.dst_blockid = cta_rank_in_cluster; + epi_load_pipeline_params.producer_arv_count = NumEpilogueLoadThreads; + epi_load_pipeline_params.consumer_arv_count = NumEpilogueThreads; + epi_load_pipeline_params.transaction_bytes = + CollectiveEpilogue::TmaTransactionBytes; + epi_load_pipeline_params.initializing_warp = 4; + EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, + epi_load_pipeline_params); + + // Epilogue Store pipeline + typename EpiStorePipeline::Params epi_store_pipeline_params; + epi_store_pipeline_params.always_wait = true; + EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); + + // Load order barrier + typename LoadOrderBarrier::Params load_order_barrier_params; + load_order_barrier_params.group_id = + (warp_category == WarpCategory::MainloopLoad) ? 0 : 1; + load_order_barrier_params.group_size = 1; + load_order_barrier_params.initializing_warp = 5; + LoadOrderBarrier load_order_barrier(shared_storage.pipelines.load_order, + load_order_barrier_params); + + EpiLoadPipelineState epi_load_pipe_consumer_state; + EpiLoadPipelineState epi_load_pipe_producer_state = + cutlass::make_producer_start_state(); + + // epilogue store pipe is producer-only (consumer is TMA unit, waits via + // scoreboarding) + EpiStorePipelineState epi_store_pipe_producer_state = + cutlass::make_producer_start_state(); + + // CLC pipeline + // Operates Scheduling Warp <--> All Warps + typename CLCPipeline::Params clc_pipeline_params; + if (WarpCategory::Sched == warp_category) { + clc_pipeline_params.role = CLCPipeline::ThreadCategory::ProducerConsumer; + } else { + clc_pipeline_params.role = CLCPipeline::ThreadCategory::Consumer; + } + clc_pipeline_params.producer_blockid = 0; + clc_pipeline_params.producer_arv_count = 1; + clc_pipeline_params.consumer_arv_count = + NumSchedThreads + + cluster_size * + (NumMainloopLoadThreads + NumMainloopLoadBThreads + + NumEpilogueThreads + NumMMAThreads + NumTransformationThreads); + if (is_epi_load_needed) { + clc_pipeline_params.consumer_arv_count += + cluster_size * NumEpilogueLoadThreads; + } + clc_pipeline_params.transaction_bytes = CLCResponseSize; + clc_pipeline_params.initializing_warp = 1; + CLCPipeline clc_pipeline(shared_storage.pipelines.clc, clc_pipeline_params, + cluster_shape); + + CLCPipelineState clc_pipeline_consumer_state; + CLCPipelineState clc_pipeline_producer_state = + cutlass::make_producer_start_state(); + + // CLC throttle pipeline + typename CLCThrottlePipeline::Params clc_throttle_pipeline_params; + if (WarpCategory::MainloopLoad == warp_category) { + clc_throttle_pipeline_params.role = + CLCThrottlePipeline::ThreadCategory::Producer; + } + if (WarpCategory::Sched == warp_category) { + clc_throttle_pipeline_params.role = + CLCThrottlePipeline::ThreadCategory::Consumer; + } + clc_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads; + clc_throttle_pipeline_params.consumer_arv_count = NumSchedThreads; + clc_throttle_pipeline_params.dst_blockid = 0; + clc_throttle_pipeline_params.initializing_warp = 3; + CLCThrottlePipeline clc_throttle_pipeline( + shared_storage.pipelines.clc_throttle, clc_throttle_pipeline_params); + CLCThrottlePipelineState clc_pipe_throttle_consumer_state; + CLCThrottlePipelineState clc_pipe_throttle_producer_state = + cutlass::make_producer_start_state(); + + // Tmem allocator + TmemAllocator tmem_allocator{}; + + // Sync allocation status between transform, MMA, and epilogue warps within + // CTA + arch::NamedBarrier tmem_allocation_result_barrier( + NumTransformationThreads + NumMMAThreads + NumEpilogueThreads, + cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); + // Sync deallocation status between MMA warps of peer CTAs + arch::ClusterBarrier& tmem_deallocation_result_barrier = + shared_storage.pipelines.tmem_dealloc; + [[maybe_unused]] uint32_t dealloc_barrier_phase = 0; + if (WarpCategory::MMA == warp_category && has_mma_peer_cta && + lane_predicate) { + tmem_deallocation_result_barrier.init(NumMMAThreads); + } + + // Initialize smem barrier for prologue throttling. Epilogue warps are + // stalled until the prologue finishes. + arch::ClusterBarrier& epilogue_throttle_barrier = + shared_storage.pipelines.epilogue_throttle; + if (WarpCategory::MMA == warp_category && lane_predicate) { + epilogue_throttle_barrier.init( + NumMMAThreads + (is_first_cta_in_cluster ? NumSchedThreads : 0) + + NumMainloopLoadThreads + NumMainloopLoadBThreads + + (is_epi_load_needed ? NumEpilogueLoadThreads : 0) + + NumTransformationThreads); + } + + // We need this to guarantee that the Pipeline init is visible + // To all producers and consumer threadblocks in the cluster + pipeline_init_arrive_relaxed(cluster_size); + + dim3 block_id_in_cluster = cute::block_id_in_cluster(); + + // Calculate mask after cluster barrier arrival + load2transform_pipeline.init_masks(cluster_shape, block_id_in_cluster, + cutlass::McastDirection::kRow); + load2mma_pipeline.init_masks(cluster_shape, cutlass::McastDirection::kCol); + transform2mma_pipeline.init_masks(cluster_shape); + mma2accum_pipeline.init_masks(cluster_shape); + + // TileID scheduler + TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, + block_id_in_cluster); + typename TileScheduler::WorkTileInfo work_tile_info = + scheduler.initial_work_tile_info(cluster_shape); + + auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info); + + // Allocate accumulators + auto acc_shape = collective_mainloop.partition_accumulator_shape(); + auto bulk_tmem = TiledMma::make_fragment_C( + append(acc_shape, Int{})); + + // Tile transform inputs now to get the k tile count + auto transform_inputs = collective_mainloop.transform_init( + params.mainloop, problem_shape_MNKL, bulk_tmem, + shared_storage.tensors.mainloop); + Tensor gA_mkl = get<0>(transform_inputs); + + // Synchronization call. Blocks wait until barriers are initialized in + // shared memory. + pipeline_init_wait(cluster_size); + + if (is_participant.main_load) { + // Ensure that the prefetched kernel does not touch + // unflushed global memory prior to this instruction + cutlass::arch::wait_on_dependent_grids(); + + bool do_load_order_arrive = is_epi_load_needed; + auto load_inputs = collective_mainloop.load_init( + problem_shape_MNKL, params.mainloop, shared_storage.tensors.mainloop); + + // Signal the epilogue warps to proceed once the prologue is complete + epilogue_throttle_barrier.arrive(); + bool requires_clc_query = true; + + do { + cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info); + auto k_tile_iter = + scheduler.get_k_tile_iterator(work_tile_info, problem_shape_MNKL, + CtaShape_MNK{}, shape<3>(gA_mkl)); + auto k_tile_count = TileScheduler::get_work_k_tile_count( + work_tile_info, problem_shape_MNKL, CtaShape_MNK{}); + auto k_tile_prologue = + min(Load2TransformPipeline::Stages, k_tile_count); + + if (is_participant.main_loadA) { + if constexpr (IsSchedDynamicPersistent) { + if (is_first_cta_in_cluster && requires_clc_query) { + clc_throttle_pipeline.producer_acquire( + clc_pipe_throttle_producer_state); + clc_throttle_pipeline.producer_commit( + clc_pipe_throttle_producer_state); + ++clc_pipe_throttle_producer_state; + } + } + } + + if (lane_predicate) { + if (is_participant.main_loadA) { + auto [load2transform_pipeline_producer_state_next, + k_tile_iter_next] = + collective_mainloop.load_A( + params.mainloop, load2transform_pipeline, + load2transform_pipeline_producer_state, load_inputs, + cta_coord_mnkl, k_tile_iter, k_tile_prologue); + load2transform_pipeline_producer_state = + load2transform_pipeline_producer_state_next; + + if (do_load_order_arrive) { + load_order_barrier.arrive(); + do_load_order_arrive = false; + } + + auto [load2transform_pipeline_producer_state_next_, unused_] = + collective_mainloop.load_A( + params.mainloop, load2transform_pipeline, + load2transform_pipeline_producer_state, load_inputs, + cta_coord_mnkl, k_tile_iter_next, + k_tile_count - k_tile_prologue); + load2transform_pipeline_producer_state = + load2transform_pipeline_producer_state_next_; + } + + if (is_participant.main_loadB) { + auto [load2mma_pipeline_producer_state_next, k_tile_iter_next] = + collective_mainloop.load_B(params.mainloop, load2mma_pipeline, + load2mma_pipeline_producer_state, + load_inputs, cta_coord_mnkl, + k_tile_iter, k_tile_prologue); + load2mma_pipeline_producer_state = + load2mma_pipeline_producer_state_next; + + auto [load2mma_pipeline_producer_state_next_, unused_] = + collective_mainloop.load_B(params.mainloop, load2mma_pipeline, + load2mma_pipeline_producer_state, + load_inputs, cta_coord_mnkl, + k_tile_iter_next, + k_tile_count - k_tile_prologue); + load2mma_pipeline_producer_state = + load2mma_pipeline_producer_state_next_; + } + } + + // Sync warp to prevent non-participating threads entering next wave + // early + __syncwarp(); + + // Fetch next work tile + auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work( + work_tile_info, clc_pipeline, clc_pipeline_consumer_state); + + requires_clc_query = increment_pipe; + if (increment_pipe) { + ++clc_pipeline_consumer_state; + } + work_tile_info = next_work_tile_info; + } while (work_tile_info.is_valid()); + + if (is_participant.main_loadA) { + if (lane_predicate) { + load2transform_pipeline.producer_tail( + load2transform_pipeline_producer_state); + } + } + if (is_participant.main_loadB) { + if (lane_predicate) { + load2mma_pipeline.producer_tail(load2mma_pipeline_producer_state); + } + } + + } + + else if (is_participant.sched) { + // Signal the epilogue warps to proceed once the prologue is complete + epilogue_throttle_barrier.arrive(); + + if constexpr (IsSchedDynamicPersistent) { + // Whether a new CLC query must be performed. + // See comment below where this variable is updated for a description of + // why this variable is needed. + bool requires_clc_query = true; + + cutlass::arch::wait_on_dependent_grids(); + + do { + if (requires_clc_query) { + // Throttle CLC query to mitigate workload imbalance caused by skews + // among persistent workers. + clc_throttle_pipeline.consumer_wait( + clc_pipe_throttle_consumer_state); + clc_throttle_pipeline.consumer_release( + clc_pipe_throttle_consumer_state); + ++clc_pipe_throttle_consumer_state; + + // Query next clcID and update producer state + clc_pipeline_producer_state = scheduler.advance_to_next_work( + clc_pipeline, clc_pipeline_producer_state); + } + + // Fetch next work tile + auto [next_work_tile_info, increment_pipe] = + scheduler.fetch_next_work(work_tile_info, clc_pipeline, + clc_pipeline_consumer_state); + + // Only perform a new CLC query if we consumed a new CLC query result + // in `fetch_next_work`. An example of a case in which CLC + // `fetch_next_work` does not consume a new CLC query response is when + // processing stream-K units. The current stream-K scheduler uses + // single WorkTileInfo to track multiple (potentially-partial) tiles + // to be computed via stream-K. In this case, `fetch_next_work` simply + // performs in-place updates on the existing WorkTileInfo, rather than + // consuming a CLC query response. + requires_clc_query = increment_pipe; + if (increment_pipe) { + ++clc_pipeline_consumer_state; + } + + work_tile_info = next_work_tile_info; + } while (work_tile_info.is_valid()); + clc_pipeline.producer_tail(clc_pipeline_producer_state); + } + } + + else if (is_participant.transformation) { + // Signal the epilogue warps to proceed once the prologue is complete + epilogue_throttle_barrier.arrive(); + + // Wait for tmem allocation + tmem_allocation_result_barrier.arrive_and_wait_unaligned(); + + do { + auto k_tile_count = TileScheduler::get_work_k_tile_count( + work_tile_info, problem_shape_MNKL, CtaShape_MNK{}); + auto k_tile_start = + TileScheduler::get_work_k_tile_start(work_tile_info); + auto k_tile_iter = cute::make_coord_iterator( + idx2crd(k_tile_start, shape<3>(gA_mkl)), shape<3>(gA_mkl)); + auto [load2transform_pipeline_consumer_state_next, + transform2mma_pipeline_producer_state_next] = + collective_mainloop.transform( + load2transform_pipeline, load2transform_pipeline_consumer_state, + transform2mma_pipeline, transform2mma_pipeline_producer_state, + bulk_tmem, transform_inputs, k_tile_iter, k_tile_count); + transform2mma_pipeline_producer_state = + transform2mma_pipeline_producer_state_next; + load2transform_pipeline_consumer_state = + load2transform_pipeline_consumer_state_next; + + // Fetch next work tile + auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work( + work_tile_info, clc_pipeline, clc_pipeline_consumer_state); + work_tile_info = next_work_tile_info; + + if (increment_pipe) { + ++clc_pipeline_consumer_state; + } + } while (work_tile_info.is_valid()); + + transform2mma_pipeline.producer_tail( + transform2mma_pipeline_producer_state); + } + + else if (is_participant.mma) { + // Tmem allocation sequence + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + tmem_allocation_result_barrier.arrive(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + + auto mma_input_operands = collective_mainloop.mma_init( + bulk_tmem, shared_storage.tensors.mainloop); + + // Signal the epilogue warps to proceed once the prologue is complete + epilogue_throttle_barrier.arrive(); + + do { + auto k_tile_count = TileScheduler::get_work_k_tile_count( + work_tile_info, problem_shape_MNKL, CtaShape_MNK{}); + // Fetch next work tile + auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work( + work_tile_info, clc_pipeline, clc_pipeline_consumer_state); + work_tile_info = next_work_tile_info; + + if (increment_pipe) { + ++clc_pipeline_consumer_state; + } + + if (is_mma_leader_cta) { + auto [load2mma_pipeline_consumer_state_next, + transform2mma_pipeline_consumer_state_next, + mma2accum_pipeline_producer_state_next] = + collective_mainloop.mma( + load2mma_pipeline, load2mma_pipeline_consumer_state, + transform2mma_pipeline, transform2mma_pipeline_consumer_state, + mma2accum_pipeline, mma2accum_pipeline_producer_state, + bulk_tmem, mma_input_operands, k_tile_count); + // Advance the mm2accum pipe + load2mma_pipeline_consumer_state = + load2mma_pipeline_consumer_state_next; + transform2mma_pipeline_consumer_state = + transform2mma_pipeline_consumer_state_next; + mma2accum_pipeline_producer_state = + mma2accum_pipeline_producer_state_next; + } + } while (work_tile_info.is_valid()); + + // leader MMA waits for leader + peer epilogues to release accumulator + // stage + if (is_mma_leader_cta) { + mma2accum_pipeline.producer_tail(mma2accum_pipeline_producer_state); + } + + // Hint on an early release of global memory resources. + // The timing of calling this function only influences performance, + // not functional correctness. + cutlass::arch::launch_dependent_grids(); + + // Signal to peer MMA that entire tmem allocation can be deallocated + if constexpr (has_mma_peer_cta) { + // Leader does wait + arrive, follower does arrive + wait + tmem_deallocation_result_barrier.arrive(mma_peer_cta_rank, + not is_mma_leader_cta); + tmem_deallocation_result_barrier.wait(dealloc_barrier_phase); + tmem_deallocation_result_barrier.arrive(mma_peer_cta_rank, + is_mma_leader_cta); + } + + // Free entire tmem allocation + tmem_allocator.free(tmem_base_ptr, + TmemAllocator::Sm100TmemCapacityColumns); + } + + else if (is_participant.epi_load) { + // Ensure that the prefetched kernel does not touch + // unflushed global memory prior to this instruction + cutlass::arch::wait_on_dependent_grids(); + + bool do_load_order_wait = true; + bool do_tail_load = false; + + // Signal the epilogue warps to proceed once the prologue is complete + epilogue_throttle_barrier.arrive(); + + do { + bool compute_epilogue = + TileScheduler::compute_epilogue(work_tile_info, params.scheduler); + // Get current work tile and fetch next work tile + auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work( + work_tile_info, clc_pipeline, clc_pipeline_consumer_state); + work_tile_info = next_work_tile_info; + + if (increment_pipe) { + ++clc_pipeline_consumer_state; + } + + if (compute_epilogue) { + if (do_load_order_wait) { + load_order_barrier.wait(); + do_load_order_wait = false; + } + + epi_load_pipe_producer_state = collective_epilogue.load( + epi_load_pipeline, epi_load_pipe_producer_state, + problem_shape_MNKL, CtaShape_MNK{}, cta_coord_mnkl, TileShape{}, + TiledMma{}, shared_storage.tensors.epilogue); + + do_tail_load = true; + } + + // Calculate the cta coordinates of the next work tile + cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info); + } while (work_tile_info.is_valid()); + + // Only perform a tail load if one of the work units processed performed + // an epilogue load. An example of a case in which a tail load should not + // be performed is in split-K if a cluster is only assigned non-final + // splits (for which the cluster does not compute the epilogue). + if (do_tail_load) { + collective_epilogue.load_tail( + epi_load_pipeline, epi_load_pipe_producer_state, epi_store_pipeline, + epi_store_pipe_producer_state); + } + } + + else if (is_participant.epilogue) { + // Throttle the epilogue warps to improve prologue performance + static constexpr int epilogue_throttle_phase_bit = 0; + epilogue_throttle_barrier.wait(epilogue_throttle_phase_bit); + + // Wait for tmem allocation + tmem_allocation_result_barrier.arrive_and_wait_unaligned(); + + auto accum_inputs = collective_mainloop.accum_init( + bulk_tmem, typename CollectiveEpilogue::CopyOpT2R{}, + typename CollectiveEpilogue::EpilogueTile{}); + bool do_tail_store = false; + do { + // Fetch next work tile + auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work( + work_tile_info, clc_pipeline, clc_pipeline_consumer_state); + + if (increment_pipe) { + ++clc_pipeline_consumer_state; + } + + auto k_tile_count = TileScheduler::get_work_k_tile_count( + work_tile_info, problem_shape_MNKL, CtaShape_MNK{}); + mma2accum_pipeline.consumer_wait(mma2accum_pipeline_consumer_state); + + // Accumulators + Tensor accumulators = + bulk_tmem(_, _, _, + mma2accum_pipeline_consumer_state + .index()); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N) + + mma2accum_pipeline_consumer_state = scheduler.template fixup( + TiledMma{}, work_tile_info, accumulators, mma2accum_pipeline, + mma2accum_pipeline_consumer_state, + typename CollectiveEpilogue::CopyOpT2R{}); + + // + // Epilogue and write to gD + // + if (scheduler.compute_epilogue(work_tile_info)) { + auto [load_state_next, store_state_next, + mma2accum_pipeline_state_next] = + collective_epilogue.store( + epi_load_pipeline, epi_load_pipe_consumer_state, + epi_store_pipeline, epi_store_pipe_producer_state, + mma2accum_pipeline, mma2accum_pipeline_consumer_state, + problem_shape_MNKL, CtaShape_MNK{}, cta_coord_mnkl, + TileShape{}, TiledMma{}, accumulators, + shared_storage.tensors.epilogue); + epi_load_pipe_consumer_state = load_state_next; + epi_store_pipe_producer_state = store_state_next; + do_tail_store = true; + + // Advance the mma2accum pipe + mma2accum_pipeline_consumer_state = mma2accum_pipeline_state_next; + } + + work_tile_info = next_work_tile_info; + cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info); + } while (work_tile_info.is_valid()); + + // Only perform a tail load if one of the work units processed performed + // an epilogue load. An example of a case in which a tail load should not + // be performed is in split-K if a cluster is only assigned non-final + // splits (for which the cluster does not compute the epilogue). + if (do_tail_store) { + collective_epilogue.store_tail( + epi_load_pipeline, epi_load_pipe_consumer_state, epi_store_pipeline, + epi_store_pipe_producer_state, CtaShape_MNK{}); + } + } else { + } + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::gemm::kernel diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_mainloop.cuh b/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_mainloop.cuh new file mode 100644 index 0000000000..cd58e86999 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_mainloop.cuh @@ -0,0 +1,1283 @@ +// Swordfish prefill mainloop, an sm100 tcgen05 mixed-input collective reading +// the Swordfish packed-weight ABI v1 directly. +// +// This is a fork of CUTLASS 4.4.2's +// include/cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp +// (BSD-3-Clause, Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES), +// following the Machete precedent of forking a stock collective. Line +// references below are against that file. The pipeline structure +// (Load2Transform / Load2Mma / Transform2Mma / Mma2Accum), the B (activation) +// path, the scale TMA path, and the MMA stage are stock. Two things change. +// +// 1. B-operand TMA over the packed ABI (stock lines 310-313, 460-467, +// 529-535, 735-814, 883-973). The quantized operand (riding the swapped +// A slot) is +// presented to TMA as a dense byte tensor (256, 8, KB, NB); each (nb, kb) +// ABI block is a linear 2048 B run (invariant I3), and the Marlin in-tile +// permutation is invisible to a byte copy (I4). The input smem staging is a +// plain byte buffer; the canonical-layout input atom machinery is deleted. +// +// 2. The Transform stage (stock lines 975-1144). Instead of +// MixedInputUtils::dequantize_A_kblock_for_transform over canonical int4, +// each transform thread consumes packed words in Marlin tile order (lane T +// <-> word 4T of each 512 B sub-tile, the contract proven by the decode +// path, swordfish_decode.cuh) and dequantizes with the marlin u4b8 LOP3 +// idiom, applies group scales, and writes K-major bf16 into the +// tcgen05-descriptor-legal compute smem buffer (invariant I5, four 32-bit +// lane-local stores per packed word). Because the stores are free-form smem +// writes, this fork uses the SS (A-from-SMEM) UMMA atom, i.e. the TiledMma +// of KernelTmaWarpSpecialized1SmMixedInputSmemSm100. (The stock Smem +// schedule does not compile in CUTLASS 4.4.2, since its transform_init calls +// make_tmem_copy on an SS descriptor fragment; this fork replaces that code +// entirely.) +// +// Scope: 1-SM or 2-SM (cta_group::2) MMA, u4b8 weights, bf16 +// activations/scales, ConvertAndScale without zero points, group_size 64 or +// 128, L = 1, N and K multiples of 128 (the ABI tail policy rejects rests). +#pragma once + +#include + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/numeric_conversion.h" +#include "cutlass/detail/sm100_tmem_helper.hpp" +#include "cutlass/detail/cluster.hpp" +#include "cutlass/detail/collective/mixed_input_utils.hpp" +#include "cutlass/detail/sm100_mixed_dtype_blockwise_layout.hpp" +#include "cutlass/detail/blockwise_scale_layout.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/atom/copy_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/arch/mma_sm100.hpp" +#include "cutlass/trace.h" +#include "cutlass/kernel_hardware_info.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { +using namespace cute; + +namespace swordfish_detail { + +// u4b8 -> bf16x2 pair dequant, the marlin LOP3 idiom +// (csrc/libtorch_stable/quantization/marlin/dequant.h, +// dequant; re-stated here so this header only depends on +// CUTLASS/CUDA). Consumes nibbles {0,4} into frag[0] and {1,5} into frag[1]; +// feed `q >> 8` for nibbles {2,6}/{3,7}. Under the Marlin pack_idx interleave +// {0,2,4,6,1,3,5,7} this yields, for a word held by lane T (column c = T/4, +// k-quad t = T%4): +// dequant(q): frag[0] = (k=2t, 2t+1), frag[1] = (k=2t+8, 2t+9) @ col +// c dequant(q >> 8): same k positions @ +// col c+8 +__device__ __forceinline__ void dequant_u4b8_bf16x2(uint32_t q, + __nv_bfloat162 frag[2]) { + static constexpr uint32_t kMask = 0x000f000f; + static constexpr uint32_t kEx = 0x43004300; + static constexpr uint32_t kSub = + 0x43084308; // {136, 136}: (v | 0x4300) - 136 = v - 8 + static constexpr uint32_t kImmLut = (0xf0 & 0xcc) | 0xaa; + uint32_t lo, hi; + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(lo) + : "r"(q), "n"(kMask), "n"(kEx), "n"(kImmLut)); + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(hi) + : "r"(q >> 4), "n"(kMask), "n"(kEx), "n"(kImmLut)); + frag[0] = __hsub2(reinterpret_cast<__nv_bfloat162 const&>(lo), + reinterpret_cast<__nv_bfloat162 const&>(kSub)); + frag[1] = __hsub2(reinterpret_cast<__nv_bfloat162 const&>(hi), + reinterpret_cast<__nv_bfloat162 const&>(kSub)); +} + +// u4b8 -> f16x2 pair dequant (marlin dequant). Same fragment +// semantics as the bf16 helper; the high nibble plane rides an embedded +// x16 exponent folded out by the fma. +__device__ __forceinline__ void dequant_u4b8_f16x2(uint32_t q, + __half2 frag[2]) { + static constexpr uint32_t kLo = 0x000f000f; + static constexpr uint32_t kHi = 0x00f000f0; + static constexpr uint32_t kEx = 0x64006400; + static constexpr uint32_t kSub = 0x64086408; + static constexpr uint32_t kMul = 0x2c002c00; + static constexpr uint32_t kAdd = 0xd480d480; + static constexpr uint32_t kImmLut = (0xf0 & 0xcc) | 0xaa; + uint32_t lo, hi; + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(lo) + : "r"(q), "n"(kLo), "n"(kEx), "n"(kImmLut)); + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(hi) + : "r"(q), "n"(kHi), "n"(kEx), "n"(kImmLut)); + frag[0] = __hsub2(reinterpret_cast<__half2 const&>(lo), + reinterpret_cast<__half2 const&>(kSub)); + frag[1] = __hfma2(reinterpret_cast<__half2 const&>(hi), + reinterpret_cast<__half2 const&>(kMul), + reinterpret_cast<__half2 const&>(kAdd)); +} + +// u8b128 -> f16x2 pair dequant (marlin dequant). +__device__ __forceinline__ void dequant_u8b128_f16x2(uint32_t q, + __half2 frag[2]) { + static constexpr uint32_t kMask01 = 0x5250; + static constexpr uint32_t kMask23 = 0x5351; + static constexpr uint32_t kBase = 0x64646464; + static constexpr uint32_t kSub = 0x64806480; + const uint32_t lo = __byte_perm(q, kBase, kMask01); + const uint32_t hi = __byte_perm(q, kBase, kMask23); + frag[0] = __hsub2(reinterpret_cast<__half2 const&>(lo), + reinterpret_cast<__half2 const&>(kSub)); + frag[1] = __hsub2(reinterpret_cast<__half2 const&>(hi), + reinterpret_cast<__half2 const&>(kSub)); +} + +// u8b128 -> bf16x2 pair dequant (marlin dequant, the +// FasterTransformer fp32-bias idiom). Under the 8-bit pack interleave +// {0,2,1,3} one word yields the same fragment pair as the 4-bit dequant of +// one nibble plane: frag[0] = perm values (0,1), frag[1] = values (2,3). +__device__ __forceinline__ void dequant_u8b128_bf16x2(uint32_t q, + __nv_bfloat162 frag[2]) { + float f[4]; + uint32_t* fc = reinterpret_cast(f); + static constexpr uint32_t kBase = 0x4B000000; + fc[0] = __byte_perm(q, kBase, 0x7650); + fc[1] = __byte_perm(q, kBase, 0x7652); + fc[2] = __byte_perm(q, kBase, 0x7651); + fc[3] = __byte_perm(q, kBase, 0x7653); + f[0] -= 8388736.f; + f[1] -= 8388736.f; + f[2] -= 8388736.f; + f[3] -= 8388736.f; + uint32_t* out = reinterpret_cast(frag); + out[0] = __byte_perm(fc[0], fc[1], 0x7632); + out[1] = __byte_perm(fc[2], fc[3], 0x7632); +} + +} // namespace swordfish_detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Same template signature as the stock CollectiveMma specialization (stock +// lines 62-103) so the instantiation site can mirror the stock builder; a +// standalone struct (the kernel layer dispatches on DispatchPolicy::Schedule +// only). +template +struct SwordfishMainloopSm100MixedInput { + public: + // + // Type Aliases (stock lines 105-263, SwapAB plumbing removed: the quantized + // operand is ALWAYS the A slot here) + // + using ConversionMode = cutlass::detail::ConversionMode; + using AtomThrShapeMNK = + Shape(typename TiledMma_::ThrLayoutVMNK{})), _1, _1>; + using DispatchPolicy = MainloopSm100TmaUmmaWarpSpecializedMixedInput< + Load2TransformPipelineStageCount_, Transform2MmaPipelineStageCount_, + SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_, + ClusterShape>; + using TileShape = TileShape_; + using TiledMma = TiledMma_; + using KernelSchedule = typename DispatchPolicy::Schedule; + static constexpr bool IsDynamicCluster = not cute::is_static_v; + static_assert(!IsDynamicCluster, + "swordfish prefill v1 requires a static cluster"); + static_assert(cute::is_same_v> || + cute::is_same_v>, + "swordfish prefill supports 1-SM (cluster 1x1x1) or 2-SM " + "(cluster 2x1x1) only"); + using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{})); + // Number of CTAs the 2-SM MMA atom spans in the tile-M (weight-N) axis: + // 1 for the 1-SM atom, 2 for the 2-SM atom. Drives the per-CTA A split. + static constexpr int kAtomCtasM = size<0>(AtomThrShapeMNK{}); + + using ElementAOptionalTuple = ElementAOptionalTuple_; + using ElementBOptionalTuple = ElementBOptionalTuple_; + static_assert(cute::is_tuple::value && + !cute::is_tuple::value, + "swordfish prefill: quantized operand must ride the A slot " + "(pass {ElementA, ElementScale} for A)"); + + using ElementA = detail::deduce_mixed_width_dtype_t<0, ElementAOptionalTuple>; + using ElementB = detail::deduce_mixed_width_dtype_t<0, ElementBOptionalTuple>; + static constexpr bool IsATransformed = true; + using ElementScale = + detail::deduce_mixed_width_dtype_t<1, ElementAOptionalTuple>; + using ElementZero = + detail::deduce_mixed_width_dtype_t<2, ElementAOptionalTuple>; + using NonVoidElementScale = + cute::conditional_t, float, ElementScale>; + using NonVoidElementZero = + cute::conditional_t, float, ElementZero>; + + // The packed operand is staged as raw bytes. + static_assert(cute::is_same_v, + "swordfish prefill: pass uint8_t as the (packed) A element"); + // Zero-point checkpoints (AWQ/HQQ) pass a third element in the A tuple. + // The zero tensor holds prescaled (8 - zp) * scale rows, scale-shaped and + // scale-typed, so it rides the scale TMA machinery verbatim and the + // transform's scaling multiply becomes an fma. + static constexpr bool HasZp = !cute::is_void_v; + static_assert(!HasZp || cute::is_same_v, + "swordfish prefill zero points are scale-typed (8 - zp) * s"); + + using StrideA = cute::remove_cvref_t(StridePairA_{}))>; + using LayoutScale = cute::remove_cvref_t(StridePairA_{}))>; + using InternalStrideA = cute::remove_pointer_t; + using StrideB = StrideB_; + using InternalStrideB = cute::remove_pointer_t; + + using CtaShapeA_MK = decltype(partition_shape_A( + TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{})))); + using CtaShapeB_NK = decltype(partition_shape_B( + TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{})))); + + using ElementAMma = typename TiledMma::ValTypeA; + using ElementBMma = typename TiledMma::ValTypeB; + static constexpr bool kActF16 = cute::is_same_v; + static_assert(kActF16 || cute::is_same_v, + "swordfish prefill dequantizes to fp16 or bf16"); + static_assert(cute::is_same_v, + "swordfish prefill group scales match the activation dtype"); + // Register pair type matching ElementAMma for the transform. + using Elem2 = cute::conditional_t; + + using ElementAccumulator = typename TiledMma::ValTypeC; + + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using GmemTiledCopyScale = GmemTiledCopyA_; + + using SmemLayoutAtomsA = SmemLayoutAtomsA_; + using SmemLayoutAtomsB = SmemLayoutAtomsB_; + using CopyAtomsA = CopyAtomsA_; + using CopyAtomsB = CopyAtomsB_; + + using SmemLayoutAtomACompute = typename SmemLayoutAtomsA::ComputeLayoutAtom; + using SmemLayoutAtomB = typename SmemLayoutAtomsB::InputLayoutAtom; + + using TmaElementA = uint8_t; + + using ArchTag = typename DispatchPolicy::ArchTag; + + using Load2TransformPipeline = cutlass::PipelineTmaTransformAsync< + DispatchPolicy::Load2TransformPipelineStageCount, AtomThrShapeMNK>; + using Load2TransformPipelineState = + typename Load2TransformPipeline::PipelineState; + + using Load2MmaPipeline = cutlass::PipelineTmaUmmaAsync< + DispatchPolicy::Load2TransformPipelineStageCount, ClusterShape, + AtomThrShapeMNK>; + using Load2MmaPipelineState = typename Load2MmaPipeline::PipelineState; + + using Transform2MmaPipeline = cutlass::PipelineUmmaConsumerAsync< + DispatchPolicy::Transform2MmaPipelineStageCount, AtomThrShapeMNK>; + using Transform2MmaPipelineState = + typename Transform2MmaPipeline::PipelineState; + + using Mma2AccumPipeline = cutlass::PipelineUmmaAsync< + DispatchPolicy::Schedule::AccumulatorPipelineStageCount, AtomThrShapeMNK>; + using Mma2AccumPipelineState = typename Mma2AccumPipeline::PipelineState; + + // ---- scale layout machinery (stock lines 265-278, unchanged) -------------- + static constexpr int ScaleGranularityMN = size<0, 0>(LayoutScale{}); + static constexpr int ScaleGranularityK = size<1, 0>(LayoutScale{}); + static_assert(ScaleGranularityMN == 1, "swordfish scales are per-column"); + static_assert(ScaleGranularityK == 32 || ScaleGranularityK == 64 || + ScaleGranularityK == 128, + "swordfish prefill supports group sizes 32, 64 and 128"); + // At granularity 32 a k64 sub-block spans two scale groups, so the + // transform keeps one register set per 32-row group of its share. + static constexpr int kScaleGroupsPerWarp = ScaleGranularityK == 32 ? 2 : 1; + using ScaleConfig = + cutlass::detail::Sm100MixedInputBlockwiseScaleConfig; + + using ScaleTileShape = + decltype(make_shape(size<0>(TileShape{}), size<2>(TileShape{}))); + using SmemLayoutAtomScaleFull = + decltype(ScaleConfig::smem_atom_layout_scale(ScaleTileShape{})); + using SmemLayoutAtomScale = + decltype(slice(make_coord(make_coord(_, 0), make_coord(_, 0)), + SmemLayoutAtomScaleFull{})); + + static_assert(cute::rank(SmemLayoutAtomB{}) == 2, + "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, + "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, + "SmemLayoutAtom must evenly divide tile shape."); + + // Thread counts (stock lines 292-294) + // Parametric transform width. The forked kernel layer derives its warp + // layout from this. 128 means one warp per packed block; 256 (one per half + // block) measured slightly worse on B200, so the transform is not the + // bottleneck there. + static constexpr uint32_t NumTransformationThreads = 128; + static constexpr uint32_t NumAccumThreads = 128; + + constexpr static int AccumulatorPipelineStageCount = + DispatchPolicy::Schedule::AccumulatorPipelineStageCount; + constexpr static int StagesPerTile = size<2>(CtaShapeA_MK{}); + + // ---- Swordfish geometry + // ---------------------------------------------------- ABI v1 constants + // (swordfish_types.cuh; restated to keep this header dependent only on + // CUTLASS/CUDA). + static constexpr int kBlockN = 64; // columns per packed block + static constexpr int kBlockK = 64; // K rows per packed block + static_assert(WBits == 4 || WBits == 8, "swordfish weights are 4 or 8 bit"); + static constexpr int kSubTileBytes = WBits == 8 ? 1024 : 512; // 16x64 tile + static constexpr int kBlockBytes = 4 * kSubTileBytes; + static constexpr int kBlockRows = kBlockBytes / 256; // TMA inner rows + + // Full MMA-tile weight-N drives the packed-A gmem tiler so the tile + // coordinate matches the scheduler. Per-CTA weight-N drives the smem + // buffers, since each CTA of a 2-SM pair stages only its own half. + static constexpr int TileN_Weights_Full = size<0>(TileShape{}); + static constexpr int CtaTileN_Weights = size<0>(CtaShape_MNK{}); // per-CTA + static constexpr int CtaTileK = size<2>(CtaShape_MNK{}); + static_assert( + CtaTileN_Weights == 128 && CtaTileK == 128, + "swordfish prefill is tuned/asserted for a 128x128x128 CTA tile"); + static constexpr int kBlocksPerTileN = + CtaTileN_Weights / kBlockN; // 2 (per CTA) + static constexpr int kBlocksPerTileK = CtaTileK / kBlockK; // 2 + static constexpr int kStageBytes = + kBlocksPerTileN * kBlocksPerTileK * kBlockBytes; // 8192 + + // CHANGE (1): input staging is the packed byte stream, laid out + // (byte-lo, byte-hi, kb, nb, PIPE). Modes 0-1 pre-split the 2048 B block run + // so every TMA box extent is <= 256. Sized/tiled PER-CTA (128 weight-cols); + // in 2-SM the load loop maps the cluster's weight-N tile coord to this CTA's + // 128-col half via kAtomCtasM*coord + atom_half. + using SmemLayoutA = Layout< + Shape<_256, Int, Int, Int, + Int>, + Stride<_1, _256, Int, Int, + Int>>; + using SwordfishTilerA = + Shape<_256, Int, Int, Int>; + + // B (activations) staging: stock (lines 320-323). + using SmemLayoutB = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomB{}, + append(CtaShapeB_NK{}, + Int{}), + (cute::conditional_t(), + Step<_2, _1, _3>, Step<_1, _2, _3>>{}))); + + // CHANGE (2): the compute buffer keeps the stock tcgen05-descriptor-legal + // core-matrix layout, but is built flat-first so the transform can address + // it by logical (n, k, stage) with a memory function IDENTICAL (by + // construction) to the MMA-facing grouped view. + // Per-CTA shape. Each CTA dequants only its own weight-N half. + using SmemLayoutAComputeMK = decltype(tile_to_shape( + SmemLayoutAtomACompute{}, + make_shape(size<0>(CtaShape_MNK{}), size<2>(CtaShape_MNK{}), + Int{}), + Step<_1, _2, _3>{})); + // Same regroup tile_to_mma_shape applies + // (cute/atom/mma_traits_sm100.hpp:116). + using SmemLayoutACompute = decltype(tiled_divide( + SmemLayoutAComputeMK{}, product_each(shape<0>(CtaShapeA_MK{})))); + + // Closed-form of the compute buffer's memory function (element offset for + // logical (n, k, stage)), used by the transform's store path so each 32-bit + // store is one ADD instead of a hierarchical cute crd2idx evaluation. For + // the Sw<3,4,3> (8, 64):(64, 1) 16-bit-flagged atom tiled M-first + // (Step<_1,_2,_3>): + // pre = (n%8)*64 + (k%64) + (n/8)*512 + (k/64)*8192 + stage*16384 + // off = pre XOR ((n%8) << 3) + // The swizzle of the smem_ptr_flag_bits<16> atom acts in the BYTE domain + // (element bits [3..6) ^= element bits [6..9) = n%8), matching the UMMA + // SWIZZLE_128B descriptor; evaluating the flagged ComposedLayout directly + // applies the swizzle in the wrong unit, so the static_asserts below pin + // the formula against the position-independent view (what the numerics + // verified end-to-end). + CUTLASS_HOST_DEVICE + static constexpr int compute_elem_offset(int n, int k, int stage) { + int const pre = (n & 7) * 64 + (k & 63) + (n >> 3) * 512 + (k >> 6) * 8192 + + stage * 16384; + return pre ^ ((n & 7) << 3); + } + using SmemLayoutAComputePosInd = + decltype(as_position_independent_swizzle_tensor( + make_tensor( + make_smem_ptr(static_cast(nullptr)), + SmemLayoutAComputeMK{})) + .layout()); + static_assert(SmemLayoutAComputePosInd{}(0, 0, 0) == + compute_elem_offset(0, 0, 0)); + static_assert(SmemLayoutAComputePosInd{}(1, 0, 0) == + compute_elem_offset(1, 0, 0)); + static_assert(SmemLayoutAComputePosInd{}(9, 8, 0) == + compute_elem_offset(9, 8, 0)); + static_assert(SmemLayoutAComputePosInd{}(3, 17, 1) == + compute_elem_offset(3, 17, 1)); + static_assert(SmemLayoutAComputePosInd{}(77, 101, 1) == + compute_elem_offset(77, 101, 1)); + static_assert(SmemLayoutAComputePosInd{}(127, 127, 0) == + compute_elem_offset(127, 127, 0)); + static_assert(SmemLayoutAComputePosInd{}(64, 64, 0) == + compute_elem_offset(64, 64, 0)); + static_assert(SmemLayoutAComputePosInd{}(15, 40, 0) == + compute_elem_offset(15, 40, 0)); + + // Scale staging: stock (lines 325-328). + using SmemLayoutScale = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomScale{}, + append(CtaShapeA_MK{}, + Int{}), + Step<_1, _2, _3>{})); + static constexpr int kScalesPerStage = cosize(take<0, 3>(SmemLayoutScale{})); + static constexpr int kScaleKGroupsPerStage = CtaTileK / ScaleGranularityK; + static_assert(kScalesPerStage == CtaTileN_Weights * kScaleKGroupsPerStage, + "unexpected scale smem layout"); + + static_assert(DispatchPolicy::Load2TransformPipelineStageCount >= 2 && + DispatchPolicy::Transform2MmaPipelineStageCount >= 2, + "Specialization requires Stages set to value 2 or more."); + // SS MMA only: A operand consumed from SMEM through a descriptor. + static_assert( + cute::is_base_of::value && + cute::is_base_of::value, + "swordfish prefill requires an SS UMMA atom (A and B from SMEM)"); + static_assert(cute::is_same_v, + "swordfish prefill v1 uses plain (non-multicast) TMA for the " + "packed operand"); + + static constexpr ConversionMode KernelConversionMode = + ConversionMode::ConvertAndScale; + static constexpr bool ModeHasScales = true; + + static constexpr size_t SmemAlignmentA = 1024; + static constexpr size_t SmemAlignmentB = + cutlass::detail::alignment_for_swizzle(SmemLayoutB{}); + + struct PipelineStorage { + using Load2TransformPipelineStorage = + typename Load2TransformPipeline::SharedStorage; + alignas(16) Load2TransformPipelineStorage load2transform_pipeline; + using Load2MmaPipelineStorage = typename Load2MmaPipeline::SharedStorage; + alignas(16) Load2MmaPipelineStorage load2mma_pipeline; + using Transform2MmaPipelineStorage = + typename Transform2MmaPipeline::SharedStorage; + alignas(16) Transform2MmaPipelineStorage transform2mma_pipeline; + using Mma2AccumPipelineStorage = typename Mma2AccumPipeline::SharedStorage; + alignas(16) Mma2AccumPipelineStorage mma2accum_pipeline; + }; + + struct SharedStorage { + struct TensorStorage : cute::aligned_struct<128, _0> { + struct TensorStorageUntransformed { + alignas(1024) + cute::ArrayEngine> smem_A; + alignas(1024) + cute::ArrayEngine> smem_B; + cute::ArrayEngine> + smem_scale; + cute::ArrayEngine : 1> + smem_zero; + }; + struct TensorStorageTransformed { + alignas(1024) cute::ArrayEngine< + ElementAMma, cute::cosize_v> smem_ACompute; + }; + TensorStorageUntransformed input; + TensorStorageTransformed compute; + } tensors; + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + + // Per-stage mbarrier transaction bytes: packed bytes + scales. Only TMA + // bytes count here (the transform's smem stores arrive through the + // transform2mma pipeline and not this barrier, per the stock collective's + // mbarrier caution). + static constexpr uint32_t kScaleTxBytes = cutlass::bits_to_bytes( + kScalesPerStage * cute::sizeof_bits_v); + static constexpr uint32_t TmaTransactionBytes_A = + kStageBytes + kScaleTxBytes * (HasZp ? 2 : 1); + // AtomThrShape-scaled as in stock. The cta_group::2 TMA loads both CTAs' + // B halves with one instruction and its arrival, covering both halves' + // bytes, lands on the MMA leader's barrier. In 1-SM this reduces to one + // tile. + static constexpr uint32_t TmaTransactionBytes_B = cutlass::bits_to_bytes( + size(AtomThrShapeMNK{}) * cosize(take<0, 3>(SmemLayoutB{})) * + cute::sizeof_bits_v); + static constexpr uint32_t TmaTransactionBytes = + TmaTransactionBytes_A + TmaTransactionBytes_B; + + // Host side kernel arguments (stock lines 424-432; ptr_A is the packed ABI + // tensor base; dA is carried for interface parity but the packed layout is + // derived from the problem shape). + struct Arguments { + ElementA const* ptr_A{nullptr}; + StrideA dA{}; + ElementB const* ptr_B{nullptr}; + StrideB dB{}; + ElementScale const* ptr_S{nullptr}; + LayoutScale layout_S{}; + ElementZero const* ptr_Z{nullptr}; // scale-shaped (8 - zp) * scale rows + }; + + // Device side kernel params + struct Params { + using ClusterLayout_VMNK = + decltype(tiled_divide(make_layout(ClusterShape{}), + make_tile(typename TiledMma::AtomThrID{}))); + + // CHANGE (1): TMA over the packed byte tensor (256, 8, KB, NB). + using GmemLayoutAPacked = + Layout, int32_t, int32_t>, + Stride<_1, _256, Int, int64_t>>; + using TMA_A = decltype(make_tma_copy( + GmemTiledCopyA{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + GmemLayoutAPacked{}), + SmemLayoutA{}(_, _, _, _, cute::Int<0>{}))); + + using TMA_B = decltype(make_tma_atom_B_sm100( + GmemTiledCopyB{}, + make_tensor(static_cast(nullptr), + repeat_like(StrideB{}, int32_t(0)), StrideB{}), + SmemLayoutB{}(_, _, _, cute::Int<0>{}), TileShape{}, TiledMma{}, + ClusterLayout_VMNK{})); + + using TMA_Scale = decltype(make_tma_atom_A_sm100( + GmemTiledCopyScale{}, + make_tensor(static_cast(nullptr), + LayoutScale{}), + SmemLayoutScale{}(_, _, _, cute::Int<0>{}), TileShape{}, TiledMma{}, + ClusterLayout_VMNK{})); + + TMA_A tma_load_a; + TMA_B tma_load_b; + TMA_Scale tma_load_scale; + TMA_Scale tma_load_zero; // constructed over ptr_Z (layout shared with S) + uint32_t tma_transaction_bytes{TmaTransactionBytes}; + int32_t blocks_k{0}; // KB = K / 64 + int32_t blocks_n{0}; // NB = N_weights / 64 + }; + + CUTLASS_DEVICE + SwordfishMainloopSm100MixedInput(Params const& params, ClusterShape, + uint32_t block_rank_in_cluster) + : observed_tma_load_a_(¶ms.tma_load_a), + observed_tma_load_b_(¶ms.tma_load_b), + block_rank_in_cluster_(block_rank_in_cluster) {} + + template + static constexpr Params to_underlying_arguments( + ProblemShape const& problem_shape, Arguments const& args, void* workspace, + cutlass::KernelHardwareInfo const& hw_info = + cutlass::KernelHardwareInfo{}) { + (void)workspace; + (void)hw_info; + + auto problem_shape_MNKL = append<4>(problem_shape, 1); + auto [M, N, K, L] = problem_shape_MNKL; // M = weight N (swapped operands) + + int32_t const blocks_k = int32_t(K / kBlockK); + int32_t const blocks_n = int32_t(M / kBlockN); + + // Packed operand as a dense byte tensor (ABI invariant I3). + auto gA_layout = + make_layout(make_shape(_256{}, Int{}, blocks_k, blocks_n), + make_stride(_1{}, _256{}, Int{}, + int64_t(kBlockBytes) * blocks_k)); + Tensor tensor_a = make_tensor( + make_gmem_ptr(reinterpret_cast(args.ptr_A)), gA_layout); + typename Params::TMA_A tma_load_a = make_tma_copy( + GmemTiledCopyA{}, tensor_a, SmemLayoutA{}(_, _, _, _, cute::Int<0>{})); + + Tensor tensor_b = + make_tensor(args.ptr_B, make_layout(make_shape(N, K, L), args.dB)); + auto cluster_layout_vmnk = tiled_divide( + make_layout(ClusterShape{}), make_tile(typename TiledMma::AtomThrID{})); + + typename Params::TMA_B tma_load_b = make_tma_atom_B_sm100( + GmemTiledCopyB{}, tensor_b, SmemLayoutB{}(_, _, _, cute::Int<0>{}), + TileShape{}, TiledMma{}, cluster_layout_vmnk); + + Tensor tensor_scale = + make_tensor(detail::get_logical_ptr(args.ptr_S), args.layout_S); + typename Params::TMA_Scale tma_load_scale = + make_tma_atom_A_sm100(GmemTiledCopyScale{}, tensor_scale, + SmemLayoutScale{}(_, _, _, cute::Int<0>{}), + TileShape{}, TiledMma{}, cluster_layout_vmnk); + + Tensor tensor_zero = make_tensor( + detail::get_logical_ptr( + reinterpret_cast(args.ptr_Z)), + args.layout_S); + typename Params::TMA_Scale tma_load_zero = + make_tma_atom_A_sm100(GmemTiledCopyScale{}, tensor_zero, + SmemLayoutScale{}(_, _, _, cute::Int<0>{}), + TileShape{}, TiledMma{}, cluster_layout_vmnk); + + return {tma_load_a, tma_load_b, tma_load_scale, tma_load_zero, + TmaTransactionBytes, blocks_k, blocks_n}; + } + + template + static bool can_implement(ProblemShape const& problem_shape, + Arguments const& args) { + auto problem_shape_MNKL = append<4>(problem_shape, 1); + auto [M, N, K, L] = problem_shape_MNKL; + + bool implementable = true; + // ABI v1 tail policy: reject non-multiples (M here = weight N). + implementable &= (M % CtaTileN_Weights == 0); + implementable &= (K % CtaTileK == 0); + implementable &= (L == 1); + implementable &= (args.ptr_S != nullptr); + implementable &= ((args.ptr_Z != nullptr) == HasZp); + + constexpr int tma_alignment_bits_B = + cutlass::detail::get_input_alignment_bits(); + constexpr int min_tma_aligned_elements_B = + tma_alignment_bits_B / cutlass::sizeof_bits::value; + implementable &= + cutlass::detail::check_alignment( + cute::make_shape(N, K, L), StrideB{}); + + if (!implementable) { + CUTLASS_TRACE_HOST( + " CAN IMPLEMENT: swordfish prefill shape/argument requirements not " + "met.\n"); + } + return implementable; + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best + /// performance + CUTLASS_DEVICE static void prefetch_tma_descriptors(Params const& params) { + cute::prefetch_tma_descriptor(params.tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_scale.get_tma_descriptor()); + if constexpr (HasZp) { + cute::prefetch_tma_descriptor(params.tma_load_zero.get_tma_descriptor()); + } + } + + /// Construct A Single Stage's Accumulator Shape + CUTLASS_DEVICE auto partition_accumulator_shape() { + return partition_shape_C(TiledMma{}, take<0, 2>(TileShape{})); + } + + /// Produce the inputs to the transform threads by loading inputs from gmem -> + /// smem. (stock lines 735-814; the A copy now walks the packed block tensor) + template + CUTLASS_DEVICE auto load_A( + Params const& params, Load2TransformPipeline load2xform_pipeline, + Load2TransformPipelineState load2xform_pipeline_state, + cute::tuple> const& load_inputs, + TileCoordMNKL const& cta_coord_mnkl, KTileIterator k_tile_iter, + int k_tile_count) { + auto [unused_gA, unused_gB, tAgA_nk, tBgB_nkl, tAsA, tBsB, mcast_mask_a, + mcast_mask_b, extra_input_partitions] = load_inputs; + + // tAgA_nk : ((TMA), 1, 1, K_TILES, N_TILES) over the packed byte tensor, + // tiled PER-CTA (128 weight-cols). Unlike the scales (multicast, routed + // through cta_mma.partition_A which applies the atom split), the packed A + // is DISJOINT across the 2-SM pair: each CTA loads its own 128-col half. + // So index the CTA's own 128-tile: AtomThrID*(base MMA tile) + atom half. + // Reduces to get<0>(cta_coord_mnkl) in the 1-SM case (AtomThrID size 1). + constexpr int kAtom = size(typename TiledMma::AtomThrID{}); + const int a_ntile = + kAtom * (get<0>(cta_coord_mnkl) / kAtom) + int(block_rank_in_cluster_); + Tensor tAgA = tAgA_nk(_, _0{}, _0{}, _, a_ntile); + + uint32_t skip_wait = (k_tile_count <= 0); + auto load2xform_pipeline_flag = load2xform_pipeline.producer_try_acquire( + load2xform_pipeline_state, skip_wait); + + using BarrierType = typename Load2TransformPipeline::ProducerBarrierType; + + CUTLASS_PRAGMA_NO_UNROLL + for (; k_tile_count > 0; --k_tile_count) { + load2xform_pipeline.producer_acquire(load2xform_pipeline_state, + load2xform_pipeline_flag); + + int tile_A_write_stage = load2xform_pipeline_state.index(); + BarrierType* load2xform_tma_barrier = + load2xform_pipeline.producer_get_barrier(load2xform_pipeline_state); + + ++load2xform_pipeline_state; + skip_wait = (k_tile_count <= 1); + load2xform_pipeline_flag = load2xform_pipeline.producer_try_acquire( + load2xform_pipeline_state, skip_wait); + + // TMA load: 4 packed blocks (2 nb x 2 kb) = one k tile of this CTA's + // weight stripe, plus the k tile's group scales. + copy(observed_tma_load_a_->with(*load2xform_tma_barrier, mcast_mask_a), + tAgA(_, *k_tile_iter), tAsA(_, tile_A_write_stage)); + + auto tSgS_mkl = get<0>(extra_input_partitions); + auto tSgS = tSgS_mkl( + _, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, + get<3>(cta_coord_mnkl)); + auto tSsS = get<1>(extra_input_partitions); + copy(params.tma_load_scale.with(*load2xform_tma_barrier, mcast_mask_a), + tSgS(_, *k_tile_iter), tSsS(_, tile_A_write_stage)); + + if constexpr (HasZp) { + auto tZgZ_mkl = get<2>(extra_input_partitions); + auto tZgZ = tZgZ_mkl( + _, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, + get<3>(cta_coord_mnkl)); + auto tZsZ = get<3>(extra_input_partitions); + copy(params.tma_load_zero.with(*load2xform_tma_barrier, mcast_mask_a), + tZgZ(_, *k_tile_iter), tZsZ(_, tile_A_write_stage)); + } + + ++k_tile_iter; + } + + return cute::make_tuple(load2xform_pipeline_state, k_tile_iter); + } + + /// (stock lines 816-876, unchanged: activations) + template + CUTLASS_DEVICE auto load_B( + Params const& params, Load2MmaPipeline load2mma_pipeline, + Load2MmaPipelineState load2mma_pipeline_state, + cute::tuple> const& load_inputs, + TileCoordMNKL const& cta_coord_mnkl, KTileIterator k_tile_iter, + int k_tile_count) { + auto [unused_gA, unused_gB, tAgA_nk, tBgB_nkl, tAsA, tBsB, mcast_mask_a, + mcast_mask_b, extra_input_partitions] = load_inputs; + + Tensor tBgB = + tBgB_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)); + + uint32_t skip_wait = (k_tile_count <= 0); + auto load2mma_pipeline_flag = load2mma_pipeline.producer_try_acquire( + load2mma_pipeline_state, skip_wait); + + using BarrierType = typename Load2TransformPipeline::ProducerBarrierType; + + CUTLASS_PRAGMA_NO_UNROLL + for (; k_tile_count > 0; --k_tile_count) { + load2mma_pipeline.producer_acquire(load2mma_pipeline_state, + load2mma_pipeline_flag); + + int tile_B_write_stage = load2mma_pipeline_state.index(); + BarrierType* load2mma_tma_barrier = + load2mma_pipeline.producer_get_barrier(load2mma_pipeline_state); + + ++load2mma_pipeline_state; + skip_wait = (k_tile_count <= 1); + load2mma_pipeline_flag = load2mma_pipeline.producer_try_acquire( + load2mma_pipeline_state, skip_wait); + + copy(observed_tma_load_b_->with(*load2mma_tma_barrier, mcast_mask_b), + tBgB(_, *k_tile_iter), tBsB(_, tile_B_write_stage)); + + ++k_tile_iter; + } + + return cute::make_tuple(load2mma_pipeline_state, k_tile_iter); + } + + /// Set up the data needed by this collective for load. + /// (stock lines 883-973; A partitioning swapped for the packed byte tensor) + template + CUTLASS_DEVICE auto load_init(ProblemShape_MNKL const& problem_shape_MNKL, + Params const& params, + TensorStorage& shared_storage) const { + auto [M, N, K, L] = problem_shape_MNKL; + + auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL); + + ThrMMA cta_mma = + TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{})); + + // ---- A: packed byte tensor, tiled by (256, 8, kb-per-tile, nb-per-tile). + Tensor mA_packed = observed_tma_load_a_->get_tma_tensor( + make_shape(_256{}, _8{}, params.blocks_k, params.blocks_n)); + Tensor gA_packed = flat_divide(mA_packed, SwordfishTilerA{}); + // (256, 8, KBt, NBt, 1, 1, K_TILES, N_TILES) + + Tensor sA = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), + SmemLayoutA{}); + Tensor sB = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), + SmemLayoutB{}); + + Layout cta_layout_mnk = make_layout(ClusterShape{}); + Layout cta_layout_vmnk = + tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{})); + auto cta_coord_vmnk = + cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_); + + auto [tAgA_nk, tAsA] = + tma_partition(*observed_tma_load_a_, Int<0>{}, Layout<_1>{}, + group_modes<0, 4>(sA), group_modes<0, 4>(gA_packed)); + + Tensor tCgB_nkl = cta_mma.partition_B(gB_nkl); + auto [tBgB_nkl, tBsB] = + tma_partition(*observed_tma_load_b_, get<1>(cta_coord_vmnk), + make_layout(size<1>(cta_layout_vmnk)), + group_modes<0, 3>(sB), group_modes<0, 3>(tCgB_nkl)); + + uint16_t mcast_mask_a = 0; + uint16_t mcast_mask_b = + create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + + // ---- scales: stock path (lines 925-947). + Tensor mS_mkl = params.tma_load_scale.get_tma_tensor(shape(LayoutScale{})); + Tensor gS_mkl = local_tile(mS_mkl, TileShape{}, make_coord(_, _, _), + Step<_1, cute::Underscore, _1>{}); + Tensor sS = + make_tensor(make_smem_ptr(shared_storage.input.smem_scale.begin()), + SmemLayoutScale{}); + Tensor tCgS_mkl = cta_mma.partition_A(gS_mkl); + + auto [tSgS_mkl, tSsS] = + tma_partition(params.tma_load_scale, get<2>(cta_coord_vmnk), + make_layout(size<2>(cta_layout_vmnk)), + group_modes<0, 3>(sS), group_modes<0, 3>(tCgS_mkl)); + + // Zero rows share the scale layout; the partitions are layout-only and + // never copied from when HasZp is false. + Tensor mZ_mkl = params.tma_load_zero.get_tma_tensor(shape(LayoutScale{})); + Tensor gZ_mkl = local_tile(mZ_mkl, TileShape{}, make_coord(_, _, _), + Step<_1, cute::Underscore, _1>{}); + Tensor sZ = + make_tensor(make_smem_ptr(shared_storage.input.smem_zero.begin()), + SmemLayoutScale{}); + Tensor tCgZ_mkl = cta_mma.partition_A(gZ_mkl); + + auto [tZgZ_mkl, tZsZ] = + tma_partition(params.tma_load_zero, get<2>(cta_coord_vmnk), + make_layout(size<2>(cta_layout_vmnk)), + group_modes<0, 3>(sZ), group_modes<0, 3>(tCgZ_mkl)); + + return cute::make_tuple(gA_mkl, gB_nkl, // for scheduler (shapes only) + tAgA_nk, tBgB_nkl, tAsA, + tBsB, // for input tensor values + mcast_mask_a, mcast_mask_b, // multicast masks + cute::make_tuple(tSgS_mkl, tSsS, tZgZ_mkl, tZsZ)); + } + + /// CHANGE (2): the Transform stage (stock lines 975-1059). Consumes the + /// packed bytes in Marlin tile order, dequantizes via the marlin u4b8 LOP3 + /// sequences, applies group scales, and writes K-major bf16 into the + /// tcgen05 compute buffer. + /// + /// Thread assignment (128 threads, 4 warps; ABI read contract): + /// warp w <-> packed block (nb = w>>1, kb = w&1) of the k tile + /// lane T <-> words [4T, 4T+4) of each of the block's 4 sub-tiles + /// (one 16 B vector load per sub-tile) + /// word 4T+j -> dequants to columns {16j+c, 16j+8+c} (c = T/4) at + /// k rows 16s + 2t + {0,1,8,9} (t = T%4) of the sub-tile. + template + CUTLASS_DEVICE auto transform( + Load2TransformPipeline load2transform_pipeline, + Load2TransformPipelineState load2transform_pipeline_consumer_state, + Transform2MmaPipeline transform2mma_pipeline, + Transform2MmaPipelineState transform2mma_pipeline_producer_state, + Accumulator accumulators, + cute::tuple + input_operands, + KTileIterator k_tile_iter, int k_tile_count) { + cutlass::arch::NamedBarrier transform_bar( + NumTransformationThreads, + cutlass::arch::ReservedNamedBarriers::TransformBarrier); + + auto [unused_gA, sA, sACompute, sS, sZ] = input_operands; + uint8_t const* smem_a_base = + reinterpret_cast(raw_pointer_cast(sA.data())); + NonVoidElementScale const* smem_s_base = raw_pointer_cast(sS.data()); + NonVoidElementScale const* smem_z_base = raw_pointer_cast(sZ.data()); + ElementAMma* smem_c_base = raw_pointer_cast(sACompute.data()); + + const int tid = threadIdx.x % NumTransformationThreads; + const int lane = tid % 32; + const int warp = tid / 32; + // Parametric width: kWarpsPerBlock warps share a packed block, each + // covering kSubTilesPerWarp of its 4 sub-tiles (4 warps -> whole block, + // 8 warps -> half each). + constexpr int kWarpsPerBlock = int(NumTransformationThreads) / 128; + constexpr int kSubTilesPerWarp = 4 / kWarpsPerBlock; + const int blk_id = warp / kWarpsPerBlock; + const int nbi = blk_id >> 1; // n64 block within the 128-col stripe + const int kbi = blk_id & 1; // k64 block within the k tile + const int sh = (warp % kWarpsPerBlock) * kSubTilesPerWarp; + const int c = lane >> 2; // column octet within each n16 group + const int t = lane & 3; // k-quad + + // This warp's k rows fall in [kbi*64, kbi*64+64), one scale k-group for + // group sizes 64 and 128 and two consecutive groups at 32. + const int scale_kg = (kbi * kBlockK) / ScaleGranularityK; + + // Per-thread compute-buffer addressing (see compute_elem_offset). The + // word (s, j) covers columns n = nbi*64 + 16j + 8b + c and rows + // k = kbi*64 + 16s + 2t + 8p, and n%8 == c for all of them, so the + // swizzle XOR mask (c << 3) is a per-thread constant confined to the + // (k%64) bits: + // off = A(j,b) + (((16s + 8p) ^ (c<<3)) + 2t) + // A = c*64 + (nbi*8 + 2j + b)*512 + kbi*8192 + stage*16384 + // Everything except the stage term is loop-invariant. + int klow[4][2]; + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < 4; s++) { + klow[s][0] = ((16 * s) ^ (c << 3)) + 2 * t; + klow[s][1] = ((16 * s + 8) ^ (c << 3)) + 2 * t; + } + ElementAMma* const thread_c_base = smem_c_base + c * 64 + kbi * 8192; + + uint32_t skip_wait = (k_tile_count <= 0); + auto load2transform_flag = load2transform_pipeline.consumer_try_wait( + load2transform_pipeline_consumer_state, skip_wait); + auto transform2mma_flag = transform2mma_pipeline.producer_try_acquire( + transform2mma_pipeline_producer_state, skip_wait); + + CUTLASS_PRAGMA_NO_UNROLL + for (; k_tile_count > 0; --k_tile_count) { + load2transform_pipeline.consumer_wait( + load2transform_pipeline_consumer_state, load2transform_flag); + transform2mma_pipeline.producer_acquire( + transform2mma_pipeline_producer_state, transform2mma_flag); + + int load2transform_consumer_index = + load2transform_pipeline_consumer_state.index(); + int transform2mma_producer_index = + transform2mma_pipeline_producer_state.index(); + + auto curr_load2transform_pipeline_consumer_state = + load2transform_pipeline_consumer_state; + auto curr_transform2mma_pipeline_producer_state = + transform2mma_pipeline_producer_state; + + // ---- pull this thread's packed words and scales into registers + // (this warp's share of the block: sub-tiles [sh, sh+kSubTilesPerWarp)) + uint32_t w[kSubTilesPerWarp][WBits == 8 ? 8 : 4]; + uint8_t const* blk = smem_a_base + + load2transform_consumer_index * kStageBytes + + (nbi * kBlocksPerTileK + kbi) * kBlockBytes + + lane * (WBits == 8 ? 32 : 16); + CUTLASS_PRAGMA_UNROLL + for (int si = 0; si < kSubTilesPerWarp; si++) { + uint4 v = + *reinterpret_cast(blk + (sh + si) * kSubTileBytes); + w[si][0] = v.x; + w[si][1] = v.y; + w[si][2] = v.z; + w[si][3] = v.w; + if constexpr (WBits == 8) { + uint4 v1 = *reinterpret_cast( + blk + (sh + si) * kSubTileBytes + 16); + w[si][4] = v1.x; + w[si][5] = v1.y; + w[si][6] = v1.z; + w[si][7] = v1.w; + } + } + + auto bcast2 = [](NonVoidElementScale v) { + if constexpr (kActF16) { + return __half2half2(reinterpret_cast<__half const&>(v)); + } else { + return __bfloat162bfloat162( + reinterpret_cast<__nv_bfloat16 const&>(v)); + } + }; + // [group of the warp's share][j][column octet 0/1] broadcast pairs + Elem2 sreg[kScaleGroupsPerWarp][4][2]; + Elem2 zreg[kScaleGroupsPerWarp][4][2]; + CUTLASS_PRAGMA_UNROLL + for (int kg = 0; kg < kScaleGroupsPerWarp; kg++) { + NonVoidElementScale const* srow = + smem_s_base + load2transform_consumer_index * kScalesPerStage + + (scale_kg + kg) * CtaTileN_Weights + nbi * kBlockN; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < 4; j++) { + sreg[kg][j][0] = bcast2(srow[16 * j + c]); + sreg[kg][j][1] = bcast2(srow[16 * j + 8 + c]); + } + if constexpr (HasZp) { + NonVoidElementScale const* zrow = + smem_z_base + load2transform_consumer_index * kScalesPerStage + + (scale_kg + kg) * CtaTileN_Weights + nbi * kBlockN; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < 4; j++) { + zreg[kg][j][0] = bcast2(zrow[16 * j + c]); + zreg[kg][j][1] = bcast2(zrow[16 * j + 8 + c]); + } + } + } + + // Loads from SMEM are done. Signal the mainloop load as early as possible + transform_bar.sync(); + load2transform_pipeline.consumer_release( + curr_load2transform_pipeline_consumer_state); + + // ---- dequant + scale + I5 stores (4x 32-bit per packed word) + ElementAMma* const stage_base = + thread_c_base + transform2mma_producer_index * 16384; + CUTLASS_PRAGMA_UNROLL + for (int si = 0; si < kSubTilesPerWarp; si++) { + const int s = sh + si; + // Sub-tile s covers k rows [16s, 16s+16); group index within the + // warp's register set (always 0 at granularity 64 and 128). + const int kg = kScaleGroupsPerWarp == 1 ? 0 : (s / 2) % 2; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < 4; j++) { + ElementAMma* const row0 = stage_base + (nbi * 8 + 2 * j) * 512; + ElementAMma* const row1 = row0 + 512; + Elem2 f0[2], f1[2]; + if constexpr (WBits == 8 && kActF16) { + swordfish_detail::dequant_u8b128_f16x2(w[si][2 * j], f0); + swordfish_detail::dequant_u8b128_f16x2(w[si][2 * j + 1], f1); + } else if constexpr (WBits == 8) { + swordfish_detail::dequant_u8b128_bf16x2(w[si][2 * j], f0); + swordfish_detail::dequant_u8b128_bf16x2(w[si][2 * j + 1], f1); + } else if constexpr (kActF16) { + swordfish_detail::dequant_u4b8_f16x2(w[si][j], f0); + swordfish_detail::dequant_u4b8_f16x2(w[si][j] >> 8, f1); + } else { + swordfish_detail::dequant_u4b8_bf16x2(w[si][j], f0); + swordfish_detail::dequant_u4b8_bf16x2(w[si][j] >> 8, f1); + } + if constexpr (HasZp) { + f0[0] = __hfma2(f0[0], sreg[kg][j][0], zreg[kg][j][0]); + f0[1] = __hfma2(f0[1], sreg[kg][j][0], zreg[kg][j][0]); + f1[0] = __hfma2(f1[0], sreg[kg][j][1], zreg[kg][j][1]); + f1[1] = __hfma2(f1[1], sreg[kg][j][1], zreg[kg][j][1]); + } else { + f0[0] = __hmul2(f0[0], sreg[kg][j][0]); + f0[1] = __hmul2(f0[1], sreg[kg][j][0]); + f1[0] = __hmul2(f1[0], sreg[kg][j][1]); + f1[1] = __hmul2(f1[1], sreg[kg][j][1]); + } + *reinterpret_cast(row0 + klow[s][0]) = + reinterpret_cast(f0[0]); + *reinterpret_cast(row0 + klow[s][1]) = + reinterpret_cast(f0[1]); + *reinterpret_cast(row1 + klow[s][0]) = + reinterpret_cast(f1[0]); + *reinterpret_cast(row1 + klow[s][1]) = + reinterpret_cast(f1[1]); + } + } + + // Publish the generic STS to the async proxy the tcgen05 MMA reads + // through, then let the MMA know this stage is transformed. + cutlass::arch::fence_view_async_shared(); + transform2mma_pipeline.producer_commit( + curr_transform2mma_pipeline_producer_state); + + ++load2transform_pipeline_consumer_state; + ++transform2mma_pipeline_producer_state; + + skip_wait = (k_tile_count <= 1); + load2transform_flag = load2transform_pipeline.consumer_try_wait( + load2transform_pipeline_consumer_state, skip_wait); + transform2mma_flag = transform2mma_pipeline.producer_try_acquire( + transform2mma_pipeline_producer_state, skip_wait); + } + return cute::make_tuple(load2transform_pipeline_consumer_state, + transform2mma_pipeline_producer_state); + } + + /// (stock lines 1061-1144, replaced. No tiled-copy plumbing, the transform + /// addresses the byte staging and the compute buffer directly) + template + CUTLASS_DEVICE auto transform_init( + Params const& params, ProblemShape_MNKL const& problem_shape_MNKL, + Accumulator accumulators, TensorStorage& shared_storage) { + auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL); + + Tensor sA = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), + SmemLayoutA{}); + Tensor sACompute = as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), + SmemLayoutAComputeMK{})); + Tensor sS = + make_tensor(make_smem_ptr(shared_storage.input.smem_scale.begin()), + SmemLayoutScale{}); + Tensor sZ = + make_tensor(make_smem_ptr(shared_storage.input.smem_zero.begin()), + SmemLayoutScale{}); + + return cute::make_tuple(gA_mkl, sA, sACompute, sS, sZ); + } + + /// Perform a collective-scoped matrix multiply-accumulate: stock (lines + /// 1146-1232), SS descriptor path. + template + CUTLASS_DEVICE auto mma( + Load2MmaPipeline load2mma_pipeline, + Load2MmaPipelineState load2mma_pipeline_consumer_state, + Transform2MmaPipeline transform2mma_pipeline, + Transform2MmaPipelineState transform2mma_pipeline_consumer_state, + Mma2AccumPipeline mma2accum_pipeline, + Mma2AccumPipelineState mma2accum_pipeline_producer_state, + cute::Tensor const& accumulators, + cute::tuple const& input_operands, int k_tile_count) { + TiledMma tiled_mma; + + auto curr_load2mma_pipeline_consumer_state = + load2mma_pipeline_consumer_state; + auto next_load2mma_pipeline_consumer_state = + load2mma_pipeline_consumer_state; + + auto curr_transform2mma_pipeline_consumer_state = + transform2mma_pipeline_consumer_state; + auto next_transform2mma_pipeline_consumer_state = + transform2mma_pipeline_consumer_state; + + uint32_t skip_wait = (k_tile_count <= 0); + auto transform2mma_flag = transform2mma_pipeline.consumer_try_wait( + next_transform2mma_pipeline_consumer_state, skip_wait); + auto load2mma_flag = load2mma_pipeline.consumer_try_wait( + next_load2mma_pipeline_consumer_state, skip_wait); + ++next_transform2mma_pipeline_consumer_state; + ++next_load2mma_pipeline_consumer_state; + + auto const [tCrA, tCrB] = input_operands; + + mma2accum_pipeline.producer_acquire(mma2accum_pipeline_producer_state); + + int mma2accum_pipeline_producer_state_index = + mma2accum_pipeline_producer_state.index(); + auto tCtC = accumulators(_, _, _, mma2accum_pipeline_producer_state_index); + auto curr_mma2accum_pipeline_producer_state = + mma2accum_pipeline_producer_state; + ++mma2accum_pipeline_producer_state; + + tiled_mma.accumulate_ = UMMA::ScaleOut::Zero; + + CUTLASS_PRAGMA_NO_UNROLL + for (; k_tile_count > 0; --k_tile_count) { + load2mma_pipeline.consumer_wait(curr_load2mma_pipeline_consumer_state, + load2mma_flag); + transform2mma_pipeline.consumer_wait( + curr_transform2mma_pipeline_consumer_state, transform2mma_flag); + + int load2mma_pipeline_consumer_state_index = + curr_load2mma_pipeline_consumer_state.index(); + int transform2mma_pipeline_consumer_state_index = + curr_transform2mma_pipeline_consumer_state.index(); + + auto tCrA0 = tCrA(_, _, _, transform2mma_pipeline_consumer_state_index); + auto tCrB0 = tCrB(_, _, _, load2mma_pipeline_consumer_state_index); + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); k_block++) { + cute::gemm(tiled_mma, tCrA0(_, _, k_block), tCrB0(_, _, k_block), tCtC); + tiled_mma.accumulate_ = UMMA::ScaleOut::One; + } + + load2mma_pipeline.consumer_release(curr_load2mma_pipeline_consumer_state); + transform2mma_pipeline.consumer_release( + curr_transform2mma_pipeline_consumer_state); + + skip_wait = (k_tile_count <= 1); + load2mma_flag = load2mma_pipeline.consumer_try_wait( + next_load2mma_pipeline_consumer_state, skip_wait); + transform2mma_flag = transform2mma_pipeline.consumer_try_wait( + next_transform2mma_pipeline_consumer_state, skip_wait); + + curr_load2mma_pipeline_consumer_state = + next_load2mma_pipeline_consumer_state; + curr_transform2mma_pipeline_consumer_state = + next_transform2mma_pipeline_consumer_state; + + ++next_load2mma_pipeline_consumer_state; + ++next_transform2mma_pipeline_consumer_state; + } + + mma2accum_pipeline.producer_commit(curr_mma2accum_pipeline_producer_state); + + return cute::make_tuple(curr_load2mma_pipeline_consumer_state, + curr_transform2mma_pipeline_consumer_state, + mma2accum_pipeline_producer_state); + } + + /// (stock lines 1234-1255; SS branch only) + template + CUTLASS_DEVICE auto mma_init( + cute::Tensor const& accumulators, + TensorStorage& shared_storage) const { + TiledMma tiled_mma; + + Tensor sACompute = + make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), + SmemLayoutACompute{}); + Tensor tCrA = tiled_mma.make_fragment_A(sACompute); + + Tensor sB = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), + SmemLayoutB{}); + Tensor tCrB = tiled_mma.make_fragment_B(sB); + return cute::make_tuple(tCrA, tCrB); + } + + template + CUTLASS_DEVICE auto accum_init( + cute::Tensor const& accumulators, + TmemCopyAtom tmem_cp_atom, EpilogueTile epilogue_tile) { + return accumulators; + } + + private: + /// gA_mkl slot 0 is only consumed for its SHAPE (the kernel derives the + /// k-tile count from shape<3>); return an identity tensor with the stock + /// tiled shape. gB_nkl is the stock activation tiling. + template + CUTLASS_DEVICE constexpr auto tile_input_tensors( + Params const& params, ProblemShape_MNKL const& problem_shape_MNKL) const { + using X = cute::Underscore; + auto [M, N, K, L] = problem_shape_MNKL; + + auto gA_mkl = make_identity_tensor( + make_shape(size<0>(TileShape{}), size<2>(TileShape{}), + ceil_div(M, size<0>(TileShape{})), + ceil_div(K, size<2>(TileShape{})), L)); + + Tensor mB_nkl = observed_tma_load_b_->get_tma_tensor(make_shape(N, K, L)); + Tensor gB_nkl = + local_tile(mB_nkl, TileShape{}, make_coord(_, _, _), Step{}); + + return cute::make_tuple(gA_mkl, gB_nkl); + } + + typename Params::TMA_A const* observed_tma_load_a_ = nullptr; + typename Params::TMA_B const* observed_tma_load_b_ = nullptr; + uint32_t block_rank_in_cluster_ = 0; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_prepack.cu b/csrc/libtorch_stable/quantization/swordfish/swordfish_prepack.cu new file mode 100644 index 0000000000..f7059d6e31 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_prepack.cu @@ -0,0 +1,98 @@ +// Packs a GPTQ int4 weight into the Swordfish ABI v1 block-linear layout +//. Reuses the Marlin repack kernel for the in-tile +// permutation, which ABI v1 adopts verbatim, then re-tiles its flat +// [K/16, N*2] int32 layout into (NB, KB, 512) int32 blocks. Runs once at +// weight load. + +#include + +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" +#include "swordfish_abi.cuh" + +// Defined in quantization/marlin/gptq_marlin_repack.cu (same extension). +torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight, + torch::stable::Tensor& perm, + int64_t size_k, int64_t size_n, + int64_t num_bits, bool is_a_8bit); + +namespace swordfish { + +namespace { + +// One CTA per output block: 512 threads gather the block's int32 words from +// the flat Marlin layout (two words each at 8-bit). Writes are perfectly +// coalesced; reads are tile-sized runs from 4 Marlin rows. +template +__global__ void retile_marlin_to_blocks(const int32_t* __restrict__ marlin, + int32_t* __restrict__ out, + int64_t num_kb, int64_t size_n) { + constexpr int kWords = 4 * kTileInt32; + const int64_t nb = blockIdx.x; + const int64_t kb = blockIdx.y; + for (int w = threadIdx.x; w < kWords; w += kBlockInt32) { + const int64_t dst = (nb * num_kb + kb) * kWords + w; + out[dst] = marlin[marlin_word_index(nb, kb, w, size_n)]; + } +} + +} // namespace + +torch::stable::Tensor swordfish_prepack_B( + torch::stable::Tensor& b_q_weight, + std::optional const& perm, int64_t size_k, + int64_t size_n, int64_t num_bits) { + STD_TORCH_CHECK(shape_ok(size_k, size_n), "swordfish ABI v1 requires K % ", + kBlockK, " == 0 and N % ", kBlockN, " == 0; got K=", size_k, + " N=", size_n, " (v1 tail policy: reject)"); + STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, + "swordfish supports 4-bit and 8-bit weights"); + + const int32_t device_index = b_q_weight.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); + + // Stage 1: Marlin in-tile permutation. A non-empty perm applies the + // act_order row sort (g_idx argsort) during the repack. + torch::stable::Tensor perm_t = + perm.has_value() + ? *perm + : torch::stable::empty({0}, torch::headeronly::ScalarType::Int, + std::nullopt, b_q_weight.device()); + torch::stable::Tensor marlin_flat = + gptq_marlin_repack(b_q_weight, perm_t, size_k, size_n, num_bits, + /*is_a_8bit=*/false); + + // Stage 2: re-tile to (NB, KB, 512|1024) int32 blocks. + const int64_t nb = num_blocks_n(size_n); + const int64_t kb = num_blocks_k(size_k); + const int64_t words = num_bits == 8 ? kBlockInt32_8 : kBlockInt32; + torch::stable::Tensor out = + torch::stable::empty({nb, kb, words}, torch::headeronly::ScalarType::Int, + std::nullopt, b_q_weight.device()); + + dim3 grid(nb, kb); + if (num_bits == 8) { + retile_marlin_to_blocks<256><<>>( + reinterpret_cast(marlin_flat.const_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), kb, size_n); + } else { + retile_marlin_to_blocks<128><<>>( + reinterpret_cast(marlin_flat.const_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), kb, size_n); + } + + return out; +} + +} // namespace swordfish + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("swordfish_prepack_B", TORCH_BOX(&swordfish::swordfish_prepack_B)); +} diff --git a/csrc/libtorch_stable/quantization/swordfish/swordfish_types.cuh b/csrc/libtorch_stable/quantization/swordfish/swordfish_types.cuh new file mode 100644 index 0000000000..691506d347 --- /dev/null +++ b/csrc/libtorch_stable/quantization/swordfish/swordfish_types.cuh @@ -0,0 +1,54 @@ +// Swordfish packed-weight ABI v1 constants and metadata. The ABI encodes +// layout, quant metadata, alignment, and versioning only, never launch +// geometry. +#pragma once + +#include + +namespace swordfish { + +// ---- ABI v1 layout constants ------------------------------------------------ +inline constexpr int kAbiVersion = 1; +inline constexpr int kTileLayoutId = + 1; // marlin16x64 / pack_idx{0,2,4,6,1,3,5,7} + // in (NB, KB, 2048B) block-linear order + +inline constexpr int kBlockN = 64; // columns per packed block +inline constexpr int kBlockK = 64; // K rows per packed block +inline constexpr int kMarlinTileK = 16; +inline constexpr int kMarlinTileN = 64; +inline constexpr int kSubTileBytes = 512; // one marlin 16x64 int4 tile +inline constexpr int kTilesPerBlock = kBlockK / kMarlinTileK; // 4 +inline constexpr int kBlockBytes = kTilesPerBlock * kSubTileBytes; // 2048 +inline constexpr int kBlockInt32 = kBlockBytes / 4; // 512 + +// 8-bit blocks double every byte figure; the tile permutation is the same. +inline constexpr int kSubTileBytes8 = 2 * kSubTileBytes; +inline constexpr int kBlockBytes8 = 2 * kBlockBytes; +inline constexpr int kBlockInt32_8 = 2 * kBlockInt32; + +// ---- schemes (metadata enums; v1 supports exactly one of each) -------------- +enum class WeightScheme : uint8_t { kU4B8 = 1, kU8B128 = 2 }; +enum class ScaleScheme : uint8_t { + kGroupContiguous = 1 +}; // [groups][N] fp16/bf16 +enum class ZpScheme : uint8_t { kNone = 0 }; + +// feature_bits (all zero in v1) +inline constexpr uint32_t kFeatHasZp = 1u << 0; +inline constexpr uint32_t kFeatActOrder = 1u << 1; +inline constexpr uint32_t kFeatPaddedTail = 1u << 2; + +// ---- shape rules (the v1 tail policy rejects non-multiples) +// ------------------ +__host__ __device__ inline constexpr bool shape_ok(int64_t k, int64_t n) { + return k > 0 && n > 0 && (k % kBlockK) == 0 && (n % kBlockN) == 0; +} +__host__ __device__ inline constexpr int64_t num_blocks_n(int64_t n) { + return n / kBlockN; +} +__host__ __device__ inline constexpr int64_t num_blocks_k(int64_t k) { + return k / kBlockK; +} + +} // namespace swordfish diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 032567677a..d675a0a6b0 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -108,6 +108,47 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "marlin_int4_fp8_preprocess(Tensor qweight, " "Tensor? qzeros_or_none, bool inplace) -> Tensor"); // conditionally compiled so impl registrations are in source file + + // swordfish (sm100/sm110 w4a16): pack a GPTQ int4 weight into the + // Swordfish ABI v1 block-linear layout. + ops.def( + "swordfish_prepack_B(Tensor b_q_weight, Tensor? perm, SymInt size_k, " + "SymInt size_n, int num_bits) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // swordfish w4a16 decode GEMM over the ABI v1 packed weight. group_zps + // holds prescaled (8 - zp) * scale per group when present (AWQ/HQQ). + ops.def( + "swordfish_mm(Tensor a, Tensor b_packed, Tensor group_scales, " + "Tensor? group_zps, Tensor? perm, int num_bits, int group_size, " + "SymInt size_k, SymInt size_n) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // swordfish dequant to dense fp16/bf16, optionally out-major transposed + // and expert-stacked, for the dense tiers. + ops.def( + "swordfish_dequant_dense(Tensor b_packed, Tensor group_scales, " + "Tensor? group_zps, int num_bits, int group_size, SymInt size_k, " + "SymInt size_n, bool transpose) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // swordfish fused-MoE decode GEMM: one persistent Stream-K launch over + // per-expert ABI v1 weights, token-sorted by moe_align_block_size(16). + ops.def( + "swordfish_moe_mm(Tensor a, Tensor b_packed, Tensor group_scales, " + "Tensor sorted_token_ids, Tensor expert_ids, " + "Tensor num_tokens_post_padded, Tensor? topk_weights, " + "int moe_block_size, int top_k, bool mul_topk_weights, int num_bits, " + "int group_size, SymInt size_k, SymInt size_n) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // swordfish w4a16 prefill GEMM (sm100 tcgen05 mixed-input mainloop fork) + // over the same ABI v1 packed weight. + ops.def( + "swordfish_prefill_mm(Tensor a, Tensor b_packed, Tensor group_scales, " + "Tensor? group_zps, int num_bits, int group_size, SymInt size_k, " + "SymInt size_n) -> Tensor"); + // conditionally compiled so impl registrations are in source file #endif #ifndef USE_ROCM diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 97fe239ac2..97d6cdebf8 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -159,7 +159,7 @@ cuda-bindings==13.0.3 # via torch cuda-pathfinder==1.3.3 # via cuda-bindings -cuda-toolkit==13.0.2 +cuda-toolkit==13.0.3.0 # via torch cupy-cuda12x==13.6.0 # via ray @@ -603,7 +603,7 @@ numpy==2.2.6 # tritonclient # vocos # xgrammar -nvidia-cublas==13.1.0.3 +nvidia-cublas==13.1.1.3 # via # cuda-toolkit # nvidia-cudnn-cu13 @@ -611,10 +611,12 @@ nvidia-cublas==13.1.0.3 nvidia-cuda-cupti==13.0.85 # via cuda-toolkit nvidia-cuda-nvrtc==13.0.88 - # via cuda-toolkit + # via + # cuda-toolkit + # nvidia-cublas nvidia-cuda-runtime==13.0.96 # via cuda-toolkit -nvidia-cudnn-cu13==9.19.0.56 +nvidia-cudnn-cu13==9.20.0.48 # via torch nvidia-cufft==12.0.0.61 # via cuda-toolkit @@ -628,9 +630,9 @@ nvidia-cusparse==12.6.3.3 # via # cuda-toolkit # nvidia-cusolver -nvidia-cusparselt-cu13==0.8.0 +nvidia-cusparselt-cu13==0.8.1 # via torch -nvidia-nccl-cu13==2.28.9 +nvidia-nccl-cu13==2.29.7 # via torch nvidia-nvjitlink==13.0.88 # via diff --git a/tests/kernels/quantization/test_swordfish_mm.py b/tests/kernels/quantization/test_swordfish_mm.py new file mode 100644 index 0000000000..1b4b729792 --- /dev/null +++ b/tests/kernels/quantization/test_swordfish_mm.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Swordfish decode GEMM correctness: swordfish_mm vs a dequantized reference. + +The oracle is A @ w_ref where w_ref is the fp reconstruction of the quantized +weight (from swordfish_quantize). Tolerances mirror test_machete_mm. +""" + +import pytest +import torch + +from aphrodite import _custom_ops as ops +from aphrodite.model_executor.layers.quantization.utils.swordfish_utils_test import ( + swordfish_quantize, + swordfish_quantize_act_order, + swordfish_quantize_awq, +) +from aphrodite.platforms import current_platform +from aphrodite.scalar_type import scalar_types +from tests.kernels.utils import opcheck + +if not current_platform.is_cuda(): + pytest.skip(reason="swordfish requires CUDA", allow_module_level=True) +if not current_platform.has_device_capability(100): + pytest.skip(reason="swordfish requires sm100 family", allow_module_level=True) + +DEVICE = "cuda" +QT = scalar_types.uint4b8 + +# (M, K, N) decode sweep including regime boundaries, plus K/N variety +MNK = [ + (1, 256, 128), + (52, 1024, 256), + (60, 11008, 256), + (2, 256, 128), + (4, 512, 256), + (8, 4096, 4096), + (16, 4096, 4096), + (17, 512, 128), + (32, 4096, 512), + (33, 256, 256), +] +GROUPS = [-1, 32, 64, 128] +DTYPES = [torch.float16, torch.bfloat16] + + +@pytest.mark.parametrize("mnk", MNK) +@pytest.mark.parametrize("group", GROUPS) +@pytest.mark.parametrize("dtype", DTYPES) +def test_swordfish_mm_correct(mnk, group, dtype): + m, k, n = mnk + if group != -1 and k % group != 0: + pytest.skip("k not divisible by group") + + torch.manual_seed(k * 13 + n + m) + w = torch.randn((k, n), dtype=dtype, device=DEVICE) / (k**0.5) + w_ref, packed, scales = swordfish_quantize(w, QT, group) + + a = torch.randn((m, k), dtype=dtype, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + + out = ops.swordfish_mm(a, packed, scales, group, k, n) + + assert out.shape == (m, n) + assert out.dtype == dtype + # machete-style tolerance + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=5e-2 if dtype == torch.float16 else 8e-2) + + +def test_swordfish_mm_opcheck(): + # M > 96 with fp16 exercises the deterministic decode epilogue. The + # atomic window's summation order varies run to run, which opcheck's + # trace comparison flags, and fp16 never routes to the prefill kernel. + m, k, n = 128, 512, 128 + torch.manual_seed(1) + w = torch.randn((k, n), dtype=torch.float16, device=DEVICE) / (k**0.5) + _, packed, scales = swordfish_quantize(w, QT, 128) + a = torch.randn((m, k), dtype=torch.float16, device=DEVICE) + opcheck(torch.ops._C.swordfish_mm, (a, packed, scales, None, None, 4, 128, k, n)) + + +@pytest.mark.parametrize("dtype", DTYPES) +def test_swordfish_mm_determinism(dtype): + # M > 96 exercises the deterministic paths, fp16 through the + # smem-reduction decode epilogue and bf16 through the tcgen05 prefill. + # The window below uses the atomic epilogue, which is not run-stable by + # design. + m, k, n = 128, 1024, 256 + torch.manual_seed(9) + w = torch.randn((k, n), dtype=dtype, device=DEVICE) / (k**0.5) + _, packed, scales = swordfish_quantize(w, QT, 128) + a = torch.randn((m, k), dtype=dtype, device=DEVICE) + o0 = ops.swordfish_mm(a, packed, scales, 128, k, n) + for _ in range(5): + assert torch.equal(o0, ops.swordfish_mm(a, packed, scales, 128, k, n)) + + +@pytest.mark.parametrize("mnk", MNK) +@pytest.mark.parametrize("group", [64, 128]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_swordfish_mm_awq_correct(mnk, group, dtype): + m, k, n = mnk + if k % group != 0: + pytest.skip("k not divisible by group") + + torch.manual_seed(k * 7 + n + m) + w = torch.randn((k, n), dtype=dtype, device=DEVICE) / (k**0.5) + w_ref, packed, scales, zps_neg = swordfish_quantize_awq(w, group) + + a = torch.randn((m, k), dtype=dtype, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + + out = ops.swordfish_mm(a, packed, scales, group, k, n, group_zps=zps_neg) + + assert out.shape == (m, n) + assert out.dtype == dtype + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=5e-2 if dtype == torch.float16 else 8e-2) + + +def test_swordfish_mm_awq_large_m(): + # M > 96 with bf16 routes through the prefill zero-point row. + m, k, n = 256, 1024, 256 + torch.manual_seed(3) + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, scales, zps_neg = swordfish_quantize_awq(w, 128) + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + out = ops.swordfish_mm(a, packed, scales, 128, k, n, group_zps=zps_neg) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=8e-2) + + +@pytest.mark.parametrize("mnk", MNK) +@pytest.mark.parametrize("group", GROUPS) +@pytest.mark.parametrize("dtype", DTYPES) +def test_swordfish_mm_8bit_correct(mnk, group, dtype): + m, k, n = mnk + if group != -1 and k % group != 0: + pytest.skip("k not divisible by group") + + torch.manual_seed(k * 3 + n + m) + w = torch.randn((k, n), dtype=dtype, device=DEVICE) / (k**0.5) + w_ref, packed, scales = swordfish_quantize(w, scalar_types.uint8b128, group) + + a = torch.randn((m, k), dtype=dtype, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + + out = ops.swordfish_mm(a, packed, scales, group, k, n, num_bits=8) + + assert out.shape == (m, n) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=5e-2 if dtype == torch.float16 else 8e-2) + + +def test_swordfish_mm_8bit_large_m(): + # M > 96 with bf16 routes through the 8-bit prefill mainloop. + m, k, n = 256, 1024, 256 + torch.manual_seed(5) + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, scales = swordfish_quantize(w, scalar_types.uint8b128, 128) + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + out = ops.swordfish_mm(a, packed, scales, 128, k, n, num_bits=8) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=8e-2) + + +@pytest.mark.parametrize( + "mnk", [(1, 512, 256), (8, 512, 256), (16, 4096, 512), (33, 2048, 1024), (256, 2048, 512), (512, 1024, 512)] +) +@pytest.mark.parametrize("bits", [4, 8]) +def test_swordfish_mm_act_order(mnk, bits): + # The row sort realigns group boundaries, so the kernel runs the plain + # grouped path; only the activation columns carry the permutation. + m, k, n = mnk + qt = scalar_types.uint4b8 if bits == 4 else scalar_types.uint8b128 + torch.manual_seed(k + n + m + bits) + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, scales, sort_indices = swordfish_quantize_act_order(w, qt, 128) + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + out = ops.swordfish_mm(a, packed, scales, 128, k, n, num_bits=bits, perm=sort_indices) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=8e-2) + + +@pytest.mark.parametrize("bits", [4, 8]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_swordfish_mm_dense_tier(bits, dtype, monkeypatch): + # Force the dequant + cuBLAS tier regardless of device size. + monkeypatch.setenv("APHRODITE_SWORDFISH_DENSE_M", "64") + m, k, n = 128, 1024, 512 + qt = scalar_types.uint4b8 if bits == 4 else scalar_types.uint8b128 + torch.manual_seed(bits + m) + w = torch.randn((k, n), dtype=dtype, device=DEVICE) / (k**0.5) + w_ref, packed, scales = swordfish_quantize(w, qt, 128) + a = torch.randn((m, k), dtype=dtype, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + out = ops.swordfish_mm(a, packed, scales, 128, k, n, num_bits=bits) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=5e-2 if dtype == torch.float16 else 8e-2) + + +def test_swordfish_mm_dense_tier_awq(monkeypatch): + monkeypatch.setenv("APHRODITE_SWORDFISH_DENSE_M", "64") + m, k, n = 128, 1024, 512 + torch.manual_seed(2) + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, scales, zps_neg = swordfish_quantize_awq(w, 128) + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + out = ops.swordfish_mm(a, packed, scales, 128, k, n, group_zps=zps_neg) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=8e-2) + + +def test_swordfish_mm_dense_tier_act_order(monkeypatch): + # The dense tier scatters weight rows through the sort instead of + # permuting the activations. + monkeypatch.setenv("APHRODITE_SWORDFISH_DENSE_M", "64") + m, k, n = 128, 1024, 512 + torch.manual_seed(7) + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, scales, sort_indices = swordfish_quantize_act_order(w, scalar_types.uint4b8, 128) + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + out = ops.swordfish_mm(a, packed, scales, 128, k, n, perm=sort_indices) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=8e-2) + + +@pytest.mark.parametrize("m", [1, 24, 64, 256]) +def test_swordfish_mm_channelwise_replication(m): + # The python layer replicates a channelwise scale row to group 128 and + # passes group -1: the op runs the grouped tiers over the replicated + # rows and native channelwise decode over row 0. All presentations must + # agree with the reference. + k, n = 1024, 512 + torch.manual_seed(4) + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, scales = swordfish_quantize(w, QT, -1) + rep = scales.expand(k // 128, n).contiguous() + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + o_cw = ops.swordfish_mm(a, packed, scales, -1, k, n) + o_rep = ops.swordfish_mm(a, packed, rep, -1, k, n) + o_g = ops.swordfish_mm(a, packed, rep, 128, k, n) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + torch.testing.assert_close(o_g.to(torch.float32), ref, rtol=1e-1, atol=8e-2) + torch.testing.assert_close(o_rep.to(torch.float32), ref, rtol=1e-1, atol=8e-2) + torch.testing.assert_close(o_cw.to(torch.float32), ref, rtol=1e-1, atol=8e-2) diff --git a/tests/kernels/quantization/test_swordfish_moe.py b/tests/kernels/quantization/test_swordfish_moe.py new file mode 100644 index 0000000000..cee04a0a01 --- /dev/null +++ b/tests/kernels/quantization/test_swordfish_moe.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""swordfish_moe_mm correctness: per-expert dense references over a random +router assignment, exercised through the real moe_align_block_size buffers.""" + +import pytest +import torch + +from aphrodite import _custom_ops as ops +from aphrodite.model_executor.layers.fused_moe.moe_align_block_size import ( + moe_align_block_size, +) +from aphrodite.model_executor.layers.quantization.utils.swordfish_utils_test import ( + swordfish_quantize, +) +from aphrodite.platforms import current_platform +from aphrodite.scalar_type import scalar_types + +if not current_platform.is_cuda(): + pytest.skip(reason="swordfish requires CUDA", allow_module_level=True) +if not current_platform.has_device_capability(100): + pytest.skip(reason="swordfish requires sm100 family", allow_module_level=True) + +DEVICE = "cuda" + + +def _moe_ref(a, w_refs, topk_ids, topk_weights, top_k, mul_topk_weights): + m, k = a.shape + n = w_refs[0].shape[1] + out = torch.zeros((m * top_k, n), dtype=torch.float32, device=a.device) + for t in range(m): + for slot in range(top_k): + e = int(topk_ids[t, slot]) + v = a[t].float() @ w_refs[e].float() + if mul_topk_weights: + v = v * topk_weights[t, slot] + out[t * top_k + slot] = v + return out + + +@pytest.mark.parametrize("m", [1, 7, 33, 128]) +@pytest.mark.parametrize("bits", [4, 8]) +@pytest.mark.parametrize("mul_topk", [False, True]) +@pytest.mark.parametrize("block", [16, 32]) +def test_swordfish_moe_mm(m, bits, mul_topk, block): + e, top_k, k, n, group = 8, 2, 512, 256, 128 + qt = scalar_types.uint4b8 if bits == 4 else scalar_types.uint8b128 + torch.manual_seed(m * 31 + bits + mul_topk) + + w_refs, packs, scales = [], [], [] + for _ in range(e): + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, s = swordfish_quantize(w, qt, group) + w_refs.append(w_ref) + packs.append(packed) + scales.append(s) + b_packed = torch.stack(packs).contiguous() + group_scales = torch.stack(scales).contiguous() + + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + router = torch.randn((m, e), device=DEVICE) + topk_weights, topk_ids = torch.topk(torch.softmax(router, -1), top_k) + topk_weights = topk_weights.float().contiguous() + topk_ids = topk_ids.to(torch.int32).contiguous() + + sorted_ids, expert_ids, num_post_padded = moe_align_block_size(topk_ids, block, e) + + out = ops.swordfish_moe_mm( + a, + b_packed, + group_scales, + sorted_ids, + expert_ids, + num_post_padded, + topk_weights if mul_topk else None, + block, + top_k, + mul_topk, + bits, + group, + k, + n, + ) + ref = _moe_ref(a, w_refs, topk_ids, topk_weights, top_k, mul_topk) + + assert out.shape == (m * top_k, n) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=8e-2) + + +def test_swordfish_moe_mm_top1_w2_style(): + # The w2 GEMM runs with top_k=1 over the [M * topk, N] intermediate. + e, k, n, group = 4, 256, 128, 128 + torch.manual_seed(11) + w_refs, packs, scales = [], [], [] + for _ in range(e): + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, s = swordfish_quantize(w, scalar_types.uint4b8, group) + w_refs.append(w_ref) + packs.append(packed) + scales.append(s) + b_packed = torch.stack(packs).contiguous() + group_scales = torch.stack(scales).contiguous() + + m = 48 + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + ids = torch.randint(0, e, (m, 1), dtype=torch.int32, device=DEVICE) + weights = torch.rand((m, 1), dtype=torch.float32, device=DEVICE) + sorted_ids, expert_ids, num_post_padded = moe_align_block_size(ids, 16, e) + + out = ops.swordfish_moe_mm( + a, + b_packed, + group_scales, + sorted_ids, + expert_ids, + num_post_padded, + weights, + 16, + 1, + True, + 4, + group, + k, + n, + ) + ref = _moe_ref(a, w_refs, ids, weights, 1, True) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=8e-2) diff --git a/tests/kernels/quantization/test_swordfish_prefill_mm.py b/tests/kernels/quantization/test_swordfish_prefill_mm.py new file mode 100644 index 0000000000..63ded71875 --- /dev/null +++ b/tests/kernels/quantization/test_swordfish_prefill_mm.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Swordfish prefill GEMM correctness: swordfish_prefill_mm (sm100 tcgen05 +mixed-input mainloop fork) vs a dequantized reference. + +Same oracle as test_swordfish_mm (A @ w_ref from swordfish_quantize, +machete-style tolerances) at prefill-sized M. v1 scope: bf16, group_size 128, +K % 128 == 0, N % 128 == 0. +""" + +import pytest +import torch + +from aphrodite import _custom_ops as ops +from aphrodite.model_executor.layers.quantization.utils.swordfish_utils_test import ( + swordfish_quantize, + swordfish_quantize_awq, +) +from aphrodite.platforms import current_platform +from aphrodite.scalar_type import scalar_types +from tests.kernels.utils import opcheck + +if not current_platform.is_cuda(): + pytest.skip(reason="swordfish requires CUDA", allow_module_level=True) +if not current_platform.has_device_capability(100): + pytest.skip(reason="swordfish requires sm100 family", allow_module_level=True) + +DEVICE = "cuda" +QT = scalar_types.uint4b8 +GROUP = 128 + +# (M, K, N) with prefill-sized M plus non-multiple-of-tile M tails +MNK = [ + (256, 4096, 4096), + (1024, 4096, 4096), + (1024, 4096, 11008), + (333, 2048, 4096), + (128, 256, 128), +] + + +@pytest.mark.parametrize("mnk", MNK) +def test_swordfish_prefill_mm_correct(mnk): + m, k, n = mnk + + torch.manual_seed(k * 13 + n + m) + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, scales = swordfish_quantize(w, QT, GROUP) + + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + + out = ops.swordfish_prefill_mm(a, packed, scales, GROUP, k, n) + + assert out.shape == (m, n) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=5e-2) + + +def test_swordfish_prefill_mm_opcheck(): + m, k, n = 256, 512, 256 + torch.manual_seed(7) + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + _, packed, scales = swordfish_quantize(w, QT, GROUP) + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + opcheck( + torch.ops._C.swordfish_prefill_mm, + (a, packed, scales, None, 4, GROUP, k, n), + ) + + +def test_swordfish_prefill_mm_rejects_bad_args(): + k, n = 512, 256 + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) + _, packed, scales = swordfish_quantize(w, QT, GROUP) + a = torch.randn((16, k), dtype=torch.bfloat16, device=DEVICE) + with pytest.raises(Exception, match="dtype"): + ops.swordfish_prefill_mm(a.to(torch.float16), packed, scales, GROUP, k, n) + with pytest.raises(Exception, match="group"): + ops.swordfish_prefill_mm(a, packed, scales, 48, k, n) + + +@pytest.mark.parametrize("mnk", [(256, 512, 256), (128, 2048, 1024), (512, 11008, 2048)]) +def test_swordfish_prefill_mm_awq_correct(mnk): + m, k, n = mnk + torch.manual_seed(m + k + n) + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, scales, zps_neg = swordfish_quantize_awq(w, GROUP) + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + out = ops.swordfish_prefill_mm(a, packed, scales, GROUP, k, n, group_zps=zps_neg) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=8e-2) + + +@pytest.mark.parametrize("mnk", [(256, 512, 256), (128, 2048, 1024), (512, 11008, 2048)]) +def test_swordfish_prefill_mm_8bit_correct(mnk): + m, k, n = mnk + torch.manual_seed(m + k + n + 1) + w = torch.randn((k, n), dtype=torch.bfloat16, device=DEVICE) / (k**0.5) + w_ref, packed, scales = swordfish_quantize(w, scalar_types.uint8b128, GROUP) + a = torch.randn((m, k), dtype=torch.bfloat16, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + out = ops.swordfish_prefill_mm(a, packed, scales, GROUP, k, n, num_bits=8) + torch.testing.assert_close(out.to(torch.float32), ref, rtol=1e-1, atol=8e-2) + + +@pytest.mark.parametrize("mnk", [(256, 512, 256), (128, 2048, 1024)]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("group", [32, 64, 128]) +def test_swordfish_prefill_mm_dtypes_groups(mnk, dtype, group): + m, k, n = mnk + torch.manual_seed(m + k + n + group) + w = torch.randn((k, n), dtype=dtype, device=DEVICE) / (k**0.5) + w_ref, packed, scales = swordfish_quantize(w, QT, group) + a = torch.randn((m, k), dtype=dtype, device=DEVICE) + ref = a.to(torch.float32) @ w_ref.to(torch.float32) + out = ops.swordfish_prefill_mm(a, packed, scales, group, k, n) + torch.testing.assert_close( + out.to(torch.float32), + ref, + rtol=1e-1, + atol=5e-2 if dtype == torch.float16 else 8e-2, + ) diff --git a/tests/kernels/quantization/test_swordfish_prepack.py b/tests/kernels/quantization/test_swordfish_prepack.py new file mode 100644 index 0000000000..aa307dc710 --- /dev/null +++ b/tests/kernels/quantization/test_swordfish_prepack.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Swordfish ABI v1 prepack: bit-exact contract tests. + +The CUDA op `swordfish_prepack_B` must reproduce the pure-torch reference +(`swordfish_pack_weights_ref`) bit-exactly — the ABI is frozen against these +tests (csrc/libtorch_stable/quantization/swordfish/docs/abi-design.md). +""" + +import pytest +import torch + +from aphrodite import _custom_ops as ops +from aphrodite.model_executor.layers.quantization.utils.marlin_utils_test import ( + get_weight_perm, + marlin_permute_weights, +) +from aphrodite.model_executor.layers.quantization.utils.quant_utils import pack_rows +from aphrodite.model_executor.layers.quantization.utils.swordfish_utils_test import ( + SWORDFISH_BLOCK_INT32, + swordfish_pack_weights_ref, + swordfish_shape_ok, +) +from aphrodite.platforms import current_platform +from tests.kernels.utils import opcheck + +if not current_platform.is_cuda(): + pytest.skip(reason="swordfish requires CUDA", allow_module_level=True) + +if not current_platform.has_device_capability(100): + pytest.skip( + reason="swordfish requires Blackwell (sm100 family)", + allow_module_level=True, + ) + +# (K, N) coverage includes the minimum, odd tile multiples, and typical +# model shapes +SHAPES = [ + (64, 64), + (192, 320), + (256, 128), + (2048, 4096), + (4096, 4096), +] + +DEVICE = "cuda" + + +def _random_q4(size_k: int, size_n: int, seed: int) -> torch.Tensor: + g = torch.Generator(device="cpu").manual_seed(seed) + return torch.randint(0, 16, (size_k, size_n), generator=g, dtype=torch.int32) + + +def _gptq_pack(q_w: torch.Tensor, size_k: int, size_n: int) -> torch.Tensor: + return pack_rows(q_w, 4, size_k, size_n) + + +@pytest.mark.parametrize("shape", SHAPES) +def test_prepack_bit_exact(shape): + size_k, size_n = shape + q_w = _random_q4(size_k, size_n, seed=size_k * 31 + size_n) + + ref = swordfish_pack_weights_ref(q_w, size_k, size_n) + + gptq = _gptq_pack(q_w, size_k, size_n).to(DEVICE) + got = ops.swordfish_prepack_B(gptq, size_k, size_n) + + assert got.dtype == torch.int32 + assert tuple(got.shape) == (size_n // 64, size_k // 64, SWORDFISH_BLOCK_INT32) + assert torch.equal(got.cpu(), ref.cpu()), ( + f"prepack mismatch at shape {shape}: first diff word {torch.nonzero(got.cpu() != ref.cpu())[0].tolist()}" + ) + + +@pytest.mark.parametrize("shape", [(256, 128), (2048, 4096)]) +def test_prepack_roundtrip(shape): + """Unpack the packed ABI back to the original int4 codes via the + index-tracked inverse of the reference permutation.""" + size_k, size_n = shape + q_w = _random_q4(size_k, size_n, seed=7) + + packed = swordfish_pack_weights_ref(q_w, size_k, size_n) + + # Rebuild the flat Marlin layout from blocks (inverse of stage 2). + nb, kb = size_n // 64, size_k // 64 + marlin_flat = packed.reshape(nb, kb, 4, 128).permute(1, 2, 0, 3).reshape(size_k // 16, size_n * 2) + + # Index-tracked inverse of stage 1, the permutation run on an index grid. + perm = get_weight_perm(num_bits=4) + idx = torch.arange(size_k * size_n, dtype=torch.int64).reshape(size_k, size_n) + idx_perm = marlin_permute_weights(idx, size_k, size_n, perm) # (K/16, N*16) + + # Nibble-unpack marlin_flat in pack order and scatter to original slots. + words = marlin_flat.reshape(-1).to(torch.int64) + nibbles = torch.stack([(words >> (4 * i)) & 0xF for i in range(8)], dim=1) + recovered = torch.empty(size_k * size_n, dtype=torch.int32) + recovered[idx_perm.reshape(-1)] = nibbles.reshape(-1).to(torch.int32) + assert torch.equal(recovered.reshape(size_k, size_n), q_w) + + +@pytest.mark.parametrize("shape", [(96, 64), (64, 96), (100, 100)]) +def test_prepack_rejects_bad_shapes(shape): + size_k, size_n = shape + assert not swordfish_shape_ok(size_k, size_n) + # build a plausibly-shaped gptq tensor; op must reject on k/n args + gptq = torch.zeros((max(size_k // 8, 1), size_n), dtype=torch.int32, device=DEVICE) + with pytest.raises(RuntimeError): + ops.swordfish_prepack_B(gptq, size_k, size_n) + + +def test_prepack_opcheck(): + size_k, size_n = 256, 128 + q_w = _random_q4(size_k, size_n, seed=3) + gptq = _gptq_pack(q_w, size_k, size_n).to(DEVICE) + opcheck( + torch.ops._C.swordfish_prepack_B, + (gptq, None, size_k, size_n, 4), + ) + + +@pytest.mark.parametrize("shape", [(64, 64), (256, 128), (2048, 4096)]) +def test_prepack_bit_exact_8bit(shape): + size_k, size_n = shape + g = torch.Generator(device="cpu").manual_seed(size_k + size_n) + q_w = torch.randint(0, 256, (size_k, size_n), generator=g, dtype=torch.int32) + ref = swordfish_pack_weights_ref(q_w, size_k, size_n, num_bits=8) + gptq = pack_rows(q_w, 8, size_k, size_n).to(DEVICE) + got = ops.swordfish_prepack_B(gptq, size_k, size_n, num_bits=8) + assert got.shape == (size_n // 64, size_k // 64, 1024) + assert torch.equal(got.cpu(), ref)