diff --git a/aphrodite/config/model.py b/aphrodite/config/model.py index a8a4d6363e..b81e201e1e 100644 --- a/aphrodite/config/model.py +++ b/aphrodite/config/model.py @@ -219,7 +219,8 @@ class ModelConfig: determine the data type of the weights.""" quantization_config: dict[str, Any] | QuantizationConfigArgs | None = None """User-facing quantization configuration. Carries per-layer-kind specs - (linear, moe) and ignore patterns; see :class:`QuantizationConfigArgs`. + (linear, moe), ignore patterns, and ordered precision overrides; see + :class:`QuantizationConfigArgs`. Auto-populated from the matching online shorthand when `quantization` is one of the values in `ONLINE_QUANT_SHORTHAND_NAMES`.""" allow_deprecated_quantization: bool = False diff --git a/aphrodite/config/quantization.py b/aphrodite/config/quantization.py index 5d62c12184..0c1bf05cec 100644 --- a/aphrodite/config/quantization.py +++ b/aphrodite/config/quantization.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# mypy: disable-error-code=call-arg -from typing import Annotated, Any +from typing import Annotated, Any, Literal from pydantic import Field, GetPydanticSchema, ValidationInfo, field_validator from pydantic_core import core_schema @@ -17,6 +18,10 @@ kFp8StaticTensorSym, kInt8StaticChannelSym, kMxfp4Dynamic, + kMxfp6E2m3Dynamic, + kMxfp6E2m3Static, + kMxfp6E3m2Dynamic, + kMxfp6E3m2Static, kMxfp8Dynamic, kNvfp4Static, ) @@ -31,6 +36,10 @@ "fp8_per_block_dynamic": kFp8Dynamic128Sym, "mxfp8": kMxfp8Dynamic, "mxfp4": kMxfp4Dynamic, + "mxfp6_e2m3": kMxfp6E2m3Static, + "mxfp6_e2m3_dynamic": kMxfp6E2m3Dynamic, + "mxfp6_e3m2": kMxfp6E3m2Static, + "mxfp6_e3m2_dynamic": kMxfp6E3m2Dynamic, "int8_per_channel_static": kInt8StaticChannelSym, } @@ -69,6 +78,32 @@ class QuantSpec: """Activation quantization key, or a name from QUANT_KEY_NAMES.""" +def _coerce_override_weight(v: Any) -> Any: + if v in (None, "bf16"): + return v + return _coerce_quant_key(v) + + +OverrideWeightField = Annotated[ + QuantKey | Literal["bf16"] | None, + GetPydanticSchema(lambda _src, _handler: core_schema.no_info_plain_validator_function(_coerce_override_weight)), +] + + +@config +class QuantOverride: + """Ordered module-level override for online quantization.""" + + pattern: str = "" + """Exact module prefix or ``re:`` regular expression.""" + + weight: OverrideWeightField = None + """Replacement weight format; ``bf16`` leaves the module unquantized.""" + + activation: QuantKeyField = None + """Replacement activation format; omitted fields inherit the base spec.""" + + @config class QuantizationConfigArgs: """User-facing quantization configuration. @@ -86,6 +121,9 @@ class QuantizationConfigArgs: ignore: list[str] = Field(default_factory=list) """Layers to skip quantization for.""" + overrides: list[QuantOverride] = Field(default_factory=list) + """Ordered module precision overrides. Later matching rules win.""" + @field_validator("linear", "moe", mode="before") @classmethod def _coerce_spec(cls, v: Any, info: ValidationInfo) -> Any: @@ -122,6 +160,16 @@ def _coerce_spec(cls, v: Any, info: ValidationInfo) -> Any: linear=QuantSpec(weight=kMxfp8Dynamic), moe=QuantSpec(weight=kMxfp8Dynamic), ), + "mxfp6": QuantizationConfigArgs( + linear=QuantSpec(weight=kMxfp6E2m3Static, activation=kMxfp8Dynamic), + moe=QuantSpec(weight=kMxfp6E2m3Static, activation=kMxfp8Dynamic), + overrides=[ + QuantOverride( + pattern=r"re:(^|.*\.)(gate|router|shared_expert_gate|lm_head)$", + weight="bf16", + ) + ], + ), # INT8 weight-only on MoE; linear stays unquantized (no `linear` field). "int8_per_channel_weight_only": QuantizationConfigArgs( moe=QuantSpec(weight=kInt8StaticChannelSym), @@ -177,4 +225,5 @@ def resolve_quantization_config( linear=quantization_config.linear or base.linear, moe=quantization_config.moe or base.moe, ignore=quantization_config.ignore or base.ignore, + overrides=[*base.overrides, *quantization_config.overrides], ) diff --git a/aphrodite/engine/arg_utils.py b/aphrodite/engine/arg_utils.py index beb1ff190b..9e8cb679f8 100644 --- a/aphrodite/engine/arg_utils.py +++ b/aphrodite/engine/arg_utils.py @@ -513,7 +513,7 @@ class EngineArgs: quantization: QuantizationMethods | str | None = ModelConfig.quantization quantization_config: "dict[str, Any] | QuantizationConfigArgs | None" = None """User-facing quantization configuration. Carries per-layer-kind - QuantSpecs (linear, moe) and ignore patterns; see + QuantSpecs (linear, moe), ignore patterns, and ordered precision overrides; see :class:`QuantizationConfigArgs`. Auto-populated from the matching online shorthand when `quantization` is one of the values in `ONLINE_QUANT_SHORTHAND_NAMES`.""" diff --git a/aphrodite/model_executor/kernels/linear/mxfp6/SOURCE.md b/aphrodite/model_executor/kernels/linear/mxfp6/SOURCE.md new file mode 100644 index 0000000000..9b9f455697 --- /dev/null +++ b/aphrodite/model_executor/kernels/linear/mxfp6/SOURCE.md @@ -0,0 +1,18 @@ + + + +# Vendored CuTe DSL source + +`cutedsl_kernel.py` is derived from NVIDIA CUTLASS commit +`f94ec46f4f63f96003d6cfdf2014731e7672c281`: + +`examples/python/CuTeDSL/cute/blackwell/kernel/blockscaled_gemm/dense_blockscaled_gemm_persistent.py` + +Sonar integrates the kernel with its packed-weight ABI and Torch custom +operator in `cutedsl.py`. The upstream command-line reference helpers remain +in the vendored file to make future CUTLASS updates easier. + +`cutedsl_grouped_kernel.py` derives from the grouped block-scaled GEMM example +at the same CUTLASS commit. It ports mixed-width operand and FP6 TMA-unpack +support from the dense implementation. NVIDIA's grouped example only accepted +same-width operands at that revision. diff --git a/aphrodite/model_executor/kernels/linear/mxfp6/__init__.py b/aphrodite/model_executor/kernels/linear/mxfp6/__init__.py new file mode 100644 index 0000000000..9f092ae0a3 --- /dev/null +++ b/aphrodite/model_executor/kernels/linear/mxfp6/__init__.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .base import Mxfp6LinearKernel, Mxfp6LinearLayerConfig +from .cutedsl import CutedslMxfp6LinearKernel +from .cutedsl_grouped import cutedsl_grouped_mxfp6_gemm + +__all__ = [ + "CutedslMxfp6LinearKernel", + "cutedsl_grouped_mxfp6_gemm", + "Mxfp6LinearKernel", + "Mxfp6LinearLayerConfig", +] diff --git a/aphrodite/model_executor/kernels/linear/mxfp6/base.py b/aphrodite/model_executor/kernels/linear/mxfp6/base.py new file mode 100644 index 0000000000..8b761a32e6 --- /dev/null +++ b/aphrodite/model_executor/kernels/linear/mxfp6/base.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Literal + +import torch + + +@dataclass(frozen=True) +class Mxfp6LinearLayerConfig: + weight_format: Literal["e2m3", "e3m2"] = "e2m3" + activation_format: Literal["mxfp8", "mxfp6_e2m3", "mxfp6_e3m2"] = "mxfp8" + + +class Mxfp6LinearKernel(ABC): + def __init__(self, config: Mxfp6LinearLayerConfig) -> None: + supported, reason = self.is_supported() + if not supported: + raise ValueError(reason) + self.config = config + + @classmethod + @abstractmethod + def is_supported(cls) -> tuple[bool, str | None]: ... + + @classmethod + @abstractmethod + def can_implement_shape(cls, n: int, k: int) -> tuple[bool, str | None]: ... + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ... + + @abstractmethod + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: ... diff --git a/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl.py b/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl.py new file mode 100644 index 0000000000..fb6a35f8e6 --- /dev/null +++ b/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Native Blackwell MXFP8 x MXFP6 linear kernel.""" + +from functools import lru_cache +from typing import Literal + +import torch +from torch.nn.parameter import Parameter + +from aphrodite.model_executor.layers.quantization.utils.mxfp8_utils import swizzle_mxfp8_scale +from aphrodite.platforms import current_platform +from aphrodite.utils.import_utils import has_cutedsl + +from .base import Mxfp6LinearKernel + + +@lru_cache(maxsize=16) +def _compile_gemm(activation_format: str, weight_format: str, output_dtype: torch.dtype): + import cuda.bindings.driver as cuda + import cutlass + from cutlass import utils + + from .cutedsl_kernel import ( + Sm100BlockScaledPersistentDenseGemmKernel, + scaled_mm, + ) + + weight_dtype = cutlass.Float6E2M3FN if weight_format == "e2m3" else cutlass.Float6E3M2FN + activation_dtype = { + "mxfp8": cutlass.Float8E4M3FN, + "mxfp6_e2m3": cutlass.Float6E2M3FN, + "mxfp6_e3m2": cutlass.Float6E3M2FN, + }[activation_format] + out_dtype = cutlass.BFloat16 if output_dtype == torch.bfloat16 else cutlass.Float16 + cluster = (1, 1) + gemm = Sm100BlockScaledPersistentDenseGemmKernel(32, (128, 128), cluster) + max_clusters = utils.HardwareInfo().get_max_active_clusters(1) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return scaled_mm( + gemm, + activation_dtype, + weight_dtype, + out_dtype, + cutlass.Float8E8M0FNU, + "k", + "k", + "n", + max_clusters, + stream, + ) + + +@torch.library.custom_op("aphrodite::cutedsl_mxfp6_gemm", mutates_args={"out"}) +def _cutedsl_mxfp6_gemm( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + out: torch.Tensor, + activation_format: str, + weight_format: str, +) -> None: + """Launch the CuTe DSL kernel behind an opaque Torch operator boundary.""" + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import make_ptr + + from aphrodite.model_executor.layers.quantization.utils.mxfp6_online_utils import quantize_mxfp6_cuda + from aphrodite.model_executor.layers.quantization.utils.mxfp8_utils import mxfp8_e4m3_quantize + + m, k = x.shape + n = out.shape[1] + if activation_format == "mxfp8": + x_q, x_scale = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True) + activation_dtype = cutlass.Float8E4M3FN + else: + activation_encoding: Literal["e2m3", "e3m2"] = "e2m3" if activation_format == "mxfp6_e2m3" else "e3m2" + x_q, x_scale = quantize_mxfp6_cuda(x, activation_encoding) + x_scale = swizzle_mxfp8_scale(x_scale, M=m, K=k) + activation_dtype = cutlass.Float6E2M3FN if activation_encoding == "e2m3" else cutlass.Float6E3M2FN + weight_dtype = cutlass.Float6E2M3FN if weight_format == "e2m3" else cutlass.Float6E3M2FN + out_dtype = cutlass.BFloat16 if out.dtype == torch.bfloat16 else cutlass.Float16 + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled = _compile_gemm(activation_format, weight_format, out.dtype) + compiled( + make_ptr( + activation_dtype, + x_q.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + make_ptr( + weight_dtype, + weight.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + make_ptr( + cutlass.Float8E8M0FNU, + x_scale.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + make_ptr( + cutlass.Float8E8M0FNU, + weight_scale.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + make_ptr( + out_dtype, + out.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + (m, n, k, 1), + stream, + ) + + +class CutedslMxfp6LinearKernel(Mxfp6LinearKernel): + """Thor-native tcgen05 MXFP6 GEMM through NVIDIA CUTLASS DSL.""" + + @classmethod + def is_supported(cls) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "MXFP6 requires CUDA" + capability = current_platform.get_device_capability() + if capability is None or capability.to_int() != 110: + return False, "the initial native MXFP6 kernel requires SM110" + if not has_cutedsl(): + return False, "MXFP6 requires nvidia-cutlass-dsl" + return True, None + + @classmethod + def can_implement_shape(cls, n: int, k: int) -> tuple[bool, str | None]: + if n < 128 or n % 128: + return False, "output width must be a multiple of 128" + if k < 128 or k % 128: + return False, "input width must be a multiple of 128" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + n = layer.mxfp6_logical_n + k = layer.mxfp6_logical_k + supported, reason = self.can_implement_shape(n, k) + if not supported: + raise ValueError(reason) + scales = swizzle_mxfp8_scale(layer.weight_scale.data, M=n, K=k) + layer.weight = Parameter(layer.weight.data.contiguous(), requires_grad=False) + layer.weight_scale = Parameter(scales.contiguous(), requires_grad=False) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + n = layer.mxfp6_logical_n + k = layer.mxfp6_logical_k + input_shape = x.shape + x_2d = x.reshape(-1, k) + m = x_2d.shape[0] + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + + _cutedsl_mxfp6_gemm( + x_2d, + layer.weight, + layer.weight_scale, + out, + self.config.activation_format, + self.config.weight_format, + ) + if bias is not None: + out.add_(bias) + return out.view(*input_shape[:-1], n) diff --git a/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_grouped.py b/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_grouped.py new file mode 100644 index 0000000000..488577071a --- /dev/null +++ b/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_grouped.py @@ -0,0 +1,285 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Native SM110 grouped MXFP8 x MXFP6 matrix multiplication.""" + +from collections.abc import Callable +from dataclasses import dataclass +from functools import lru_cache + +import torch + +from aphrodite.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, +) + + +def _grouped_swizzle_scales( + scales: torch.Tensor, + expert_offsets: torch.Tensor, + num_experts: int, + k: int, +) -> torch.Tensor: + """Swizzle row-major scales into one fixed-size slot per expert.""" + rows = scales.shape[0] + num_k_tiles = (k + 127) // 128 + slot_size = ((rows + 127) // 128) * num_k_tiles * 512 + result = torch.zeros( + (num_experts, slot_size), + dtype=torch.uint8, + device=scales.device, + ) + if rows == 0: + return result + + row = torch.arange(rows, device=scales.device) + expert = torch.bucketize(row, expert_offsets[1:], right=True) + local_row = row - expert_offsets[expert] + scale_col = torch.arange(k // 32, device=scales.device) + + mt = local_row[:, None] // 128 + group4 = local_row[:, None] % 128 // 32 + row32 = local_row[:, None] % 32 + kt = scale_col[None, :] // 4 + col4 = scale_col[None, :] % 4 + index = (((mt * num_k_tiles + kt) * 32 + row32) * 4 + group4) * 4 + col4 + result[expert[:, None], index] = scales + return result + + +def _grouped_swizzle_activation_scales( + scales: torch.Tensor, + expert_offsets: torch.Tensor, + num_experts: int, + k: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Swizzle scales into compact, variable-size expert slots. + + The allocation has one 128-row alignment tile per expert plus enough tiles + for all routed rows. It avoids reserving the worst-case row count for every + expert while keeping all layout calculations on the GPU. + """ + rows = scales.shape[0] + num_k_tiles = (k + 127) // 128 + tile_size = num_k_tiles * 512 + counts = expert_offsets[1:] - expert_offsets[:-1] + tiles = (counts + 127) // 128 + tile_bases = torch.cat((torch.zeros_like(tiles[:1]), tiles.cumsum(0)[:-1])) + # sum(ceil(count / 128)) <= ceil(total / 128) + E - 1. + max_tiles = (rows + 127) // 128 + max(num_experts - 1, 0) + result = torch.zeros(max_tiles * tile_size, dtype=torch.uint8, device=scales.device) + if rows == 0: + return result, tile_bases * tile_size + + row = torch.arange(rows, device=scales.device) + expert = torch.bucketize(row, expert_offsets[1:], right=True) + local_row = row - expert_offsets[expert] + scale_col = torch.arange(k // 32, device=scales.device) + mt = local_row[:, None] // 128 + group4 = local_row[:, None] % 128 // 32 + row32 = local_row[:, None] % 32 + kt = scale_col[None, :] // 4 + col4 = scale_col[None, :] % 4 + local_index = (((mt * num_k_tiles + kt) * 32 + row32) * 4 + group4) * 4 + col4 + index = tile_bases[expert, None] * tile_size + local_index + result[index] = scales + return result, tile_bases * tile_size + + +@dataclass +class _CompiledGroupedGemm: + fn: Callable[..., object] + initial: tuple[object, object, object, object, object] + tensormap: object + backing_tensors: tuple[torch.Tensor, ...] + max_active_clusters: int + + +@lru_cache(maxsize=16) +def _compile_grouped( + num_experts: int, + activation_format: str, + weight_format: str, + output_dtype: torch.dtype, +) -> _CompiledGroupedGemm: + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + import cutlass.torch as cutlass_torch + from cutlass import utils + + from .cutedsl_grouped_kernel import ( + Sm100GroupedBlockScaledGemmKernel, + create_tensor_and_stride, + ) + + a_dtype = { + "mxfp8": cutlass.Float8E4M3FN, + "mxfp6_e2m3": cutlass.Float6E2M3FN, + "mxfp6_e3m2": cutlass.Float6E3M2FN, + }[activation_format] + b_dtype = cutlass.Float6E2M3FN if weight_format == "e2m3" else cutlass.Float6E3M2FN + c_dtype = cutlass.BFloat16 if output_dtype == torch.bfloat16 else cutlass.Float16 + sf_dtype = cutlass.Float8E8M0FNU + + # The initial tensors carry only type and layout information. + initial_a = create_tensor_and_stride(1, 128, 128, False, a_dtype)[1:3] + initial_b = create_tensor_and_stride(1, 128, 128, False, b_dtype)[1:3] + initial_c = create_tensor_and_stride(1, 128, 128, False, c_dtype)[1:3] + initial_sfa = create_tensor_and_stride(1, 128, 128, False, sf_dtype)[1:3] + initial_sfb = create_tensor_and_stride(1, 128, 128, False, sf_dtype)[1:3] + initial_pairs = (initial_a, initial_b, initial_c, initial_sfa, initial_sfb) + initial = tuple(pair[1] for pair in initial_pairs) + + dummy_shapes, shapes_backing = cutlass_torch.cute_tensor_like( + torch.empty((num_experts, 4), dtype=torch.int32), + cutlass.Int32, + is_dynamic_layout=False, + assumed_align=16, + ) + dummy_strides, strides_backing = cutlass_torch.cute_tensor_like( + torch.empty((num_experts, 3, 2), dtype=torch.int32), + cutlass.Int32, + is_dynamic_layout=False, + assumed_align=16, + ) + dummy_ptrs, ptrs_backing = cutlass_torch.cute_tensor_like( + torch.empty((num_experts, 3), dtype=torch.int64), + cutlass.Int64, + is_dynamic_layout=False, + assumed_align=16, + ) + dummy_sf_ptrs, sf_ptrs_backing = cutlass_torch.cute_tensor_like( + torch.empty((num_experts, 2), dtype=torch.int64), + cutlass.Int64, + is_dynamic_layout=False, + assumed_align=16, + ) + + hardware = utils.HardwareInfo() + sm_count = hardware.get_max_active_clusters(1) + max_active_clusters = hardware.get_max_active_clusters(1) + tensormap_shape = ( + sm_count, + Sm100GroupedBlockScaledGemmKernel.num_tensormaps, + Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8, + ) + tensormap, tensormap_backing = cutlass_torch.cute_tensor_like( + torch.empty(tensormap_shape, dtype=torch.int64), + cutlass.Int64, + is_dynamic_layout=False, + ) + kernel = Sm100GroupedBlockScaledGemmKernel(32, (128, 128), (1, 1)) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + fn = cute.compile( + kernel, + *initial, + num_experts, + dummy_shapes, + dummy_strides, + dummy_ptrs, + dummy_sf_ptrs, + max_active_clusters, + tensormap, + max_active_clusters, + stream, + options="--opt-level 2", + ) + backing = tuple(pair[0] for pair in initial_pairs) + ( + shapes_backing, + strides_backing, + ptrs_backing, + sf_ptrs_backing, + tensormap_backing, + ) + return _CompiledGroupedGemm( + fn, + initial, + tensormap, + backing, + max_active_clusters, + ) + + +@torch.library.custom_op( + "aphrodite::cutedsl_grouped_mxfp6_gemm", + mutates_args={"out"}, +) +def cutedsl_grouped_mxfp6_gemm( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + expert_offsets: torch.Tensor, + out: torch.Tensor, + logical_n: int, + logical_k: int, + activation_format: str, + weight_format: str, +) -> None: + """Quantize grouped activations and execute one persistent grouped GEMM.""" + import cuda.bindings.driver as cuda + from cutlass.cute.runtime import from_dlpack + + from aphrodite.model_executor.layers.quantization.utils.mxfp6_online_utils import ( + quantize_mxfp6_cuda, + ) + + num_experts = weight.shape[0] + if activation_format == "mxfp8": + x_q, x_scale = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=False) + else: + encoding = "e2m3" if activation_format == "mxfp6_e2m3" else "e3m2" + x_q, x_scale = quantize_mxfp6_cuda(x, encoding) # type: ignore[arg-type] + x_scale, x_scale_offsets = _grouped_swizzle_activation_scales( + x_scale, + expert_offsets, + num_experts, + logical_k, + ) + + counts = expert_offsets[1:] - expert_offsets[:-1] + ones = torch.ones_like(counts) + shapes = torch.stack( + (counts, torch.full_like(counts, logical_n), torch.full_like(counts, logical_k), ones), + dim=1, + ).to(torch.int32) + strides = torch.empty((num_experts, 3, 2), dtype=torch.int32, device=x.device) + strides[:, 0, 0].fill_(logical_k) + strides[:, 0, 1].fill_(1) + strides[:, 1, 0].fill_(logical_k) + strides[:, 1, 1].fill_(1) + strides[:, 2, 0].fill_(logical_n) + strides[:, 2, 1].fill_(1) + + expert = torch.arange(num_experts, dtype=torch.int64, device=x.device) + ptrs = torch.stack( + ( + x_q.data_ptr() + expert_offsets[:-1] * x_q.stride(0) * x_q.element_size(), + weight.data_ptr() + expert * weight.stride(0) * weight.element_size(), + out.data_ptr() + expert_offsets[:-1] * out.stride(0) * out.element_size(), + ), + dim=1, + ) + sf_ptrs = torch.stack( + ( + x_scale.data_ptr() + x_scale_offsets, + weight_scale.data_ptr() + expert * weight_scale.stride(0) * weight_scale.element_size(), + ), + dim=1, + ) + + compiled = _compile_grouped( + num_experts, + activation_format, + weight_format, + out.dtype, + ) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled.fn( + *compiled.initial, + from_dlpack(shapes, assumed_align=16), + from_dlpack(strides, assumed_align=16), + from_dlpack(ptrs, assumed_align=16), + from_dlpack(sf_ptrs, assumed_align=16), + compiled.tensormap, + stream, + ) diff --git a/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_grouped_kernel.py b/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_grouped_kernel.py new file mode 100644 index 0000000000..2024009ef8 --- /dev/null +++ b/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_grouped_kernel.py @@ -0,0 +1,3022 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa +# mypy: disable-error-code="name-defined, return-value" + +# 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. + +import argparse +import functools +from typing import List, Type, Tuple, Union +from inspect import isclass + +import torch +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.runtime import from_dlpack + +""" +This example provides an experimental implementation of the SM100 grouped blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases. + +A grouped blockscaled GEMM example for the NVIDIA Blackwell SM100 architecture using CUTE DSL + +This example demonstrates an implementation of grouped blockscaled GEMM using a TMA plus Blackwell SM100 TensorCore +warp-specialized persistent kernel. +The grouped GEMM workload computes a batch of GEMM operations with distinct problem sizes. Pointers to matrices +in global memory are passed to the kernel in an array (also held in global memory). Similarly, problem shapes and +strides are also stored in arrays in GMEM. + +This differs from "Batched Array" GEMM since the size of each GEMM problem in the grouped GEMM concept may be distinct. + +To run this example: + +.. code-block:: bash + + python examples/blackwell/grouped_blockscaled_gemm.py \ + --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 128,128 --cluster_shape_mn 1,1 \ + --problem_sizes_mnkl "(8192,1280,32,1),(32,384,1536,1),(640,1280,32,1),(640,160,32,1)" \ + --num_groups 4 + +The above example command makes 4 groups of different m, n, k sizes. The Blackwell tcgen05 MMA tile shape +is specified as (128, 64) and the cluster shape is (1,1). The input, mma accumulator and output data type +are set as fp16, fp32 and fp16, respectively. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/grouped_blockscaled_gemm.py \ + --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 128,128 --cluster_shape_mn 1,1 \ + --problem_sizes_mnkl "(8192,1280,32,1),(32,384,1536,1),(640,1280,32,1),(640,160,32,1)" \ + --num_groups 4 + --warmup_iterations 1 --iterations 10 --skip_ref_check + +Constraints: +* Supported input data types: mxf8, mxf4, nvf4 + see detailed valid dtype combinations in below Sm100GroupedBlockScaledGemmKernel class documentation +* A/B tensors must have the same data type, mixed data type is not supported (e.g., mxf8 x mxf4) +* Mma tiler M must be 128 or 256(use_2cta_instrs) +* Mma tiler N must be 128 or 256 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors +* Cluster shape M must be multiple of 2 if Mma tiler M is 256(use_2cta_instrs) +* The l mode(aka, batch size) for each group must be 1. +* The majorness for A, B and C must be the same across all groups. +* The contiguous dimension of A/B/C tensors in each group must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 16 and 32 for Float8 and Float4, respectively. +""" + + +class Sm100GroupedBlockScaledGemmKernel: + """This example demonstrates an implementation of grouped blockscaled GEMM using a TMA plus Blackwell SM100 TensorCore + warp-specialized persistent kernel. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + + :note: In current version, A and B tensors must have the same data type + - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported + + :note: Supported combinations of A/B data types, SF data typs and SF vector size: + - MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32 + - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32 + - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16 + + :note: Supported accumulator data types: + - Float32 + + :note: Supported C data types: + - Float32 + - Float16/BFloat16 + - Float8E4M3FN/Float8E5M2 + :note: Constraints: + - MMA tiler M must be 128 or 256 (use_2cta_instrs) + - MMA tiler N must be 128/256 + - Cluster shape M must be multiple of 2 if Mma tiler M is 256 + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + - Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors + """ + + def __init__( + self, + sf_vec_size: int, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ): + """Initializes the configuration for a Blackwell grouped blockscaled GEMM kernel. + + Besides configurations for dense persistent blockscaled GEMM, there is an extra config specific to grouped blockscaled GEMM: + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :param mma_tiler_mn: tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: tuple[int, int] + :param cluster_shape_mn: tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: tuple[int, int] + """ + self.acc_dtype = cutlass.Float32 + self.sf_vec_size = sf_vec_size + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler = (*mma_tiler_mn, 1) + + self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + + self.tensormap_update_mode = utils.TensorMapUpdateMode.SMEM + + self.occupancy = 1 + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.threads_per_cta = 32 * len((self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)) + # Set barrier for epilogue sync and tmem ptr sync + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=32 * len(self.epilog_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + # Barrier used by MMA/TMA warps to signal A/B tensormap initialization completion + self.tensormap_ab_init_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=64, + ) + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + + # Set up configurations that dependent on gemm inputs. + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B/SFA/SFB + - Computing epilogue subtile + - Setting up A/B/SFA/SFB/C stage counts in shared memory + - Computing A/B/SFA/SFB/C shared memory layout + - Checking reserved smem bytes size capacity for mbar, tensor memory management and tensormap updates utilization + """ + # Compute mma instruction shapes + # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K) + self.mma_inst_shape_mn = ( + self.mma_tiler[0], + self.mma_tiler[1], + ) + # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_inst_shape_mn[0], + self.mma_inst_shape_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_tile_shape_mnk = tuple(x * y for x, y in zip(self.cta_tile_shape_mnk, (*self.cluster_shape_mn, 1))) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Compute epilogue subtile + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.smem_alloc_a_dtype, + self.smem_alloc_b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.sf_dtype, + self.sf_vec_size, + self.smem_capacity, + self.occupancy, + ) + + # Compute A/B/SFA/SFB/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.smem_alloc_a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.smem_alloc_b_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + + mbar_smem_bytes = self._get_mbar_smem_bytes( + num_acc_stage=self.num_acc_stage, + num_ab_stage=self.num_ab_stage, + num_c_stage=self.num_c_stage, + ) + + # Use utils.TensorMapUpdateMode.SMEM by default + tensormap_smem_bytes = ( + Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap * Sm100GroupedBlockScaledGemmKernel.num_tensormaps + ) + if ( + mbar_smem_bytes + tensormap_smem_bytes + Sm100GroupedBlockScaledGemmKernel.tensor_memory_management_bytes + > self.reserved_smem_bytes + ): + raise ValueError( + f"smem consumption for mbar and tensormap {mbar_smem_bytes + tensormap_smem_bytes} exceeds the " + f"reserved smem bytes {self.reserved_smem_bytes}" + ) + + @cute.jit + def __call__( + self, + initial_a: cute.Tensor, + initial_b: cute.Tensor, + initial_c: cute.Tensor, + initial_sfa: cute.Tensor, + initial_sfb: cute.Tensor, + group_count: cutlass.Constexpr[int], + problem_shape_mnkl: cute.Tensor, + strides_abc: cute.Tensor, + tensor_address_abc: cute.Tensor, + tensor_address_sfasfb: cute.Tensor, + total_num_clusters: cutlass.Constexpr[int], + tensormap_cute_tensor: cute.Tensor, + max_active_clusters: cutlass.Constexpr[int], + stream: cuda.CUstream, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + For grouped GEMM, tensor shapes, tensor strides, and tensor address are all provided + by different tensors in global memory. The "initial" tensors only carry data type and + majorness information. + + :param initial_a: Initial tensor A, used for data type and majorness information. + :type initial_a: cute.Tensor + :param initial_b: Initial tensor B, used for data type and majorness information. + :type initial_b: cute.Tensor + :param initial_c: Initial tensor C, used for data type and majorness information. + :type initial_c: cute.Tensor + :param initial_sfa: Initial tensor SFA, used for data type and majorness information. + :type initial_sfa: cute.Tensor + :param initial_sfb: Initial tensor SFB, used for data type and majorness information. + :type initial_sfb: cute.Tensor + :param group_count: The number of GEMM groups. + :type group_count: cutlass.Constexpr[int] + :param problem_shape_mnkl: Tensor containing the (M, N, K, L) shape for each group. + :type problem_shape_mnkl: cute.Tensor + :param strides_abc: Tensor containing the strides for A, B, and C for each group. + :type strides_abc: cute.Tensor + :param tensor_address_abc: Tensor containing the base addresses for A, B, and C for each group. + :type tensor_address_abc: cute.Tensor + :param tensor_address_sfasfb: Tensor containing the base addresses for SFA and SFB for each group. + :type tensor_address_sfasfb: cute.Tensor + :param total_num_clusters: Total number of clusters needed for all groups. + :type total_num_clusters: cutlass.Constexpr[int] + :param tensormap_cute_tensor: Tensor for storing tensormaps. + :type tensormap_cute_tensor: cute.Tensor + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr[int] + :param stream: CUDA stream for asynchronous execution. + :type stream: cuda.CUstream + :raises TypeError: If A and B data types do not match. + """ + self.a_dtype = initial_a.element_type + self.b_dtype = initial_b.element_type + self.sf_dtype = initial_sfa.element_type + self.c_dtype = initial_c.element_type + self.mxf8f6f4 = self.a_dtype.width != self.b_dtype.width or self.a_dtype.width == 6 or self.b_dtype.width == 6 + self.smem_alloc_a_dtype = cutlass.Int8 if self.mxf8f6f4 and self.a_dtype.width < 8 else self.a_dtype + self.smem_alloc_b_dtype = cutlass.Int8 if self.mxf8f6f4 and self.b_dtype.width < 8 else self.b_dtype + self.is_nvfp4_output = self.c_dtype is cutlass.Float4E2M1FN + self.a_major_mode = utils.LayoutEnum.from_tensor(initial_a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(initial_b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(initial_c) + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_a.shape, self.sf_vec_size) + initial_sfa = cute.make_tensor(initial_sfa.iterator, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_b.shape, self.sf_vec_size) + initial_sfb = cute.make_tensor(initial_sfb.iterator, sfb_layout) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + initial_a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=self.smem_alloc_a_dtype if self.mxf8f6f4 and self.a_dtype.width < 8 else None, + ) + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + initial_b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=self.smem_alloc_b_dtype if self.mxf8f6f4 and self.b_dtype.width < 8 else None, + ) + + # Setup TMA load for SFA + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + initial_sfa, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Setup TMA load for SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id) + sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + initial_sfb, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Int16, + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size) * atom_thr_size + + # Setup TMA store for C + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + initial_c, + epi_smem_layout, + self.epi_tile, + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + total_num_clusters, self.cluster_shape_mn, max_active_clusters + ) + + self.buffer_align_bytes = 1024 + self.size_tensormap_in_i64 = ( + Sm100GroupedBlockScaledGemmKernel.num_tensormaps + * Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap + // 8 + ) + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + tensormap_buffer: cute.struct.MemRange[cutlass.Int64, self.size_tensormap_in_i64] + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[self.smem_alloc_a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[self.smem_alloc_b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sSFA: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sSFB: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tiled_mma_sfb, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + tma_atom_c, + tma_tensor_c, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + group_count, + problem_shape_mnkl, + strides_abc, + tensor_address_abc, + tensor_address_sfasfb, + tensormap_cute_tensor, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + group_count: cutlass.Constexpr, + problem_sizes_mnkl: cute.Tensor, + strides_abc: cute.Tensor, + ptrs_abc: cute.Tensor, + ptrs_sfasfb: cute.Tensor, + tensormaps: cute.Tensor, + ): + """ + GPU device kernel performing the grouped GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + if warp_idx == self.tma_warp_id: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_a) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_b) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_sfa) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_sfb) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(cta_rank_in_cluster) + # coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: tensormap buffer, a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + tensormap_smem_ptr = storage.tensormap_buffer.data_ptr() + tensormap_a_smem_ptr = tensormap_smem_ptr + tensormap_b_smem_ptr = tensormap_a_smem_ptr + Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8 + tensormap_sfa_smem_ptr = tensormap_b_smem_ptr + Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8 + tensormap_sfb_smem_ptr = tensormap_sfa_smem_ptr + Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8 + tensormap_c_smem_ptr = tensormap_sfb_smem_ptr + Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8 + + tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar.ptr + tmem_holding_buf_ptr = storage.tmem_holding_buf.ptr + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Tensor memory dealloc barrier init + if use_2cta_instrs: + if warp_idx == self.tma_warp_id: + num_tmem_dealloc_threads = 32 + with cute.arch.elect_one(): + cute.arch.mbarrier_init(tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/SFA/SFB/C + # + sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner) + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + # (MMA, MMA_M, MMA_K, STAGE) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + # (MMA, MMA_N, MMA_K, STAGE) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + + # + # Compute multicast mask for A/B/SFA/SFB buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) + # (bM, bK, RestM, RestK, RestL) + gSFA_mkl = cute.local_tile(mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + # (bN, bK, RestN, RestK, RestL) + gSFB_nkl = cute.local_tile(mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # TMA load scaled factor A partition_S/D + sfa_cta_layout = a_cta_layout + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfa, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # TMA load scaled factor B partition_S/D + sfb_cta_layout = cute.make_layout(cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Get tensormap buffer address + # + grid_dim = cute.arch.grid_dim() + tensormap_workspace_idx = bidz * grid_dim[1] * grid_dim[0] + bidy * grid_dim[0] + bidx + + tensormap_manager = utils.TensorMapManager( + utils.TensorMapUpdateMode.SMEM, + Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap, + ) + tensormap_a_gmem_ptr = tensormap_manager.get_tensormap_ptr( + tensormaps[(tensormap_workspace_idx, 0, None)].iterator + ) + tensormap_b_gmem_ptr = tensormap_manager.get_tensormap_ptr( + tensormaps[(tensormap_workspace_idx, 1, None)].iterator + ) + tensormap_sfa_gmem_ptr = tensormap_manager.get_tensormap_ptr( + tensormaps[(tensormap_workspace_idx, 2, None)].iterator + ) + tensormap_sfb_gmem_ptr = tensormap_manager.get_tensormap_ptr( + tensormaps[(tensormap_workspace_idx, 3, None)].iterator + ) + tensormap_c_gmem_ptr = tensormap_manager.get_tensormap_ptr( + tensormaps[(tensormap_workspace_idx, 4, None)].iterator + ) + + # + # Persistent tile scheduling loop + # + # When the problem shapes are on device, we launch one CTA per SM. + # The if condition later prevents the warps from extra CTAs from doing any work. + tile_sched = utils.StaticPersistentGroupTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + grid_dim, + self.cluster_tile_shape_mnk, + utils.create_initial_search_state(), + group_count, + problem_sizes_mnkl, + ) + initial_work_tile_info = tile_sched.initial_work_tile_info() + + # + # Specialized TMA load warp + # + if warp_idx == self.tma_warp_id and initial_work_tile_info.is_valid_tile: + # + # Persistent tile scheduling loop + # + work_tile = initial_work_tile_info + + tensormap_init_done = cutlass.Boolean(False) + # group index of last tile + last_group_idx = cutlass.Int32(-1) + + ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage) + + while work_tile.is_valid_tile: + grouped_gemm_cta_tile_info = work_tile.group_search_result + cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k + cur_group_idx = grouped_gemm_cta_tile_info.group_idx + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 + # Do not load any data if cur_k_tile_cnt is 0 + if not is_k_tile_cnt_zero: + is_group_changed = cur_group_idx != last_group_idx + # skip tensormap update if we're working on the same group + if is_group_changed: + real_tensor_a = self.make_tensor_abc_for_tensormap_update( + cur_group_idx, + self.a_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 0, # 0 for tensor A + ) + real_tensor_b = self.make_tensor_abc_for_tensormap_update( + cur_group_idx, + self.b_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 1, # 1 for tensor B + ) + real_tensor_sfa = self.make_tensor_sfasfb_for_tensormap_update( + cur_group_idx, + self.sf_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + ptrs_sfasfb, + 0, # 0 for tensor SFA + ) + real_tensor_sfb = self.make_tensor_sfasfb_for_tensormap_update( + cur_group_idx, + self.sf_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + ptrs_sfasfb, + 1, # 1 for tensor SFB + ) + if not tensormap_init_done: + # wait tensormap initialization complete + self.tensormap_ab_init_barrier.arrive_and_wait() + tensormap_init_done = True + + tensormap_manager.update_tensormap( + ( + real_tensor_a, + real_tensor_b, + real_tensor_sfa, + real_tensor_sfb, + ), + (tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb), + ( + tensormap_a_gmem_ptr, + tensormap_b_gmem_ptr, + tensormap_sfa_gmem_ptr, + tensormap_sfb_gmem_ptr, + ), + self.tma_warp_id, + ( + tensormap_a_smem_ptr, + tensormap_b_smem_ptr, + tensormap_sfa_smem_ptr, + tensormap_sfb_smem_ptr, + ), + ) + + mma_tile_coord_mnl = ( + grouped_gemm_cta_tile_info.cta_tile_idx_m // cute.size(tiled_mma.thr_id.shape), + grouped_gemm_cta_tile_info.cta_tile_idx_n, + 0, + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + + # ((atom_v, rest_v), RestK) + tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgSFB_slice = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < cur_k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + + if is_group_changed: + tensormap_manager.fence_tensormap_update(tensormap_a_gmem_ptr) + tensormap_manager.fence_tensormap_update(tensormap_b_gmem_ptr) + tensormap_manager.fence_tensormap_update(tensormap_sfa_gmem_ptr) + tensormap_manager.fence_tensormap_update(tensormap_sfb_gmem_ptr) + # + # Tma load loop + # + for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + + # TMA load A/B/SFA/SFB + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=a_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_a_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=b_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_b_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfa_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_sfa_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfb_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_sfb_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < cur_k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + else: + if not tensormap_init_done: + # wait tensormap initialization complete + self.tensormap_ab_init_barrier.arrive_and_wait() + tensormap_init_done = True + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + last_group_idx = cur_group_idx + + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id and initial_work_tile_info.is_valid_tile: + # + # Initialize tensormaps for A, B, SFA and SFB + # + tensormap_manager.init_tensormap_from_atom(tma_atom_a, tensormap_a_smem_ptr, self.mma_warp_id) + tensormap_manager.init_tensormap_from_atom(tma_atom_b, tensormap_b_smem_ptr, self.mma_warp_id) + tensormap_manager.init_tensormap_from_atom(tma_atom_sfa, tensormap_sfa_smem_ptr, self.mma_warp_id) + tensormap_manager.init_tensormap_from_atom(tma_atom_sfb, tensormap_sfb_smem_ptr, self.mma_warp_id) + # indicate tensormap initialization has finished + self.tensormap_ab_init_barrier.arrive_and_wait() + + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + self.tmem_alloc_barrier.arrive_and_wait() + + # + # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor + # + # Make accumulator tmem tensor + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + self.acc_dtype, + alignment=16, + ptr_to_buffer_holding_addr=tmem_holding_buf_ptr, + ) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # Make SFA tmem tensor + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), + dtype=self.sf_dtype, + ) + # (MMA, MMA_M, MMA_K) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + # Make SFB tmem tensor + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) + + tcgen05.find_tmem_tensor_col_offset(tCtSFA), + dtype=self.sf_dtype, + ) + # (MMA, MMA_N, MMA_K) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + # + # Partition for S2T copy of SFA/SFB + # + tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = self.mainloop_s2t_copy_and_partition( + sSFA, tCtSFA + ) + tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = self.mainloop_s2t_copy_and_partition( + sSFB, tCtSFB + ) + + # + # Persistent tile scheduling loop + # + work_tile = initial_work_tile_info + + ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage) + acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) + while work_tile.is_valid_tile: + cur_group_idx = work_tile.group_search_result.group_idx + problem_shape_k = work_tile.group_search_result.problem_shape_k + + # MMA warp is only interested in number of tiles along K dimension + cur_k_tile_cnt = (problem_shape_k + self.cluster_tile_shape_mnk[2] - 1) // self.cluster_tile_shape_mnk[ + 2 + ] + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 + + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < cur_k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + # + # Wait for accumulator buffer empty + # + if is_leader_cta and not is_k_tile_cnt_zero: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(cur_k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + + # Copy SFA/SFB from smem to tmem + s2t_stage_coord = ( + None, + None, + None, + None, + ab_consumer_state.index, + ) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t_staged, + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t_staged, + tCtSFB_compact_s2t, + ) + + # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = ( + None, + None, + kblock_idx, + ab_consumer_state.index, + ) + + # Set SFA/SFB tensor to tiled_mma + sf_kblock_coord = (None, None, kblock_idx) + tiled_mma.set( + tcgen05.Field.SFA, + tCtSFA[sf_kblock_coord].iterator, + ) + tiled_mma.set( + tcgen05.Field.SFB, + tCtSFB[sf_kblock_coord].iterator, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord], + tCrB[kblock_coord], + tCtAcc, + ) + + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < cur_k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + # + # Async arrive accumulator buffer full + # + if not is_k_tile_cnt_zero: + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id and initial_work_tile_info.is_valid_tile: + # initialize tensorap for C + tensormap_manager.init_tensormap_from_atom( + tma_atom_c, + tensormap_c_smem_ptr, + self.epilog_warp_id[0], + ) + # + # Alloc tensor memory buffer + # + if warp_idx == self.epilog_warp_id[0]: + cute.arch.alloc_tmem( + self.num_tmem_alloc_cols, + tmem_holding_buf_ptr, + is_two_cta=use_2cta_instrs, + ) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + self.tmem_alloc_barrier.arrive_and_wait() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + self.acc_dtype, + alignment=16, + ptr_to_buffer_holding_addr=tmem_holding_buf_ptr, + ) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + ### Start from here + # + # Partition for epilogue + # + epi_tidx = tidx + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(tiled_copy_t2r, tTR_rC, epi_tidx, sC) + tma_atom_c, bSG_sC, bSG_gC_partitioned = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC, epi_tile, sC + ) + + # + # Persistent tile scheduling loop + # + work_tile = initial_work_tile_info + + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) + + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + # group index to start searching + last_group_idx = cutlass.Int32(-1) + + while work_tile.is_valid_tile: + grouped_gemm_cta_tile_info = work_tile.group_search_result + cur_group_idx = grouped_gemm_cta_tile_info.group_idx + cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 + is_group_changed = cur_group_idx != last_group_idx + + # We still need to store 0s when k_tile_cnt is 0 + if is_group_changed: + # construct tensor c based on real shape, stride information + real_tensor_c = self.make_tensor_abc_for_tensormap_update( + cur_group_idx, + self.c_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 2, # 2 for tensor C + ) + tensormap_manager.update_tensormap( + ((real_tensor_c),), + ((tma_atom_c),), + ((tensormap_c_gmem_ptr),), + self.epilog_warp_id[0], + (tensormap_c_smem_ptr,), + ) + + mma_tile_coord_mnl = ( + grouped_gemm_cta_tile_info.cta_tile_idx_m // cute.size(tiled_mma.thr_id.shape), + grouped_gemm_cta_tile_info.cta_tile_idx_n, + 0, + ) + + # + # Slice to per mma tile index + # + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + + # + # Wait for accumulator buffer full + # + if not is_k_tile_cnt_zero: + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + if is_group_changed: + if warp_idx == self.epilog_warp_id[0]: + tensormap_manager.fence_tensormap_update(tensormap_c_gmem_ptr) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in range(subtile_cnt): + if not is_k_tile_cnt_zero: + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + tRS_rC.store(acc_vec.to(self.c_dtype)) + else: + if cutlass.const_expr(self.is_nvfp4_output): + zeros_i8 = cute.make_rmem_tensor( + cute.recast_layout( + cutlass.Int8.width, + self.c_dtype.width, + tRS_rC.layout, + ), + cutlass.Int8, + ) + zeros_i8.fill(0) + tRS_rC.store(cute.recast_tensor(zeros_i8, self.c_dtype).load()) + else: + tRS_rC.fill(0) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_c_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + # + # Async arrive accumulator buffer empty + # + if not is_k_tile_cnt_zero: + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + last_group_idx = cur_group_idx + + # + # Dealloc the tensor memory buffer + # + if warp_idx == self.epilog_warp_id[0]: + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=use_2cta_instrs) + self.epilog_sync_barrier.arrive_and_wait() + if warp_idx == self.epilog_warp_id[0]: + if use_2cta_instrs: + cute.arch.mbarrier_arrive(tmem_dealloc_mbar_ptr, cta_rank_in_cluster ^ 1) + cute.arch.mbarrier_wait(tmem_dealloc_mbar_ptr, 0) + cute.arch.dealloc_tmem(acc_tmem_ptr, self.num_tmem_alloc_cols, is_two_cta=use_2cta_instrs) + # + # Wait for C store complete + # + c_pipeline.producer_tail() + + @cute.jit + def make_tensor_abc_for_tensormap_update( + self, + group_idx: cutlass.Int32, + dtype: Type[cutlass.Numeric], + problem_shape_mnk: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32], + strides_abc: cute.Tensor, + tensor_address_abc: cute.Tensor, + tensor_index: int, + ): + """Extract stride and tensor address for a given group and construct a global tensor for A, B or C. + + This function is used within the kernel to dynamically create a CUTE tensor + representing A, B, or C for the current group being processed, using the + group-specific address, shape, and stride information. + + :param group_idx: The index of the current group within the grouped GEMM. + :type group_idx: cutlass.Int32 + :param dtype: The data type of the tensor elements (e.g., cutlass.Float16). + :type dtype: Type[cutlass.Numeric] + :param problem_shape_mnk: The (M, N, K) problem shape for the current group. + :type problem_shape_mnk: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32] + :param strides_abc: Tensor containing strides for A, B, C for all groups. Layout: (group_count, 3, 2). + :type strides_abc: cute.Tensor + :param tensor_address_abc: Tensor containing global memory addresses for A, B, C for all groups. Layout: (group_count, 3). + :type tensor_address_abc: cute.Tensor + :param tensor_index: Specifies which tensor to create: 0 for A, 1 for B, 2 for C. + :type tensor_index: int + :return: A CUTE tensor representing the requested global memory tensor (A, B, or C) for the specified group. + :rtype: cute.Tensor + :raises TypeError: If the provided dtype is not a subclass of cutlass.Numeric. + """ + ptr_i64 = tensor_address_abc[(group_idx, tensor_index)] + if cutlass.const_expr(not isclass(dtype) or not issubclass(dtype, cutlass.Numeric)): + raise TypeError(f"dtype must be a type of cutlass.Numeric, got {type(dtype)}") + tensor_gmem_ptr = cute.make_ptr(dtype, ptr_i64, cute.AddressSpace.gmem, assumed_align=16) + + strides_tensor_gmem = strides_abc[(group_idx, tensor_index, None)] + strides_tensor_reg = cute.make_rmem_tensor( + cute.make_layout(2), + strides_abc.element_type, + ) + cute.autovec_copy(strides_tensor_gmem, strides_tensor_reg) + stride_mn = strides_tensor_reg[0] + stride_k = strides_tensor_reg[1] + c1 = cutlass.Int32(1) + c0 = cutlass.Int32(0) + + if cutlass.const_expr(tensor_index == 0): # tensor A + m = problem_shape_mnk[0] + k = problem_shape_mnk[2] + return cute.make_tensor( + tensor_gmem_ptr, + cute.make_layout((m, k, c1), stride=(stride_mn, stride_k, c0)), + ) + elif cutlass.const_expr(tensor_index == 1): # tensor B + n = problem_shape_mnk[1] + k = problem_shape_mnk[2] + return cute.make_tensor( + tensor_gmem_ptr, + cute.make_layout((n, k, c1), stride=(stride_mn, stride_k, c0)), + ) + else: # tensor C + m = problem_shape_mnk[0] + n = problem_shape_mnk[1] + return cute.make_tensor( + tensor_gmem_ptr, + cute.make_layout((m, n, c1), stride=(stride_mn, stride_k, c0)), + ) + + @cute.jit + def make_tensor_sfasfb_for_tensormap_update( + self, + group_idx: cutlass.Int32, + dtype: Type[cutlass.Numeric], + problem_shape_mnk: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32], + tensor_address_sfasfb: cute.Tensor, + tensor_index: int, + ): + """Extract tensor address for a given group and construct a global tensor for SFA or SFB. + + This function is used within the kernel to dynamically create a CUTE tensor + representing SFA or SFB for the current group being processed, using the + group-specific address, shape information. + + :param group_idx: The index of the current group within the grouped GEMM. + :type group_idx: cutlass.Int32 + :param dtype: The data type of the tensor elements (e.g., cutlass.Float16). + :type dtype: Type[cutlass.Numeric] + :param problem_shape_mnk: The (M, N, K) problem shape for the current group. + :type problem_shape_mnk: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32] + :param tensor_address_sfasfb: Tensor containing global memory addresses for SFA, SFB for all groups. Layout: (group_count, 2). + :type tensor_address_sfasfb: cute.Tensor + :param tensor_index: Specifies which tensor to create: 0 for SFA, 1 for SFB. + :type tensor_index: int + :return: A CUTE tensor representing the requested global memory tensor (SFA, SFB) for the specified group. + :rtype: cute.Tensor + :raises TypeError: If the provided dtype is not a subclass of cutlass.Numeric. + """ + ptr_i64 = tensor_address_sfasfb[(group_idx, tensor_index)] + if cutlass.const_expr(not isclass(dtype) or not issubclass(dtype, cutlass.Numeric)): + raise TypeError(f"dtype must be a type of cutlass.Numeric, got {type(dtype)}") + tensor_gmem_ptr = cute.make_ptr(dtype, ptr_i64, cute.AddressSpace.gmem, assumed_align=16) + + c1 = cutlass.Int32(1) + if cutlass.const_expr(tensor_index == 0): # tensor SFA + m = problem_shape_mnk[0] + k = problem_shape_mnk[2] + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF((m, k, c1), self.sf_vec_size) + return cute.make_tensor( + tensor_gmem_ptr, + sfa_layout, + ) + else: # tensor SFB + n = problem_shape_mnk[1] + k = problem_shape_mnk[2] + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size) + return cute.make_tensor( + tensor_gmem_ptr, + sfb_layout, + ) + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination). + + :param sSF: The scale factor tensor in smem + :type sSF: cute.Tensor + :param tSF: The scale factor tensor in tmem + :type tSF: cute.Tensor + + :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: + - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) + - tCsSF_compact_s2t: The partitioned scale factor tensor in smem + - tSF_compact_s2t: The partitioned scale factor tensor in tmem + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSF_compact = cute.filter_zeros(sSF) + # (MMA, MMA_MN, MMA_K) + tCtSF_compact = cute.filter_zeros(tSF) + + # Make S2T CopyAtom and tiledCopy + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc: The accumulated tensor in register used to hold t2r results + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_mnl_epi = cute.flat_divide(gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilog_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + :type sepi: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rC: The partitioned tensor C (register source) + - tRS_sC: The partitioned tensor C (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op(self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: Union[cute.CopyAtom, cute.TiledCopy], + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Make tiledCopy for global memory store, then use it to: + partition shared memory (source) and global memory (destination) for TMA store version. + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where: + - tma_atom_c: The TMA copy atom + - bSG_sC: The partitioned shared memory tensor C + - bSG_gC: The partitioned global tensor C + :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_epi = cute.flat_divide(gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + + tma_atom_c = atom + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + return tma_atom_c, bSG_sC, bSG_gC + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + epi_tile: cute.Tile, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + smem_capacity: int, + occupancy: int, + ) -> Tuple[int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout enum of operand C. + :type c_layout: utils.LayoutEnum + :param sf_dtype: Data type of Scale factor. + :type sf_dtype: type[cutlass.Numeric] + :param sf_vec_size: Scale factor vector size. + :type sf_vec_size: int + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # ACC stages + num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 + + # Default C stages + num_c_stage = 2 + + # Calculate smem layout and size for one stage of A, B, SFA, SFB and C + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, # a tmp 1 stage is provided + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, # a tmp 1 stage is provided + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + mbar_helpers_bytes = 1024 + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B/SFA/SFB stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B/SFA/SFB stage + num_ab_stage = (smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes + # Add remaining unused smem to epilogue + num_c_stage += ( + smem_capacity - occupancy * ab_bytes_per_stage * num_ab_stage - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + + return num_acc_stage, num_ab_stage, num_c_stage + + @staticmethod + def _compute_grid( + total_num_clusters: int, + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr[int], + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """Compute tile scheduler parameters and grid shape for grouped GEMM operations. + + :param total_num_clusters: Total number of clusters to process across all groups. + :type total_num_clusters: int + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr[int] + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: tuple[utils.PersistentTileSchedulerParams, tuple[int, ...]] + """ + # Create problem shape with M, N dimensions from cluster shape + # and L dimension representing the total number of clusters. + problem_shape_ntile_mnl = ( + cluster_shape_mn[0], + cluster_shape_mn[1], + cutlass.Int32(total_num_clusters), + ) + + tile_sched_params = utils.PersistentTileSchedulerParams(problem_shape_ntile_mnl, (*cluster_shape_mn, 1)) + + grid = utils.StaticPersistentTileScheduler.get_grid_shape(tile_sched_params, max_active_clusters) + + return tile_sched_params, grid + + @staticmethod + def _get_mbar_smem_bytes(**kwargs_stages: int) -> int: + """Calculate shared memory consumption for memory barriers based on provided stages. + + Each stage requires 2 barriers, and each barrier consumes 8 bytes of shared memory. + The total consumption is the sum across all provided stages. This function calculates the total + shared memory needed for these barriers. + + :param kwargs_stages: Variable keyword arguments where each key is a stage name + (e.g., num_acc_stage, num_ab_stage) and each value is the + number of stages of that type. + :type kwargs_stages: int + :return: Total shared memory bytes required for all memory barriers. + :rtype: int + """ + num_barriers_per_stage = 2 + num_bytes_per_barrier = 8 + mbar_smem_consumption = sum( + [num_barriers_per_stage * num_bytes_per_barrier * stage for stage in kwargs_stages.values()] + ) + return mbar_smem_consumption + + @staticmethod + def is_valid_dtypes_and_scale_factor_vec_size( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + ) -> bool: + """ + Check if the dtypes and sf_vec_size are valid combinations + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size of the scale factor + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes and sf_vec_size are valid, False otherwise + :rtype: bool + """ + is_valid = True + + # Check valid ab_dtype + if ab_dtype not in { + cutlass.Float4E2M1FN, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + is_valid = False + + # Check valid sf_vec_size + if sf_vec_size not in {16, 32}: + is_valid = False + + # Check valid sf_dtype + if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: + is_valid = False + + # Check valid sf_dtype and sf_vec_size combinations + if sf_dtype == cutlass.Float8E4M3FN and sf_vec_size == 32: + is_valid = False + if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} and sf_vec_size == 16: + is_valid = False + + # Check valid c_dtype + if c_dtype not in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + is_valid = False + + return is_valid + + @staticmethod + def is_valid_layouts( + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if layouts and dtypes are valid combinations + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major dimension of the A tensor + :type a_major: str + :param b_major: The major dimension of the B tensor + :type b_major: str + :param c_major: The major dimension of the C tensor + :type c_major: str + + :return: True if the layouts are valid, False otherwise + :rtype: bool + """ + is_valid = True + + if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"): + is_valid = False + return is_valid + + @staticmethod + def is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ) -> bool: + """ + Check if the mma tiler and cluster shape are valid + + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + + :return: True if the mma tiler and cluster shape are valid, False otherwise + :rtype: bool + """ + is_valid = True + # Skip invalid mma tile shape + if mma_tiler_mn[0] not in [128, 256]: + is_valid = False + if mma_tiler_mn[1] not in [128, 256]: + is_valid = False + # Skip illegal cluster shape + if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + # Special cluster shape check for scale factor multicasts. + # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. + or cluster_shape_mn[0] > 4 + or cluster_shape_mn[1] > 4 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + @staticmethod + def is_valid_tensor_alignment( + problem_sizes_mnkl: List[Tuple[int, int, int, int]], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the tensor alignment is valid + + :param problem_sizes_mnkl: The problem shape for each group + :type problem_sizes_mnkl: List[Tuple[int, int, int, int]] + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + for m, n, k, l in problem_sizes_mnkl: + if ( + not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + return is_valid + + @staticmethod + def can_implement( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + problem_sizes_mnkl: List[Tuple[int, int, int, int]], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the gemm can be implemented + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the gemm can be implemented, False otherwise + :rtype: bool + """ + can_implement = True + # Skip unsupported types + if not Sm100GroupedBlockScaledGemmKernel.is_valid_dtypes_and_scale_factor_vec_size( + ab_dtype, sf_dtype, sf_vec_size, c_dtype + ): + can_implement = False + # Skip unsupported layouts + if not Sm100GroupedBlockScaledGemmKernel.is_valid_layouts(ab_dtype, c_dtype, a_major, b_major, c_major): + can_implement = False + # Skip invalid mma tile shape and cluster shape + if not Sm100GroupedBlockScaledGemmKernel.is_valid_mma_tiler_and_cluster_shape(mma_tiler_mn, cluster_shape_mn): + can_implement = False + # Skip illegal problem shape for load/store alignment + if not Sm100GroupedBlockScaledGemmKernel.is_valid_tensor_alignment( + problem_sizes_mnkl, ab_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + return can_implement + + # Size of smem we reserved for mbarrier, tensor memory management and tensormap update + reserved_smem_bytes = 1024 + bytes_per_tensormap = 128 + num_tensormaps = 5 + # size of smem used for tensor memory management + tensor_memory_management_bytes = 12 + + +# Create tensor and return the pointer, tensor, and stride +def create_tensor_and_stride( + l: int, + mode0: int, + mode1: int, + is_mode0_major: bool, + dtype: type[cutlass.Numeric], + is_dynamic_layout: bool = True, +) -> tuple[int, torch.Tensor, cute.Tensor, torch.Tensor, tuple[int, int]]: + """Create GPU tensor from either a new or existing CPU tensor. + + :param torch_tensor_cpu: Optional existing CPU tensor to reuse. If None, creates a new one. + :type torch_tensor_cpu: torch.Tensor, optional + """ + + # Create new CPU tensor + torch_tensor_cpu = cutlass_torch.matrix( + l, + mode0, + mode1, + is_mode0_major, + cutlass.Float32, + ) + + # Create GPU tensor from CPU tensor (new or existing) + cute_tensor, torch_tensor = cutlass_torch.cute_tensor_like( + torch_tensor_cpu, dtype, is_dynamic_layout, assumed_align=16 + ) + + # omit stride for L mode as it is always 1 + stride = (1, mode0) if is_mode0_major else (mode1, 1) + + return ( + torch_tensor.data_ptr(), + torch_tensor, + cute_tensor, + torch_tensor_cpu, + stride, + ) + + +def create_tensors_abc_for_all_groups( + problem_sizes_mnkl: List[tuple[int, int, int, int]], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, +) -> tuple[ + List[List[int]], + List[List[torch.Tensor]], + List[tuple], + List[List[tuple]], + List[List[torch.Tensor]], +]: + ref_torch_fp32_tensors_abc = [] + torch_tensors_abc = [] + cute_tensors_abc = [] + strides_abc = [] + ptrs_abc = [] + + # Iterate through all groups and create tensors for each group + for group_idx, (m, n, k, l) in enumerate(problem_sizes_mnkl): + # Create tensors A, B, C + ( + ptr_a, + torch_tensor_a, + cute_tensor_a, + ref_torch_fp32_tensor_a, + stride_mk_a, + ) = create_tensor_and_stride(l, m, k, a_major == "m", a_dtype) + + ( + ptr_b, + torch_tensor_b, + cute_tensor_b, + ref_torch_fp32_tensor_b, + stride_nk_b, + ) = create_tensor_and_stride(l, n, k, b_major == "n", b_dtype) + + ( + ptr_c, + torch_tensor_c, + cute_tensor_c, + ref_torch_fp32_tensor_c, + stride_mn_c, + ) = create_tensor_and_stride(l, m, n, c_major == "m", c_dtype) + + ref_torch_fp32_tensors_abc.append([ref_torch_fp32_tensor_a, ref_torch_fp32_tensor_b, ref_torch_fp32_tensor_c]) + + ptrs_abc.append([ptr_a, ptr_b, ptr_c]) + torch_tensors_abc.append([torch_tensor_a, torch_tensor_b, torch_tensor_c]) + strides_abc.append([stride_mk_a, stride_nk_b, stride_mn_c]) + cute_tensors_abc.append( + ( + cute_tensor_a, + cute_tensor_b, + cute_tensor_c, + ) + ) + + return ( + ptrs_abc, + torch_tensors_abc, + cute_tensors_abc, + strides_abc, + ref_torch_fp32_tensors_abc, + ) + + +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_tensor: cute.Tensor, + sf_mma_tensor: cute.Tensor, +): + """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout""" + # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l) + # group to ((32, 4, rest_m), (4, rest_k), l) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + + +# Create scale factor tensor SFA/SFB +def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = max(1, ceil_div(k, sf_vec_size)) + ref_shape = (l, mn, sf_k) + + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + ref_permute_order = (1, 2, 0) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + # Create f32 ref torch tensor (cpu) + ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + ref_shape, + torch.float32, + permute_order=ref_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=1, + max_val=3, + ), + ) + + # Create f32 cute torch tensor (cpu) + cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=0, + max_val=1, + ), + ) + + # convert ref f32 tensor to cute f32 tensor + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + from_dlpack(ref_f32_torch_tensor_cpu), + from_dlpack(cute_f32_torch_tensor_cpu), + ) + cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() + + # reshape makes memory contiguous + ref_f32_torch_tensor_cpu = ( + ref_f32_torch_tensor_cpu.permute(2, 0, 1) + .unsqueeze(-1) + .expand(l, mn, sf_k, sf_vec_size) + .reshape(l, mn, sf_k * sf_vec_size) + .permute(*ref_permute_order) + ) + # prune to mkl for reference check. + ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] + + # Create dtype cute torch tensor (cpu) + cute_tensor, cute_torch_tensor = cutlass_torch.cute_tensor_like( + cute_f32_torch_tensor_cpu, + dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + + # Convert f32 cute tensor to dtype cute tensor + cute_tensor = cutlass_torch.convert_cute_tensor( + cute_f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=True, + ) + # get pointer of the tensor + ptr = cute_torch_tensor.data_ptr() + return ref_f32_torch_tensor_cpu, ptr, cute_tensor, cute_torch_tensor + + +def create_tensors_sfasfb_for_all_groups( + problem_sizes_mnkl: List[tuple[int, int, int, int]], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, +) -> tuple[ + List[List[int]], + List[List[torch.Tensor]], + List[tuple], + List[List[torch.Tensor]], +]: + ptrs_sfasfb = [] + torch_tensors_sfasfb = [] + cute_tensors_sfasfb = [] + refs_sfasfb = [] + + # Iterate through all groups and create tensors for each group + for group_idx, (m, n, k, l) in enumerate(problem_sizes_mnkl): + sfa_ref, ptr_sfa, sfa_tensor, sfa_torch = create_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype) + sfb_ref, ptr_sfb, sfb_tensor, sfb_torch = create_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype) + ptrs_sfasfb.append([ptr_sfa, ptr_sfb]) + torch_tensors_sfasfb.append([sfa_torch, sfb_torch]) + cute_tensors_sfasfb.append( + ( + sfa_tensor, + sfb_tensor, + ) + ) + refs_sfasfb.append([sfa_ref, sfb_ref]) + + return ( + ptrs_sfasfb, + torch_tensors_sfasfb, + cute_tensors_sfasfb, + refs_sfasfb, + ) + + +def run( + num_groups: int, + problem_sizes_mnkl: List[Tuple[int, int, int, int]], + host_problem_shape_available: bool, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +): + """Run SM100 grouped blockscaledGEMM example with specified configurations. + + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False + :type use_cold_l2: bool, optional + :return: Execution time of the GEMM kernel in microseconds + :rtype: float + """ + print("Running Blackwell Grouped GEMM test with:") + print(f"{num_groups} groups") + for i, (m, n, k, l) in enumerate(problem_sizes_mnkl): + print(f"Group {i}: {m}x{n}x{k}x{l}") + print(f"A/B dtype: {a_dtype}/{b_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}") + print(f"C dtype: {c_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + + # Skip unsupported testcase + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(2025) + + # Create tensors A, B, C for all groups + ( + ptrs_abc, + torch_tensors_abc, + cute_tensors_abc, + strides_abc, + ref_f32_torch_tensors_abc, + ) = create_tensors_abc_for_all_groups( + problem_sizes_mnkl, + a_dtype, + b_dtype, + c_dtype, + a_major, + b_major, + c_major, + ) + # Create tensors SFA, SFB for all groups + ( + ptrs_sfasfb, + torch_tensors_sfasfb, + cute_tensors_sfasfb, + refs_f32_torch_tensors_sfasfb, + ) = create_tensors_sfasfb_for_all_groups( + problem_sizes_mnkl, + sf_dtype, + sf_vec_size, + ) + + # Setup inital tensors for TMA of A,B and C + alignment = 16 # 16 bytes aligned + divisibility_a = 32 if a_dtype == cutlass.Float4E2M1FN else 16 + divisibility_b = 32 if b_dtype == cutlass.Float4E2M1FN else 16 + divisibility_c = 32 if c_dtype == cutlass.Float4E2M1FN else 16 + divisibility_sf = 32 if sf_dtype == cutlass.Float4E2M1FN else 16 + + min_a_size = alignment * 8 // a_dtype.width + min_a_size *= (divisibility_a + min_a_size - 1) // min_a_size + min_b_size = 128 if b_dtype.width == 6 else alignment * 8 // b_dtype.width + min_b_size *= (divisibility_b + min_b_size - 1) // min_b_size + + min_c_size = alignment * 8 // c_dtype.width + div_mul_c = (divisibility_c + min_c_size - 1) // min_c_size + min_c_size = min_c_size * div_mul_c + + min_sf_size = alignment * 8 // sf_dtype.width + div_mul_sf = (divisibility_sf + min_sf_size - 1) // min_sf_size + min_sf_size = min_sf_size * div_mul_sf + + initial_cute_tensors_abc = [ + create_tensor_and_stride(1, min_a_size, min_a_size, a_major == "m", a_dtype)[2], + create_tensor_and_stride(1, min_b_size, min_b_size, b_major == "n", b_dtype)[2], + create_tensor_and_stride(1, min_c_size, min_c_size, c_major == "m", c_dtype)[2], + ] + initial_cute_tensors_sfasfb = [ + create_tensor_and_stride(1, min_sf_size, min_sf_size, a_major == "m", sf_dtype)[2], + create_tensor_and_stride(1, min_sf_size, min_sf_size, b_major == "n", sf_dtype)[2], + ] + + hardware_info = cutlass.utils.HardwareInfo() + sm_count = hardware_info.get_max_active_clusters(1) + max_active_clusters = hardware_info.get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) + # Prepare tensormap buffer for each SM + num_tensormap_buffers = sm_count + tensormap_shape = ( + num_tensormap_buffers, + Sm100GroupedBlockScaledGemmKernel.num_tensormaps, + Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8, + ) + tensor_of_tensormap, tensor_of_tensormap_torch = cutlass_torch.cute_tensor_like( + torch.empty(tensormap_shape, dtype=torch.int64), + cutlass.Int64, + is_dynamic_layout=False, + ) + + grouped_blockscaled_gemm = Sm100GroupedBlockScaledGemmKernel( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + ) + + # layout (num_groups, 4):(4, 1) + ( + tensor_of_dim_size_mnkl, + tensor_of_dim_size_mnkl_torch, + ) = cutlass_torch.cute_tensor_like( + torch.tensor(problem_sizes_mnkl, dtype=torch.int32), + cutlass.Int32, + is_dynamic_layout=False, + assumed_align=16, + ) + + # layout (num_groups, 3, 2):(6, 2, 1) + tensor_of_strides_abc, tensor_of_strides_abc_torch = cutlass_torch.cute_tensor_like( + torch.tensor(strides_abc, dtype=torch.int32), + cutlass.Int32, + is_dynamic_layout=False, + assumed_align=16, + ) + + # layout (num_groups,3):(3, 1) + tensor_of_ptrs_abc, tensor_of_ptrs_abc_torch = cutlass_torch.cute_tensor_like( + torch.tensor(ptrs_abc, dtype=torch.int64), + cutlass.Int64, + is_dynamic_layout=False, + assumed_align=16, + ) + + # layout (num_groups,2):(2, 1) + tensor_of_ptrs_sfasfb, tensor_of_ptrs_sfasfb_torch = cutlass_torch.cute_tensor_like( + torch.tensor(ptrs_sfasfb, dtype=torch.int64), + cutlass.Int64, + is_dynamic_layout=False, + assumed_align=16, + ) + + # Compute total number of cluster tiles we need to compute for given grouped GEMM problem + def compute_total_num_clusters( + problem_sizes_mnkl: List[tuple[int, int, int, int]], + cluster_tile_shape_mn: tuple[int, int], + ) -> int: + total_num_clusters = 0 + for m, n, _, _ in problem_sizes_mnkl: + num_clusters_mn = tuple((x + y - 1) // y for x, y in zip((m, n), cluster_tile_shape_mn)) + total_num_clusters += functools.reduce(lambda x, y: x * y, num_clusters_mn) + return total_num_clusters + + # Compute cluster tile shape + def compute_cluster_tile_shape( + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + ) -> tuple[int, int]: + cta_tile_shape_mn = [128, mma_tiler_mn[1]] + return tuple(x * y for x, y in zip(cta_tile_shape_mn, cluster_shape_mn)) + + cluster_tile_shape_mn = compute_cluster_tile_shape(mma_tiler_mn, cluster_shape_mn) + total_num_clusters = compute_total_num_clusters(problem_sizes_mnkl, cluster_tile_shape_mn) + + # Initialize Stream + current_stream = cutlass_torch.default_stream() + + # If the host problem shape is available, we will launch the grid with only + # the necessary clusters. The function compute_total_num_clusters() does that. + # If the problem shape only exists on device, we will need to launch all active + # clusters possible on a device. + if host_problem_shape_available: + print("Problem shapes available on host and device") + total_num_clusters = compute_total_num_clusters(problem_sizes_mnkl, cluster_tile_shape_mn) + else: + print("Problem shapes available only on device") + total_num_clusters = max_active_clusters + + # Compile grouped GEMM kernel + compiled_grouped_gemm = cute.compile( + grouped_blockscaled_gemm, + initial_cute_tensors_abc[0], + initial_cute_tensors_abc[1], + initial_cute_tensors_abc[2], + initial_cute_tensors_sfasfb[0], + initial_cute_tensors_sfasfb[1], + num_groups, + tensor_of_dim_size_mnkl, + tensor_of_strides_abc, + tensor_of_ptrs_abc, + tensor_of_ptrs_sfasfb, + total_num_clusters, + tensor_of_tensormap, + max_active_clusters, + current_stream, + options=f"--opt-level 2", + ) + + # reference check + if not skip_ref_check: + compiled_grouped_gemm( + initial_cute_tensors_abc[0], + initial_cute_tensors_abc[1], + initial_cute_tensors_abc[2], + initial_cute_tensors_sfasfb[0], + initial_cute_tensors_sfasfb[1], + tensor_of_dim_size_mnkl, + tensor_of_strides_abc, + tensor_of_ptrs_abc, + tensor_of_ptrs_sfasfb, + tensor_of_tensormap, + current_stream, + ) + print("Verifying results...") + + for i, ( + (a_ref, b_ref, c_ref), + (sfa_ref, sfb_ref), + (a_tensor, b_tensor, c_tensor), + (m, n, k, l), + ) in enumerate( + zip( + ref_f32_torch_tensors_abc, + refs_f32_torch_tensors_sfasfb, + cute_tensors_abc, + problem_sizes_mnkl, + ) + ): + ref_res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref) + ref_res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref) + ref = torch.einsum("mkl,nkl->mnl", ref_res_a, ref_res_b) + + print(f"checking group {i}") + c_ref_device = c_ref.cuda() + + cute.testing.convert( + c_tensor, + from_dlpack(c_ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=(1 if c_major == "n" else 0) + ), + ) + + c_ref = c_ref_device.cpu() + + if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Convert ref : f32 -> f8 -> f32 + ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, device="cuda").permute(1, 2, 0) + ref_f8 = from_dlpack(ref_f8_, assumed_align=16).mark_layout_dynamic(leading_dim=1) + ref_f8.element_type = c_dtype + ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda() + ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic(leading_dim=1) + cute.testing.convert(ref_tensor, ref_f8) + cute.testing.convert(ref_f8, ref_tensor) + ref = ref_device.cpu() + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + + def generate_tensors(): + ( + ptrs_abc_workspace, + torch_tensors_abc_workspace, + cute_tensors_abc_workspace, + strides_abc_workspace, + _, + ) = create_tensors_abc_for_all_groups( + problem_sizes_mnkl, + a_dtype, + b_dtype, + c_dtype, + a_major, + b_major, + c_major, + ) + + ( + ptrs_sfasfb_workspace, + torch_tensors_sfasfb_workspace, + cute_tensors_sfasfb_workspace, + _, + ) = create_tensors_sfasfb_for_all_groups( + problem_sizes_mnkl, + sf_dtype, + sf_vec_size, + ) + + initial_cute_tensors_abc_workspace = [ + create_tensor_and_stride(1, min_a_size, min_a_size, a_major == "m", a_dtype)[2], + create_tensor_and_stride(1, min_b_size, min_b_size, b_major == "n", b_dtype)[2], + create_tensor_and_stride(1, min_c_size, min_c_size, c_major == "m", c_dtype)[2], + ] + initial_cute_tensors_sfasfb_workspace = [ + create_tensor_and_stride(1, min_sf_size, min_sf_size, a_major == "m", sf_dtype)[2], + create_tensor_and_stride(1, min_sf_size, min_sf_size, b_major == "n", sf_dtype)[2], + ] + + # Create new tensors for this workspace + tensor_of_strides_abc_workspace, _ = cutlass_torch.cute_tensor_like( + torch.tensor(strides_abc_workspace, dtype=torch.int32), + cutlass.Int32, + is_dynamic_layout=False, + assumed_align=16, + ) + + tensor_of_ptrs_abc_workspace, _ = cutlass_torch.cute_tensor_like( + torch.tensor(ptrs_abc_workspace, dtype=torch.int64), + cutlass.Int64, + is_dynamic_layout=False, + assumed_align=16, + ) + + tensor_of_ptrs_sfasfb_workspace, _ = cutlass_torch.cute_tensor_like( + torch.tensor(ptrs_sfasfb_workspace, dtype=torch.int64), + cutlass.Int64, + is_dynamic_layout=False, + assumed_align=16, + ) + + tensormap_workspace, _ = cutlass_torch.cute_tensor_like( + torch.empty(tensormap_shape, dtype=torch.int64), + cutlass.Int64, + is_dynamic_layout=False, + ) + + args = cute.testing.JitArguments( + initial_cute_tensors_abc_workspace[0], + initial_cute_tensors_abc_workspace[1], + initial_cute_tensors_abc_workspace[2], + initial_cute_tensors_sfasfb_workspace[0], + initial_cute_tensors_sfasfb_workspace[1], + tensor_of_dim_size_mnkl, + tensor_of_strides_abc_workspace, + tensor_of_ptrs_abc_workspace, + tensor_of_ptrs_sfasfb_workspace, + tensormap_workspace, + current_stream, + ) + args.add_to_scope([torch_tensors_abc_workspace, torch_tensors_sfasfb_workspace]) + return args + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + sum( + [ + sum([torch_tensor.numel() * torch_tensor.element_size() for torch_tensor in group_tensors]) + for group_tensors in torch_tensors_abc + torch_tensors_sfasfb + ] + ) + + + # Add size of strides tensor + tensor_of_strides_abc_torch.numel() * tensor_of_strides_abc_torch.element_size() + + + # Add size of ptrs tensor A, B, C + tensor_of_ptrs_abc_torch.numel() * tensor_of_ptrs_abc_torch.element_size() + + + # Add size of ptrs tensor SFA, SFB + tensor_of_ptrs_sfasfb_torch.numel() * tensor_of_ptrs_sfasfb_torch.element_size() + + + # Add size of tensormap tensor + tensor_of_tensormap_torch.numel() * tensor_of_tensormap_torch.element_size() + ) + workspace_count = cute.testing.get_workspace_count(one_workspace_bytes, warmup_iterations, iterations) + + exec_time = cute.testing.benchmark( + compiled_grouped_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + runtime_s = exec_time / 1.0e6 + fmas = 0 + for group in range(num_groups): + [M, N, K, _] = problem_sizes_mnkl[group] + fmas += M * N * K + flop = 2 * fmas + gflop = flop / 1.0e9 + gflops = gflop / runtime_s + + print("Average Runtime : ", exec_time / 1000, "ms") + print("GFLOPS : ", gflops) + + return exec_time # Return execution time in microseconds + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError("Invalid format. Expected comma-separated integers.") + + def parse_comma_separated_tuples(s: str) -> List[tuple[int, ...]]: + if s.strip().startswith("("): + # Split on ),( to separate tuples + tuples = s.strip("()").split("),(") + result = [] + tuple_len = None + + for t in tuples: + # Parse individual tuple + nums = [int(x.strip()) for x in t.split(",")] + + # Validate tuple length consistency + if tuple_len is None: + tuple_len = len(nums) + elif len(nums) != tuple_len: + raise argparse.ArgumentTypeError("All tuples must have the same length") + + result.append(tuple(nums)) + return result + + raise argparse.ArgumentTypeError("Invalid format. Expected comma-separated integers or list of tuples") + + parser = argparse.ArgumentParser(description="Example of Grouped GEMM on Blackwell.") + parser.add_argument( + "--num_groups", + type=int, + default=2, + help="Number of groups", + ) + parser.add_argument( + "--problem_sizes_mnkl", + type=parse_comma_separated_tuples, + default=((128, 128, 128, 1), (128, 128, 128, 1)), + help="a tuple of problem sizes for each group (comma-separated tuples)", + ) + parser.add_argument( + "--mma_tiler_mn", + type=parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + parser.add_argument( + "--host_problem_shape_available", + action="store_true", + help="Enable the compute of grid based upon host problem shape", + ) + parser.add_argument( + "--cluster_shape_mn", + type=parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E8M0FNU) + parser.add_argument("--sf_vec_size", type=int, default=16) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float16) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument("--tolerance", type=float, default=1e-01, help="Tolerance for validation") + parser.add_argument("--warmup_iterations", type=int, default=0, help="Warmup iterations") + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument("--skip_ref_check", action="store_true", help="Skip reference checking") + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + + args = parser.parse_args() + + if len(args.problem_sizes_mnkl) != 0 and len(args.problem_sizes_mnkl) != args.num_groups: + parser.error("--problem_sizes_mnkl must contain exactly num_groups tuples") + + # l mode must be 1 for all groups + for _, _, _, l in args.problem_sizes_mnkl: + if l != 1: + parser.error("l must be 1 for all groups") + + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + run( + args.num_groups, + args.problem_sizes_mnkl, + args.host_problem_shape_available, + args.ab_dtype, + args.sf_dtype, + args.sf_vec_size, + args.c_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + print("PASS") diff --git a/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_kernel.py b/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_kernel.py new file mode 100644 index 0000000000..6a8809a77d --- /dev/null +++ b/aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_kernel.py @@ -0,0 +1,3124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa +# mypy: disable-error-code=name-defined + +# 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. + +from typing import Literal + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +import torch +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 +from cutlass.cute.runtime import make_ptr +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +""" +This example provides an experimental implementation of the SM100 batched dense blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases. + +A high-performance persistent batched dense blockscaled GEMM example for the NVIDIA Blackwell SM100 architecture +using CUTE DSL. +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") for MXF8 input type and can only be row-major("K") for MXF4/NVF4 input type +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") for MXF8 input type and can only be row-major("K") for MXF4/NVF4 input type +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") +- Matrix SFA layout is filled internally according to A shape and BlockScaledBasicChunk, which has M×ceil_div(K, sf_vec_size)×L elements respectively +- Matrix SFB layout is filled internally according to B shape and BlockScaledBasicChunk, which has N×ceil_div(K, sf_vec_size)×L elements respectively + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Support persistent tile scheduling to better overlap memory load/store with mma between tiles + - Support warp specialization to avoid explicit pipelining between mainloop load and mma + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: + - Load scale factor A/B from shared memory (SMEM) to tensor memory (TMEM) using tcgen05.cp instruction. + - Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Type convert C matrix to output type. + - Optionally store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations, + or directly store C matrix from registers (RMEM) to global memory (GMEM) without TMA operations. + - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor: + e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) + +SM100 tcgen05.mma.kind.block_scale instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Read scalefactor A from TMEM +- Read scalefactor B from TMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Input arguments to this example is shown below: + +.. code-block:: bash + + python examples/cute/blackwell/kernel/blockscaled_gemm/dense_blockscaled_gemm_persistent.py \ + --a_dtype Float4E2M1FN --b_dtype Float4E2M1FN \ + --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,1024,1 + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/cute/blackwell/kernel/blockscaled_gemm/dense_blockscaled_gemm_persistent.py \ + --a_dtype Float4E2M1FN --b_dtype Float4E2M1FN \ + --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,1024,1 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + + +Constraints: +* Supported input data types: mxf8, mxf4, nvf4, and mixed-precision f8f6f4 combinations + see detailed valid dtype combinations in below Sm100BlockScaledPersistentDenseGemmKernel class documentation +* A and B may use different element data types (e.g., Float8E4M3FN x Float4E2M1FN) +* Mma tiler M must be 128 or 256(use_2cta_instrs) +* Mma tiler N must be 64/128/192/256 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if Mma tiler M is 256(use_2cta_instrs) +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 16 and 32 for Float8 and Float4, respectively. +""" + + +class Sm100BlockScaledPersistentDenseGemmKernel: + """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types + and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + + :note: A and B may use different element data types (e.g. mixed f8f6f4). Each of + a_dtype / b_dtype may independently be any of the supported element types below. + + :note: Supported combinations of A/B data types, SF data typs and SF vector size: + - MXF8: A/B in {Float8E5M2, Float8E4M3FN} + SF: Float8E8M0FNU + sf_vec_size: 32 + - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32 + - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16 + - Mixed f8f6f4: A and B from different element types; + restricted to SF: Float8E8M0FNU + sf_vec_size: 32 + + :note: Supported accumulator data types: + - Float32 + + :note: Supported C data types: + - Float32 + - Float16/BFloat16 + - Float8E4M3FN/Float8E5M2 + + :note: Constraints: + - MMA tiler M must be 128 or 256 (use_2cta_instrs) + - MMA tiler N must be 64/128/192/256 + - Cluster shape M must be multiple of 2 if Mma tiler M is 256 + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + - Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors + + Example: + >>> gemm = Sm100BlockScaledPersistentDenseGemmKernel( + ... sf_vec_size=16, mma_tiler_mn=(256, 128), cluster_shape_mn=(2, 1) + ... ) + >>> gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, max_active_clusters, stream) + """ + + def __init__( + self, + sf_vec_size: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + ): + """Initializes the configuration for a Blackwell dense GEMM kernel. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator, always set to Float32 + - sf_vec_size: Scalefactor A/B vector size. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + """ + + self.acc_dtype = cutlass.Float32 + self.sf_vec_size = sf_vec_size + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler = (*mma_tiler_mn, 1) + + self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + + self.occupancy = 1 + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len((self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)) + # Set barrier id for epilogue sync and tmem ptr sync + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_warp * len(self.epilog_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=self.threads_per_warp * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B/SFA/SFB + - Computing epilogue subtile + - Setting up A/B/SFA/SFB/C stage counts in shared memory + - Computing A/B/SFA/SFB/C shared memory layout + """ + # Compute mma instruction shapes + # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K) + self.mma_inst_shape_mn = ( + self.mma_tiler[0], + self.mma_tiler[1], + ) + # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_inst_shape_mn[0], + self.mma_inst_shape_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Compute epilogue subtile + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + self.epi_tile_n = cute.size(self.epi_tile[1]) + + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.smem_alloc_a_dtype, + self.smem_alloc_b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.sf_dtype, + self.sf_vec_size, + self.smem_capacity, + self.occupancy, + ) + + # Compute A/B/SFA/SFB/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.smem_alloc_a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.smem_alloc_b_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + + # Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case + self.overlapping_accum = self.num_acc_stage == 1 + + # Compute number of TMEM columns for SFA/SFB/Accumulator + sf_atom_mn = 32 + self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k + self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k + self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols + self.num_accumulator_tmem_cols = ( + self.cta_tile_shape_mnk[1] * self.num_acc_stage + if not self.overlapping_accum + else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols + ) + + # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue + self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n + + @cute.jit + def __call__( + self, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + sfa_ptr: cute.Pointer, + sfb_ptr: cute.Pointer, + c_ptr: cute.Pointer, + layouts: cutlass.Constexpr[tuple[OperandMajorMode, OperandMajorMode, utils.LayoutEnum]], + problem_mnkl: tuple[int, int, int, int], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a_tensor: Input tensor A + :type a_tensor: cute.Tensor + :param b_tensor: Input tensor B + :type b_tensor: cute.Tensor + :param sfa_tensor: Scale factor tensor A + :type sfa_tensor: cute.Tensor + :param sfb_tensor: Scale factor tensor B + :type sfb_tensor: cute.Tensor + :param c_tensor: Output tensor C + :type c_tensor: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor + :type epilogue_op: cutlass.Constexpr + :raises TypeError: If input data types are incompatible with the MMA instruction. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: type[cutlass.Numeric] = a_ptr.value_type + self.b_dtype: type[cutlass.Numeric] = b_ptr.value_type + self.sf_dtype: type[cutlass.Numeric] = sfa_ptr.value_type + self.c_dtype: type[cutlass.Numeric] = c_ptr.value_type + self.mxf8f6f4 = self.needs_unpack_tma(self.a_dtype, self.b_dtype) + self.smem_alloc_a_dtype = cutlass.Int8 if (self.mxf8f6f4 and self.a_dtype.width < 8) else self.a_dtype + self.smem_alloc_b_dtype = cutlass.Int8 if (self.mxf8f6f4 and self.b_dtype.width < 8) else self.b_dtype + m, n, k, l = problem_mnkl + self.a_major_mode, self.b_major_mode, self.c_layout = layouts + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + a_layout = cute.make_ordered_layout((m, cute.assume(k, 32), l), order=(0, 1, 2)) + if cutlass.const_expr(self.a_major_mode == OperandMajorMode.K): + a_layout = cute.make_ordered_layout((cute.assume(m, 32), k, l), order=(1, 0, 2)) + b_layout = cute.make_ordered_layout((n, cute.assume(k, 32), l), order=(0, 1, 2)) + if cutlass.const_expr(self.b_major_mode == OperandMajorMode.K): + b_layout = cute.make_ordered_layout((cute.assume(n, 32), k, l), order=(1, 0, 2)) + c_layout = cute.make_ordered_layout((cute.assume(m, 32), n, l), order=(0, 1, 2)) + if cutlass.const_expr(self.c_layout == utils.LayoutEnum.ROW_MAJOR): + c_layout = cute.make_ordered_layout((m, cute.assume(n, 32), l), order=(1, 0, 2)) + a_tensor = cute.make_tensor(a_ptr, a_layout) + b_tensor = cute.make_tensor(b_ptr, b_layout) + c_tensor = cute.make_tensor(c_ptr, c_layout) + + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, self.sf_vec_size) + sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, self.sf_vec_size) + sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a_tensor, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=self.smem_alloc_a_dtype if (self.mxf8f6f4 and self.a_dtype.width < 8) else None, + ) + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b_tensor, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=self.smem_alloc_b_dtype if (self.mxf8f6f4 and self.b_dtype.width < 8) else None, + ) + + # Setup TMA load for SFA + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + sfa_tensor, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Setup TMA load for SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id) + sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb_tensor, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Int16, + ) + + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + x = tma_tensor_sfb.stride[0][1] + y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) + + new_shape = ( + (tma_tensor_sfb.shape[0][0], ((2, 2), y)), + tma_tensor_sfb.shape[1], + tma_tensor_sfb.shape[2], + ) + # Use right multiplication for ScaledBasis (3 * x instead of x * 3) + x_times_3 = 3 * x + new_stride = ( + (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), + tma_tensor_sfb.stride[1], + tma_tensor_sfb.stride[2], + ) + tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) + tma_tensor_sfb = cute.make_tensor(tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size) * atom_thr_size + + # Setup TMA store for C + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c_tensor, + epi_smem_layout, + self.epi_tile, + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c_tensor, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + self.buffer_align_bytes = 1024 + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.smem_alloc_a_dtype, + cute.cosize(self.a_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[ + self.smem_alloc_b_dtype, + cute.cosize(self.b_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sSFA: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sSFB: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tiled_mma_sfb, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + tma_atom_c, + tma_tensor_c, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_sfa) + cpasync.prefetch_descriptor(tma_atom_sfb) + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(cta_rank_in_cluster) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = self.threads_per_warp * len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/SFA/SFB/C + # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner) + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + # (MMA, MMA_M, MMA_K, STAGE) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + # (MMA, MMA_N, MMA_K, STAGE) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + + # + # Compute multicast mask for A/B/SFA/SFB buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) + # (bM, bK, RestM, RestK, RestL) + gSFA_mkl = cute.local_tile(mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + # (bN, bK, RestN, RestK, RestL) + gSFB_nkl = cute.local_tile( + mSFB_nkl, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # TMA load scaled factor A partition_S/D + sfa_cta_layout = a_cta_layout + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfa, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # TMA load scaled factor B partition_S/D + sfb_cta_layout = cute.make_layout(cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + if cutlass.const_expr(self.overlapping_accum): + num_acc_stage_overlapped = 2 + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage_overlapped)) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = cute.make_tensor( + tCtAcc_fake.iterator, + cute.make_layout( + tCtAcc_fake.shape, + stride=( + tCtAcc_fake.stride[0], + tCtAcc_fake.stride[1], + tCtAcc_fake.stride[2], + (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1], + ), + ), + ) + else: + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Specialized TMA load warp + # + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + + # ((atom_v, rest_v), RestK) + tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + + slice_n = mma_tile_coord_mnl[1] + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + slice_n = mma_tile_coord_mnl[1] // 2 + # ((atom_v, rest_v), RestK) + tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + + # TMA load A/B/SFA/SFB + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfa_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfb_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # Make accumulator tmem tensor + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # Make SFA tmem tensor + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + # (MMA, MMA_M, MMA_K) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + # Make SFB tmem tensor + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + # (MMA, MMA_N, MMA_K) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + # + # Partition for S2T copy of SFA/SFB + # + ( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t, + tCtSFA_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + ( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t, + tCtSFB_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage) + acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Get accumulator stage index + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_producer_state.phase ^ 1 + else: + acc_stage_index = acc_producer_state.index + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + tCtSFB_mma = tCtSFB + if cutlass.const_expr(self.cta_tile_shape_mnk[1] in {64, 192}): + # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB) + offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + + # Copy SFA/SFB from smem to tmem + s2t_stage_coord = ( + None, + None, + None, + None, + ab_consumer_state.index, + ) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t_staged, + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t_staged, + tCtSFB_compact_s2t, + ) + + # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, ab_consumer_state.index) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[tile_crd], tCtSFA], + [tCrB[tile_crd], tCtSFB_mma], + tCtAcc, + ) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # + # Partition for epilogue + # + epi_tidx = tidx + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = self.epilog_tmem_copy_and_partition(epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(tiled_copy_t2r, tTR_rC, epi_tidx, sC) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_c, tCgC, epi_tile, sC) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) + + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + + # Get accumulator stage index + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_consumer_state.phase + reverse_subtile = True if acc_stage_index == 0 else False + else: + acc_stage_index = acc_consumer_state.index + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in cutlass.range(subtile_cnt): + real_subtile_idx = subtile_idx + if cutlass.const_expr(self.overlapping_accum): + if reverse_subtile: + real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n - 1 - subtile_idx + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Async arrive accumulator buffer empty ealier when overlapping_accum is enabled + # + if cutlass.const_expr(self.overlapping_accum): + if subtile_idx == self.iter_acc_early_release_in_epilogue: + # Fence for TMEM load + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, real_subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + if cutlass.const_expr(not self.overlapping_accum): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(acc_tmem_ptr) + # + # Wait for C store complete + # + c_pipeline.producer_tail() + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination). + + :param sSF: The scale factor tensor in smem + :type sSF: cute.Tensor + :param tSF: The scale factor tensor in tmem + :type tSF: cute.Tensor + + :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: + - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) + - tCsSF_compact_s2t: The partitioned scale factor tensor in smem + - tSF_compact_s2t: The partitioned scale factor tensor in tmem + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSF_compact = cute.filter_zeros(sSF) + # (MMA, MMA_MN, MMA_K) + tCtSF_compact = cute.filter_zeros(tSF) + + # Make S2T CopyAtom and tiledCopy + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: cutlass.Boolean | bool, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc: The accumulated tensor in register used to hold t2r results + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_mnl_epi = cute.flat_divide(gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilog_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + :type sepi: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rC: The partitioned tensor C (register source) + - tRS_sC: The partitioned tensor C (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op(self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: cute.CopyAtom | cute.TiledCopy, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Make tiledCopy for global memory store, then use it to: + partition shared memory (source) and global memory (destination) for TMA store version. + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where: + - tma_atom_c: The TMA copy atom + - bSG_sC: The partitioned shared memory tensor C + - bSG_gC: The partitioned global tensor C + :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_epi = cute.flat_divide(gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + + tma_atom_c = atom + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + return tma_atom_c, bSG_sC, bSG_gC + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + epi_tile: cute.Tile, + c_dtype: type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + sf_dtype: type[cutlass.Numeric], + sf_vec_size: int, + smem_capacity: int, + occupancy: int, + ) -> tuple[int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout enum of operand C. + :type c_layout: utils.LayoutEnum + :param sf_dtype: Data type of Scale factor. + :type sf_dtype: type[cutlass.Numeric] + :param sf_vec_size: Scale factor vector size. + :type sf_vec_size: int + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # ACC stages + num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 + + # Default C stages + num_c_stage = 2 + + # Calculate smem layout and size for one stage of A, B, SFA, SFB and C + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, # a tmp 1 stage is provided + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, # a tmp 1 stage is provided + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + mbar_helpers_bytes = 1024 + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B/SFA/SFB stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B/SFA/SFB stage + num_ab_stage = (smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes + # Add remaining unused smem to epilogue + num_c_stage += ( + smem_capacity - occupancy * ab_bytes_per_stage * num_ab_stage - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + + return num_acc_stage, num_ab_stage, num_c_stage + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + :param swizzle_size: Swizzling size in the unit of cluster for improving L2 cache hit rate, defaults to 1 + :type swizzle_size: int + :param raster_order: Rasterization order of clusters ('m' or 'n'), defaults to 'm' + :type raster_order: Literal["m", "n"] + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + # Convert raster_order ("m" or "n") to raster_along_m (True or False) + raster_along_m = raster_order == "m" + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl, swizzle_size, raster_along_m + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape(tile_sched_params, max_active_clusters) + + return tile_sched_params, grid + + @staticmethod + def needs_unpack_tma( + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + ) -> bool: + """ + Decide whether TMA must use the UNPACK_U8 variant (U4_UNPACK_U8 / + U6_UNPACK_U8) for narrow-precision operands. + + Unpack is required when: + * Operand widths differ (mxf8f6f4 mixed-precision) — A and B must + share a uniform byte-per-element SMEM layout, so the narrower + operand is unpacked into 1B/elem containers in SMEM. + * Either operand is 6-bit — there is no packed U6 TMA format, + only U6_UNPACK_U8 exists. + + Otherwise (same-width and no 6-bit operand, e.g. f4xf4 / f8xf8 / + f8E4M3xf8E5M2) TMA can use the natural packed format (U4 for 4-bit, + U8 for 8-bit). + + :param a_dtype: Element data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: Element data type of the B operand + :type b_dtype: Type[cutlass.Numeric] + :return: True if UNPACK_U8 TMA format must be used, False otherwise + :rtype: bool + """ + if a_dtype.width != b_dtype.width: + return True + if a_dtype.width == 6 or b_dtype.width == 6: + return True + return False + + @staticmethod + def is_valid_dtypes_and_scale_factor_vec_size( + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: type[cutlass.Numeric], + ) -> bool: + """ + Check if the dtypes and sf_vec_size are valid combinations + + :param a_dtype: The data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operand + :type b_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size of the scale factor + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes and sf_vec_size are valid, False otherwise + :rtype: bool + """ + supported_ab_dtypes = { + cutlass.Float4E2M1FN, + cutlass.Float6E2M3FN, + cutlass.Float6E3M2FN, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + } + + # Check A/B element types + if a_dtype not in supported_ab_dtypes or b_dtype not in supported_ab_dtypes: + return False + + # Check SF element type + if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: + return False + + # sf_vec_size rules: + # * 16 is only supported for Float4E2M1FN x Float4E2M1FN (NVF4 / MXF4 fp4-pair) + # * 32 is required for every other A/B combination (MXF8, mxf8f6f4 mixed, MXF4-pair with MX scaling) + # SF dtype pairing with sf_vec_size: + # * sf_vec_size == 16 requires sf_dtype in {Float8E4M3FN (NVF4), Float8E8M0FNU (MXF4)} + # * sf_vec_size == 32 requires sf_dtype == Float8E8M0FNU (MX scaling) + both_fp4 = a_dtype is cutlass.Float4E2M1FN and b_dtype is cutlass.Float4E2M1FN + if sf_vec_size == 16: + if not both_fp4: + return False + elif sf_vec_size == 32: + if sf_dtype is not cutlass.Float8E8M0FNU: + return False + else: + return False + + # Check valid c_dtype + if c_dtype not in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + return False + + return True + + @staticmethod + def is_valid_layouts( + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + ) -> bool: + """ + Check if layouts and dtypes are valid combinations + + :param a_dtype: The data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operand + :type b_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major dimension of the A tensor + :type a_major: Literal["m", "k"] + :param b_major: The major dimension of the B tensor + :type b_major: Literal["n", "k"] + :param c_major: The major dimension of the C tensor + :type c_major: Literal["m", "n"] + + :return: True if the layouts are valid, False otherwise + :rtype: bool + """ + is_valid = True + + # FP4 operands can only be k-major (checked per operand) + if a_dtype is cutlass.Float4E2M1FN and a_major != "k": + is_valid = False + if b_dtype is cutlass.Float4E2M1FN and b_major != "k": + is_valid = False + return is_valid + + @staticmethod + def is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + ) -> bool: + """ + Check if the mma tiler and cluster shape are valid + + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + + :return: True if the mma tiler and cluster shape are valid, False otherwise + :rtype: bool + """ + is_valid = True + # Skip invalid mma tile shape + if mma_tiler_mn[0] not in [128, 256]: + is_valid = False + if mma_tiler_mn[1] not in [64, 128, 192, 256]: + is_valid = False + # Skip illegal cluster shape + if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + # Special cluster shape check for scale factor multicasts. + # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. + or cluster_shape_mn[0] > 4 + or cluster_shape_mn[1] > 4 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + @staticmethod + def is_valid_tensor_alignment( + m: int, + n: int, + k: int, + l: int, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + mma_tiler_mn: tuple[int, int], + ) -> bool: + """ + Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param a_dtype: The data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operand + :type b_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: Literal["m", "k"] + :param b_major: The major axis of the B tensor + :type b_major: Literal["n", "k"] + :param c_major: The major axis of the C tensor + :type c_major: Literal["m", "n"] + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler, + needed to verify per-CTA UNPACK alignment under 2CTA MMA. + :type mma_tiler_mn: Tuple[int, int] + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + # TMA requires the contiguous inner dimension to be a multiple of + # 16 B (= 128 bits). Work in bits so non-byte-aligned widths + # (e.g. 6-bit) are handled correctly: 16 * 8 // dtype.width is + # wrong when dtype.width does not divide 128 (it returns 21 for + # 6-bit instead of the real requirement K*6 % 128 == 0). + return (num_major_elements * dtype.width) % (16 * 8) == 0 + + def check_contigous_128_alignment(dtype, is_mode0_major, tensor_shape): + # we only need to check alignment for subbyte dtype + if dtype.width >= 8: + return True + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 128 + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(a_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(b_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + # When an operand is loaded via the UNPACK TMA variant + # (U4_UNPACK_U8 or U6_UNPACK_U8), its inner tensor dimension in bytes + # must be a multiple of 64B (4-bit) or 96B (6-bit); both work out to + # a multiple of 128 elements along the contiguous dim. The check only + # applies to sub-byte operands and only when the pair triggers UNPACK. + if Sm100BlockScaledPersistentDenseGemmKernel.needs_unpack_tma(a_dtype, b_dtype) and ( + not check_contigous_128_alignment(a_dtype, a_major == "m", (m, k, l)) + or not check_contigous_128_alignment(b_dtype, b_major == "n", (n, k, l)) + ): + is_valid = False + # Additional UNPACK constraint for any sub-byte operand on its contig + # axis: when a sub-byte A is m-major (contig=M) or a sub-byte B is + # n-major (contig=N), the MMA tile's contig dim (after 2CTA M-split, + # which splits M for both A and B on the non-multicast atom path) + # must be a multiple of 128 elements to satisfy the 64B (fp4) / + # 96B (fp6) inner-dim requirement of U4_/U6_UNPACK_U8. Observed + # failures: (128,192)/(1,1)/m-n-m (1CTA) and (256,128)/(2,2)/m-n-m + # (2CTA) trigger CUDA illegal instruction for fp6 when mma_tiler_N + # (or N/2 after 2CTA split) is not a 128-multiple; the same rule + # applies to any other sub-byte operand on its non-K contig axis. + use_2cta_instrs = mma_tiler_mn[0] == 256 + cta_div = 2 if use_2cta_instrs else 1 + if ( + Sm100BlockScaledPersistentDenseGemmKernel.needs_unpack_tma(a_dtype, b_dtype) + and a_major == "m" + and a_dtype.width < 8 + and (mma_tiler_mn[0] // cta_div) % 128 != 0 + ): + is_valid = False + if ( + Sm100BlockScaledPersistentDenseGemmKernel.needs_unpack_tma(a_dtype, b_dtype) + and b_major == "n" + and b_dtype.width < 8 + and (mma_tiler_mn[1] // cta_div) % 128 != 0 + ): + is_valid = False + return is_valid + + @staticmethod + def can_implement( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + sf_vec_size: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + ) -> bool: + """ + Check if the gemm can be implemented + + :param mnkl: The problem size as a tuple (M, N, K, L). + :type mnkl: Tuple[int, int, int, int] + :param a_dtype: The data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operand + :type b_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: Literal["m", "k"] + :param b_major: The major axis of the B tensor + :type b_major: Literal["n", "k"] + :param c_major: The major axis of the C tensor + :type c_major: Literal["m", "n"] + :param sf_vec_size: The vector size + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + :return: True if the gemm can be implemented, False otherwise + :rtype: bool + """ + # Unpack parameters + m, n, k, l = mnkl + can_implement = True + # Skip unsupported types + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size( + a_dtype, b_dtype, sf_dtype, sf_vec_size, c_dtype + ): + can_implement = False + # Skip unsupported layouts + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_layouts( + a_dtype, b_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + # Skip invalid mma tile shape and cluster shape + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn, cluster_shape_mn + ): + can_implement = False + # Skip illegal problem shape for load/store alignment + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_tensor_alignment( + m, + n, + k, + l, + a_dtype, + b_dtype, + c_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + ): + can_implement = False + return can_implement + + +# Helper function to convert scale factor tensor from MKL layout to (32, 4, restM, 4, restK, l) format +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_ptr: cute.Pointer, + sf_mma_ptr: cute.Pointer, + mn: int, + sf_k: int, + l: int, + mma_shape: tuple, +): + mma_permute_order = (3, 4, 1, 5, 2, 0) + permuted_shape = tuple(mma_shape[i] for i in mma_permute_order) + cute_layout = cute.make_ordered_layout(permuted_shape, order=(2, 1, 4, 0, 3, 5)) + + sf_ref_tensor = cute.make_tensor(sf_ref_ptr, cute.make_layout((mn, sf_k, l), stride=(sf_k, 1, mn * sf_k))) + sf_mma_tensor = cute.make_tensor(sf_mma_ptr, cute_layout) + + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + pass + + +# Helper function for ceil division +def ceil_div(a, b): + return (a + b - 1) // b + + +# Convert scale factor tensors from (m, k, l) to (32, 4, restM, 4, restK, l) format +def create_and_reorder_scale_factor_tensor(l, mn, k, sf_vec_size, sf_dtype, torch_tensor): + """ + Create the CUTE-format scale factor tensor on CUDA based on the reference tensor. + """ + sf_k = ceil_div(k, sf_vec_size) + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, # batch size + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + # Generate a random int8 tensor, then convert to float8_e4m3fn + cute_tensor = torch.ones(mma_shape, dtype=cutlass_torch.dtype(sf_dtype)).permute(3, 4, 1, 5, 2, 0) + + # Call the helper function to do layout conversion + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + make_ptr( + sf_dtype, + torch_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + make_ptr( + sf_dtype, + cute_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + mn, + sf_k, + l, + mma_shape, + ) + return cute_tensor.cuda() + + +# Compile the persistent dense blockscaled GEMM operation +def scaled_mm( + gemm_obj: Sm100BlockScaledPersistentDenseGemmKernel, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + options: str = "", +): + # Construct CuTe Pointers + a_ptr = make_ptr(a_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + b_ptr = make_ptr(b_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32) + sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32) + + a_major_mode = OperandMajorMode.K if a_major == "k" else OperandMajorMode.MN + b_major_mode = OperandMajorMode.K if b_major == "k" else OperandMajorMode.MN + c_layout = utils.LayoutEnum.ROW_MAJOR if c_major == "n" else utils.LayoutEnum.COL_MAJOR + return cute.compile( + gemm_obj, + a_ptr, + b_ptr, + sfa_ptr, + sfb_ptr, + c_ptr, + (a_major_mode, b_major_mode, c_layout), + (cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0)), + max_active_clusters, + stream, + epilogue_op, + options=options, + ) + + +def is_emulated_dtype( + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], +) -> bool: + # torch natively represents A/B only when both operands share the same + # dtype and the (dtype, sf_dtype) pair matches a supported non-emulated case. + if a_dtype != b_dtype: + return True + if c_dtype in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + }: + if a_dtype == cutlass.Float4E2M1FN and sf_dtype == cutlass.Float8E4M3FN: + return False + if a_dtype == cutlass.Float8E4M3FN and sf_dtype == cutlass.Float8E8M0FNU: + return False + + return True + + +# Convert scale factor tensor from MKL layout to blocked layout +def to_blocked(input_matrix): + rows, cols = input_matrix.shape + # Please ensure rows and cols are multiples of 128 and 4 respectively + n_row_blocks = ceil_div(rows, 128) + n_col_blocks = ceil_div(cols, 4) + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + + # Pad the input matrix if necessary + if padded_rows != rows or padded_cols != cols: + # For FP8 types, convert to float32 for padding, then convert back + original_dtype = input_matrix.dtype + input_float32 = input_matrix.to(torch.float32) + padded = torch.nn.functional.pad( + input_float32, + (0, padded_cols - cols, 0, padded_rows - rows), + mode="constant", + value=0, + ) + # Convert back to original dtype if needed + if original_dtype != input_float32.dtype: + padded = padded.to(original_dtype) + else: + padded = input_matrix + blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3) + rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16) + return rearranged.flatten() + + +# Reference implementation of the persistent dense blockscaled GEMM operation (emulated version) +def reference_scaled_mm_emulated( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + mnkl: tuple[int, int, int, int], + sf_vec_size: int, + c_dtype: type[cutlass.Numeric], +): + m, n, k, l = mnkl + sfa_expanded = torch.repeat_interleave(sfa, sf_vec_size, dim=1)[:, :k, :].to(dtype=torch.float32).cuda() + sfb_expanded = torch.repeat_interleave(sfb, sf_vec_size, dim=1)[:, :k, :].to(dtype=torch.float32).cuda() + res_a = torch.einsum("mkl,mkl->mkl", a, sfa_expanded) + res_b = torch.einsum("nkl,nkl->nkl", b, sfb_expanded) + # Cast res_a and res_b to float32 for einsum to avoid NotImplementedError on 'Byte' + ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) + c_ref = ref.to(dtype=cutlass_torch.dtype(c_dtype)) + return c_ref + + +# Reference implementation of the persistent dense blockscaled GEMM operation (non-emulated version) +def reference_scaled_mm( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + mnkl: tuple[int, int, int, int], + c_dtype: type[cutlass.Numeric], +): + m, n, k, l = mnkl + c_ref = torch.clone(c) + for l_idx in range(l): + # Convert the scale factor tensor to blocked format + scale_a = to_blocked(sfa[:, :, l_idx]) + scale_b = to_blocked(sfb[:, :, l_idx]) + # Ensure a_slice is row-major (M, K) with stride (K, 1) + a_slice = a[:, :, l_idx].contiguous() + # Ensure b_slice is row-major (N, K) so that transpose gives column-major (K, N) + b_slice = b[:, :, l_idx].contiguous() + # (m, k) @ (n, k).T -> (m, n) + res = torch._scaled_mm( + a_slice, + b_slice.transpose(0, 1), + scale_a.cuda(), + scale_b.cuda(), + bias=None, + out_dtype=c_ref.dtype, + ) + c_ref[:, :, l_idx] = res + return c_ref + + +# Construct CuTe Pointers for the persistent dense blockscaled GEMM operation (emulated version) +def construct_cute_pointers_emulated( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], +): + a_cute, _ = cutlass_torch.cute_tensor_like( + a.cpu(), + a_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + a_cute = cutlass_torch.convert_cute_tensor( + a, + a_cute, + a_dtype, + is_dynamic_layout=True, + ) + b_cute, _ = cutlass_torch.cute_tensor_like( + b.cpu(), + b_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + b_cute = cutlass_torch.convert_cute_tensor( + b, + b_cute, + b_dtype, + is_dynamic_layout=True, + ) + a_ptr = a_cute.iterator + b_ptr = b_cute.iterator + + sfa_ptr = make_ptr(sf_dtype, sfa.data_ptr(), cute.AddressSpace.gmem, assumed_align=32) + sfb_ptr = make_ptr(sf_dtype, sfb.data_ptr(), cute.AddressSpace.gmem, assumed_align=32) + c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + return a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, a_cute, b_cute + + +# Construct CuTe Pointers for the persistent dense blockscaled GEMM operation (non-emulated version) +def construct_cute_pointers( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], +): + a_ptr = make_ptr(a_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + b_ptr = make_ptr(b_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + sfa_ptr = make_ptr(sf_dtype, sfa.data_ptr(), cute.AddressSpace.gmem, assumed_align=32) + sfb_ptr = make_ptr(sf_dtype, sfb.data_ptr(), cute.AddressSpace.gmem, assumed_align=32) + c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + return a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr + + +# Use uint8 and uint32 to emulate unsupported +# dtype in torch +def prepare_tensors_emulated( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + init_normal: bool = False, + normal_mean: float = 0.0, + normal_std: float = 1.0, +): + m, n, k, l = mnkl + sf_k = ceil_div(k, sf_vec_size) + + # Create tensor SFA/SFB with values in [1, 3) + sfa = torch.randint(0, 3, (l, m, sf_k), dtype=torch.uint8).permute(1, 2, 0).to(dtype=cutlass_torch.dtype(sf_dtype)) + sfb = torch.randint(0, 3, (l, n, sf_k), dtype=torch.uint8).permute(1, 2, 0).to(dtype=cutlass_torch.dtype(sf_dtype)) + + # Create tensor A/B + if a_major == "k": + a = torch.empty((l, m, k), dtype=torch.float32, device="cuda").permute(1, 2, 0) + else: + a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(2, 1, 0) + if b_major == "k": + b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(1, 2, 0) + else: + b = torch.empty((l, k, n), dtype=torch.float32, device="cuda").permute(2, 1, 0) + + # Initialize A/B tensors with either normal distribution or random integers + for tensor in [a, b]: + if init_normal: + tensor.normal_(mean=normal_mean, std=normal_std) + else: + tensor.copy_(torch.randint(-2, 2, tensor.shape, dtype=torch.float32, device="cuda")) + + if c_major == "n": + c = torch.empty((l, m, n), dtype=cutlass_torch.dtype(c_dtype), device="cuda").permute(1, 2, 0) + else: + c = torch.empty((l, n, m), dtype=cutlass_torch.dtype(c_dtype), device="cuda").permute(2, 1, 0) + return a, b, c, sfa, sfb + + +def prepare_tensors( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + init_normal: bool = False, + normal_mean: float = 0.0, + normal_std: float = 1.0, +): + m, n, k, l = mnkl + + if a_dtype == cutlass.Float4E2M1FN and b_dtype == cutlass.Float4E2M1FN: + # Using int8 for torch.float4_e2m1fn_x2 tensor allocation + # Thus the size of k needs to be halved in this case. + k_fct = 2 + else: + k_fct = 1 + + sf_k = ceil_div(k, sf_vec_size) + + # Create tensor SFA/SFB + sfa = torch.randint(0, 3, (l, m, sf_k), dtype=torch.uint8).permute(1, 2, 0).to(dtype=cutlass_torch.dtype(sf_dtype)) + sfb = torch.randint(0, 3, (l, n, sf_k), dtype=torch.uint8).permute(1, 2, 0).to(dtype=cutlass_torch.dtype(sf_dtype)) + + # Create tensor A/B + if a_major == "k": + a = torch.empty((l, m, k // k_fct), dtype=torch.int8, device="cuda").permute(1, 2, 0) + else: + a = torch.empty((l, k, m), dtype=torch.int8, device="cuda").permute(2, 1, 0) + if b_major == "k": + b = torch.empty((l, n, k // k_fct), dtype=torch.int8, device="cuda").permute(1, 2, 0) + else: + b = torch.empty((l, k, n), dtype=torch.int8, device="cuda").permute(2, 1, 0) + + # Initialize A/B tensors with random integers + # Note: int8 types always use random init (normal distribution not supported), + # consistent with use_normal_init pattern in generate_tensors + for tensor in [a, b]: + tensor.copy_(torch.randint(-2, 2, tensor.shape, dtype=torch.int8, device="cuda")) + + # Create and initialize tensor C + if c_major == "n": + c_shape, c_perm = (l, m, n), (1, 2, 0) + else: + c_shape, c_perm = (l, n, m), (2, 1, 0) + + c = torch.randint(-2, 2, c_shape, dtype=cutlass_torch.dtype(c_dtype), device="cuda").permute(c_perm) + + if a_dtype == cutlass.Float4E2M1FN and b_dtype == cutlass.Float4E2M1FN: + a = a.view(dtype=torch.float4_e2m1fn_x2) + b = b.view(dtype=torch.float4_e2m1fn_x2) + else: + a = a.to(dtype=cutlass_torch.dtype(a_dtype)) + b = b.to(dtype=cutlass_torch.dtype(b_dtype)) + + c = c.to(dtype=cutlass_torch.dtype(c_dtype)) + return a, b, c, sfa, sfb + + +# This will show how to convert torch tensor +# and pass to CuTe kernel +def run_scaled_mm( + gemm_obj: Sm100BlockScaledPersistentDenseGemmKernel, + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + init_normal: bool = False, + normal_mean: float = 0.0, + normal_std: float = 1.0, + **kwargs, +): + """Execute a persistent batched dense blockscaled GEMM operation on Blackwell architecture with performance benchmarking (non-emulated dtypes). + + This function prepares input tensors, configures and launches the persistent GEMM kernel, + optionally performs reference validation, and benchmarks the execution performance. + + :param gemm_obj: A gemm object which is created and passed along to be used + :type gemm_obj: A gemm_obj of Sm100BlockScaledPersistentDenseGemmKernel + :param mnkl: Problem size (M, N, K, L) + :type mnkl: Tuple[int, int, int, int] + :param a_dtype: Data type for input tensor A + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: Data type for input tensor B + :type b_dtype: Type[cutlass.Numeric] + :param sf_dtype: Data type for scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: Vector size for scale factor tensor + :type sf_vec_size: int + :param c_dtype: Data type for output tensor C + :type c_dtype: Type[cutlass.Numeric] + :param a_major/b_major/c_major: Memory layout of tensor A/B/C + :type a_major/b_major/c_major: Literal["m", "k", "n"] + :param mma_tiler_mn: MMA tiling size. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster shape. + :type cluster_shape_mn: Tuple[int, int] + :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 + :type tolerance: float, optional + :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0 + :type warmup_iterations: int, optional + :param iterations: Number of benchmark iterations to run, defaults to 1 + :type iterations: int, optional + :param skip_ref_check: Whether to skip reference result validation, defaults to False + :type skip_ref_check: bool, optional + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False + :type use_cold_l2: bool, optional + :param init_normal: Whether to initialize tensors using normal distribution + instead of uniform random, defaults to False. + :type init_normal: bool, optional + :param normal_mean: Mean for normal distribution initialization, defaults to 0.0. + :type normal_mean: float, optional + :param normal_std: Standard deviation for normal distribution initialization, + defaults to 1.0. + :type normal_std: float, optional + :raises RuntimeError: If CUDA GPU is not available + :raises ValueError: If the configuration is invalid or unsupported by the kernel + :return: Execution time of the GEMM kernel + :rtype: float + """ + print(f"Running {gemm_obj.__class__.__name__} test with:") + print(f"mnkl: {mnkl}") + print(f"A dtype: {a_dtype}, B dtype: {b_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}") + print(f"C dtype: {c_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + + # Unpack parameters + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + + # Check if configuration can be implemented + max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) + + # Compile gemm kernel with fake tensors + compiled_gemm = scaled_mm( + gemm_obj, + a_dtype, + b_dtype, + c_dtype, + sf_dtype, + a_major, + b_major, + c_major, + max_active_clusters, + current_stream, + ) + + # Create Torch Tensors for A, scale factor A, B, scale factor B, C + a, b, c, sfa, sfb = prepare_tensors( + mnkl, + a_dtype, + b_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + init_normal=init_normal, + normal_mean=normal_mean, + normal_std=normal_std, + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype, sfa) + sfb_reordered = create_and_reorder_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype, sfb) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr = construct_cute_pointers( + a, + b, + sfa_reordered, + sfb_reordered, + c, + a_dtype, + b_dtype, + sf_dtype, + c_dtype, + ) + + # Compute reference result + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_gemm(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream) + c_ref = reference_scaled_mm(a, b, sfa, sfb, c, (m, n, k, l), c_dtype) + if c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Rtol=0.001 and atol=0.1 are not supported for bitwise comparison of + # low dimensional floats. Please use rtol=0.0 and atol=0.0. + tolerance = 0.0 + torch.testing.assert_close(c, c_ref, atol=tolerance, rtol=tolerance) + + def generate_inputs(): + a, b, c, sfa, sfb = prepare_tensors( + mnkl, + a_dtype, + b_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + init_normal=init_normal, + normal_mean=normal_mean, + normal_std=normal_std, + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype, sfa) + sfb_reordered = create_and_reorder_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype, sfb) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr = construct_cute_pointers( + a, + b, + sfa_reordered, + sfb_reordered, + c, + a_dtype, + b_dtype, + sf_dtype, + c_dtype, + ) + jit_args = cutlass.testing.JitArguments(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream) + # Keep references to external variables (e.g., Torch tensors when taking a view) + jit_args.add_to_scope([a, b, sfa_reordered, sfb_reordered, c]) + return jit_args + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a.numel() * a.element_size() + + b.numel() * b.element_size() + + sfa.numel() * sfa.element_size() + + sfb.numel() * sfb.element_size() + + c.numel() * c.element_size() + ) + workspace_count = cutlass.testing.get_workspace_count(one_workspace_bytes, warmup_iterations, iterations) + + exec_time = cutlass.testing.benchmark( + compiled_gemm, + workspace_generator=generate_inputs, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + return exec_time # Return execution time in microseconds + + +# This is to compatible with the other narrow +# precision combinations are not supported in either +# torch or dlpack. For example, Float4E2M1FN with Float8E8M0FNU. +def run_scaled_mm_with_emulated_dtype( + gemm_obj: Sm100BlockScaledPersistentDenseGemmKernel, + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + init_normal: bool = False, + normal_mean: float = 0.0, + normal_std: float = 1.0, + **kwargs, +): + """Execute a persistent batched dense blockscaled GEMM operation on Blackwell architecture with performance benchmarking (emulated dtypes). + + This function prepares input tensors, configures and launches the persistent GEMM kernel, + optionally performs reference validation, and benchmarks the execution performance. + + :param gemm_obj: A gemm object which is created and passed along to be used + :type gemm_obj: A gemm_obj of Sm100BlockScaledPersistentDenseGemmKernel + :param mnkl: Problem size (M, N, K, L) + :type mnkl: Tuple[int, int, int, int] + :param a_dtype: Data type for input tensor A + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: Data type for input tensor B + :type b_dtype: Type[cutlass.Numeric] + :param sf_dtype: Data type for scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: Vector size for scale factor tensor + :type sf_vec_size: int + :param c_dtype: Data type for output tensor C + :type c_dtype: Type[cutlass.Numeric] + :param a_major/b_major/c_major: Memory layout of tensor A/B/C + :type a_major/b_major/c_major: Literal["m", "n","k"] + :param mma_tiler_mn: MMA tiling size. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster shape. + :type cluster_shape_mn: Tuple[int, int] + :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 + :type tolerance: float, optional + :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0 + :type warmup_iterations: int, optional + :param iterations: Number of benchmark iterations to run, defaults to 1 + :type iterations: int, optional + :param skip_ref_check: Whether to skip reference result validation, defaults to False + :type skip_ref_check: bool, optional + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False + :type use_cold_l2: bool, optional + :param init_normal: Whether to initialize tensors using normal distribution + instead of uniform random, defaults to False. + :type init_normal: bool, optional + :param normal_mean: Mean for normal distribution initialization, defaults to 0.0. + :type normal_mean: float, optional + :param normal_std: Standard deviation for normal distribution initialization, + defaults to 1.0. + :type normal_std: float, optional + :raises RuntimeError: If CUDA GPU is not available + :raises ValueError: If the configuration is invalid or unsupported by the kernel + :return: Execution time of the GEMM kernel + :rtype: float + """ + print(f"Running {gemm_obj.__class__.__name__} test (Emulated) with:") + print(f"mnkl: {mnkl}") + print(f"A dtype: {a_dtype}, B dtype: {b_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}") + print(f"C dtype: {c_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + + # Unpack parameters + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + + # Check if configuration can be implemented + max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) + + # Compile gemm kernel with fake tensors + compiled_gemm = scaled_mm( + gemm_obj, + a_dtype, + b_dtype, + c_dtype, + sf_dtype, + a_major, + b_major, + c_major, + max_active_clusters, + current_stream, + ) + + # Create Torch Tensors for A, scale factor A, B, scale factor B, C + a, b, c, sfa, sfb = prepare_tensors_emulated( + mnkl, + a_dtype, + b_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + init_normal=init_normal, + normal_mean=normal_mean, + normal_std=normal_std, + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype, sfa) + sfb_reordered = create_and_reorder_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype, sfb) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, a_cute, b_cute = construct_cute_pointers_emulated( + a, + b, + sfa_reordered, + sfb_reordered, + c, + a_dtype, + b_dtype, + sf_dtype, + c_dtype, + ) + + # Compute reference result + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_gemm(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream) + c_ref = reference_scaled_mm_emulated(a, b, sfa, sfb, c, (m, n, k, l), sf_vec_size, c_dtype) + if c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Rtol=0.001 and atol=0.1 are not supported for bitwise comparison of + # low dimensional floats. Please use rtol=0.0 and atol=0.0. + tolerance = 0.0 + torch.testing.assert_close(c, c_ref, atol=tolerance, rtol=tolerance) + + def generate_inputs(): + a, b, c, sfa, sfb = prepare_tensors_emulated( + mnkl, + a_dtype, + b_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + init_normal=init_normal, + normal_mean=normal_mean, + normal_std=normal_std, + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype, sfa) + sfb_reordered = create_and_reorder_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype, sfb) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, a_cute, b_cute = construct_cute_pointers_emulated( + a, + b, + sfa_reordered, + sfb_reordered, + c, + a_dtype, + b_dtype, + sf_dtype, + c_dtype, + ) + jit_args = cutlass.testing.JitArguments(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream) + # Keep references to external variables (e.g., Torch tensors when taking a view) + jit_args.add_to_scope([a, b, sfa_reordered, sfb_reordered, c, a_cute, b_cute]) + return jit_args + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a.numel() * a.element_size() + + b.numel() * b.element_size() + + sfa.numel() * sfa.element_size() + + sfb.numel() * sfb.element_size() + + c.numel() * c.element_size() + ) + workspace_count = cutlass.testing.get_workspace_count(one_workspace_bytes, warmup_iterations, iterations) + + exec_time = cutlass.testing.benchmark( + compiled_gemm, + workspace_generator=generate_inputs, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + return exec_time # Return execution time in microseconds + + +def run( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + init_normal: bool = False, + normal_mean: float = 0.0, + normal_std: float = 1.0, + **kwargs, +): + """ + Execute the appropriate GEMM function based on dtype. + + Routes to either run_scaled_mm_with_emulated_dtype or run_scaled_mm + depending on whether the dtypes require emulation. + + :param init_normal: Whether to initialize tensors using normal distribution + instead of uniform random, defaults to False. Note: for int8/uint8 dtypes, + tensors always use random integer initialization regardless of this flag. + :type init_normal: bool, optional + :param normal_mean: Mean for normal distribution initialization, defaults to 0.0. + :type normal_mean: float, optional + :param normal_std: Standard deviation for normal distribution initialization, + defaults to 1.0. + :type normal_std: float, optional + """ + + # Unpack once for logging / error messages + m, n, k, l = mnkl + + # Configure gemm kernel + gemm = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + ) + + # Skip unsupported testcase + if not gemm.can_implement( + mnkl, + a_dtype, + b_dtype, + sf_dtype, + c_dtype, + a_major, + b_major, + c_major, + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + ): + raise cutlass.testing.CantImplementError( + "Unsupported testcase " + f"{a_dtype}, {b_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, " + f"{mma_tiler_mn}, {cluster_shape_mn}, " + f"{m}, {n}, {k}, {l}, " + f"{a_major}, {b_major}, {c_major}" + ) + + if is_emulated_dtype(a_dtype, b_dtype, sf_dtype, c_dtype): + exec_time = run_scaled_mm_with_emulated_dtype( + gemm, + mnkl, + a_dtype, + b_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + tolerance, + warmup_iterations, + iterations, + skip_ref_check, + use_cold_l2, + init_normal, + normal_mean, + normal_std, + ) + else: + exec_time = run_scaled_mm( + gemm, + mnkl, + a_dtype, + b_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + tolerance, + warmup_iterations, + iterations, + skip_ref_check, + use_cold_l2, + init_normal, + normal_mean, + normal_std, + ) + return exec_time + + +def prepare_parser(): + def parse_comma_separated_ints(s: str) -> tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError("Invalid format. Expected comma-separated integers.") + + parser = argparse.ArgumentParser(description="Example of Sm100 Dense Persistent BlockScaled GEMM.") + + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=(512, 256, 256, 1), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--mma_tiler_mn", + type=parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + parser.add_argument( + "--cluster_shape_mn", + type=parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--a_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--b_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E4M3FN) + parser.add_argument("--sf_vec_size", type=int, default=16) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float16) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument("--tolerance", type=float, default=1e-01, help="Tolerance for validation") + parser.add_argument("--warmup_iterations", type=int, default=0, help="Warmup iterations") + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument("--skip_ref_check", action="store_true", help="Skip reference checking") + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + testing.add_tensor_init_args(parser, supports_int_dtypes=False) + return parser + + +if __name__ == "__main__": + parser = prepare_parser() + args = parser.parse_args() + + testing.validate_tensor_init_args(args, parser) + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + # Execute GEMM with appropriate function based on dtype + run( + args.mnkl, + args.a_dtype, + args.b_dtype, + args.sf_dtype, + args.sf_vec_size, + args.c_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + args.init_normal, + args.normal_mean, + args.normal_std, + ) + print("PASS") diff --git a/aphrodite/model_executor/layers/fused_moe/experts/cutedsl_mxfp6_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/cutedsl_mxfp6_moe.py new file mode 100644 index 0000000000..c29e0e6fcd --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/experts/cutedsl_mxfp6_moe.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Native SM110 grouped MXFP6 experts implemented with CuTe DSL.""" + +import torch + +import aphrodite.model_executor.layers.fused_moe.modular_kernel as mk +from aphrodite.model_executor.kernels.linear.mxfp6.cutedsl_grouped import ( + cutedsl_grouped_mxfp6_gemm, +) +from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation +from aphrodite.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, +) +from aphrodite.model_executor.layers.fused_moe.moe_permute_unpermute import ( + MoEPermuteScratch, + moe_permute, + moe_permute_unpermute_supported, + moe_unpermute, +) +from aphrodite.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from aphrodite.model_executor.layers.fused_moe.utils import _resize_cache +from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp6E2m3Dynamic, + kMxfp6E2m3Static, + kMxfp6E3m2Dynamic, + kMxfp6E3m2Static, + kMxfp8Dynamic, +) +from aphrodite.platforms import current_platform + + +class CutedslMxfp6Experts(mk.FusedMoEExpertsModular): + """Standard-layout MXFP6 MoE experts for Blackwell SM110.""" + + def __init__(self, *args, weight_format: str, activation_format: str, **kwargs): + super().__init__(*args, **kwargs) + self.weight_format = weight_format + self.activation_quant_format = activation_format + self._permute_scratch: MoEPermuteScratch | None = None + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + return current_platform.is_cuda() and current_platform.is_device_capability_family(110) + + @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 { + kMxfp6E2m3Static, + kMxfp6E3m2Static, + } and activation_key in { + kMxfp8Dynamic, + kMxfp6E2m3Dynamic, + kMxfp6E3m2Dynamic, + } + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SILU + + @staticmethod + def _supports_parallel_config(config: FusedMoEParallelConfig) -> bool: + # The standard path handles TP. EP requires an expert-map-aware grouped + # metadata builder and is deliberately rejected for the first backend. + return config.ep_size == 1 + + @staticmethod + def _supports_shape(hidden_dim: int) -> bool: + return hidden_dim >= 128 and hidden_dim % 128 == 0 + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + 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, ...]]: + del global_num_experts, local_num_experts, expert_tokens_meta + activation_out_dim = self.adjust_N_for_activation(N, activation) + return ( + (M * topk, max(N, K)), + (M * topk, max(activation_out_dim, K)), + (M, K), + ) + + def _get_permute_scratch(self) -> MoEPermuteScratch | None: + if self._permute_scratch is None and moe_permute_unpermute_supported(): + self._permute_scratch = MoEPermuteScratch( + max_num_tokens=self.moe_config.max_num_tokens, + topk=self.moe_config.experts_per_token, + num_experts=self.moe_config.num_experts, + num_local_experts=self.moe_config.num_local_experts, + device=torch.device(self.moe_config.device), + ) + return self._permute_scratch + + 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, + ) -> None: + del a1q_scale, a2_scale, expert_tokens_meta + if apply_router_weight_on_input: + if topk_ids.shape[1] != 1: + raise ValueError("router-weight-on-input requires top-k=1") + hidden_states = hidden_states * topk_weights.to(hidden_states.dtype) + + m, k = hidden_states.shape + topk = topk_ids.shape[1] + n = self.moe_config.intermediate_size_per_partition + local_experts = w1.shape[0] + num_experts = global_num_experts if expert_map is None else expert_map.numel() + permuted = _resize_cache(workspace2, (m * topk, k)) + permuted, _, offsets, inv_perm, _ = moe_permute( + hidden_states, + None, + topk_ids, + num_experts, + local_experts, + expert_map, + permuted_hidden_states=permuted, + scratch=self._get_permute_scratch(), + ) + + gate_up = _resize_cache(workspace13, (m * topk, 2 * n)) + cutedsl_grouped_mxfp6_gemm( + permuted, + w1, + self.w1_scale, + offsets, + gate_up, + 2 * n, + k, + self.activation_quant_format, + self.weight_format, + ) + activated = _resize_cache(workspace2, (m * topk, n)) + self.activation(activation, activated, gate_up) + + down = _resize_cache(workspace13, (m * topk, k)) + cutedsl_grouped_mxfp6_gemm( + activated, + w2, + self.w2_scale, + offsets, + down, + k, + n, + self.activation_quant_format, + self.weight_format, + ) + moe_unpermute( + out=output, + permuted_hidden_states=down, + topk_weights=None if apply_router_weight_on_input else topk_weights, + inv_permuted_idx=inv_perm, + expert_first_token_offset=offsets, + ) diff --git a/aphrodite/model_executor/layers/quantization/__init__.py b/aphrodite/model_executor/layers/quantization/__init__.py index 70644f6695..648647f23a 100644 --- a/aphrodite/model_executor/layers/quantization/__init__.py +++ b/aphrodite/model_executor/layers/quantization/__init__.py @@ -45,6 +45,7 @@ "int8_per_channel_weight_only", "nvfp4_per_token", "mxfp8", + "mxfp6", ] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) diff --git a/aphrodite/model_executor/layers/quantization/online/base.py b/aphrodite/model_executor/layers/quantization/online/base.py index 7af9b350c7..2def40d7ae 100644 --- a/aphrodite/model_executor/layers/quantization/online/base.py +++ b/aphrodite/model_executor/layers/quantization/online/base.py @@ -36,6 +36,10 @@ from aphrodite.model_executor.layers.quantization.online.int8 import ( Int8OnlineMoEMethod, ) +from aphrodite.model_executor.layers.quantization.online.mxfp6 import ( + Mxfp6OnlineLinearMethod, + Mxfp6OnlineMoEMethod, +) from aphrodite.model_executor.layers.quantization.online.mxfp8 import ( Mxfp8OnlineLinearMethod, Mxfp8OnlineMoEMethod, @@ -49,6 +53,8 @@ kFp8StaticChannelSym, kFp8StaticTensorSym, kInt8StaticChannelSym, + kMxfp6E2m3Static, + kMxfp6E3m2Static, kMxfp8Dynamic, kNvfp4Static, ) @@ -64,6 +70,8 @@ kFp8Static128BlockSym: Fp8PerBlockOnlineLinearMethod, kFp8StaticChannelSym: Fp8PtpcOnlineLinearMethod, kMxfp8Dynamic: Mxfp8OnlineLinearMethod, + kMxfp6E2m3Static: Mxfp6OnlineLinearMethod, + kMxfp6E3m2Static: Mxfp6OnlineLinearMethod, } _ONLINE_MOE_METHODS: dict[QuantKey, type] = { @@ -73,6 +81,8 @@ kMxfp8Dynamic: Mxfp8OnlineMoEMethod, kInt8StaticChannelSym: Int8OnlineMoEMethod, kNvfp4Static: Nvfp4OnlineMoEMethod, + kMxfp6E2m3Static: Mxfp6OnlineMoEMethod, + kMxfp6E3m2Static: Mxfp6OnlineMoEMethod, } @@ -94,6 +104,40 @@ def __init__( self.args = args self.ignored_layers: list[str] = args.ignore + def _resolve_spec(self, prefix: str, base: QuantSpec | None) -> QuantSpec | None: + """Resolve ordered precision rules for a module. + + ``ignore`` remains an unconditional compatibility escape hatch. The + newer overrides are evaluated in declaration order, so a later rule + can refine (or undo) an earlier broad rule. + """ + if should_ignore_layer( + prefix, + ignore=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return None + + spec = base + for override in self.args.overrides: + if not should_ignore_layer( + prefix, + ignore=[override.pattern], + fused_mapping=self.packed_modules_mapping, + ): + continue + if override.weight == "bf16": + spec = None + continue + + inherited_weight = spec.weight if spec is not None else None + inherited_activation = spec.activation if spec is not None else None + spec = QuantSpec( + weight=override.weight or inherited_weight, + activation=override.activation or inherited_activation, + ) # type: ignore[call-arg] + return spec + @classmethod def get_name(cls) -> QuantizationMethods: return "online" @@ -135,6 +179,38 @@ def _dispatch( f"weight={spec.weight} is not supported; supported weight " f"keys: {sorted(str(k) for k in table)}" ) + if cls is Mxfp6OnlineLinearMethod: + return cls(spec) + if cls is Mxfp6OnlineMoEMethod: + from aphrodite.model_executor.kernels.linear.mxfp6 import CutedslMxfp6LinearKernel + from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation + + supported, reason = CutedslMxfp6LinearKernel.is_supported() + moe = layer.moe_config + if supported and moe.activation != MoEActivation.SILU: + supported = False + reason = f"activation {moe.activation} is not supported" + elif supported and moe.moe_parallel_config.ep_size != 1: + supported = False + reason = "expert parallel execution is not supported yet" + elif supported and moe.has_bias: + supported = False + reason = "biased experts are not supported yet" + elif supported and moe.swiglu_limit is not None: + supported = False + reason = "SwiGLU clamping is not supported yet" + elif supported and ( + moe.hidden_dim < 128 + or moe.hidden_dim % 128 + or moe.intermediate_size_per_partition < 128 + or moe.intermediate_size_per_partition % 128 + ): + supported = False + reason = "hidden and intermediate dimensions must be multiples of 128" + if not supported: + logger.warning_once("Keeping MXFP6-targeted MoE layers in BF16 because %s", reason) + return None + return cls(layer=layer, spec=spec) # Online method classes pick their own activation format internally. # Per-class activation overrides are not yet wired through; reject # explicit overrides until the relevant method class opts in. @@ -148,21 +224,17 @@ def _dispatch( def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> "QuantizeMethodBase | None": if isinstance(layer, LinearBase): - if should_ignore_layer( - prefix, - ignore=self.ignored_layers, - fused_mapping=self.packed_modules_mapping, - ): - return UnquantizedLinearMethod() - method = self._dispatch(self.args.linear, _ONLINE_LINEAR_METHODS, layer) + method = self._dispatch( + self._resolve_spec(prefix, self.args.linear), + _ONLINE_LINEAR_METHODS, + layer, + ) return method if method is not None else UnquantizedLinearMethod() elif isinstance(layer, RoutedExperts): - if should_ignore_layer( - prefix, - ignore=self.ignored_layers, - fused_mapping=self.packed_modules_mapping, - ): - return UnquantizedFusedMoEMethod(layer.moe_config) - method = self._dispatch(self.args.moe, _ONLINE_MOE_METHODS, layer) + method = self._dispatch( + self._resolve_spec(prefix, self.args.moe), + _ONLINE_MOE_METHODS, + layer, + ) return method if method is not None else UnquantizedFusedMoEMethod(layer.moe_config) return None diff --git a/aphrodite/model_executor/layers/quantization/online/mxfp6.py b/aphrodite/model_executor/layers/quantization/online/mxfp6.py new file mode 100644 index 0000000000..f7c5a2cdcb --- /dev/null +++ b/aphrodite/model_executor/layers/quantization/online/mxfp6.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Online OCP MXFP6 quantization for Blackwell SM110.""" + +from typing import Literal + +import torch +from torch.nn import Module + +from aphrodite.config.quantization import QuantSpec +from aphrodite.logger import init_logger +from aphrodite.model_executor.kernels.linear.mxfp6 import ( + CutedslMxfp6LinearKernel, + Mxfp6LinearLayerConfig, +) +from aphrodite.model_executor.layers.linear import UnquantizedLinearMethod +from aphrodite.model_executor.layers.quantization.online.fp8 import ( + _Fp8OnlineLinearBase, +) +from aphrodite.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, +) +from aphrodite.model_executor.layers.quantization.utils.mxfp6_online_utils import ( + quantize_mxfp6_cuda, +) +from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp6E2m3Dynamic, + kMxfp6E2m3Static, + kMxfp6E3m2Dynamic, + kMxfp6E3m2Static, + kMxfp8Dynamic, +) +from aphrodite.model_executor.utils import replace_parameter + +logger = init_logger(__name__) + + +def _swizzle_mxfp6_expert_scales(scales: torch.Tensor, m: int, k: int) -> torch.Tensor: + """Vectorized F8_128x4 scale swizzle for ``[expert, M, K/32]``.""" + experts = scales.shape[0] + num_m_tiles = (m + 127) // 128 + num_k_tiles = (k + 127) // 128 + padded = torch.zeros( + (experts, num_m_tiles * 128, num_k_tiles * 4), + dtype=scales.dtype, + device=scales.device, + ) + padded[:, :m, : k // 32] = scales + return padded.view(experts, num_m_tiles, 4, 32, num_k_tiles, 4).transpose(2, 4).contiguous().view(experts, -1) + + +class Mxfp6OnlineLinearMethod(_Fp8OnlineLinearBase): + def __init__(self, spec: QuantSpec): + super().__init__() + if spec.weight in (kMxfp6E2m3Static, kMxfp6E2m3Dynamic): + weight_format: Literal["e2m3", "e3m2"] = "e2m3" + elif spec.weight in (kMxfp6E3m2Static, kMxfp6E3m2Dynamic): + weight_format = "e3m2" + else: + raise ValueError(f"unsupported MXFP6 weight format: {spec.weight}") + + if spec.activation == kMxfp8Dynamic: + activation_format: Literal["mxfp8", "mxfp6_e2m3", "mxfp6_e3m2"] = "mxfp8" + elif spec.activation == kMxfp6E2m3Dynamic: + activation_format = "mxfp6_e2m3" + elif spec.activation == kMxfp6E3m2Dynamic: + activation_format = "mxfp6_e3m2" + else: + raise ValueError("MXFP6 requires dynamic MXFP8 or MXFP6 activations") + self.kernel = CutedslMxfp6LinearKernel(Mxfp6LinearLayerConfig(weight_format, activation_format)) + self.unquantized = UnquantizedLinearMethod() + self.fallback_reason: str | None = None + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + super().create_weights( + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ) + n = sum(output_partition_sizes) + k = input_size_per_partition + supported, reason = self.kernel.is_supported() + if supported: + supported, reason = self.kernel.can_implement_shape(n, k) + if not supported: + self.fallback_reason = reason + logger.warning_once( + "Keeping an MXFP6-targeted linear layer in %s because %s", + params_dtype, + reason, + ) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + if self.fallback_reason is not None: + layer._already_called_process_weights_after_loading = True + return + + logger.info_once("Converting full-precision weights to MXFP6 one layer at a time") + n, k = layer.weight.shape + packed, scales = quantize_mxfp6_cuda( + layer.weight.contiguous(), + self.kernel.config.weight_format, + ) + layer.mxfp6_logical_n = n + layer.mxfp6_logical_k = k + replace_parameter(layer, "weight", packed.data) + replace_parameter(layer, "weight_scale", scales.data) + self.kernel.process_weights_after_loading(layer) + layer._already_called_process_weights_after_loading = True + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if self.fallback_reason is not None: + return self.unquantized.apply(layer, x, bias) + return self.kernel.apply_weights(layer, x, bias) + + +class Mxfp6OnlineMoEMethod(OnlineMoEMethodBase): + """Quantize BF16 MoE weights to packed MXFP6 while each layer loads.""" + + def __init__(self, *, layer: torch.nn.Module, spec: QuantSpec): + super().__init__(layer.moe_config) + if spec.weight == kMxfp6E2m3Static: + weight_format: Literal["e2m3", "e3m2"] = "e2m3" + elif spec.weight == kMxfp6E3m2Static: + weight_format = "e3m2" + else: + raise ValueError(f"unsupported MXFP6 MoE weight format: {spec.weight}") + if spec.activation == kMxfp8Dynamic: + activation_format: Literal["mxfp8", "mxfp6_e2m3", "mxfp6_e3m2"] = "mxfp8" + elif spec.activation == kMxfp6E2m3Dynamic: + activation_format = "mxfp6_e2m3" + elif spec.activation == kMxfp6E3m2Dynamic: + activation_format = "mxfp6_e3m2" + else: + raise ValueError("MXFP6 MoE requires dynamic MXFP8 or MXFP6 activations") + self.weight_format = weight_format + self.activation_format = activation_format + + def create_weights( + self, + layer: Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + if hidden_size % 128 or intermediate_size_per_partition % 128: + raise ValueError("native MXFP6 MoE dimensions must be divisible by 128") + super().create_weights( + layer, + num_experts, + hidden_size, + intermediate_size_per_partition, + params_dtype, + **extra_weight_attrs, + ) + + def get_fused_moe_quant_config(self, layer: torch.nn.Module): + del layer + return getattr(self, "moe_quant_config", None) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + from aphrodite.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, + ) + from aphrodite.model_executor.layers.fused_moe.config import FusedMoEQuantConfig + from aphrodite.model_executor.layers.fused_moe.experts.cutedsl_mxfp6_moe import ( + CutedslMxfp6Experts, + ) + + logger.info_once("Converting full-precision weights to MXFP6 one layer at a time") + w13_k = layer.w13_weight.shape[2] + w13, s13 = quantize_mxfp6_cuda(layer.w13_weight.contiguous(), self.weight_format) + s13 = _swizzle_mxfp6_expert_scales(s13, w13.shape[1], w13_k) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w13_weight_scale", s13) + + # Release the full-precision gate/up tensor before converting the down + # projection. This bounds conversion memory to one expert tensor. + w2_k = layer.w2_weight.shape[2] + w2, s2 = quantize_mxfp6_cuda(layer.w2_weight.contiguous(), self.weight_format) + s2 = _swizzle_mxfp6_expert_scales(s2, w2.shape[1], w2_k) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w2_weight_scale", s2) + + act_dtype = self.activation_format + weight_dtype = f"mxfp6_{self.weight_format}" + quant_config = FusedMoEQuantConfig.make( + act_dtype, + weight_dtype=weight_dtype, + block_shape=[1, 32], + w1_scale=s13, + w2_scale=s2, + is_scale_swizzled=True, + ) + prepare_finalize = maybe_make_prepare_finalize( + moe=self.moe, + quant_config=quant_config, + routing_tables=layer._expert_routing_tables(), + allow_new_interface=True, + ) + assert prepare_finalize is not None + experts = CutedslMxfp6Experts( + moe_config=self.moe, + quant_config=quant_config, + weight_format=self.weight_format, + activation_format=self.activation_format, + ) + import aphrodite.model_executor.layers.fused_moe.modular_kernel as mk + + self.moe_quant_config = quant_config + self.moe_kernel = mk.FusedMoEKernel(prepare_finalize, experts) + layer._already_called_process_weights_after_loading = True diff --git a/aphrodite/model_executor/layers/quantization/utils/mxfp6_online_utils.py b/aphrodite/model_executor/layers/quantization/utils/mxfp6_online_utils.py new file mode 100644 index 0000000000..ede4dda4ed --- /dev/null +++ b/aphrodite/model_executor/layers/quantization/utils/mxfp6_online_utils.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reference utilities for online OCP MXFP6 quantization. + +The production SM110 converter is implemented in CUDA. These routines define +the packed-weight ABI and provide a device-independent correctness oracle. +""" + +from functools import lru_cache +from typing import Literal + +import torch + +from aphrodite.model_executor.layers.quantization.utils.mxfp8_utils import ( + swizzle_mxfp8_scale, +) + +MXFP6_BLOCK_SIZE = 32 +MXFP6_PACK_INPUT = 4 +MXFP6_PACK_BYTES = 3 +Mxfp6Format = Literal["e2m3", "e3m2"] + + +@lru_cache +def _mxfp6_quantize_triton_kernel(): + from aphrodite.triton_utils import tl, triton + + @triton.jit + def kernel( + x_ptr, + packed_ptr, + scale_ptr, + k: tl.constexpr, + mantissa_bits: tl.constexpr, + exponent_bias: tl.constexpr, + max_value: tl.constexpr, + min_normal: tl.constexpr, + subnormal_step: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + row = tl.program_id(0) + tile_k = tl.program_id(1) + offsets = tile_k * BLOCK_K + tl.arange(0, BLOCK_K) + mask = offsets < k + values = tl.load(x_ptr + row * k + offsets, mask=mask, other=0.0).to(tl.float32) + magnitude = tl.reshape(tl.abs(values), (BLOCK_K // 32, 32)) + amax = tl.max(magnitude, axis=1) + scale_exp = tl.ceil(tl.log2(amax / max_value)) + scale_exp = tl.maximum(-127.0, tl.minimum(127.0, scale_exp)) + scale_exp = tl.where(amax == 0.0, -127.0, scale_exp) + scaled = magnitude * tl.reshape(tl.exp2(-scale_exp), (BLOCK_K // 32, 1)) + scaled = tl.reshape(scaled, (BLOCK_K,)) + values = tl.reshape(values, (BLOCK_K,)) + + mantissa_scale: tl.constexpr = 1 << mantissa_bits + normal_exp = tl.floor(tl.log2(tl.maximum(scaled, min_normal))) + exponent = normal_exp.to(tl.int32) + exponent_bias + normal_mantissa = tl.extra.cuda.libdevice.rint((scaled * tl.exp2(-normal_exp) - 1.0) * mantissa_scale).to( + tl.int32 + ) + carry = normal_mantissa == (1 << mantissa_bits) + exponent += carry.to(tl.int32) + normal_mantissa = tl.where(carry, 0, normal_mantissa) + normal_code = (exponent << mantissa_bits) | normal_mantissa + subnormal_code = tl.extra.cuda.libdevice.rint(scaled / subnormal_step).to(tl.int32) + code = tl.where(scaled < min_normal, subnormal_code, normal_code) + code = tl.maximum(0, tl.minimum(31, code)) + code |= tl.where(values < 0.0, 32, 0) + + grouped = tl.reshape(code, (BLOCK_K // 4, 4)) + shifts = tl.reshape(tl.arange(0, 4) * 6, (1, 4)) + words = tl.sum( + grouped << shifts, + axis=1, + ) + byte_offsets = row * (k * 3 // 4) + tile_k * (BLOCK_K * 3 // 4) + tl.arange(0, BLOCK_K // 4) * 3 + byte_mask = byte_offsets < (row + 1) * (k * 3 // 4) + tl.store(packed_ptr + byte_offsets, words & 0xFF, mask=byte_mask) + tl.store(packed_ptr + byte_offsets + 1, (words >> 8) & 0xFF, mask=byte_mask) + tl.store(packed_ptr + byte_offsets + 2, (words >> 16) & 0xFF, mask=byte_mask) + scale_offsets = row * (k // 32) + tile_k * (BLOCK_K // 32) + tl.arange(0, BLOCK_K // 32) + scale_mask = scale_offsets < (row + 1) * (k // 32) + tl.store(scale_ptr + scale_offsets, scale_exp + 127.0, mask=scale_mask) + + return kernel + + +def quantize_mxfp6_cuda( + x: torch.Tensor, + fmt: Mxfp6Format = "e2m3", +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused CUDA conversion to packed OCP MXFP6 and row-major E8M0 scales.""" + if not x.is_cuda or x.ndim not in (2, 3): + raise ValueError("MXFP6 CUDA quantization expects a 2D/3D CUDA tensor") + if x.shape[-1] % MXFP6_BLOCK_SIZE: + raise ValueError("MXFP6 requires K to be divisible by 32") + from aphrodite.triton_utils import triton + + k = x.shape[-1] + rows = x.numel() // k + packed = torch.empty((*x.shape[:-1], k * 3 // 4), dtype=torch.uint8, device=x.device) + scales = torch.empty((*x.shape[:-1], k // 32), dtype=torch.uint8, device=x.device) + if fmt == "e2m3": + mantissa_bits, exponent_bias, max_value = 3, 1, 7.5 + min_normal, subnormal_step = 1.0, 0.125 + elif fmt == "e3m2": + mantissa_bits, exponent_bias, max_value = 2, 3, 28.0 + min_normal, subnormal_step = 0.25, 0.0625 + else: + raise ValueError(f"unsupported MXFP6 format: {fmt}") + kernel = _mxfp6_quantize_triton_kernel() + block_k = 256 + kernel[(rows, triton.cdiv(k, block_k))]( + x, + packed, + scales, + k=k, + mantissa_bits=mantissa_bits, + exponent_bias=exponent_bias, + max_value=max_value, + min_normal=min_normal, + subnormal_step=subnormal_step, + BLOCK_K=block_k, + ) + return packed, scales + + +def _positive_codebook(fmt: Mxfp6Format, device: torch.device) -> torch.Tensor: + if fmt == "e2m3": + exponent_bits, mantissa_bits, bias = 2, 3, 1 + elif fmt == "e3m2": + exponent_bits, mantissa_bits, bias = 3, 2, 3 + else: + raise ValueError(f"unsupported MXFP6 format: {fmt}") + + codes = torch.arange(1 << (exponent_bits + mantissa_bits), device=device) + exponent = codes >> mantissa_bits + mantissa = codes & ((1 << mantissa_bits) - 1) + normal = (1.0 + mantissa.float() / (1 << mantissa_bits)) * torch.exp2(exponent.float() - bias) + subnormal = mantissa.float() * 2.0 ** (1 - bias - mantissa_bits) + return torch.where(exponent == 0, subnormal, normal) + + +def _quantize_codes(x: torch.Tensor, fmt: Mxfp6Format) -> torch.Tensor: + """Quantize values already divided by their block scale to raw FP6 codes.""" + values = _positive_codebook(fmt, x.device) + magnitude = x.abs().float().clamp(max=values[-1]) + upper = torch.searchsorted(values, magnitude).clamp(max=values.numel() - 1) + lower = (upper - 1).clamp(min=0) + lower_dist = magnitude - values[lower] + upper_dist = values[upper] - magnitude + choose_upper = upper_dist < lower_dist + tied = upper_dist == lower_dist + # IEEE round-to-nearest-even selects the candidate with an even LSB. + choose_upper |= tied & ((upper & 1) == 0) + positive = torch.where(choose_upper, upper, lower).to(torch.uint8) + sign = torch.signbit(x).to(torch.uint8) << 5 + # OCP finite FP6 preserves signed zero. + return positive | sign + + +def pack_mxfp6_codes(codes: torch.Tensor) -> torch.Tensor: + """Pack four consecutive six-bit codes into three little-endian bytes.""" + if codes.shape[-1] % MXFP6_PACK_INPUT != 0: + raise ValueError("MXFP6 packing requires the last dimension to be divisible by 4") + grouped = codes.to(torch.int32).view(*codes.shape[:-1], -1, 4) + words = grouped[..., 0] | grouped[..., 1] << 6 | grouped[..., 2] << 12 | grouped[..., 3] << 18 + return ( + torch.stack( + ( + words & 0xFF, + words >> 8 & 0xFF, + words >> 16 & 0xFF, + ), + dim=-1, + ) + .flatten(-2) + .to(torch.uint8) + ) + + +def unpack_mxfp6_codes(packed: torch.Tensor) -> torch.Tensor: + """Unpack the canonical three-byte/four-value MXFP6 representation.""" + if packed.shape[-1] % MXFP6_PACK_BYTES != 0: + raise ValueError("packed MXFP6 data must contain complete three-byte groups") + grouped = packed.to(torch.int32).view(*packed.shape[:-1], -1, 3) + words = grouped[..., 0] | grouped[..., 1] << 8 | grouped[..., 2] << 16 + return ( + torch.stack( + tuple(words >> shift & 0x3F for shift in (0, 6, 12, 18)), + dim=-1, + ) + .flatten(-2) + .to(torch.uint8) + ) + + +def dequantize_mxfp6_reference( + packed: torch.Tensor, + scales: torch.Tensor, + fmt: Mxfp6Format, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + codes = unpack_mxfp6_codes(packed) + values = _positive_codebook(fmt, packed.device)[(codes & 0x1F).long()] + values = torch.where((codes & 0x20) != 0, -values, values) + blocked = values.view(*values.shape[:-1], -1, MXFP6_BLOCK_SIZE) + descale = torch.exp2(scales.float() - 127.0).unsqueeze(-1) + return (blocked * descale).flatten(-2).to(dtype) + + +def quantize_mxfp6_reference( + x: torch.Tensor, + fmt: Mxfp6Format = "e2m3", + *, + swizzle_scales: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2D/3D tensor using OCP block-32 MXFP6 semantics.""" + if x.ndim not in (2, 3): + raise ValueError(f"MXFP6 reference quantization expects a 2D or 3D tensor, got {x.ndim}D") + if x.shape[-1] % MXFP6_BLOCK_SIZE != 0: + raise ValueError("MXFP6 requires K to be divisible by 32") + + codebook = _positive_codebook(fmt, x.device) + blocked = x.float().view(*x.shape[:-1], -1, MXFP6_BLOCK_SIZE) + amax = blocked.abs().amax(dim=-1) + safe_amax = amax.clamp_min(torch.finfo(torch.float32).tiny) + scale_exp = torch.ceil(torch.log2(safe_amax / codebook[-1])).clamp(-127, 127) + scale_exp = torch.where(amax == 0, torch.full_like(scale_exp, -127), scale_exp) + scales = (scale_exp + 127).to(torch.uint8) + scaled = blocked / torch.exp2(scale_exp).unsqueeze(-1) + packed = pack_mxfp6_codes(_quantize_codes(scaled.flatten(-2), fmt)) + + if swizzle_scales: + if x.ndim == 2: + scales = swizzle_mxfp8_scale(scales, x.shape[0], x.shape[1]) + else: + scales = torch.cat([swizzle_mxfp8_scale(s, x.shape[1], x.shape[2]) for s in scales]) + return packed, scales diff --git a/aphrodite/model_executor/layers/quantization/utils/quant_utils.py b/aphrodite/model_executor/layers/quantization/utils/quant_utils.py index 2b710f9330..76c129a2bd 100644 --- a/aphrodite/model_executor/layers/quantization/utils/quant_utils.py +++ b/aphrodite/model_executor/layers/quantization/utils/quant_utils.py @@ -165,6 +165,29 @@ def __str__(self): kMxfp4StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) kMxfp4Static = QuantKey(FP4_DTYPE, scale=kMxfp4StaticGroupScale, symmetric=True) +kMxfp6StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) +kMxfp6DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) +kMxfp6E2m3Static = QuantKey( + scalar_types.float6_e2m3f, + scale=kMxfp6StaticGroupScale, + symmetric=True, +) +kMxfp6E2m3Dynamic = QuantKey( + scalar_types.float6_e2m3f, + scale=kMxfp6DynamicGroupScale, + symmetric=True, +) +kMxfp6E3m2Static = QuantKey( + scalar_types.float6_e3m2f, + scale=kMxfp6StaticGroupScale, + symmetric=True, +) +kMxfp6E3m2Dynamic = QuantKey( + scalar_types.float6_e3m2f, + scale=kMxfp6DynamicGroupScale, + symmetric=True, +) + # TODO: convert this to use SCALAR_TYPE. This is not right. kInt4StaticGroupScale = ScaleDesc(torch.float16, True, GroupShape(1, -1)) kInt4Static = QuantKey(INT4_DTYPE, scale=kInt4StaticGroupScale, symmetric=True) diff --git a/docs/astro.config.mjs b/docs/astro.config.mjs index a2289f1bec..6fc3bed162 100644 --- a/docs/astro.config.mjs +++ b/docs/astro.config.mjs @@ -71,6 +71,7 @@ export default defineConfig({ { label: 'Reasoning and tool parsers', slug: 'features/reasoning-and-tools' }, { label: 'Multimodal inputs', slug: 'features/multimodal' }, { label: 'FP8 vision attention', slug: 'features/fp8-vit-attention' }, + { label: 'Online MXFP6', slug: 'features/mxfp6' }, { label: 'Observability', slug: 'features/observability' }, ], }, diff --git a/docs/src/content/docs/features/mxfp6.md b/docs/src/content/docs/features/mxfp6.md new file mode 100644 index 0000000000..0887cf15a7 --- /dev/null +++ b/docs/src/content/docs/features/mxfp6.md @@ -0,0 +1,76 @@ +--- +title: Online MXFP6 +description: Convert BF16 checkpoints to native MXFP6 while Sonar loads them on NVIDIA Thor. +--- + +Sonar can convert BF16 or FP16 weights to OCP MXFP6 while it loads a model. The initial native backend targets NVIDIA Thor GPUs with compute capability 11.0. It does not write a converted checkpoint to disk. + +## Start a model + +Use `--quantization mxfp6` to select E2M3 weights and dynamic MXFP8 activations: + +```sh +aphrodite run Qwen/Qwen3.5-4B --quantization mxfp6 +``` + +The loader converts one layer at a time. It replaces each full-precision tensor with packed six-bit data and an E8M0 scale for each block of 32 values. This limits temporary memory use. The model does not have to fit in host memory twice. + +Sonar keeps router gates and shared-expert gates in BF16 by default. These small tensors are sensitive to quantization and do not account for much model memory. + +## Select the FP6 encoding + +The shorthand uses E2M3 weights. Use `--quantization online` with a JSON configuration to select E3M2: + +```sh +aphrodite run Qwen/Qwen3.5-4B \ + --quantization online \ + --quantization-config '{ + "linear":{"weight":"mxfp6_e3m2","activation":"mxfp8"}, + "moe":{"weight":"mxfp6_e3m2","activation":"mxfp8"} + }' +``` + +E2M3 has more mantissa precision. E3M2 has a larger finite range. Start with E2M3 unless evaluation shows that the model needs the E3M2 range. + +## Use W6A6 + +The native backend also accepts MXFP6 activations. This mode reduces activation precision and adds conversion work: + +```sh +aphrodite run Qwen/Qwen3.5-4B \ + --quantization online \ + --quantization-config '{ + "linear":{"weight":"mxfp6_e2m3","activation":"mxfp6_e2m3_dynamic"}, + "moe":{"weight":"mxfp6_e2m3","activation":"mxfp6_e2m3_dynamic"} + }' +``` + +Benchmark W6A6 against the default W6A8 mode on the target workload. W6A8 is faster on the current Thor kernel. + +## Keep selected modules in BF16 + +Use ordered overrides to control precision. Sonar applies the last matching rule: + +```sh +aphrodite run Qwen/Qwen3.5-4B \ + --quantization mxfp6 \ + --quantization-config '{ + "overrides":[ + {"pattern":"re:.*\\.lm_head$","weight":"bf16"}, + {"pattern":"model.layers.0.mlp.down_proj","weight":"bf16"} + ] + }' +``` + +An exact module prefix matches that module. A pattern that starts with `re:` is a regular expression. Later rules can refine an earlier broad rule. + +## Current limits + +- Native execution requires CUDA, compute capability 11.0, and `nvidia-cutlass-dsl`. +- Dense input and output dimensions must be multiples of 128. +- MoE hidden and intermediate dimensions must be multiples of 128. +- The native MoE path supports gated SiLU without expert bias or SwiGLU clamping. +- Expert parallel execution stays in BF16. Tensor parallel execution is supported. +- Unsupported layers stay in BF16 and produce a startup warning. + +Evaluate model quality before production use. Online conversion has no calibration dataset and can affect different models in different ways. diff --git a/docs/src/content/docs/reference/models.md b/docs/src/content/docs/reference/models.md index bd4d521bc2..8df36ba3b9 100644 --- a/docs/src/content/docs/reference/models.md +++ b/docs/src/content/docs/reference/models.md @@ -7,7 +7,7 @@ description: Model architectures that the current Sonar source tree registers. Sonar reads the `architectures` field from a model's `config.json` file. The value must match an architecture in these tables. A compatible Transformers implementation can also use the Transformers backend. -The current registry contains **371 architecture entries** in the user-facing categories below. +The current registry contains **374 architecture entries** in the user-facing categories below. ## Text generation @@ -80,7 +80,7 @@ The current registry contains **371 architecture entries** in the user-facing ca | `InternLM3ForCausalLM` | `llama:LlamaForCausalLM` | | `Jais2ForCausalLM` | `jais2:Jais2ForCausalLM` | | `JambaForCausalLM` | `jamba:JambaForCausalLM` | -| `KimiLinearForCausalLM` | `kimi_linear:KimiLinearForCausalLM` | +| `KimiLinearForCausalLM` | `aphrodite.models.kimi_k3:KimiLinearForCausalLM` | | `LLaMAForCausalLM` | `llama:LlamaForCausalLM` | | `LagunaForCausalLM` | `laguna:LagunaForCausalLM` | | `Lfm2ForCausalLM` | `lfm2:Lfm2ForCausalLM` | @@ -293,6 +293,7 @@ The current registry contains **371 architecture entries** in the user-facing ca | `KeyeForConditionalGeneration` | `keye:KeyeForConditionalGeneration` | | `KeyeVL1_5ForConditionalGeneration` | `keye_vl1_5:KeyeVL1_5ForConditionalGeneration` | | `KimiK25ForConditionalGeneration` | `kimi_k25:KimiK25ForConditionalGeneration` | +| `KimiK3ForConditionalGeneration` | `aphrodite.models.kimi_k3:KimiK3ForConditionalGeneration` | | `KimiVLForConditionalGeneration` | `kimi_vl:KimiVLForConditionalGeneration` | | `Lfm2VlForConditionalGeneration` | `lfm2_vl:Lfm2VLForConditionalGeneration` | | `LightOnOCRForConditionalGeneration` | `lightonocr:LightOnOCRForConditionalGeneration` | @@ -394,6 +395,8 @@ The current registry contains **371 architecture entries** in the user-facing ca | `GlmOcrMTPModel` | `glm_ocr_mtp:GlmOcrMTP` | | `HYV3MTPModel` | `hy_v3_mtp:HYV3MTP` | | `InklingMTPModel` | `aphrodite.models.inkling:InklingMTP` | +| `K3DSparkModel` | `aphrodite.models.kimi_k3.nvidia.dspark_mla:K3DSparkForCausalLM` | +| `KimiK3MTPModel` | `aphrodite.models.kimi_k3:KimiK3MTP` | | `LlamaForCausalLMEagle3` | `llama_eagle3:Eagle3LlamaForCausalLM` | | `LongCatFlashMTPModel` | `longcat_flash_mtp:LongCatFlashMTP` | | `MedusaModel` | `medusa:Medusa` | diff --git a/docs/src/content/docs/reference/quantization.md b/docs/src/content/docs/reference/quantization.md index df981831e4..98288200a7 100644 --- a/docs/src/content/docs/reference/quantization.md +++ b/docs/src/content/docs/reference/quantization.md @@ -41,6 +41,7 @@ A check mark means that the platform class declares the method. NVIDIA CUDA acce | `int8_per_channel_weight_only` | SM 7.5 | — | — | | `nvfp4_per_token` | SM 7.5 | — | — | | `mxfp8` | SM 8.0 | ✓ | ✓ | +| `mxfp6` | SM 11.0 | — | — | ## Read the table diff --git a/docs/src/content/docs/reference/server-arguments.md b/docs/src/content/docs/reference/server-arguments.md index 95ce1a0853..549fec5a02 100644 --- a/docs/src/content/docs/reference/server-arguments.md +++ b/docs/src/content/docs/reference/server-arguments.md @@ -83,6 +83,14 @@ Select GDN prefill backend. Choices: `flashinfer`, `triton`, `cutedsl`. +### `--kda-prefill-backend` + +`--kda-prefill-backend KDA_PREFILL_BACKEND` + +Select KDA prefill backend. + +Choices: `auto`, `triton`, `flashkda`. + ### `--enable-log-requests` `--enable-log-requests, --no-enable-log-requests` @@ -779,7 +787,8 @@ determine the data type of the weights. `--quantization-config QUANTIZATION_CONFIG` User-facing quantization configuration. Carries per-layer-kind specs -(linear, moe) and ignore patterns; see :class:`QuantizationConfigArgs`. +(linear, moe), ignore patterns, and ordered precision overrides; see +:class:`QuantizationConfigArgs`. Auto-populated from the matching online shorthand when `quantization` is one of the values in `ONLINE_QUANT_SHORTHAND_NAMES`. Should either be a valid JSON string or JSON keys passed individually. @@ -2998,7 +3007,7 @@ provided. If using `ngram` method, the related configuration `prompt_lookup_max` and `prompt_lookup_min` should be considered. -Choices: `bailing_hybrid_mtp`, `custom_class`, `deepseek_mtp`, `dflash`, `draft_model`, `dspark`, `eagle`, `eagle3`, `ernie_mtp`, `exaone4_5_mtp`, `exaone_moe_mtp`, `extract_hidden_states`, `gemma4_mtp`, `glm4_moe_lite_mtp`, `glm4_moe_mtp`, `glm_ocr_mtp`, `hy_v3_mtp`, `inkling_mtp`, `longcat_flash_mtp`, `medusa`, `mimo_mtp`, `mimo_v2_mtp`, `minimax_m3_mtp`, `mlp_speculator`, `mtp`, `nemotron_h_mtp`, `ngram`, `ngram_gpu`, `pangu_ultra_moe_mtp`, `qwen3_5_mtp`, `qwen3_next_mtp`, `step3p5_mtp`, `suffix`, `None`. +Choices: `bailing_hybrid_mtp`, `custom_class`, `deepseek_mtp`, `dflash`, `draft_model`, `dspark`, `eagle`, `eagle3`, `ernie_mtp`, `exaone4_5_mtp`, `exaone_moe_mtp`, `extract_hidden_states`, `gemma4_mtp`, `glm4_moe_lite_mtp`, `glm4_moe_mtp`, `glm_ocr_mtp`, `hy_v3_mtp`, `inkling_mtp`, `kimi_k3_mtp`, `longcat_flash_mtp`, `medusa`, `mimo_mtp`, `mimo_v2_mtp`, `minimax_m3_mtp`, `mlp_speculator`, `mtp`, `nemotron_h_mtp`, `ngram`, `ngram_gpu`, `pangu_ultra_moe_mtp`, `qwen3_5_mtp`, `qwen3_next_mtp`, `step3p5_mtp`, `suffix`, `None`. ### `--spec-model` diff --git a/pyproject.toml b/pyproject.toml index 5da866190b..7acfc9b871 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -141,7 +141,10 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "docs/governance/process.md", "docs/assets/contributing/aphrodite_bench_serve_timeline.html", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "rust/src/text/src/output/decoded.rs", - "rust/src/tokenizer/src/incremental.rs"] + "rust/src/tokenizer/src/incremental.rs", + # Vendored NVIDIA CUTLASS reference kernel. Keep identifiers aligned with upstream. + "aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_kernel.py", + "aphrodite/model_executor/kernels/linear/mxfp6/cutedsl_grouped_kernel.py"] ignore-hidden = false [tool.typos.default] diff --git a/tests/quantization/test_mxfp6_online_utils.py b/tests/quantization/test_mxfp6_online_utils.py new file mode 100644 index 0000000000..ee29165d50 --- /dev/null +++ b/tests/quantization/test_mxfp6_online_utils.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from aphrodite.model_executor.layers.quantization.utils.mxfp6_online_utils import ( + dequantize_mxfp6_reference, + pack_mxfp6_codes, + quantize_mxfp6_reference, + unpack_mxfp6_codes, +) + + +def test_mxfp6_known_pack_order(): + codes = torch.tensor([[0, 1, 2, 3]], dtype=torch.uint8) + assert pack_mxfp6_codes(codes).tolist() == [[64, 32, 12]] + assert torch.equal(unpack_mxfp6_codes(pack_mxfp6_codes(codes)), codes) + + +@pytest.mark.parametrize("fmt,max_value", [("e2m3", 7.5), ("e3m2", 28.0)]) +def test_mxfp6_all_codes_round_trip(fmt, max_value): + codes = torch.arange(64, dtype=torch.uint8).view(1, 64) + packed = pack_mxfp6_codes(codes) + scales = torch.full((1, 2), 127, dtype=torch.uint8) + values = dequantize_mxfp6_reference(packed, scales, fmt) + requantized, requant_scales = quantize_mxfp6_reference(values, fmt) + assert torch.equal(requant_scales, scales) + assert torch.equal(requantized, packed) + assert values.abs().max() == max_value + + +@pytest.mark.parametrize("fmt", ["e2m3", "e3m2"]) +def test_mxfp6_block_scaling_and_signed_zero(fmt): + x = torch.zeros((2, 32), dtype=torch.bfloat16) + x[0, 0] = 1024 + x[0, 1] = -1024 + x[1, 0] = -0.0 + packed, scales = quantize_mxfp6_reference(x, fmt) + restored = dequantize_mxfp6_reference(packed, scales, fmt, torch.bfloat16) + assert restored[0, 0] == 1024 + assert restored[0, 1] == -1024 + assert torch.signbit(restored[1, 0]) + + +def test_mxfp6_rejects_unaligned_k(): + with pytest.raises(ValueError, match="divisible by 32"): + quantize_mxfp6_reference(torch.zeros((2, 31))) diff --git a/tests/quantization/test_quantization_config_args.py b/tests/quantization/test_quantization_config_args.py index d776e3f1a8..225dda81c8 100644 --- a/tests/quantization/test_quantization_config_args.py +++ b/tests/quantization/test_quantization_config_args.py @@ -7,6 +7,7 @@ from aphrodite.config.quantization import ( QUANT_KEY_NAMES, QuantizationConfigArgs, + QuantOverride, QuantSpec, resolve_quantization_config, ) @@ -16,6 +17,8 @@ kFp8Static128BlockSym, kFp8StaticTensorSym, kInt8StaticChannelSym, + kMxfp6E2m3Dynamic, + kMxfp6E2m3Static, kMxfp8Dynamic, ) @@ -106,6 +109,44 @@ def test_resolve_merges_explicit_over_shorthand(): assert args.moe == QuantSpec(weight=kFp8StaticTensorSym) +def test_mxfp6_shorthand_uses_w6a8_and_preserves_routers(): + args = resolve_quantization_config("mxfp6", None) + expected = QuantSpec(weight=kMxfp6E2m3Static, activation=kMxfp8Dynamic) + assert args.linear == expected + assert args.moe == expected + assert args.overrides == [ + QuantOverride( + pattern=r"re:(^|.*\.)(gate|router|shared_expert_gate|lm_head)$", + weight="bf16", + ) + ] + + +def test_mxfp6_w6a6_and_ordered_overrides(): + args = resolve_quantization_config( + "mxfp6", + { + "linear": { + "weight": "mxfp6_e2m3", + "activation": "mxfp6_e2m3_dynamic", + }, + "overrides": [ + {"pattern": "re:.*\\.gate$", "weight": "mxfp6_e2m3"}, + {"pattern": "model.layers.0.mlp.gate", "weight": "bf16"}, + ], + }, + ) + assert args.linear == QuantSpec( + weight=kMxfp6E2m3Static, + activation=kMxfp6E2m3Dynamic, + ) + assert [rule.weight for rule in args.overrides] == [ + "bf16", + kMxfp6E2m3Static, + "bf16", + ] + + def test_resolve_rejects_quantization_config_with_non_shorthand_quant(): # If --quantization names something other than an online shorthand, # quantization_config is not allowed via this path (checkpoint quant diff --git a/tools/docs/generate_user_reference.py b/tools/docs/generate_user_reference.py index 7c82266b5b..0bafd8198d 100644 --- a/tools/docs/generate_user_reference.py +++ b/tools/docs/generate_user_reference.py @@ -213,6 +213,7 @@ def quantization_reference() -> str: declared = { platform: set(literal_assignment(path, "supported_quantization")) for platform, path in PLATFORMS.items() } + method_capability_overrides = {"mxfp6": 110} lines = [ frontmatter( "Quantization support", @@ -229,6 +230,7 @@ def quantization_reference() -> str: for method in methods: config_class = config_classes.get(method, "OnlineQuantizationConfig") _, minimum = capabilities.get(config_class, (None, None)) + minimum = method_capability_overrides.get(method, minimum) cuda = f"SM {minimum // 10}.{minimum % 10}" if minimum else "Any supported CUDA GPU" rocm = "✓" if method in declared["AMD ROCm"] else "—" xpu = "✓" if method in declared["Intel XPU"] else "—"