diff --git a/build/torch-cuda/_ops.py b/build/torch-cuda/_ops.py index 1b2f73243ef0cb3bd877aaf2022b702c090f9806..083e0f6afca8bf025a1a5753e285596a246958d4 100644 --- a/build/torch-cuda/_ops.py +++ b/build/torch-cuda/_ops.py @@ -4,7 +4,13 @@ def get_backend() -> str: """Detect the backend by inspecting torch.""" import torch - if hasattr(torch, "neuron"): + if hasattr(torch.backends, "tpu"): + # torch_tpu sets torch.backends.tpu when it is imported (via + # torch's device-backend autoload), regardless of whether TPU + # hardware is present — analogous to torch.version.cuda being + # set on CUDA builds without a GPU. + return "tpu" + elif hasattr(torch, "neuron"): # Needs to be sorted before specific Torch builds, since Neuron # extension can be loaded into e.g. CUDA Torch builds. return "neuron" @@ -22,7 +28,7 @@ def get_backend() -> str: def _find_ops_name() -> str: kernel_name = "sonic_moe" - unique_id = "86f75d9" + unique_id = "83d1d6e" backend = get_backend() return f"_{kernel_name}_{backend}_{unique_id}" diff --git a/build/torch-cuda/_ops_compat.py b/build/torch-cuda/_ops_compat.py deleted file mode 100644 index f4d00b106c8d2e016802140898f645d53e1c3237..0000000000000000000000000000000000000000 --- a/build/torch-cuda/_ops_compat.py +++ /dev/null @@ -1,10 +0,0 @@ -"""Compatibility helpers for op namespacing in source and built layouts.""" - -try: - from ._ops import add_op_namespace_prefix as _generated_add_op_namespace_prefix -except ImportError: - def _generated_add_op_namespace_prefix(name: str) -> str: - return name if "::" in name else f"sonicmoe::{name}" - -def add_op_namespace_prefix(name: str) -> str: - return _generated_add_op_namespace_prefix(name) diff --git a/build/torch-cuda/functional/__init__.py b/build/torch-cuda/functional/__init__.py index f626f330c53392b605d893994bcfaa07328375f4..41ef76af89b656874881815493c4b0c5748143ca 100644 --- a/build/torch-cuda/functional/__init__.py +++ b/build/torch-cuda/functional/__init__.py @@ -11,13 +11,11 @@ from ..quack.gemm_interface import gemm, gemm_dgated, gemm_gated from ..enums import ActivationType, is_glu from .backward import ( _down_projection_backward_act, - _down_projection_backward_weight, _token_broadcast_backward, _topk_softmax_bwd, _up_projection_backward_act, - _up_projection_backward_weight, ) -from .forward import _down_projection_forward, _router_forward, _topk_softmax_fwd, _up_projection_forward +from .forward import _router_forward, _topk_softmax_fwd from .triton_kernels import TC_topk_router_metadata_triton, general_routing_router_metadata_triton @@ -107,17 +105,21 @@ class _UpProjection(torch.autograd.Function): else None ) - _up_projection_forward( - x=x, - w1=w1, - h=h, - a=a, - b1=b1, - expert_frequency_offset=expert_frequency_offset, - x_gather_idx=x_gather_idx, - activation_type=activation_type.value, - is_inference_mode_enabled=is_inference_mode_enabled, - concat_layout=concat_layout, + assert activation_type.value in ( + "swiglu", + "geglu", + ), f"QuACK gemm_gated only supports glu activations, got {activation_type.value}" + gemm_gated( + x, + w1.permute(2, 1, 0), + activation=activation_type.value, + cu_seqlens_m=expert_frequency_offset, + A_idx=x_gather_idx, + preact_out=h, + postact_out=a, + store_preact=(not is_inference_mode_enabled), + bias=b1, + concat_layout=(("B", "bias") if b1 is not None else ("B",)) if concat_layout else None, ) ctx.T = T @@ -182,14 +184,15 @@ class _UpProjection(torch.autograd.Function): concat_layout=concat_layout, ) - _up_projection_backward_weight( - x=x, - dw1=dw1, - dh=dh, - expert_frequency_offset=expert_frequency_offset, - x_gather_idx=x_gather_idx, - is_glu_activation=is_glu_activation, - concat_layout=concat_layout, + gemm( + x.T, + dh, + out=dw1.permute(2, 1, 0), + cu_seqlens_k=expert_frequency_offset, + A_idx=x_gather_idx, + batch_idx_permute=None, + dynamic_scheduler=False, + concat_layout=(("out",) if concat_layout else None), ) dx_reduced = torch.empty(T, H, dtype=dh.dtype, device=dh.device) @@ -231,13 +234,7 @@ class _DownProjection(torch.autograd.Function): y = torch.empty(TK, H, dtype=a.dtype, device=a.device) - _down_projection_forward( - w2=w2, - a=a, - y=y, - b2=b2, - expert_frequency_offset=expert_frequency_offset, - ) + gemm(a, w2.permute(2, 1, 0), out=y, cu_seqlens_m=expert_frequency_offset, bias=b2) o = torch.empty(T, H, device=a.device, dtype=a.dtype) topk_scores = topk_scores.view(-1) @@ -266,6 +263,7 @@ class _DownProjection(torch.autograd.Function): expert_frequency_offset, x_gather_idx, s_scatter_idx, + s_reverse_scatter_idx, ) return o @@ -285,6 +283,7 @@ class _DownProjection(torch.autograd.Function): expert_frequency_offset, x_gather_idx, s_scatter_idx, + s_reverse_scatter_idx, ) = ctx.saved_tensors dw2 = torch.empty_like(w2) @@ -313,12 +312,14 @@ class _DownProjection(torch.autograd.Function): activation_type=activation_type.value, ) - _down_projection_backward_weight( - dout=dout, - a_prime=a_prime, - dw2=dw2, - expert_frequency_offset=expert_frequency_offset, - x_gather_idx=x_gather_idx, + gemm( + dout.T, + a_prime, + out=dw2.permute(2, 0, 1), + cu_seqlens_k=expert_frequency_offset, + A_idx=x_gather_idx, + batch_idx_permute=None, + dynamic_scheduler=False, ) # TC top-K routing @@ -369,7 +370,6 @@ def moe_TC_softmax_topk_layer( if type(activation_type) == str: activation_type = ActivationType(activation_type) - assert not torch.compiler.is_compiling() assert is_glu(activation_type), "QuACK GEMM does not support non GLU activation yet" a, h = _UpProjection.apply( @@ -467,7 +467,6 @@ def moe_general_routing_inputs( num_activated_expert_per_token_offset, ) - assert not torch.compiler.is_compiling() assert is_glu(activation_type), "QuACK GEMM does not support non GLU activation yet" a, h = _UpProjection.apply( diff --git a/build/torch-cuda/functional/backward.py b/build/torch-cuda/functional/backward.py index fe2b94302add991473cbdb799d8204d20659f6af..55f3e2e41e4c9319552e37be6f6f980c13f2a02d 100644 --- a/build/torch-cuda/functional/backward.py +++ b/build/torch-cuda/functional/backward.py @@ -11,7 +11,7 @@ import triton import triton.language as tl from ..quack.gemm_interface import gemm, gemm_dgated -from .._ops_compat import add_op_namespace_prefix +from .._ops import add_op_namespace_prefix from ..utils import get_powers_of_2 from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton @@ -208,35 +208,6 @@ def _up_projection_backward_act( _up_projection_backward_act.compile_cache = {} -@torch.library.custom_op(add_op_namespace_prefix("_up_projection_backward_weight"), mutates_args={"dw1"}) -def _up_projection_backward_weight( - x: torch.Tensor, - dw1: torch.Tensor, - dh: torch.Tensor, - expert_frequency_offset: torch.Tensor, - x_gather_idx: torch.Tensor, - is_glu_activation: bool, - concat_layout: bool = False, -) -> None: - I, H, E = dw1.size() - if is_glu_activation: - I //= 2 - - gemm( - x.T, - dh, - out=dw1.permute(2, 1, 0), - cu_seqlens_k=expert_frequency_offset, - A_idx=x_gather_idx, - batch_idx_permute=None, - dynamic_scheduler=False, - concat_layout=(("out",) if concat_layout else None), - ) - - -_up_projection_backward_weight.compile_cache = {} - - @torch.library.custom_op(add_op_namespace_prefix("_down_projection_backward_act"), mutates_args={"dh", "ds", "db2", "a_prime"}) def _down_projection_backward_act( dout: torch.Tensor, @@ -272,8 +243,6 @@ def _down_projection_backward_act( A_idx=x_gather_idx, dynamic_scheduler=False, ) - ds[s_scatter_idx] = ds_scattered - if db2 is None: ds[s_scatter_idx] = ds_scattered else: @@ -314,28 +283,6 @@ def _down_projection_backward_act( _down_projection_backward_act.compile_cache = {} -@torch.library.custom_op(add_op_namespace_prefix("_down_projection_backward_weight"), mutates_args={"dw2"}) -def _down_projection_backward_weight( - dout: torch.Tensor, - a_prime: torch.Tensor, - dw2: torch.Tensor, - expert_frequency_offset: torch.Tensor, - x_gather_idx: torch.Tensor, -) -> None: - gemm( - dout.T, - a_prime, - out=dw2.permute(2, 0, 1), - cu_seqlens_k=expert_frequency_offset, - A_idx=x_gather_idx, - batch_idx_permute=None, - dynamic_scheduler=False, - ) - - -_down_projection_backward_weight.compile_cache = {} - - @torch.library.custom_op(add_op_namespace_prefix("_token_broadcast_backward"), mutates_args={"dx_reduced"}) def _token_broadcast_backward( dx_reduced: torch.Tensor, diff --git a/build/torch-cuda/functional/forward.py b/build/torch-cuda/functional/forward.py index 2cc4e78001e2e70d61b07c95aab386f40bb4dcdf..3629d941c0c2dfb26cef3330a0bcc36fdcf3a8a0 100644 --- a/build/torch-cuda/functional/forward.py +++ b/build/torch-cuda/functional/forward.py @@ -9,9 +9,8 @@ import triton import triton.language as tl from cutlass.cute.runtime import from_dlpack from ..quack.cute_dsl_utils import torch2cute_dtype_map -from ..quack.gemm_interface import gemm, gemm_gated -from .._ops_compat import add_op_namespace_prefix +from .._ops import add_op_namespace_prefix from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton from .topk import Softmax_Over_TopK, TopK_Over_Softmax @@ -62,54 +61,6 @@ def _topk_fwd( _topk_fwd.compile_cache = {} -@torch.library.custom_op(add_op_namespace_prefix("_up_projection_forward"), mutates_args={"h", "a"}) -def _up_projection_forward( - x: torch.Tensor, - w1: torch.Tensor, - h: torch.Tensor, - a: torch.Tensor, - b1: torch.Tensor | None, - expert_frequency_offset: torch.Tensor, - x_gather_idx: torch.Tensor, - activation_type: str, - is_inference_mode_enabled: bool = False, - concat_layout: bool = False, -) -> None: - assert activation_type in ( - "swiglu", - "geglu", - ), f"QuACK gemm_gated only supports glu activations, got {activation_type}" - gemm_gated( - x, - w1.permute(2, 1, 0), - activation=activation_type, - cu_seqlens_m=expert_frequency_offset, - A_idx=x_gather_idx, - preact_out=h, - postact_out=a, - store_preact=(not is_inference_mode_enabled), - bias=b1, - concat_layout=(("B", "bias") if b1 is not None else ("B",)) if concat_layout else None, - ) - - -_up_projection_forward.compile_cache = {} - - -@torch.library.custom_op(add_op_namespace_prefix("_down_projection_forward"), mutates_args={"y"}) -def _down_projection_forward( - w2: torch.Tensor, - a: torch.Tensor, - y: torch.Tensor, - b2: torch.Tensor | None, - expert_frequency_offset: torch.Tensor, -) -> None: - gemm(a, w2.permute(2, 1, 0), out=y, cu_seqlens_m=expert_frequency_offset, bias=b2) - - -_down_projection_forward.compile_cache = {} - - @torch.library.custom_op(add_op_namespace_prefix("_router_forward"), mutates_args={"o"}) def _router_forward( y: torch.Tensor, diff --git a/build/torch-cuda/functional/reduction_over_k_gather.py b/build/torch-cuda/functional/reduction_over_k_gather.py index c5e41f1ca56ff288cd9d1391660b3685c9cb1c68..eebf3e0f847807def84d5bc01bfd52280a10261d 100644 --- a/build/torch-cuda/functional/reduction_over_k_gather.py +++ b/build/torch-cuda/functional/reduction_over_k_gather.py @@ -71,12 +71,12 @@ def token_gather_sum_kernel( ): # 1D tiling over T only pid_t = tl.program_id(axis=0) - t_idx = pid_t.to(tl.uint32) + t_idx = pid_t.to(tl.int64) # Load segment starts and ends for this token if is_varlen_K: - Ms = tl.load(M_offset_ptr + t_idx).to(tl.uint32) - Me = tl.load(M_offset_ptr + t_idx + 1).to(tl.uint32) + Ms = tl.load(M_offset_ptr + t_idx).to(tl.int64) + Me = tl.load(M_offset_ptr + t_idx + 1).to(tl.int64) K_this_token = Me - Ms # actual K for this token else: Ms = MAX_K * t_idx @@ -84,7 +84,7 @@ def token_gather_sum_kernel( # Outer loop over H tiles for h_tile in tl.static_range(triton.cdiv(H, BLOCK_H)): - h_idx = (h_tile * BLOCK_H + tl.arange(0, BLOCK_H)).to(tl.uint32) # [BLOCK_H] + h_idx = (h_tile * BLOCK_H + tl.arange(0, BLOCK_H)).to(tl.int64) # [BLOCK_H] m_h = h_idx < H # Initialize accumulator for this H tile @@ -94,7 +94,7 @@ def token_gather_sum_kernel( for k_tile in tl.range(tl.cdiv(K_this_token, BLOCK_K)): k_offset = k_tile * BLOCK_K - k_idx = (k_offset + tl.arange(0, BLOCK_K)).to(tl.uint32) # [BLOCK_K] + k_idx = (k_offset + tl.arange(0, BLOCK_K)).to(tl.int64) # [BLOCK_K] # Mask for valid K indices m_k = k_idx < K_this_token # [BLOCK_K] @@ -103,7 +103,7 @@ def token_gather_sum_kernel( m_abs = Ms + k_idx # [BLOCK_K] # Gather permuted indices - perm_idx = tl.load(M_perm_ptr + m_abs, mask=m_k, other=0).to(tl.uint32) # [BLOCK_K] + perm_idx = tl.load(M_perm_ptr + m_abs, mask=m_k, other=0).to(tl.int64) # [BLOCK_K] # Load x values: [BLOCK_K, BLOCK_H] x_ptrs = x_ptr + perm_idx[:, None] * stride_xM + h_idx[None, :] * stride_xH diff --git a/build/torch-cuda/functional/tile_scheduler.py b/build/torch-cuda/functional/tile_scheduler.py deleted file mode 100644 index f9d9dd8101bd994c43465bb239a4759be78510b0..0000000000000000000000000000000000000000 --- a/build/torch-cuda/functional/tile_scheduler.py +++ /dev/null @@ -1,91 +0,0 @@ -# ******************************************************************************** -# Copyright (c) 2025, Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao -# ******************************************************************************** - -from __future__ import annotations - -import cutlass -import cutlass.cute as cute -from cutlass import Boolean, Int32, const_expr -from ..quack.pipeline import PipelineStateWAdvance -from ..quack.tile_scheduler import TileScheduler, VarlenMTileScheduler - - -class SonicMoETileScheduler(TileScheduler): - @staticmethod - @cute.jit - def create( - params: TileScheduler.Params, - tile_count: cute.Tensor | None = None, - scheduler_pipeline: cutlass.pipeline.PipelineAsync | None = None, - is_scheduler_warp: bool | Boolean = False, - *, - loc=None, - ip=None, - ) -> SonicMoETileScheduler: - """is_scheduler_warp should only be true for one warp in the whole cluster""" - stages = 0 - if const_expr(not params.is_persistent): - cidx, cidy, _ = cute.arch.cluster_idx() - cdimx, _, _ = cute.arch.cluster_dim() - cluster_id = cidx + cidy * cdimx - current_work_linear_idx = Int32(cluster_id) - else: - _, _, bidz = cute.arch.block_idx() - current_work_linear_idx = Int32(bidz) - if const_expr(params.tile_count_semaphore is not None): - assert tile_count is not None - assert scheduler_pipeline is not None - stages = const_expr(cute.size(tile_count)) - return SonicMoETileScheduler( - current_work_linear_idx, - Int32(0), # num_tiles_executed - tile_count, - scheduler_pipeline, - PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(1 if is_scheduler_warp else 0)), - params, - loc=loc, - ip=ip, - ) - - def prefetch_next_work(self, *, advance_count: int = 1, loc=None, ip=None): - old_current_work_linear_idx = self._current_work_linear_idx - if const_expr(self.params.is_persistent): - num_persistent_clusters = cute.arch.grid_dim()[2] - self._current_work_linear_idx += advance_count * Int32(num_persistent_clusters) - future_tile_coord_mnkl = self.get_current_work() - self._current_work_linear_idx = old_current_work_linear_idx - return future_tile_coord_mnkl - - -class SonicMoEVarlenMTileScheduler(VarlenMTileScheduler, SonicMoETileScheduler): - @staticmethod - @cute.jit - def create( - params: VarlenMTileScheduler.Params, - tile_count: cute.Tensor | None = None, - scheduler_pipeline: cutlass.pipeline.PipelineAsync | None = None, - is_scheduler_warp: bool | Boolean = False, - *, - loc=None, - ip=None, - ) -> SonicMoEVarlenMTileScheduler: - stages = 0 - _, _, bidz = cute.arch.block_idx() - current_work_linear_idx = Int32(bidz) - if const_expr(params.tile_count_semaphore is not None): - assert tile_count is not None - assert scheduler_pipeline is not None - stages = const_expr(cute.size(tile_count)) - return SonicMoEVarlenMTileScheduler( - current_work_linear_idx, - Int32(0), # num_tiles_executed - Int32(0), # current_batch_idx - Int32(0), # num_work_idx_before_cur_batch - tile_count, - scheduler_pipeline, - PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(1 if is_scheduler_warp else 0)), - params, - loc=loc, - ip=ip, - ) diff --git a/build/torch-cuda/functional/triton_kernels/__init__.py b/build/torch-cuda/functional/triton_kernels/__init__.py index e2bc28024bdd255f439c223e426640c18b22c370..fa3c7e2caec69047b447e115399159778ea4c4f9 100644 --- a/build/torch-cuda/functional/triton_kernels/__init__.py +++ b/build/torch-cuda/functional/triton_kernels/__init__.py @@ -1,10 +1,13 @@ +# ******************************************************************************** +# Copyright (c) 2026, Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao +# ******************************************************************************** import math import torch +from ..._ops import add_op_namespace_prefix import triton import triton.language as tl -from ..._ops_compat import add_op_namespace_prefix from .bitmatrix import _bitmatrix_metadata_compute_stage1, _bitmatrix_metadata_compute_stage2, _keyed_add diff --git a/build/torch-cuda/metadata.json b/build/torch-cuda/metadata.json index 5702f687d61fdc4b87afc39d820e545f5f633d2e..a90821abd5ad72fee0aae1194025a5e4037843df 100644 --- a/build/torch-cuda/metadata.json +++ b/build/torch-cuda/metadata.json @@ -1,7 +1,7 @@ { "name": "sonic-moe", - "id": "_sonic_moe_cuda_86f75d9", - "version": 1, + "id": "_sonic_moe_cuda_83d1d6e", + "version": 2, "license": "Apache-2.0", "python-depends": [ "tvm-ffi", @@ -9,5 +9,120 @@ ], "backend": { "type": "cuda" + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "DDepLW9l03NRbdlAuR2+UkWgJuQN/8ymgKZTMmdV0Dc=", + "_ops.py": "9EVL/sWTDKjpk2E9q9TV3k0BIZR8kzK88mQRF7lwkm8=", + "enums.py": "eeBuPJFCiiGCsL7OmYWCMfdvj4iQE4irrXun8niAUgg=", + "functional/__init__.py": "FDQaNUpQGY+3YeKXK9PNmtwKjXtn16GxsGdyFiQul08=", + "functional/backward.py": "W10eywFlVh1mMO+9oktuSWK7HllWZxzaByzQfXt1jM4=", + "functional/forward.py": "Tq1hGu14TvCXMptVrIZJQLgnAp0wSVzATEv9T5nc35w=", + "functional/reduction_over_k_gather.py": "fBYTaHbHTXpl7i2dtJ1InNIh0PIpxfE+6ImvPoNcOaU=", + "functional/topk.py": "Gsq7DJpeAuL0co/Zm9tSpyP0t9tSEwxSlRyjrTyCVhg=", + "functional/triton_kernels/__init__.py": "19C1f3DyiRGJfW3TcfHZH9x4q2vhqzAj2YkK/if9TUU=", + "functional/triton_kernels/bitmatrix.py": "kJ+VIr07n0e48+XQHSgQbLnmcRgFZ91M3NIUxNN4hwI=", + "jit.py": "pko7Lkttkab+0YDUkXbAAN3I/yMvTnw7d2jJE6zml+g=", + "moe.py": "7UtJ7Oe+401gE+pn0axai3MWrVipfJl9pJc72mSwB/g=", + "quack/__init__.py": "BsCyHg2lDdBg8lhuijV2CjN2CFJnQQuovcg5v4gM+Xo=", + "quack/_ops_compat.py": "Df2expY3Aqaob3ZuykC0qsSWm4DuUUSPUvI9J9gY9Rg=", + "quack/activation.py": "YIsLHpQauAokG2r/V9ToZnFzkCDll3JKU4IhTCsK8Hw=", + "quack/autotuner.py": "nbdngdlWYF19B4R6z6N7fgnDoO40MFghpeHlcKe2Yv8=", + "quack/bench/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "quack/bench/bench_utils.py": "OZfHYnYliR7+o1paEdakdQ8YS22ySnOwPjwxuwfXpxE=", + "quack/blockscaled/__init__.py": "8aZJFIdaeT2ND6x77yrR/kAFVlSAkjx/nEwgPXrSv5s=", + "quack/blockscaled/quantize.py": "EPiyWvbHg1PC0UI5IRX3CqHTNk6HnjzhkUY+vkwgNSU=", + "quack/blockscaled/utils.py": "IE/QeL4IurRCCLLF/m/PKrZjuJ/Tp1j5jcgZggxy2eM=", + "quack/broadcast_utils.py": "HVUVszpqpbhv3ZrfCMSQVtqmjfRJJljYH+3J0WuSaMs=", + "quack/cache/__init__.py": "9qjakMy4yFmpI7XdQz4NDUyqfCti6xi0tkZFPnp9KVU=", + "quack/cache/_pool_preload.py": "qGKHxfFndDBCYyHUjSRDC7A+conMOQTqsjPfX5OCZc0=", + "quack/cache/async_compile.py": "BzJjQ5YswF4YAreXFDvbI05pQHzdtE5zxSdN4M7pOp4=", + "quack/cache/jit.py": "vXgLNNeAMrEY8a1fpigiiCUWZGCTbrtxMYWFrHIffkE=", + "quack/compile_utils.py": "d8AzPa2ONwKF3DhsB1E5orFzrNmJJinrYLJ0kstFq6k=", + "quack/complex.py": "WmYNde9V3ZDLaOt7ttCnjH4Dxahan4+PvLfoHHHU+jo=", + "quack/copy_utils.py": "0e1AtA1GZjFdc9Ts0mNT+gCftK0PVbDDkzOkXiBgSHE=", + "quack/cross_entropy.py": "OL8SJ435IOxTxL7L/C6vl/+3+BQWbIgMUt86rM8hJ84=", + "quack/cute_dsl_utils.py": "9lUc4T4L+Nm05V8doN+iiUqBFa2yBUSjYaVUseXCfH4=", + "quack/dsl/__init__.py": "v0w1PKdAXyF8yOzc/OwvR5+4M5rcdni62pCtVWpj2bI=", + "quack/dsl/cute_dsl_ptxas.py": "jbam6fUwzm4X5CZc29LevlWXBAzbedWu3j6hK/2lEmA=", + "quack/dsl/cute_tensor.py": "l1Z8XNzBSQuHwRKR6gX5+Se1DR39NoVR/JEV2IJbveo=", + "quack/dsl/cute_tensor_indexing.py": "1QnPrnekkECXkgHC5NWY6JxyDmTTf5kgClCIYwxMfG0=", + "quack/dsl/smem_struct.py": "piMjPM+wy8w/Ig6IZI/AOoVpeKsKSEes7xPxY2eiWCI=", + "quack/dsl/torch_library_op.py": "QfwtFO4P/3FSivWuY1LH2dL2xiead+m2VpK3iBKjKp0=", + "quack/epi_composable.py": "iCA8mYu9CLXNVvWmdP7847Pk+51tUYv6U6lTb5ptgfQ=", + "quack/epi_ops.py": "5+MCkOUqeLxFrzevU5OH4KXgxJJTKC+10jUHkxiDplA=", + "quack/epi_utils.py": "pdXiffZgdAbSA2WM5jkOIWvrPTBGqsvIkKrIffiAXFM=", + "quack/fast_math.py": "Nvhr8+bgjPAS+e/XKzlxO0Xzqa3gi8vRxV3gl82ECJI=", + "quack/gemm.py": "mAVLwuD+tC7O+NGhLzGaYQhsoWm4lo8h7TWEsqyJosA=", + "quack/gemm_act.py": "OIIB130+VHkMvyfP/kXCXzeNr7BkK68pIhxY2YwTqXM=", + "quack/gemm_base.py": "9RUR26bnwPd5LkM0PDo5201x0QdxgLO6VRZNCxBuDR8=", + "quack/gemm_config.py": "NinZfsUbbcaFyMR+yzuBtSfF2Q6TJPd3AoxAr89I5AQ=", + "quack/gemm_dact.py": "HV3dHIx4Er+QogUnp/lvBbkOAsXm2+z2YRiuOHyO5gQ=", + "quack/gemm_default_epi.py": "YFh3eGsC5wPVHiw//IoNQKiDH0u1ioT/5QPwtasRats=", + "quack/gemm_interface.py": "XTqyTBBZxOFEnE9/DCJjAsUQ3BBICY9nFgOSGjc8OZ4=", + "quack/gemm_norm_act.py": "u5CVGXIN4g1Z6pPD4fHKQkjDl1V/hdH1S3eWBlq8Hf0=", + "quack/gemm_sm100.py": "6vehfja8xPV1A78HBC6kUjHfOkUGGVtTiTLW7YVlr48=", + "quack/gemm_sm120.py": "aCy6GOMLgyvuFq+L4mRPXMXuKVcobswO6FSwLWhfx+k=", + "quack/gemm_sm80.py": "mGyJSOf7HSTJbOQawSJiu53hfKNg7HvSPzTVDLSwB4c=", + "quack/gemm_sm90.py": "38ONe3cM8g8hAEx4JomzbbkJ5lTo92UpbLNoUMn9EQE=", + "quack/gemm_sq_reduce.py": "N6JefW8NwxiMjFs6JssVu97Ysn9n4nErQkW2J6VMZHI=", + "quack/gemm_symmetric.py": "PRVtXM+nCya7nyiDH5ozvndGnfLX5TkM9kXYnSuJ5sg=", + "quack/gemm_tvm_ffi_utils.py": "76C1usihWxUD0lRi82AUO8OVjWVW5ypxIRelaFSAzTo=", + "quack/jax_utils.py": "z9CKuM2HbYXkgib8+mVNEvI5WF2HxBl8cgTCmiosTus=", + "quack/layout_utils.py": "I6MaaKCC67T+Yh/d6QdU4xbiqfPQ0ZJRMJ7/eUGFlFg=", + "quack/linear.py": "DK6ZcxbVB5PNnvRhUilds7BbRRdqovwNXJTPUOfUo2A=", + "quack/linear_cross_entropy.py": "UmkMdBB+8Tp4C4ikwjM/AGG/Ql+lWoqgmbADEHZ4I6o=", + "quack/mlp.py": "LKrK0VX3MToBM98oELVUMQJ3Zmgt93QNyZ0u7xdaFNE=", + "quack/nvmmh_heuristic.py": "mxImbqtdhY4NxYS6kZRL41zgG8oKa94uo5yLUtbLs2E=", + "quack/pipeline.py": "veqXF6OrbggTrcjhw4E/vDc7ufja27tifRDSuok4XHM=", + "quack/reduce.py": "Ul6PlgNnYQvm4dq948stCPrBdKA0Ct+jjgSeZFbxdNk=", + "quack/reduction_base.py": "cJc9fa+DsOFigD7d7bgNWR03gXkyXNzoSXwjPJxKTl0=", + "quack/rms_final_reduce.py": "9X2rSFH25BjNdLFDaJKmZeq/i1/gnmYhoYtCyPmAhrQ=", + "quack/rmsnorm.py": "yFyoI/p3ZyZqvOaDgN/+a3PbK1kJZvnYxjgLxAorgN4=", + "quack/rmsnorm_config.py": "HDOdXv5Ej2egQVQDVk4k1ynXUE0D0HpPaN9O7Qd9EK8=", + "quack/rotary.py": "A0FaRr7LKdWw+kjeHEg79I4eYb9+MbGKSMPuawocGBw=", + "quack/rounding.py": "jI6fP/qCjSMh9s8PtjPDYfTPs200C3xx9zKll1hlodU=", + "quack/sm100_utils.py": "+AkAPWzg+2cXZMwTijIVTvWWwZmk5m3bAsOJQTffzDo=", + "quack/sm80_utils.py": "0yOgRslAIq8B7MS8tIkgZuilOMIkYnJDUl7c2gjdkzg=", + "quack/sm90_utils.py": "gX9BE07yDSBMnrFdIJY0EXvyQuutTiXlCT8j5mHDLh0=", + "quack/softmax.py": "YQ4R6wVI2J/+FsjuTVGQ2Jj4shl4KU/YrHgFZmQ99do=", + "quack/softmax_jax.py": "VZ1phS40aMQDyx4QgsrWPxbkbdLnv3jNEc/+j2VqNN0=", + "quack/sort/bitonic_sort.py": "XjGKnZ51NI2N2BoHBgxz2Ne2OvYfnuOTgiOSnyjN08Q=", + "quack/sort/generate_sorting_networks.py": "al5BepGSLhn/bBPg64p2KpW7itx9gaCIGxhpqIceKB8=", + "quack/sort/sorting_networks.py": "okUmplKNQMAu8/lHClRiQf31X3tIe7MR07KybQJfH8Y=", + "quack/sort/utils.py": "lifuGThMk4D1VxfAUYDxrrUzHwxLGgLdoinpE8UVN3g=", + "quack/spec/__init__.py": "gZsivx8yPk0Tk7GqzHpOtrjN2b3dSB0/MNKVt0/e8qs=", + "quack/spec/mma.py": "3SSQiMkkncb3vuaoc9Dkr687vIBxEf7icnsdWJv5q58=", + "quack/spec/smem.py": "McZhRq5NJlHEFye+TQZpvDvnarRgLFKh2Hk6JslyOUA=", + "quack/spec/tensor_spec.py": "2d30W/TI7kwvc7tTw6/q1JU35/k0tJGD81JAXeX/P2A=", + "quack/spec/tma.py": "sz94fm7k50nGQ2jr7ubCzeZBX/YWZRmIerwp8P4+EFs=", + "quack/spec/tmem.py": "7YVy9F+Xn3iYEx2aDKiVW1xSIO8GMw+nyrwIdFgnrxs=", + "quack/sync/__init__.py": "5dkW0RJn9GodDYJeXxyQJfNaMBkXczp5BjkXeQK4/yg=", + "quack/sync/barrier.py": "sTHk/GcIC+OcgF6SK+YxSo7R4mdCB1n/6i42fdx+iVg=", + "quack/tensormap_manager.py": "twLLdCS6s8+BCYIh4XYdPtE3XX4OAASHklnsGssznPg=", + "quack/testing/__init__.py": "F4fOLg+7e5MRZUU+2kKdpR+XykrUhy3hrkrXke139+0=", + "quack/testing/pytest_plugin.py": "drWmLh35ffAa9J2VjGOiyhfWYBOcyr54HRNb1UkGTHo=", + "quack/testing/trace.py": "/pVNHOyiv1mgnvGKWy7RmreheJzUZx229+fMJoQZJo0=", + "quack/tile_scheduler.py": "VftFntcYMwF3t6op2H4arOSjmcQzvSqBp5YXjDSS5a0=", + "quack/topk.py": "md6UD7+W1OlZ3K0UDwhfT5fIT3kBrL3ccPYDuEcIedM=", + "quack/trace.py": "nmHaefNLu7k9MjI+LPYmGe69WU1UumN6D7IDAtX9sms=", + "quack/transform/__init__.py": "JhfsXbnxVdojWTz6VbGHCTkCANeR1b9Js2wDyldnO+c=", + "quack/transform/hadamard.py": "fehxPxaWlhWr1LHgBJijuFHQ2Q9houh/ZLEeO1WOvtQ=", + "quack/utils.py": "/FsIscjaxmOSj/g6frh4gLgB6iuIZ6e3EeGNunKWBuw=", + "quack/varlen_utils.py": "UtdJmeCWbAMeJhTYNF78cjN9CH3fEpNQvicoF1YhGL8=", + "utils.py": "mggbSpSYFQIJTyYgccngut/SeK2AqJfSXL1/CMZ8VyM=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "sha": "a6564d1f481adcbd3273099c0f4432e7b833e846", + "dirty": false + }, + "kernel": { + "sha": "83d1d6e670b94603c6bdda6f039416ce988d7683", + "dirty": false + } } } \ No newline at end of file diff --git a/build/torch-cuda/metadata.json.sigstore b/build/torch-cuda/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..861b9e9d137e267faaddd5171422f838275f403a --- /dev/null +++ b/build/torch-cuda/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtCgAwIBAgIUFiUTErfHXgcT0avZDUrtqVpjPVMwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODA0MDgyNzA5WhcNMjYwODA0MDgzNzA5WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE8E1Q2yBXTAmi4adPxCkrbQcPf17/xBZ3tfE6qg9RHWEdmi7WM5D7S43lqmnObiB4CkefF+OIxcSlHV8UJ2opOKOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU36UElV8bsJe0ibWvj7caucWF9xgwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoODNkMWQ2ZTY3MGI5NDYwM2M2YmRkYTZmMDM5NDE2Y2U5ODhkNzY4MzATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoODNkMWQ2ZTY3MGI5NDYwM2M2YmRkYTZmMDM5NDE2Y2U5ODhkNzY4MzAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoODNkMWQ2ZTY3MGI5NDYwM2M2YmRkYTZmMDM5NDE2Y2U5ODhkNzY4MzAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDgzZDFkNmU2NzBiOTQ2MDNjNmJkZGE2ZjAzOTQxNmNlOTg4ZDc2ODMwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA4OTE2ODY1NTMvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn8viPTkAAAQDAEYwRAIgfLsNELCVlyuHUG5H2cJ2naaV5F0SYQZv0YdXyC9K6asCIGpaO2l3hSUXIiLRBq6dsIXZHZyZqZ2/6nNPdklmZMlzMAoGCCqGSM49BAMDA2gAMGUCMHeyUvclqJRFc6yuhmZtURySzkJ6RpnXasTQ5hXcx4H/3BVH1hHtFeXi4epa0GNcLgIxAPspbTGLLNMVY+7Rhv7SaSsX5AGzaWVABXmQPDy7S7T4IJArASfsqBm3NmO0SR0vtA=="}, "tlogEntries":[{"logIndex":"2339551903", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785832029", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQDitpUL1h/KusQ9OMEtNPCsHo084FfnQmZK0WwFpKnoZAIhAPTvy8x20w6lcWucTqcZVLYevLSY3i2S/rLsZ34hFk44"}, "inclusionProof":{"logIndex":"2217647641", "rootHash":"nQRMpxeZMxZKoBM0Aid5ajFNCBrokJCPxAZe9n5BiME=", "treeSize":"2217647645", "hashes":["48NZ8qo8UwRJEAZ3CLfS1KfxJZxkxn+iOuNafjEb1/M=", "u2Qapxkuy0/IKrnn3gJDePaVd9YyBiTAyigk5uwxM8Y=", "yzLVy2ZAAv+TndVc45A90uQ1k7RULtWa9nqp7rNIt9U=", "ZeI1SGr9GP/SskvVce8yHk+6DRrxLrWqJpQnRQRaKpc=", "nQWu7U6/pF5XpZlDhkQlNjCedXmCTB3GYqdStVFhcGo=", "WeMerEQh8gnrs+xRVnIEQXfykDf4KiXgRHX8oVHQYIs=", "2uIbGD5iEQJHWVjTX0311tWS7CFZVaKdy3hZGbyhI8s=", "DLtcvkp4dRTrpmSEDyxo9Y5x1mCGHjNNX5a+Pu2gajA=", "7VIK52c6IvZNnE+Au/wuGXo7lX61HQyDCibo+Dfcvs8=", "jDU3DoSktlu0G+qce/KFrCSxbMWDf2MyMt8j1NVTV/g=", "p0bIaoDFhBhr45Lurz3ap+CC2ZtzymaJJb6Ij0aF6Xg=", "7xWj0qvDr9hJ4vooBulXteOWWqidK4mqZKahpqVHaoE=", "r3u4sx0tkA1pkDYoU+5zLDv1az999M9yZP3h7rXTwYE=", "8cBkID7yRXftjQqclK9Jj0hUDPpp7HMqoKXxzdkwOzg=", "i5Zl8FZrDwxCDv2e2DNO2M8JvpR/c11ElvCZS53/teA=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2217647645\nnQRMpxeZMxZKoBM0Aid5ajFNCBrokJCPxAZe9n5BiME=\n\n— rekor.sigstore.dev wNI9ajBFAiEAm8BuxbwwhrjVkIerX6mU3KfouLswOpqLlyxRei9yV58CIEW4vOQQYXUx761ALJD54e5Q5qfizDBuosCfSV/l8+Jc\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI0ZWYzZTQwM2RjMGJlM2Q1N2YzNWYyYjNkNzIwNTMwNDZkZmZlYjczZWQ2YmE3ODE4NGMwYjA0Njc5MzllNzU2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJSElhM3JzWDI4a1ZOKzFPVDQyUDhLNk82anY3Wk9iSW9neXg1M2t5a0xFWkFpRUF5K1QxY0VyTjJ3M2xhRWpwaXRZYi9JNi9KMEhlOFNiQVN6UjBVVFdEYk84PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SRFowRjNTVUpCWjBsVlJtbFZWRVZ5WmtoWVoyTlVNR0YyV2tSVmNuUnhWbkJxVUZaTmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFUVRCTlJHZDVUbnBCTlZkb1kwNU5hbGwzVDBSQk1FMUVaM3BPZWtFMVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVU0UlRGUk1ubENXRlJCYldrMFlXUlFlRU5yY21KUlkxQm1NVGN2ZUVKYU0zUm1SVFlLY1djNVVraFhSV1J0YVRkWFRUVkVOMU0wTTJ4eGJXNVBZbWxDTkVOclpXWkdLMDlKZUdOVGJFaFdPRlZLTW05d1QwdFBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlV6TmxWRkNteFdPR0p6U21Vd2FXSlhkbW8zWTJGMVkxZEdPWGhuZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVBSRTVyVFZkUk1scFVXVE5OUjBrMVRrUlpkMDB5VFRKWmJWSnJDbGxVV20xTlJFMDFUa1JGTWxreVZUVlBSR2hyVG5wWk5FMTZRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDlFVG10TlYxRXlXbFJaTTAxSFNUVk9SRmwzVFRKTk1sbHRVbXRaVkZwdFRVUk5OVTVFUlRJS1dUSlZOVTlFYUd0T2VsazBUWHBCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDlFVG10TlYxRXlXbFJaTTAxSFNUVk9SRmwzVFRKTk1sbHRVbXNLV1ZSYWJVMUVUVFZPUkVVeVdUSlZOVTlFYUd0T2VsazBUWHBCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUm5lbHBFUm1zS1RtMVZNazU2UW1sUFZGRXlUVVJPYWs1dFNtdGFSMFV5V21wQmVrOVVVWGhPYlU1c1QxUm5ORnBFWXpKUFJFMTNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFMFQxUkZNazlFV1RGT1ZFMTJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDRkbWxRVkd0QlFVRlJSQXBCUlZsM1VrRkpaMlpNYzA1RlRFTldiSGwxU0ZWSE5VZ3lZMG95Ym1GaFZqVkdNRk5aVVZwMk1GbGtXSGxET1VzMllYTkRTVWR3WVU4eWJETm9VMVZZQ2tscFRGSkNjVFprYzBsWVdraGFlVnB4V2pJdk5tNU9VR1JyYkcxYVRXeDZUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlaMEZOUjFWRFRVaGxlVlYyWTJ3S2NVcFNSbU0yZVhWb2JWcDBWVko1VTNwclNqWlNjRzVZWVhOVVVUVm9XR040TkVndk0wSldTREZvU0hSR1pWaHBOR1Z3WVRCSFRtTk1aMGw0UVZCemNBcGlWRWRNVEU1TlZsa3JOMUpvZGpkVFlWTnpXRFZCUjNwaFYxWkJRbGh0VVZCRWVUZFROMVEwU1VwQmNrRlRabk54UW0welRtMVBNRk5TTUhaMFFUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgx31mO/vQ8UTM6GrWl7ysWYCg9VM1q5BDPZqvFVfIaYMCFQD1qIDvOHE1Cgo864r6GhDv1PSMORgPMjAyNjA4MDQwODI3MDlaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgwNDA4MjcwOVowLwYJKoZIhvcNAQkEMSIEIHmfCoWI9C2TufxbIehLRJjjeOkcBfMZEcgFW3CNowVHMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMHTZ3i8JqZ77m5Vccy3iUVKabZVryj/JRhzGaPfzud9V06kWLkG/Y8pmaHEOZXR4tAIxAOK0xDAb+/Kj7Vn16OaMNEQGjedN9J3fDf2Nyf+/TvnAlwVVpbX++c717VMpDOG6pw=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"TvPkA9wL49V/NfKz1yBTBG3/63Pta6eBhMCwRnk551Y="}, "signature":"MEUCIHIa3rsX28kVN+1OT42P8K6O6jv7ZObIogyx53kykLEZAiEAy+T1cErN2w3laEjpitYb/I6/J0He8SbASzR0UTWDbO8="}} \ No newline at end of file diff --git a/build/torch-cuda/quack/__init__.py b/build/torch-cuda/quack/__init__.py index 56d29ecc4e4e18eacd629562efd03cb41980949d..673349cc413317b49a1a02224df8134e32a61afa 100644 --- a/build/torch-cuda/quack/__init__.py +++ b/build/torch-cuda/quack/__init__.py @@ -1,8 +1,17 @@ -__version__ = "0.3.11" +__version__ = "0.6.1" import os +from . import dsl as _quack_dsl # noqa: F401 + if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None: - from . import cute_dsl_ptxas # noqa: F401 + from .dsl import cute_dsl_ptxas as _cute_dsl_ptxas + + # Patch before importing any modules that instantiate CuTeDSL. The patch + # forces PTX dumping so the CUDA library loader can replace CUTLASS DSL's + # embedded ptxas-library cubin with one assembled by system ptxas. + _cute_dsl_ptxas.patch() - cute_dsl_ptxas.patch() +# Pythonic CuTe tensor indexing (`:` / `...` sugar) is installed as a side effect +# of importing `quack.dsl`, which imports `quack.dsl.cute_tensor_indexing` and +# monkey-patches CuTe's tensor classes process-wide. diff --git a/build/torch-cuda/quack/_compile_worker.py b/build/torch-cuda/quack/_compile_worker.py deleted file mode 100644 index 05fb5c1f0a26b1cfbd0814efe92a376b8467d63c..0000000000000000000000000000000000000000 --- a/build/torch-cuda/quack/_compile_worker.py +++ /dev/null @@ -1,102 +0,0 @@ -# Copyright (c) 2025, Tri Dao. -# Persistent subprocess worker for parallel autotuning pre-compilation. -# Receives length-prefixed pickled tasks on stdin, creates FakeTensors -# matching the parent's tensor metadata, and compiles with COMPILE_ONLY=True. -# Stays alive to process multiple configs (amortizes import overhead). - -import importlib -import pickle -import struct -import sys - -import torch -from torch._subclasses.fake_tensor import FakeTensorMode - -from . import cache_utils - -cache_utils.COMPILE_ONLY = True - -_dtype_map = { - "torch.float16": torch.float16, - "torch.bfloat16": torch.bfloat16, - "torch.float32": torch.float32, - "torch.float64": torch.float64, - "torch.int32": torch.int32, - "torch.int64": torch.int64, - "torch.int8": torch.int8, - "torch.uint8": torch.uint8, - "torch.bool": torch.bool, -} - - -def _make_fake_tensor(meta): - shape = meta["shape"] - stride = meta["stride"] - dtype = _dtype_map[meta["dtype"]] - return torch.empty_strided(shape, stride, dtype=dtype, device="cuda") - - -def _recv(stream): - """Read a length-prefixed pickled message. Returns None on EOF.""" - header = stream.read(4) - if len(header) < 4: - return None - length = struct.unpack(" str: - return add_op_namespace_prefix(f"quack__{name}") + + +# For quack we need to prefix the function name because some names +# overlap between quack and sonic-moe itself. Name the function the +# same as the function it is wrapping for the prefix check to be +# happy. +def add_op_namespace_prefix(name: str) -> str: + return _add_op_namespace_prefix(f"quack__{name}") diff --git a/build/torch-cuda/quack/activation.py b/build/torch-cuda/quack/activation.py index c689f0b7f3a0c7ca792c5089df9772f2b865b2d4..2f48657eba368fa035ce43b7407ef43111a10d36 100644 --- a/build/torch-cuda/quack/activation.py +++ b/build/torch-cuda/quack/activation.py @@ -4,10 +4,12 @@ import math from typing import Tuple from functools import partial +import cutlass import cutlass.cute as cute from cutlass import Float32, Boolean, const_expr -from cutlass.cutlass_dsl import T, dsl_user_op -from cutlass._mlir.dialects import llvm, nvvm +from cutlass.cutlass_dsl import dsl_user_op +from cutlass._mlir.dialects import nvvm +from cutlass._mlir_helpers import math as mlir_math F32_or_F32x2 = Float32 | Tuple[Float32, Float32] @@ -21,30 +23,61 @@ sub_packed_f32x2 = partial( @dsl_user_op -def tanh(a: float | Float32, *, loc=None, ip=None) -> Float32: - return Float32( - llvm.inline_asm( - T.f32(), - [Float32(a).ir_value(loc=loc, ip=ip)], - "tanh.approx.f32 $0, $1;", - "=f,f", - has_side_effects=False, - is_align_stack=False, +def tanh(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: + if const_expr(not isinstance(x, tuple)): + return cute.math.tanh(x, fastmath=True, loc=loc, ip=ip) + else: + return ( + cute.math.tanh(x[0], fastmath=True, loc=loc, ip=ip), + cute.math.tanh(x[1], fastmath=True, loc=loc, ip=ip), ) - ) @dsl_user_op -def sigmoid(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: +def dtanh( + x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None +) -> Tuple[F32_or_F32x2, F32_or_F32x2]: + if const_expr(not isinstance(x, tuple)): + tanh_x = tanh(x, loc=loc, ip=ip) + dx = dout * (1.0 - tanh_x * tanh_x) + return dx, tanh_x + else: + tanh_x = tanh(x, loc=loc, ip=ip) + sech2_x = cute.arch.fma_packed_f32x2(tanh_x, (-tanh_x[0], -tanh_x[1]), (1.0, 1.0)) + dx = cute.arch.mul_packed_f32x2(dout, sech2_x) + return dx, tanh_x + + +@dsl_user_op +def sigmoid_tanh(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: if const_expr(not isinstance(x, tuple)): # return 0.5 + 0.5 * cute.math.tanh(0.5 * x, fastmath=True) return 0.5 + 0.5 * tanh(0.5 * x) else: x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x) - tanh_x_half = (tanh(x_half[0]), tanh(x_half[1])) + tanh_x_half = tanh(x_half) return cute.arch.fma_packed_f32x2(tanh_x_half, (0.5, 0.5), (0.5, 0.5)) +@dsl_user_op +def sigmoid(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: + log2_e = math.log2(math.e) + if const_expr(not isinstance(x, tuple)): + exp_neg_x = cute.math.exp2(x * (-log2_e), fastmath=True, loc=loc, ip=ip) + return mlir_math.rcp(exp_neg_x + 1.0, approx=True, ftz=True, loc=loc, ip=ip) + else: + neg_x = cute.arch.mul_packed_f32x2(x, (-log2_e, -log2_e)) + exp_neg_x = ( + cute.math.exp2(neg_x[0], fastmath=True, loc=loc, ip=ip), + cute.math.exp2(neg_x[1], fastmath=True, loc=loc, ip=ip), + ) + denom = cute.arch.add_packed_f32x2(exp_neg_x, (1.0, 1.0)) + return ( + mlir_math.rcp(denom[0], approx=True, ftz=True, loc=loc, ip=ip), + mlir_math.rcp(denom[1], approx=True, ftz=True, loc=loc, ip=ip), + ) + + @dsl_user_op def dsigmoid_from_output(out: Float32, dout: Float32, *, loc=None, ip=None) -> Float32: # return dout * out * (1.0 - out) @@ -54,9 +87,9 @@ def dsigmoid_from_output(out: Float32, dout: Float32, *, loc=None, ip=None) -> F @dsl_user_op def relu(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: if const_expr(not isinstance(x, tuple)): - return cute.arch.fmax(x, Float32(0.0)) + return cutlass.max(x, Float32(0.0), loc=loc, ip=ip) else: - return cute.arch.fmax(x[0], Float32(0.0)), cute.arch.fmax(x[1], Float32(0.0)) + return relu(x[0], loc=loc, ip=ip), relu(x[1], loc=loc, ip=ip) @dsl_user_op @@ -66,7 +99,7 @@ def drelu( ) -> Tuple[F32_or_F32x2, F32_or_F32x2]: if const_expr(not isinstance(x, tuple)): x_pos = Boolean(x > 0) - return dout if x_pos else Float32(0.0), cute.arch.fmax(x, Float32(0.0)) + return dout if x_pos else Float32(0.0), relu(x, loc=loc, ip=ip) else: x0_pos = Boolean(x[0] > 0) x1_pos = Boolean(x[1] > 0) @@ -77,9 +110,9 @@ def drelu( @dsl_user_op def relu_sq(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: if const_expr(not isinstance(x, tuple)): - return cute.arch.fmax(x, Float32(0.0)) * x + return relu(x, loc=loc, ip=ip) * x else: - relu_x = (cute.arch.fmax(x[0], Float32(0.0)), cute.arch.fmax(x[1], Float32(0.0))) + relu_x = relu(x, loc=loc, ip=ip) return cute.arch.mul_packed_f32x2(relu_x, x) @@ -117,19 +150,17 @@ def gelu_tanh_approx(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: sqrt_2_over_pi = math.sqrt(2 / math.pi) # ~0.797885 sqrt_2_over_pi_coeff = 0.044715 * sqrt_2_over_pi # ~0.0356774 if const_expr(not isinstance(x, tuple)): - return 0.5 * ( - x - # Currently cute.math.tanh(x, fastmath=True) generates very slow code - # * (1 + cute.math.tanh(x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * (x * x)), fastmath=True)) - * (1.0 + tanh(x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * (x * x)))) - ) + x_sq = x * x + z = x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * x_sq) + tanh_z = tanh(z) + return 0.5 * (x * tanh_z + x) else: x_sq = cute.arch.mul_packed_f32x2(x, x) x_sq_scaled = cute.arch.fma_packed_f32x2( x_sq, (sqrt_2_over_pi_coeff, sqrt_2_over_pi_coeff), (sqrt_2_over_pi, sqrt_2_over_pi) ) z = cute.arch.mul_packed_f32x2(x, x_sq_scaled) - tanh_z = (tanh(z[0]), tanh(z[1])) + tanh_z = tanh(z) x_tanh_z = cute.arch.fma_packed_f32x2(tanh_z, x, x) return cute.arch.mul_packed_f32x2((0.5, 0.5), x_tanh_z) @@ -162,11 +193,16 @@ def dgelu_tanh_approx( # Compute gradient # sech^2(z) = 1 - tanh^2(z) - sech2_z = 1 - tanh_z * tanh_z + # Keep this as a multiply-add expression so vectorize=True lowers to + # FFMA2 like the explicit F32x2 path; `1.0 - tanh_z * tanh_z` costs + # an extra FADD2/FMUL2 pair per vector. + sech2_z = tanh_z * (-tanh_z) + 1.0 # dz/dx = c1 + 3 * c2 * x^2 dz_dx = sqrt_2_over_pi + sqrt_2_over_pi_coeff_3 * x_sq # d/dx[gelu(x)] = 0.5 * (1 + tanh(z)) + 0.5 * x * sech^2(z) * dz/dx - dgelu = half_tanh_z_plus_one + x * (0.5 * (sech2_z * dz_dx)) + sech2_dz_dx = sech2_z * dz_dx + x_sech2_dz_dx = x * sech2_dz_dx + dgelu = x_sech2_dz_dx * 0.5 + half_tanh_z_plus_one dx = dout * dgelu return dx, gelu_out @@ -177,7 +213,7 @@ def dgelu_tanh_approx( x_sq, (sqrt_2_over_pi_coeff, sqrt_2_over_pi_coeff), (sqrt_2_over_pi, sqrt_2_over_pi) ) z = cute.arch.mul_packed_f32x2(x, x_sq_scaled) - tanh_z = (tanh(z[0]), tanh(z[1])) + tanh_z = tanh(z) half_tanh_z_plus_one = cute.arch.fma_packed_f32x2(tanh_z, (0.5, 0.5), (0.5, 0.5)) gelu_out = cute.arch.mul_packed_f32x2(x, half_tanh_z_plus_one) @@ -236,7 +272,18 @@ def dsoftplus_from_output(out: Float32, dout: Float32, *, loc=None, ip=None) -> @dsl_user_op -def silu(x: F32_or_F32x2, *, already_halved: bool = False, loc=None, ip=None) -> F32_or_F32x2: +def silu(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: + """ + silu(x) = x * sigmoid(x) = x * rcp(1 + exp(-x)). + """ + if const_expr(not isinstance(x, tuple)): + return x * sigmoid(x, loc=loc, ip=ip) + else: + return cute.arch.mul_packed_f32x2(x, sigmoid(x, loc=loc, ip=ip)) + + +@dsl_user_op +def silu_tanh(x: F32_or_F32x2, *, already_halved: bool = False, loc=None, ip=None) -> F32_or_F32x2: """ silu(x) = x * sigmoid(x) = x * (1 + tanh(x / 2)) / 2 = (0.5 * x) * tanh(0.5 * x) + (0.5 * x) This compiles down to 3 SASS instructions: FMUL to get 0.5 * x, MUFU.TANH, and FFMA. @@ -247,10 +294,91 @@ def silu(x: F32_or_F32x2, *, already_halved: bool = False, loc=None, ip=None) -> return x_half * tanh(x_half) + x_half else: x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x) if const_expr(not already_halved) else x - tanh_x_half = (tanh(x_half[0]), tanh(x_half[1])) + tanh_x_half = tanh(x_half) return cute.arch.fma_packed_f32x2(x_half, tanh_x_half, x_half) +@dsl_user_op +def dsilu( + x: F32_or_F32x2, + dout: F32_or_F32x2, + *, + loc=None, + ip=None, +) -> Tuple[F32_or_F32x2, F32_or_F32x2]: + """ + SiLU backward pass: computes d_silu(x) * dout and recomputes silu(x). + + d_silu(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x))). + """ + if const_expr(not isinstance(x, tuple)): + sigmoid_x = sigmoid(x, loc=loc, ip=ip) + silu_x = x * sigmoid_x + # This form vectorizes cleanly with cutlass.range(..., vectorize=True): + # FADD2 (1 - sigmoid_x), FFMA2 (silu_x * tmp + sigmoid_x), FMUL2 (* dout). + d_silu_x_dout = (sigmoid_x + silu_x * (1.0 - sigmoid_x)) * dout + return d_silu_x_dout, silu_x + else: + sigmoid_x = sigmoid(x) + silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_x) + sigmoid_x_minus_silu_x_sigmoid_x = cute.arch.fma_packed_f32x2( + sigmoid_x, (-silu_x[0], -silu_x[1]), sigmoid_x + ) + sigmoid_x_minus_silu_x_sigmoid_x_plus_silu_x = cute.arch.add_packed_f32x2( + sigmoid_x_minus_silu_x_sigmoid_x, silu_x + ) + d_silu_x_dout = cute.arch.mul_packed_f32x2( + sigmoid_x_minus_silu_x_sigmoid_x_plus_silu_x, dout + ) + return d_silu_x_dout, silu_x + + +@dsl_user_op +def dsilu_tanh( + x: F32_or_F32x2, + dout: F32_or_F32x2, + *, + already_halved: bool = False, + loc=None, + ip=None, +) -> Tuple[F32_or_F32x2, F32_or_F32x2]: + """ + SiLU backward using sigmoid(x) = 0.5 * (1 + tanh(0.5 * x)). + """ + if const_expr(not isinstance(x, tuple)): + if const_expr(not already_halved): + x_half = 0.5 * x + tanh_x_half = tanh(x_half) + sigmoid_x = 0.5 * tanh_x_half + 0.5 + silu_x = x_half * tanh_x_half + x_half + else: + tanh_x = tanh(x) + sigmoid_x = 0.5 * tanh_x + 0.5 + silu_x = x * tanh_x + x + d_silu_x_dout = (sigmoid_x + silu_x * (1.0 - sigmoid_x)) * dout + return d_silu_x_dout, silu_x + else: + if const_expr(not already_halved): + x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x) + tanh_x_half = tanh(x_half) + sigmoid_x = cute.arch.fma_packed_f32x2(tanh_x_half, (0.5, 0.5), (0.5, 0.5)) + silu_x = cute.arch.fma_packed_f32x2(x_half, tanh_x_half, x_half) + else: + tanh_x = tanh(x) + sigmoid_x = cute.arch.fma_packed_f32x2(tanh_x, (0.5, 0.5), (0.5, 0.5)) + silu_x = cute.arch.fma_packed_f32x2(x, tanh_x, x) + sigmoid_x_minus_silu_x_sigmoid_x = cute.arch.fma_packed_f32x2( + sigmoid_x, (-silu_x[0], -silu_x[1]), sigmoid_x + ) + sigmoid_x_minus_silu_x_sigmoid_x_plus_silu_x = cute.arch.add_packed_f32x2( + sigmoid_x_minus_silu_x_sigmoid_x, silu_x + ) + d_silu_x_dout = cute.arch.mul_packed_f32x2( + sigmoid_x_minus_silu_x_sigmoid_x_plus_silu_x, dout + ) + return d_silu_x_dout, silu_x + + @dsl_user_op def swiglu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: if const_expr(not isinstance(x, tuple)): @@ -259,13 +387,20 @@ def swiglu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32 return cute.arch.mul_packed_f32x2(silu(x), y) +@dsl_user_op +def swiglu_tanh(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: + if const_expr(not isinstance(x, tuple)): + return silu_tanh(x) * y + else: + return cute.arch.mul_packed_f32x2(silu_tanh(x), y) + + @dsl_user_op def dswiglu( x: F32_or_F32x2, y: F32_or_F32x2, dout: F32_or_F32x2, *, - already_halved: bool = False, loc=None, ip=None, ) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]: @@ -280,15 +415,8 @@ def dswiglu( to use FFMA instead of FADD and FMUL). """ if const_expr(not isinstance(x, tuple)): - # Compute sigmoid(x) using tanh: sigmoid(x) = 0.5 * (1 + tanh(0.5 * x)) - # FMUL, MUFU.TANH, then FFMA - if const_expr(not already_halved): - sigmoid_x = sigmoid(x) - silu_x = x * sigmoid_x # FMUL - else: - tanh_x = tanh(x) # MUFU.TANH - sigmoid_x = 0.5 * tanh_x + 0.5 # FFMA - silu_x = x * tanh_x + x # FFMA + sigmoid_x = sigmoid(x) + silu_x = x * sigmoid_x # FMUL silu_x_dout = silu_x * dout # FMUL # d_silu(x) * dout # = sigmoid_x * (1 + x * (1 - sigmoid_x)) * dout @@ -296,19 +424,65 @@ def dswiglu( # = (sigmoid_x + silu_x * (1 - sigmoid_x)) * dout # = (sigmoid_x + silu_x - silu_x * sigmoid_x) * dout # = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x * dout - d_silu_x_dout = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x_dout # FFMA, FFMA + # This form lets ptxas recover the same two packed FFMA instructions + # as the explicit F32x2 path while reusing silu_x_dout for dy. + d_silu_x_dout = (sigmoid_x + (-silu_x) * sigmoid_x) * dout + silu_x_dout dx = d_silu_x_dout * y # FMUL dy = silu_x_dout swiglu_out = silu_x * y # FMUL - # Overall it's 1 MUFU.TANH, 5 FMUL, 3 FFMA return dx, dy, swiglu_out else: # Compute sigmoid(x) and silu(x) + sigmoid_x = sigmoid(x) + silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_x) + silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout) + # d_silu(x) * dout = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x * dout + sigmoid_x_minus_silu_x_sigmoid_x = cute.arch.fma_packed_f32x2( + sigmoid_x, (-silu_x[0], -silu_x[1]), sigmoid_x + ) + d_silu_x_dout = cute.arch.fma_packed_f32x2( + sigmoid_x_minus_silu_x_sigmoid_x, dout, silu_x_dout + ) + dx = cute.arch.mul_packed_f32x2(d_silu_x_dout, y) + dy = silu_x_dout + swiglu_out = cute.arch.mul_packed_f32x2(silu_x, y) + return dx, dy, swiglu_out + + +@dsl_user_op +def dswiglu_tanh( + x: F32_or_F32x2, + y: F32_or_F32x2, + dout: F32_or_F32x2, + *, + already_halved: bool = False, + loc=None, + ip=None, +) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]: + """ + SwiGLU backward using sigmoid(x) = 0.5 * (1 + tanh(0.5 * x)). + """ + if const_expr(not isinstance(x, tuple)): + if const_expr(not already_halved): + sigmoid_x = sigmoid_tanh(x) + silu_x = x * sigmoid_x # FMUL + else: + tanh_x = tanh(x) + sigmoid_x = 0.5 * tanh_x + 0.5 + silu_x = x * tanh_x + x + silu_x_dout = silu_x * dout + d_silu_x_dout = (sigmoid_x + (-silu_x) * sigmoid_x) * dout + silu_x_dout + dx = d_silu_x_dout * y + dy = silu_x_dout + swiglu_out = silu_x * y + # Overall it's 1 MUFU.TANH, 5 FMUL, 3 FFMA + return dx, dy, swiglu_out + else: if const_expr(not already_halved): - sigmoid_x = sigmoid(x) + sigmoid_x = sigmoid_tanh(x) silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_x) else: - tanh_x = (tanh(x[0]), tanh(x[1])) + tanh_x = tanh(x) sigmoid_x = cute.arch.fma_packed_f32x2(tanh_x, (0.5, 0.5), (0.5, 0.5)) silu_x = cute.arch.fma_packed_f32x2(x, tanh_x, x) silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout) @@ -332,18 +506,31 @@ def swiglu_oai( """The swiglu variant used in gpt-oss, which has a scaling factor on x and bias of 1 to y. https://github.com/openai/gpt-oss/blob/7be9334950053a888e24887a57dac797a17d6e00/gpt_oss/torch/model.py#L249 x * sigmoid(alpha * x) * (y + 1) - Compile down to FMUL, FMUL, TANH, FFMA, FFMA """ - # Compute sigmoid(alpha * x) using tanh: sigmoid(z) = 0.5 * (1 + tanh(z/2)) + if const_expr(not isinstance(x, tuple)): + sigmoid_alpha_x = sigmoid(alpha * x) + silu_x = x * sigmoid_alpha_x + return silu_x * y + silu_x + else: + alpha_x = cute.arch.mul_packed_f32x2((alpha, alpha), x) + sigmoid_alpha_x = sigmoid(alpha_x) + silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_alpha_x) + return cute.arch.fma_packed_f32x2(silu_x, y, silu_x) + + +@dsl_user_op +def swiglu_oai_tanh( + x: F32_or_F32x2, y: F32_or_F32x2, alpha: float = 1.702, *, loc=None, ip=None +) -> F32_or_F32x2: + """Tanh-based swiglu_oai kept for SASS/accuracy comparison.""" if const_expr(not isinstance(x, tuple)): x_half = 0.5 * x - # silu_x = x_half * cute.math.tanh(alpha * x_half, fastmath=True) + x_half silu_x = x_half * tanh(alpha * x_half) + x_half return silu_x * y + silu_x else: x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x) alpha_x_half = cute.arch.mul_packed_f32x2((alpha, alpha), x_half) - tanh_alpha_x_half = (tanh(alpha_x_half[0]), tanh(alpha_x_half[1])) + tanh_alpha_x_half = tanh(alpha_x_half) silu_x = cute.arch.fma_packed_f32x2(x_half, tanh_alpha_x_half, x_half) return cute.arch.fma_packed_f32x2(silu_x, y, silu_x) @@ -361,28 +548,62 @@ def dswiglu_oai( d/dx[x * sigmoid(alpha * x)] = sigmoid(alpha * x) + alpha * x * sigmoid(alpha * x) * (1 - sigmoid(alpha * x)) """ if const_expr(not isinstance(x, tuple)): - # Compute sigmoid(alpha * x) using tanh: sigmoid(z) = 0.5 * (1 + tanh(z/2)) - alpha_x_half = (0.5 * alpha) * x # FMUL - # MUFU.TANH, then FFMA - # sigmoid_alpha_x = 0.5 + 0.5 * cute.math.tanh(alpha_x_half, fastmath=True) + sigmoid_alpha_x = sigmoid(alpha * x) + silu_x = x * sigmoid_alpha_x + silu_x_dout = silu_x * dout + # Keep this as two multiply-add expressions. With vectorize=True this + # matches the explicit F32x2 path; spelling it as (1 - sigmoid) costs + # an extra FADD2/FMUL2 pair per vector. + silu_x_minus_product = silu_x * (-sigmoid_alpha_x) + silu_x + sigmoid_plus_alpha_diff = silu_x_minus_product * alpha + sigmoid_alpha_x + d_silu_x_dout = sigmoid_plus_alpha_diff * dout + dx = d_silu_x_dout * y + d_silu_x_dout + dy = silu_x_dout + swiglu_out = silu_x * y + silu_x + return dx, dy, swiglu_out + else: + alpha_x = cute.arch.mul_packed_f32x2((alpha, alpha), x) + sigmoid_alpha_x = sigmoid(alpha_x) + silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_alpha_x) + silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout) + silu_x_minus_product = cute.arch.fma_packed_f32x2( + silu_x, (-sigmoid_alpha_x[0], -sigmoid_alpha_x[1]), silu_x + ) + sigmoid_plus_alpha_diff = cute.arch.fma_packed_f32x2( + (alpha, alpha), silu_x_minus_product, sigmoid_alpha_x + ) + d_silu_x_dout = cute.arch.mul_packed_f32x2(sigmoid_plus_alpha_diff, dout) + dx = cute.arch.fma_packed_f32x2(d_silu_x_dout, y, d_silu_x_dout) + dy = silu_x_dout + swiglu_out = cute.arch.fma_packed_f32x2(silu_x, y, silu_x) + return dx, dy, swiglu_out + + +@dsl_user_op +def dswiglu_oai_tanh( + x: F32_or_F32x2, y: F32_or_F32x2, dout: F32_or_F32x2, alpha: float = 1.702, *, loc=None, ip=None +) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]: + """Tanh-based dswiglu_oai kept for SASS/accuracy comparison.""" + if const_expr(not isinstance(x, tuple)): + alpha_x_half = (0.5 * alpha) * x sigmoid_alpha_x = 0.5 + 0.5 * tanh(alpha_x_half) - silu_x = x * sigmoid_alpha_x # FMUL - silu_x_dout = silu_x * dout # FMUL - # FFMA, FFMA, FMUL - d_silu_x_dout = (sigmoid_alpha_x + alpha * (silu_x - silu_x * sigmoid_alpha_x)) * dout - dx = d_silu_x_dout * y + d_silu_x_dout # FFMA, instead of multiply by y + 1 + silu_x = x * sigmoid_alpha_x + silu_x_dout = silu_x * dout + # Same spelling as dswiglu_oai: this preserves the packed FFMA2 chain + # under cutlass.range(..., vectorize=True). + silu_x_minus_product = silu_x * (-sigmoid_alpha_x) + silu_x + sigmoid_plus_alpha_diff = silu_x_minus_product * alpha + sigmoid_alpha_x + d_silu_x_dout = sigmoid_plus_alpha_diff * dout + dx = d_silu_x_dout * y + d_silu_x_dout dy = silu_x_dout - swiglu_out = silu_x * y + silu_x # FFMA, instead of multiply by y + 1 - # Overall it's 1 MUFU.TANH, 4 FMUL, 5 FFMA + swiglu_out = silu_x * y + silu_x return dx, dy, swiglu_out else: - # Compute sigmoid(alpha * x) alpha_x_half = cute.arch.mul_packed_f32x2(((0.5 * alpha), (0.5 * alpha)), x) - tanh_alpha_x_half = (tanh(alpha_x_half[0]), tanh(alpha_x_half[1])) + tanh_alpha_x_half = tanh(alpha_x_half) sigmoid_alpha_x = cute.arch.fma_packed_f32x2(tanh_alpha_x_half, (0.5, 0.5), (0.5, 0.5)) silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_alpha_x) silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout) - # d_silu_x_dout = (sigmoid_alpha_x + alpha * (silu_x - silu_x * sigmoid_alpha_x)) * dout silu_x_minus_product = cute.arch.fma_packed_f32x2( silu_x, (-sigmoid_alpha_x[0], -sigmoid_alpha_x[1]), silu_x ) @@ -400,10 +621,9 @@ def dswiglu_oai( def glu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: """GLU: Gated Linear Unit glu(x, y) = sigmoid(x) * y - Using tanh to compute sigmoid: sigmoid(x) = 0.5 * (1 + tanh(x/2)) """ if const_expr(not isinstance(x, tuple)): - sigmoid_x = sigmoid(x) # FMUL, MUFU.TANH, then FFMA + sigmoid_x = sigmoid(x) return sigmoid_x * y # FMUL else: sigmoid_x = sigmoid(x) @@ -423,8 +643,7 @@ def dglu( - glu_out = sigmoid(x) * y """ if const_expr(not isinstance(x, tuple)): - # Compute sigmoid(x) using tanh: sigmoid(x) = 0.5 * (1 + tanh(x/2)) - sigmoid_x = sigmoid(x) # FMUL, MUFU.TANH, then FFMA + sigmoid_x = sigmoid(x) sigmoid_x_dout = sigmoid_x * dout # FMUL glu_out = sigmoid_x * y # FMUL # dx = y * sigmoid(x) * (1 - sigmoid(x)) * dout @@ -452,9 +671,9 @@ def reglu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x reglu(x, y) = relu(x) * y = max(x, 0) * y """ if const_expr(not isinstance(x, tuple)): - return cute.arch.fmax(x, Float32(0.0)) * y + return relu(x, loc=loc, ip=ip) * y else: - relu_x = relu(x) + relu_x = relu(x, loc=loc, ip=ip) return cute.arch.mul_packed_f32x2(relu_x, y) @@ -473,15 +692,16 @@ def dreglu( """ if const_expr(not isinstance(x, tuple)): x_pos = Boolean(x > 0) - relu_x = cute.arch.fmax(x, Float32(0.0)) - dx = (dout * y) if x_pos else Float32(0.0) + relu_x = relu(x, loc=loc, ip=ip) + dout_y = dout * y + dx = dout_y if x_pos else Float32(0.0) dy = dout * relu_x reglu_out = relu_x * y return dx, dy, reglu_out else: x0_pos = Boolean(x[0] > 0) x1_pos = Boolean(x[1] > 0) - relu_x = relu(x) + relu_x = relu(x, loc=loc, ip=ip) dout_y = cute.arch.mul_packed_f32x2(dout, y) dx = ((dout_y[0] if x0_pos else Float32(0.0)), (dout_y[1] if x1_pos else Float32(0.0))) dy = cute.arch.mul_packed_f32x2(dout, relu_x) @@ -538,21 +758,28 @@ def dgeglu( act_fn_map = { None: None, "silu": silu, + "silu-tanh": silu_tanh, "relu": relu, "relu_sq": relu_sq, "gelu_tanh_approx": gelu_tanh_approx, + "tanh": tanh, } dact_fn_map = { None: None, + "silu": dsilu, + "silu-tanh": dsilu_tanh, "relu": drelu, "relu_sq": drelu_sq, "gelu_tanh_approx": dgelu_tanh_approx, + "tanh": dtanh, } gate_fn_map = { "swiglu": swiglu, + "swiglu-tanh": swiglu_tanh, "swiglu_oai": swiglu_oai, + "swiglu_oai-tanh": swiglu_oai_tanh, "reglu": reglu, "geglu": geglu, "glu": glu, @@ -560,7 +787,9 @@ gate_fn_map = { dgate_fn_map = { "swiglu": dswiglu, + "swiglu-tanh": dswiglu_tanh, "swiglu_oai": dswiglu_oai, + "swiglu_oai-tanh": dswiglu_oai_tanh, "reglu": dreglu, "geglu": dgeglu, "glu": dglu, diff --git a/build/torch-cuda/quack/autotuner.py b/build/torch-cuda/quack/autotuner.py index 5ca65c9e07dfc42075e7c9ff5b0611e5ec251779..f012efa4c9d1b2f740a23a7b89a45afb9fba1724 100644 --- a/build/torch-cuda/quack/autotuner.py +++ b/build/torch-cuda/quack/autotuner.py @@ -4,6 +4,7 @@ from __future__ import annotations import builtins import os +import sys import time import inspect import base64 @@ -12,6 +13,11 @@ import json from pathlib import Path from functools import cached_property, partial from typing import Dict, Tuple, List, Optional, Any +from .bench.bench_utils import ( + _bench_cuda_graph_l2_rotate, + _clone_l2_rotate_inputs, + _pick_l2_rotate_count, +) import torch from torch import Tensor @@ -25,29 +31,6 @@ PACKAGE_NAME = "quack" VERSION = __version__ -def _get_current_cuda_device() -> str | None: - """Return the physical CUDA device identifier for the current process. - - Maps the logical ``torch.cuda.current_device()`` index through - ``CUDA_VISIBLE_DEVICES`` (if set) so the result is valid as a - standalone ``CUDA_VISIBLE_DEVICES`` value (handles integer IDs, - GPU UUIDs, and MIG IDs). - - Returns ``None`` if CUDA is not initialized or the device cannot - be determined. - """ - if not (torch.cuda.is_available() and torch.cuda.is_initialized()): - return None - logical_device = torch.cuda.current_device() - parent_visible = os.environ.get("CUDA_VISIBLE_DEVICES") - if parent_visible is not None: - visible_devices = [d.strip() for d in parent_visible.split(",")] - if logical_device < len(visible_devices): - return visible_devices[logical_device] - return None - return str(logical_device) - - def get_home_dir(): return os.getenv(f"{PACKAGE_NAME.upper()}_HOME", Path.home()) @@ -75,6 +58,14 @@ def _base32(key): return base64.b32encode(bytes.fromhex(key)).decode("utf-8").rstrip("=") +#: How long a deferred config may wait on its pool compile before the bench +#: loop stops trusting the pool and benches it with the pool suppressed +#: (in-process compile). Guards against a wedged worker / a foreign flock +#: holder that never produces the .o; without it a permanently-"pending" +#: sha would rotate forever. Tests override this. +_POOL_WEDGE_TIMEOUT_S = 300.0 + + def _gpu_warmup(duration_ms=200): """Saturate the GPU to reach thermal steady-state before benchmarking. @@ -91,6 +82,21 @@ def _gpu_warmup(duration_ms=200): torch.cuda.synchronize() +# --------------------------------------------------------------------------- +# Candidate-config compilation +# +# There is no separate precompile phase: the bench loop in ``benchmark()`` +# (inside ``Autotuner.__call__``) runs under ``pool_scope()`` from +# quack.cache.async_compile. A config whose kernel misses the .o cache +# raises ``CompilePending`` from jit_cache after shipping the pickled +# ``_compile_*`` key to a CPU worker; the loop rotates that config to the +# back and benches whichever config is ready. Total wall stays +# max(parallel_compile, serial_bench), key discovery uses the real tensors +# in-process, and workers never launch kernels (they call the tensor-free +# ``_compile_*`` functions directly). +# --------------------------------------------------------------------------- + + class Autotuner: def __init__( self, @@ -163,146 +169,6 @@ class Autotuner: return partial(triton.testing.do_bench, warmup=5, rep=25) return self._do_bench - def _precompile(self, *args, configs, **kwargs): - """Pre-compile all configs in parallel subprocesses to populate .o cache. - - cute.compile() is not thread-safe (MLIR thread-local state) and fork after - CUDA init causes segfaults. So we spawn persistent subprocess workers: each - has its own CUDA context, creates FakeTensors matching the parent's tensor - metadata, and compiles with COMPILE_ONLY=True. Workers stay alive to amortize - import overhead across multiple configs. The parent then loads instantly from - the .o cache during benchmarking. - """ - from .cache_utils import CACHE_ENABLED - - if not CACHE_ENABLED: - return - - max_workers = min(len(configs), int(os.getenv("QUACK_COMPILE_WORKERS", "8"))) - if max_workers <= 1: - return - - # Quick check: compile first config in-process. If it loads from .o cache - # (<0.5s), the rest are likely cached too — skip spawning workers. - t_check = time.time() - try: - current = dict(kwargs, **configs[0].all_kwargs()) - self.fn(*args, **current) - except Exception: - pass - if time.time() - t_check < 0.5: - return - - verbose = os.getenv(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1" - if verbose: - print(f"Pre-compiling {len(configs)} configs with {max_workers} workers") - t0 = time.time() - - import pickle - import struct - import subprocess - import sys - - def _send(stream, msg): - data = pickle.dumps(msg) - stream.write(struct.pack(" sha + attempts = {} # id(config) -> int + deadline = {} # id(config) -> wedge deadline + spins = 0 + while queue: + config = queue.popleft() + sha = awaiting.get(id(config)) + wedged = sha is not None and time.monotonic() > deadline[id(config)] + if sha is not None and not wedged: + state, _ = pool.poll(sha) + if state == "pending": + queue.append(config) + spins += 1 + if spins >= len(queue): + time.sleep(0.05) + spins = 0 + continue + spins = 0 + n = attempts.get(id(config), 0) + 1 + attempts[id(config)] = n + try: + if wedged or n > _MAX_ATTEMPTS: + # Wedged pool: compile in-process so + # the sweep always terminates. + with suppress_pool(): + timings[config] = self._bench( + *args, config=config, **kwargs + ) + else: + timings[config] = self._bench( + *args, config=config, **kwargs + ) + except CompilePending as e: + awaiting[id(config)] = e.sha + deadline.setdefault( + id(config), + time.monotonic() + _POOL_WEDGE_TIMEOUT_S, + ) + queue.append(config) + finally: + # Free L2-cold sets before persisting the cache so the + # user's subsequent .fn(...) call has full HBM. + self._l2_cold_arg_sets = None + self._l2_cold_kwarg_sets = None bench_end = time.time() - if os.getenv(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1": + if verbose: for config, time_ in timings.items(): print(f"[{config}] -> {time_[0]:.3f}ms") + # Surface bench failures (configs returning inf timings) + # so smem-overflow / launch errors aren't silently masked. + n_failed = sum(1 for t in timings.values() if t[0] == float("inf")) + if n_failed: + print( + f"quack autotune: {n_failed}/{len(timings)} configs " + f"failed for {self.fn.__name__}{key}; " + f"set {PACKAGE_NAME.upper()}_PRINT_AUTOTUNING=1 for details", + file=sys.stderr, + ) self.bench_time = bench_end - bench_start self.cache[key] = builtins.min(timings, key=timings.get) self.configs_timings = timings diff --git a/build/torch-cuda/quack/bench/__init__.py b/build/torch-cuda/quack/bench/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch-cuda/quack/bench/bench_utils.py b/build/torch-cuda/quack/bench/bench_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..129d37121522461f3abf02c2c296535091f4a3c1 --- /dev/null +++ b/build/torch-cuda/quack/bench/bench_utils.py @@ -0,0 +1,202 @@ +"""Shared helpers for triton perf_report-based benchmarks.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +import torch +from torch import Tensor + +if TYPE_CHECKING: + import pandas as pd + from triton.testing import Benchmark + + +def run_and_print(mark, save_path=None): + """Run a triton ``Mark`` (from ``perf_report``) and print/save results. + + Each runner is expected to return a ``dict[str, Any]`` mapping stat name to + value, e.g. ``{"ms": 0.123, "GB/s": 1234}``. All providers in a benchmark + must return the same set of keys. Values are written through unchanged -- + rounding/formatting is the caller's responsibility. + + Output columns are ``x_names + [f"{line_name} ({stat})" for ...]``. + """ + benchmarks = mark.benchmarks if isinstance(mark.benchmarks, list) else [mark.benchmarks] + for bench in benchmarks: + df = _run_one(mark.fn, bench) + print(bench.plot_name + ":") + print(df.to_string()) + if save_path: + os.makedirs(save_path, exist_ok=True) + df.to_csv(os.path.join(save_path, f"{bench.plot_name}.csv"), index=False) + + +def _run_one(fn, bench: Benchmark) -> pd.DataFrame: + try: + import pandas as pd + except ImportError as e: + raise ImportError( + "pandas is required to format benchmark results. " + "Install it with `pip install pandas` or `pip install -e '.[bench]'`." + ) from e + + x_names = list(bench.x_names) + rows = [] + stat_keys = None # locked in from the first runner result + for x in bench.x_vals: + if not isinstance(x, (list, tuple)): + x = [x] * len(x_names) + x_args = dict(zip(x_names, x)) + row = list(x) + for line_val in bench.line_vals: + stats = fn(**x_args, **{bench.line_arg: line_val}, **bench.args) + if not isinstance(stats, dict): + raise TypeError(f"runner must return dict[str, Any], got {type(stats).__name__}") + if stat_keys is None: + stat_keys = list(stats.keys()) + elif list(stats.keys()) != stat_keys: + raise ValueError(f"runner returned keys {list(stats.keys())}, expected {stat_keys}") + row.extend(stats[k] for k in stat_keys) + rows.append(row) + cols = list(x_names) + [ + f"{name} ({stat})" for name in bench.line_names for stat in (stat_keys or []) + ] + return pd.DataFrame(rows, columns=cols) + + +def _bench_cuda_graph_l2_rotate( + fn, + arg_sets, + kwarg_sets, + extra_kwargs, + warmup_target_ms: float = 200.0, + n_timed_calls: int = 200, + quantiles=None, +): + """L2-cold single-replay CUDA-graph benchmark. + + Warmup is time-based: probe a single kernel launch to estimate the + per-call cost, then iterate round-robin over the pre-cloned + ``(arg_sets[i], kwarg_sets[i])`` pairs as a plain Python loop for + ``warmup_target_ms`` of wall-clock GPU work. Heavy pipelined configs + (TMA + smem_stages=3) need enough warmup to drain the pipeline-fill + phase or the timed window catches them artificially fast - a fixed + count of warmup launches underwarms heavy configs while overpaying + for cheap ones. + + The timed window is a single ``graph.replay()`` of a captured CUDA + graph whose body records ``n_timed_calls`` round-robin invocations - + no Python loop, no per-launch CPU overhead - so the measurement is + just GPU work / total recorded calls. + + Round-robin over fresh tensor sets defeats the L2-resident caching + that inflates short-kernel timing under ``triton.testing.do_bench`` + (which calls the kernel on the same single tensor each iteration). + For cache-cold production workloads the round-robin number predicts + real-world latency; the L2-hot number favours wider layouts / deeper + smem stages that don't actually win once data has to come from HBM. + + ``fn`` is called as ``fn(*arg_sets[i], **kwarg_sets[i], **extra_kwargs)`` + once per recorded launch. ``extra_kwargs`` holds the per-config kwargs + (e.g. ``{"config": }``) which don't need cloning; + keys in ``extra_kwargs`` must not overlap with ``kwarg_sets[i]``. + Returns ms/call, or a ``len(quantiles)``-list replicating that value + when ``quantiles`` is provided (for API parity with + ``triton.testing.do_bench``). + """ + n_sets = len(arg_sets) + # Round timed-call count to a multiple of n_sets for even L2 turnover. + rounds_timed = max(1, n_timed_calls // n_sets) + total_timed_calls = rounds_timed * n_sets + + # A few priming launches so the probe doesn't catch first-launch + # driver / kernel-load overhead. + for _ in range(3): + fn(*arg_sets[0], **kwarg_sets[0], **extra_kwargs) + torch.cuda.synchronize() + + # Probe a single launch to estimate per-call ms; size the warmup loop + # to hit ``warmup_target_ms`` of GPU work regardless of kernel cost. + probe_start = torch.cuda.Event(enable_timing=True) + probe_end = torch.cuda.Event(enable_timing=True) + probe_start.record() + fn(*arg_sets[0], **kwarg_sets[0], **extra_kwargs) + probe_end.record() + torch.cuda.synchronize() + est_ms = max(probe_start.elapsed_time(probe_end), 1e-3) + n_warmup_calls = max(50, int(warmup_target_ms / est_ms)) + + # Warmup: plain Python loop over rotating sets, no graph capture. + for i in range(n_warmup_calls): + idx = i % n_sets + fn(*arg_sets[idx], **kwarg_sets[idx], **extra_kwargs) + torch.cuda.synchronize() + + # Capture timed graph: a single replay covers all timed kernel launches. + timed_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(timed_graph): + for _ in range(rounds_timed): + for i in range(n_sets): + fn(*arg_sets[i], **kwarg_sets[i], **extra_kwargs) + torch.cuda.synchronize() + + start_evt = torch.cuda.Event(enable_timing=True) + end_evt = torch.cuda.Event(enable_timing=True) + start_evt.record() + timed_graph.replay() + end_evt.record() + torch.cuda.synchronize() + ms = start_evt.elapsed_time(end_evt) / total_timed_calls + + if quantiles: + return [ms for _ in quantiles] + return ms + + +def _clone_l2_rotate_inputs(args, kwargs, n_buffers: int): + """Clone tensor args AND tensor kwargs ``n_buffers`` times. + + Returns ``(arg_sets, kwarg_sets)``: ``arg_sets[i]`` is a tuple matching + ``args``' positional shape with tensors cloned to fresh memory; + ``kwarg_sets[i]`` is a dict matching ``kwargs``' keys with tensor values + cloned. Non-tensor values (ints, strings, None, dataclasses, etc.) are + shared across all sets (no clone). + + Both args and kwargs are cloned so that every recorded launch in the + L2-cold round-robin CUDA graph touches distinct GMEM addresses, + including for write-target kwargs like ``dw_partial`` / ``dx``. + """ + arg_sets = [] + kwarg_sets = [] + for _ in range(n_buffers): + arg_sets.append(tuple(a.clone() if isinstance(a, Tensor) else a for a in args)) + kwarg_sets.append({k: v.clone() if isinstance(v, Tensor) else v for k, v in kwargs.items()}) + return arg_sets, kwarg_sets + + +def _pick_l2_rotate_count( + args, kwargs, target_ratio: int = 3, min_buffers: int = 4, max_buffers: int = 16 +): + """Pick ``n_bufs`` so cloned input bytes per round exceed + ``target_ratio * L2_size`` (defeats L2 reuse), capped by HBM headroom and + [min_buffers, max_buffers]. Counts tensor bytes across both ``args`` and + ``kwargs`` so write-target kwargs (``dw_partial``, ``dx``, etc.) are + included in the L2-turnover calculation. + """ + if not torch.cuda.is_available(): + return min_buffers + tensor_bytes = sum(a.numel() * a.element_size() for a in args if isinstance(a, Tensor)) + sum( + v.numel() * v.element_size() for v in kwargs.values() if isinstance(v, Tensor) + ) + if tensor_bytes == 0: + return min_buffers + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + l2_size = props.L2_cache_size + n_by_l2 = (target_ratio * l2_size + tensor_bytes - 1) // tensor_bytes + free_bytes, _ = torch.cuda.mem_get_info() + # Leave half of free memory headroom for the kernel's own scratch + the + # user's other allocations. + n_by_mem = max(1, int(free_bytes * 0.5) // tensor_bytes) + return max(min_buffers, min(max_buffers, min(n_by_l2, n_by_mem))) diff --git a/build/torch-cuda/quack/blockscaled/__init__.py b/build/torch-cuda/quack/blockscaled/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..680d1f84740b5e684d9750ae0fa374fc9a35d910 --- /dev/null +++ b/build/torch-cuda/quack/blockscaled/__init__.py @@ -0,0 +1,31 @@ +# Copyright (c) 2026, Tri Dao. +"""Blockscaled (MXFP8 / MXFP4 / NVFP4) GEMM support. + +- :mod:`quack.blockscaled.quantize` — pure-PyTorch quantizers (ported from + torchao) with torch.compile'd fast paths. +- :mod:`quack.blockscaled.utils` — scale-factor packing/unpacking, operand + builders for tests/benchmarks, and the kernel-level compile path. + +The GEMM entry points live in :mod:`quack.gemm_interface` (pass ``(A, SFA)`` / +``(B, SFB)`` tuples). +""" + +from .quantize import ( # noqa: F401 + nvfp4_per_tensor_scale, + to_blocked, + to_mx, + to_mx_compiled, + to_mxfp4, + to_mxfp4_compiled, + to_nvfp4, + to_nvfp4_compiled, +) +from .utils import ( # noqa: F401 + BLOCKSCALED_FORMATS, + blockscaled_gemm_reference, + blockscaled_quantize, + dequant_operand, + pack_scale_2d_to_blocked_contig, + scale_blocked_for_cublas, + unpack_scale_blocked_to_2d, +) diff --git a/build/torch-cuda/quack/mx_utils.py b/build/torch-cuda/quack/blockscaled/quantize.py similarity index 100% rename from build/torch-cuda/quack/mx_utils.py rename to build/torch-cuda/quack/blockscaled/quantize.py diff --git a/build/torch-cuda/quack/blockscaled_gemm_utils.py b/build/torch-cuda/quack/blockscaled/utils.py similarity index 71% rename from build/torch-cuda/quack/blockscaled_gemm_utils.py rename to build/torch-cuda/quack/blockscaled/utils.py index cdd682fd7d409bdcb32b9c3e0a07890ad4af6561..4a4131ab397ad3aad3d2760d01ef11ff6ade0136 100644 --- a/build/torch-cuda/quack/blockscaled_gemm_utils.py +++ b/build/torch-cuda/quack/blockscaled/utils.py @@ -9,16 +9,16 @@ import torch import cutlass import cutlass.cute as cute -from .compile_utils import make_fake_tensor as fake_tensor -from .cute_dsl_utils import get_device_capacity, get_max_active_clusters -from .gemm_default_epi import GemmDefaultSm100 -from .gemm_tvm_ffi_utils import div_for_dtype, make_scheduler_args -from .mx_utils import ( +from ..compile_utils import make_fake_tensor as fake_tensor +from ..cute_dsl_utils import get_device_capacity, get_max_active_clusters +from ..gemm_default_epi import GemmDefaultSm100 +from ..gemm_tvm_ffi_utils import div_for_dtype, make_scheduler_args +from .quantize import ( to_mx_compiled, to_mxfp4_compiled, to_nvfp4_compiled, ) -from .varlen_utils import VarlenArguments +from ..varlen_utils import VarlenArguments TORCH_DTYPE_MAP = { @@ -188,7 +188,7 @@ def create_blockscaled_operand_tensor( def _pack_blockscaled_scales(ref_blocks: torch.Tensor) -> torch.Tensor: - """Rearrange (mn, sf_k, l) scales into the (l, rm, rk, 512) blocked layout.""" + """Rearrange (mn, sf_k, l) scales into the (l, rm, rk, 32, 4, 4) blocked layout.""" mn, sf_k, l = ref_blocks.shape rm = ceil_div(mn, 128) rk = ceil_div(sf_k, 4) @@ -205,7 +205,7 @@ def _pack_blockscaled_scales(ref_blocks: torch.Tensor) -> torch.Tensor: k_idx[None, :, None] // 4, l_idx[None, None, :], ] = ref_blocks - return packed_6d.view(l, rm, rk, 512) + return packed_6d def create_blockscaled_scale_tensor( @@ -237,10 +237,10 @@ def create_blockscaled_scale_tensor( def pack_scale_2d_to_blocked_contig(scale_2d: torch.Tensor) -> torch.Tensor: """Rearrange a (l, mn, sf_k) or (mn, sf_k) e8m0 scale tensor into the - contiguous (l, rm, rk, 512) blocked layout shared by the quack kernel and - cuBLAS's block-scaling. Each 512 B inner block holds one 128 MN × 4 K - swizzled tile. Pads `mn` to a multiple of 128 and `sf_k` to a multiple of - 4 with zeros.""" + contiguous (l, rm, rk, 32, 4, 4) blocked layout shared by the quack kernel + and cuBLAS's block-scaling. Each inner (32, 4, 4) atom (512 B) holds one + 128 MN × 4 K swizzled tile. Pads `mn` to a multiple of 128 and `sf_k` to a + multiple of 4 with zeros.""" if scale_2d.dim() == 2: scale_2d = scale_2d.unsqueeze(0) assert scale_2d.dim() == 3, f"expected (l, mn, sf_k), got shape {tuple(scale_2d.shape)}" @@ -260,22 +260,92 @@ def pack_scale_2d_to_blocked_contig(scale_2d: torch.Tensor) -> torch.Tensor: blocks = padded.view(l, rm, 128, rk, 4).permute(0, 1, 3, 2, 4) # split 128 into (4 outer, 32 inner), then swap to (32, 4) blocks = blocks.reshape(l, rm, rk, 4, 32, 4).transpose(3, 4).contiguous() - return blocks.view(l, rm, rk, 512).view(orig_dtype) + return blocks.view(orig_dtype) + + +def unpack_scale_blocked_to_2d(blocked: torch.Tensor, mn: int, sf_k: int) -> torch.Tensor: + """Unswizzle (l, rm, rk, 32, 4, 4) blocked scale factors to (l, mn, sf_k).""" + l, rm, rk = blocked.shape[:3] + assert tuple(blocked.shape[3:]) == (32, 4, 4) + orig_dtype = blocked.dtype + u8 = blocked.view(torch.uint8) + # (32=m%32, 4=m//32, 4=k%4) -> (4, 32, 4) -> (l, rm, rk, 128, 4) -> (l, mn_pad, sf_k_pad) + u8 = u8.transpose(3, 4).reshape(l, rm, rk, 128, 4) + u8 = u8.permute(0, 1, 3, 2, 4).reshape(l, rm * 128, rk * 4) + return u8[:, :mn, :sf_k].contiguous().view(orig_dtype) + + +def dequant_operand(x: torch.Tensor) -> torch.Tensor: + """Dequantize an operand tensor to float32 values (without scale factors). + + fp8 tensors convert directly; ``float4_e2m1fn_x2`` tensors unpack two codes + per byte (low nibble = even K, high nibble = odd K), doubling the last dim. + """ + if x.dtype == torch.float4_e2m1fn_x2: + u8 = x.view(torch.uint8) + lo = _fp4_unpacked_to_value(u8 & 0x0F) + hi = _fp4_unpacked_to_value((u8 >> 4) & 0x0F) + return torch.stack([lo, hi], dim=-1).reshape(*x.shape[:-1], x.shape[-1] * 2) + return x.float() + + +BLOCKSCALED_FORMATS = { + # format: (torch operand dtype, torch SF dtype, sf_vec_size) + "mxfp8": (torch.float8_e4m3fn, torch.float8_e8m0fnu, 32), + "mxfp4": (torch.float4_e2m1fn_x2, torch.float8_e8m0fnu, 32), + "nvfp4": (torch.float4_e2m1fn_x2, torch.float8_e4m3fn, 16), +} + + +def blockscaled_quantize( + x: torch.Tensor, format: str = "mxfp8", per_tensor_scale: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a (M, K) or (L, M, K) bf16/fp32 tensor along K for blockscaled GEMM. + + Returns ``(q, sf)`` ready to pass as an ``(A, SFA)`` / ``(B, SFB)`` tuple to + :func:`quack.gemm_interface.gemm`: + q: same leading shape as ``x``; fp8 for mxfp8 (M, K), packed fp4x2 for + mxfp4/nvfp4 (M, K/2), K-contiguous. + sf: blocked scale factors, (rm, rk, 32, 4, 4) or (L, rm, rk, 32, 4, 4). + For nvfp4, ``per_tensor_scale`` (scalar fp32) folds the global scale; pass the + product of A's and B's per-tensor scales as ``alpha`` to the GEMM. + """ + from .quantize import to_mx_compiled, to_mxfp4_compiled, to_nvfp4_compiled + + assert format in BLOCKSCALED_FORMATS, f"unknown blockscaled format: {format}" + q_dtype, sf_dtype, sf_vec = BLOCKSCALED_FORMATS[format] + assert x.shape[-1] % sf_vec == 0, f"K ({x.shape[-1]}) must be divisible by {sf_vec}" + batched = x.ndim == 3 + l, mn, k = x.shape if batched else (1, *x.shape) + x_flat = x.reshape(l * mn, k) + if format == "mxfp8": + q, sc = to_mx_compiled(x_flat, sf_vec) + elif format == "mxfp4": + q, sc = to_mxfp4_compiled(x_flat, sf_vec) + else: + q, sc, _ = to_nvfp4_compiled(x_flat, sf_vec, per_tensor_scale) + q = q.view(torch.uint8).view(q_dtype) if q_dtype == torch.float4_e2m1fn_x2 else q + q = q.reshape(*x.shape[:-1], -1) + sf = pack_scale_2d_to_blocked_contig(sc.view(l, mn, k // sf_vec)) + return q, sf if batched else sf.squeeze(0) def scale_view_for_kernel(scale_contig: torch.Tensor, mn: int, sf_k: int, l: int) -> torch.Tensor: - """Validate a (l, rm, rk, 512) scale tensor and return it unchanged. - Only the innermost 512-B tile must be contiguous (stride 1, size 512); - outer (L, rm, rk) strides are free — the kernel reads them from the - passed tensor. This lets callers pass a slice/view of a larger buffer - with no extra copy. Works for both E8M0 (MX) and E4M3 (NVFP4).""" + """Validate a (l, rm, rk, 32, 4, 4) scale tensor and return it unchanged. + Only the innermost (32, 4, 4) atom (one 512 B tile) must be contiguous + (strides (16, 4, 1)); outer (L, rm, rk) strides are free — the kernel + reads them from the passed tensor. This lets callers pass a slice/view of + a larger buffer with no extra copy. Works for both E8M0 (MX) and E4M3 + (NVFP4).""" rm = ceil_div(mn, 128) rk = ceil_div(sf_k, 4) - assert scale_contig.shape == (l, rm, rk, 512), ( - f"expected (l, rm, rk, 512) = ({l}, {rm}, {rk}, 512), got {tuple(scale_contig.shape)}" + assert scale_contig.shape == (l, rm, rk, 32, 4, 4), ( + f"expected (l, rm, rk, 32, 4, 4) = ({l}, {rm}, {rk}, 32, 4, 4), " + f"got {tuple(scale_contig.shape)}" ) - assert scale_contig.stride(-1) == 1, ( - f"innermost 512-B dim must be unit-stride, got stride {scale_contig.stride(-1)}" + assert scale_contig.stride()[-3:] == (16, 4, 1), ( + f"inner (32, 4, 4) atom must be contiguous with strides (16, 4, 1), " + f"got {scale_contig.stride()[-3:]}" ) return scale_contig @@ -283,9 +353,9 @@ def scale_view_for_kernel(scale_contig: torch.Tensor, mn: int, sf_k: int, l: int def scale_blocked_for_cublas( scale_contig: torch.Tensor, mn: int, sf_k: int, l_idx: int = 0 ) -> torch.Tensor: - """Flatten a (l, rm, rk, 512) scale tensor to the 1D swizzled layout + """Flatten a (l, rm, rk, 32, 4, 4) scale tensor to the 1D swizzled layout torch._scaled_mm expects. Uses a single l slice.""" - assert scale_contig.is_contiguous() and scale_contig.dim() == 4 + assert scale_contig.is_contiguous() and scale_contig.dim() == 6 return scale_contig[l_idx].reshape(-1) @@ -328,8 +398,8 @@ def create_blockscaled_operand_quantized( ref: (mn, k, l) float32 dequantized reference q_mkl: (mn, k, l) operand tensor in the layout the quack kernel consumes (float8_e4m3fn for fp8 formats; int8 with packed nibbles for fp4) - scale_contig: (l, rm, rk, 512) contiguous scale storage. Each 512 B - inner block is one 128 MN × 4 K swizzled tile. Byte layout matches + scale_contig: (l, rm, rk, 32, 4, 4) contiguous scale storage. Each inner + (32, 4, 4) atom (512 B) is one 128 MN × 4 K swizzled tile. Byte layout matches cuBLAS `to_blocked`. Pass directly to the quack kernel, or use `scale_blocked_for_cublas` for cuBLAS. """ @@ -403,7 +473,7 @@ def create_blockscaled_varlen_m_operands( """Generate bf16 randn + quantize for a varlen_m blockscaled GEMM. Per-expert seqlens may be arbitrary (not required to be multiples of 128). - SF is stored in dQaccum-style padded format: each expert `i`'s scales + SF is stored with tile-aligned per-batch padding: each expert `i`'s scales occupy `ceildiv(m_i, 128) * 128` rows at offset `(cu_seqlens_m[i] + i * 128) // 128 * 128` in the padded scale buffer. The kernel decodes via `VarlenManager.offset_batch_SFA` which applies the @@ -414,10 +484,14 @@ def create_blockscaled_varlen_m_operands( b_ref: (num_experts, n, k) fp32 dequantized qa: (total_m, k) 2D K-major quantized operand (fp8) or (total_m, k/2) (fp4) qb: (n, k, num_experts) 3D K-major quantized operand (fp8) or (n, k/2, num_experts) (fp4) - a_sc_contig: (1, total_padded_rm, rk, 512) — dQaccum-padded SFA. + a_sc_contig: (1, total_padded_rm, rk, 32, 4, 4) — M-padded SFA (tile-aligned per batch). total_padded_rm = ((total_m + num_experts * 128) // 128). - b_sc_contig: (num_experts, rn, rk, 512) — regular per-expert SFB. + b_sc_contig: (num_experts, rn, rk, 32, 4, 4) — regular per-expert SFB. cu_seqlens_m: (num_experts+1,) int32 + + Supports MXFP8 / MXFP4 / NVFP4; fp4 formats require b_major="k" (tcgen05 + MMA needs K-major fp4 operands). NVFP4 uses no per-tensor scale here (it + would just fold into alpha). """ assert k % sf_vec_size == 0 if seqlens_m is None: @@ -429,23 +503,31 @@ def create_blockscaled_varlen_m_operands( std = randn_std if randn_std is not None else k**-0.5 sf_k = k // sf_vec_size - if ab_dtype == cutlass.Float8E4M3FN and sf_dtype == cutlass.Float8E8M0FNU and sf_vec_size == 32: - from .mx_utils import to_mx_compiled - - to_fn = to_mx_compiled - else: - raise NotImplementedError( - f"varlen_m currently only supports MXFP8 (got ab={ab_dtype}, sf={sf_dtype}, vec={sf_vec_size}). " - "FP4 support pending." - ) - - # Quantize A: (total_m, k) bf16 -> (total_m, k) fp8 K-major. + fmt = _blockscaled_format_of(ab_dtype, sf_dtype, sf_vec_size) + if fmt != "mxfp8": + assert b_major == "k", f"{fmt} requires K-major operands, got b_major={b_major!r}" + + def quantize(x2d): + """(rows, k) bf16 -> (q, scale_2d, dequant_ref); q is fp8 (rows, k) or fp4x2 (rows, k/2).""" + if fmt == "mxfp8": + q, sc = to_mx_compiled(x2d, sf_vec_size) + vals = q.float() + else: + if fmt == "mxfp4": + q_packed, sc = to_mxfp4_compiled(x2d, sf_vec_size) + else: # nvfp4 + q_packed, sc, _ = to_nvfp4_compiled(x2d, sf_vec_size, None) + q = q_packed.view(torch.uint8).view(torch.float4_e2m1fn_x2) + vals = dequant_operand(q) + ref = vals * sc.float().repeat_interleave(sf_vec_size, dim=-1) + return q, sc, ref + + # Quantize A: (total_m, k) bf16 -> (total_m, k[/2]) K-major. # A data itself is stored packed (no per-expert padding); only SFA is padded. a_hp = (torch.randn(total_m, k, dtype=torch.bfloat16, device="cuda") * std).contiguous() - qa, sa_2d = to_fn(a_hp, sf_vec_size) # (total_m, k), (total_m, sf_k) - a_ref = qa.float() * sa_2d.float().repeat_interleave(sf_vec_size, dim=-1) + qa, sa_2d, a_ref = quantize(a_hp) # (total_m, k[/2]), (total_m, sf_k), (total_m, k) - # Build padded SFA storage (dQaccum format). Each expert's m_i rows of + # Build padded SFA storage (tile-aligned per-batch). Each expert's m_i rows of # scales are written at padded tile offset `cu_seqlens[i] // 128 + i`. # Allocation: `ceildiv(total_m, 128) + (L - 1)` tiles — proven sufficient # in AI/varlen_blockscaled_sf_layout.md (proof 2's "tighter alternative"). @@ -462,24 +544,22 @@ def create_blockscaled_varlen_m_operands( offset += m_i a_sc_contig = pack_scale_2d_to_blocked_contig(sa_2d_padded.view(1, total_padded_m, sf_k)) - # Quantize B: (num_experts, n, k) bf16 -> (n, k, num_experts). b_major selects - # k-major (stride (k, 1, n*k)) or n-major (stride (1, n, n*k)). + # Quantize B: (num_experts, n, k) bf16 -> (n, k[/2], num_experts). b_major selects + # k-major (stride (kb, 1, n*kb)) or n-major (stride (1, n, n*k), mxfp8 only). assert b_major in ("k", "n"), f"b_major must be 'k' or 'n', got {b_major!r}" b_hp = (torch.randn(num_experts, n, k, dtype=torch.bfloat16, device="cuda") * std).contiguous() - qb_flat, sb_2d = to_fn(b_hp.view(num_experts * n, k), sf_vec_size) + qb_flat, sb_2d, b_ref_flat = quantize(b_hp.view(num_experts * n, k)) + kb = qb_flat.shape[-1] # k for fp8, k/2 for packed fp4 if b_major == "k": qb = ( - qb_flat.view(num_experts, n, k).contiguous().permute(1, 2, 0) - ) # (n, k, l) stride (k, 1, n*k) + qb_flat.view(num_experts, n, kb).contiguous().permute(1, 2, 0) + ) # (n, kb, l) stride (kb, 1, n*kb) else: qb = ( - qb_flat.view(num_experts, n, k).transpose(1, 2).contiguous().permute(2, 1, 0) + qb_flat.view(num_experts, n, kb).transpose(1, 2).contiguous().permute(2, 1, 0) ) # (n, k, l) stride (1, n, n*k) - sb_2d = sb_2d.view(num_experts, n, sf_k) - b_sc_contig = pack_scale_2d_to_blocked_contig(sb_2d) - b_ref = qb_flat.float().view(num_experts, n, k) * sb_2d.float().repeat_interleave( - sf_vec_size, dim=-1 - ) + b_sc_contig = pack_scale_2d_to_blocked_contig(sb_2d.view(num_experts, n, sf_k)) + b_ref = b_ref_flat.view(num_experts, n, k) cu_seqlens_m = torch.tensor( [0] + list(itertools.accumulate(seqlens_m)), dtype=torch.int32, device="cuda" @@ -498,23 +578,34 @@ def create_blockscaled_varlen_k_operands( *, randn_std: Optional[float] = None, seqlens_k: Optional[list] = None, + sf_pad_byte: int = 0, ): """Generate bf16 randn + quantize for a varlen_k blockscaled GEMM. - Per-expert `k_i` must be a multiple of `sf_vec_size` (quantization chunk) - but NOT necessarily a multiple of `sf_vec_size * 4` (= 128 for MXFP8). - The SF buffer uses dQaccum-style K padding: each expert `i`'s scales occupy + Per-expert `k_i` is arbitrary (any positive int): neither `sf_vec_size` nor + `sf_vec_size * 4` (= 128 for MXFP8) alignment is required. A non-multiple-of-32 + `k_i` just means the expert's last scale block covers a partial chunk; the + kernel's ragged value TMA zero-fills beyond `cu_seqlens_k[i+1]`, so the tail + contributes exactly 0. + The SF buffer uses tile-aligned per-batch K padding: each expert `i`'s scales occupy `ceildiv(k_i, 128) * 128` bytes worth of K at offset `(cu_seqlens_k[i] + i * 128) // 128 * 128` (in source-K units). A and B operand data stay packed and unpadded along K — only their SF buffers pad. + SF pad regions inside each expert's last 512 B atom are loaded by the + kernel (TMA loads whole atom columns) but never consumed: the mma loop + skips the MMA instructions for pad k-blocks (one instruction per SF block + for mxfp8; see `GemmSm100.mma`), so the pad may hold arbitrary bytes — + including 0xFF (e8m0 NaN). `sf_pad_byte` sets the pad fill so tests can + poison it deliberately. + Returns (a_ref_list, b_ref_list, qa, qb, a_sc_contig, b_sc_contig, cu_seqlens_k): a_ref_list: list of per-expert (m, k_i) fp32 dequantized A. b_ref_list: list of per-expert (n, k_i) fp32 dequantized B. qa: (m, total_k) K-major fp8 (stride (total_k, 1)). qb: (n, total_k) K-major fp8 (stride (total_k, 1)). - a_sc_contig: (1, rm, total_padded_rk, 512) dQaccum-padded SFA. - b_sc_contig: (1, rn, total_padded_rk, 512) dQaccum-padded SFB. + a_sc_contig: (1, rm, total_padded_rk, 32, 4, 4) K-padded SFA (tile-aligned per batch). + b_sc_contig: (1, rn, total_padded_rk, 32, 4, 4) K-padded SFB (tile-aligned per batch). cu_seqlens_k: (num_experts+1,) int32. """ if not ( @@ -530,30 +621,37 @@ def create_blockscaled_varlen_k_operands( f"seqlens_k length {len(seqlens_k)} != num_experts {num_experts}" ) for i, k_i in enumerate(seqlens_k): - assert k_i % sf_vec_size == 0, ( - f"seqlens_k[{i}]={k_i} must be divisible by sf_vec_size={sf_vec_size}" - ) + assert k_i > 0, f"seqlens_k[{i}]={k_i} must be positive" total_k = int(sum(seqlens_k)) std = randn_std if randn_std is not None else (max(seqlens_k)) ** -0.5 - sf_k_total = total_k // sf_vec_size - from .mx_utils import to_mx_compiled + from .quantize import to_mx_compiled + + def quantize(mn, k_i): + # The quantizer reshapes K into sf_vec_size chunks, so zero-pad k_i up to a + # multiple of it; zeros never raise a chunk amax, so the real elements + # quantize identically. Values are sliced back to k_i; scales keep the + # ceil(k_i / sf_vec_size) blocks (the last one covers a partial chunk). + k_q = (k_i + sf_vec_size - 1) // sf_vec_size * sf_vec_size + hp = torch.zeros(mn, k_q, dtype=torch.bfloat16, device="cuda") + hp[:, :k_i] = torch.randn(mn, k_i, dtype=torch.bfloat16, device="cuda") * std + q, sc = to_mx_compiled(hp, sf_vec_size) + q = q[:, :k_i] + ref = q.float() * sc.float().repeat_interleave(sf_vec_size, dim=-1)[:, :k_i] + return q, sc, ref a_q_list, a_sc_list, a_ref_list = [], [], [] b_q_list, b_sc_list, b_ref_list = [], [], [] for k_i in seqlens_k: - # A slice: (m, k_i) bf16 -> fp8, scales (m, k_i // sf_vec_size). - a_hp = (torch.randn(m, k_i, dtype=torch.bfloat16, device="cuda") * std).contiguous() - a_q, a_sc = to_mx_compiled(a_hp, sf_vec_size) + a_q, a_sc, a_ref = quantize(m, k_i) a_q_list.append(a_q) a_sc_list.append(a_sc) - a_ref_list.append(a_q.float() * a_sc.float().repeat_interleave(sf_vec_size, dim=-1)) + a_ref_list.append(a_ref) - b_hp = (torch.randn(n, k_i, dtype=torch.bfloat16, device="cuda") * std).contiguous() - b_q, b_sc = to_mx_compiled(b_hp, sf_vec_size) + b_q, b_sc, b_ref = quantize(n, k_i) b_q_list.append(b_q) b_sc_list.append(b_sc) - b_ref_list.append(b_q.float() * b_sc.float().repeat_interleave(sf_vec_size, dim=-1)) + b_ref_list.append(b_ref) # Pack operand data along K: (m, total_k), (n, total_k). varlen_k's # ragged TMA descriptors are built for MN-major operands (stride 1 on @@ -572,11 +670,15 @@ def create_blockscaled_varlen_k_operands( total_padded_rk = (total_k + tile - 1) // tile + (num_experts - 1) total_padded_k = total_padded_rk * tile total_padded_sf_k = total_padded_k // sf_vec_size - sa_2d_padded = torch.zeros(m, total_padded_sf_k, dtype=a_sc_list[0].dtype, device="cuda") - sb_2d_padded = torch.zeros(n, total_padded_sf_k, dtype=b_sc_list[0].dtype, device="cuda") + sa_2d_padded = torch.full( + (m, total_padded_sf_k), sf_pad_byte, dtype=torch.uint8, device="cuda" + ).view(a_sc_list[0].dtype) + sb_2d_padded = torch.full( + (n, total_padded_sf_k), sf_pad_byte, dtype=torch.uint8, device="cuda" + ).view(b_sc_list[0].dtype) k_offset = 0 for i, k_i in enumerate(seqlens_k): - sf_k_i = k_i // sf_vec_size + sf_k_i = (k_i + sf_vec_size - 1) // sf_vec_size k_offset_padded = (k_offset // tile + i) * tile sf_k_offset_padded = k_offset_padded // sf_vec_size sa_2d_padded[:, sf_k_offset_padded : sf_k_offset_padded + sf_k_i] = a_sc_list[i] @@ -635,16 +737,20 @@ def compile_blockscaled_gemm_tvm_ffi( ) stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - from .gemm_tvm_ffi_utils import make_fake_varlen_args + from ..gemm_tvm_ffi_utils import make_fake_varlen_args varlen_args_fake = make_fake_varlen_args(varlen_m, varlen_k, False, None) or VarlenArguments() # Fake operand tensors with sym_ints (varlen-aware shapes). if varlen_m: total_m_sym = cute.sym_int() - n_sym, k_sym, l_sym = cute.sym_int(), cute.sym_int(), cute.sym_int() - # Detect each operand's leading (stride-1) dim so m-major A / n-major B - # are accepted for varlen_m (MXFP8 only — fp4 is rejected upstream). + n_sym, l_sym = cute.sym_int(), cute.sym_int() + # Sub-byte (fp4) operands need the contiguous K extent statically divisible + # by the packing factor; harmless for 8-bit dtypes. + k_sym = cute.sym_int(divisibility=div_for_dtype(ab_dtype) if ab_dtype.width < 8 else 1) + # Detect B's leading (stride-1) dim so n-major B is accepted for varlen_m + # (mxfp8 only; fp4 is always K-major). A must be K-major for varlen_m — + # the public API enforces this (see quack/gemm.py) for all dtypes. fake_mA = fake_tensor( ab_dtype, (total_m_sym, k_sym), @@ -709,7 +815,7 @@ def compile_blockscaled_gemm_tvm_ffi( varlen_args, stream, ): - gemm(a, b, d, None, compile_epi_args, scheduler_args, varlen_args, stream, sfa, sfb, None) + gemm(a, b, d, None, compile_epi_args, scheduler_args, varlen_args, stream, sfa, sfb) compiled = cute.compile( runner, diff --git a/build/torch-cuda/quack/broadcast_utils.py b/build/torch-cuda/quack/broadcast_utils.py index e7a1efc55f8a341024cc2f859fb977130cab5bf5..c16285c93dfb1d56c78797a200c2cfc1e324095e 100644 --- a/build/torch-cuda/quack/broadcast_utils.py +++ b/build/torch-cuda/quack/broadcast_utils.py @@ -5,18 +5,17 @@ import cutlass import cutlass.cute as cute from cutlass import Float32, const_expr -from .layout_utils import make_acc_tensor_mn_view +from . import layout_utils @cute.jit def vec_op(tCrC: cute.Tensor, tCrVec: cute.Tensor, op: Callable, is_colvec: bool) -> None: if const_expr(tCrC.element_type != Float32): # Convert to f32 - tCrC_f32 = cute.make_rmem_tensor(tCrC.shape, Float32) - tCrC_f32.store(tCrC.load().to(Float32)) + tCrC_f32 = tCrC.to(Float32) else: tCrC_f32 = tCrC # this happens to work for frgA layout too, not just acc layout - tCrC_f32_mn = make_acc_tensor_mn_view(tCrC_f32) + tCrC_f32_mn = layout_utils.reshape_acc_to_mn(tCrC_f32) if const_expr(is_colvec): assert cute.size(tCrC_f32_mn, mode=[0]) == cute.size(tCrVec) for r in cutlass.range(cute.size(tCrC_f32_mn, mode=[0]), unroll_full=True): diff --git a/build/torch-cuda/quack/cache/__init__.py b/build/torch-cuda/quack/cache/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bab6cdb033710214b32020008cc81acbb9236e62 --- /dev/null +++ b/build/torch-cuda/quack/cache/__init__.py @@ -0,0 +1,74 @@ +# Copyright (c) 2025-2026, Tri Dao. +"""Persistent kernel-cache utilities for QuACK. + +Public API +---------- + +Persistent ``.o`` cache: +* :func:`jit_cache` — decorator that wraps a kernel-compile function with + in-memory + persistent ``.o`` caching (see :mod:`quack.cache.jit`). +* :data:`CACHE_ENABLED`, :data:`CACHE_DIR`, :data:`EXTRA_SOURCE_DIRS` — + static-config flags. +* :class:`FileLock`, :func:`get_cache_path`, :class:`CacheInfo` — + supporting types. + +Async compilation (see :mod:`quack.cache.async_compile`): +* :class:`CompilePending` — raised by ``jit_cache`` on a cold miss while a + compile pool is active; the caller defers and retries once the ``.o`` + lands. +* :func:`pool_scope` — activate a compile pool for a scoped block (used by + the autotuner's bench loop). + +CRITICAL ORDERING: the static-config flags below MUST be defined before the +``from quack.cache.jit import ...`` block. ``quack/cache/jit.py`` does +``import quack.cache as _state`` at its module top; Python returns the +partially-initialized package object, and lookups inside ``jit_cache``'s +wrapper rely on these names already existing at that checkpoint. Reordering +the imports here, even via an auto-formatter, will break the first kernel +compile with ``AttributeError``. The defensive unit tests in +``tests/test_cache.py`` exercise this path end-to-end. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import List, Optional + +CACHE_ENABLED: bool = os.getenv("QUACK_CACHE_ENABLED", "1") == "1" +CACHE_DIR: Optional[str] = os.getenv("QUACK_CACHE_DIR", None) + +#: Downstream projects can append directories here to include their sources +#: in the cache fingerprint. Must be set before the first jit_cache call. +EXTRA_SOURCE_DIRS: List[Path] = [] + + +# --------------------------------------------------------------------------- +# Public API surface. Imported AFTER the flags are defined. +# --------------------------------------------------------------------------- + +from .jit import ( # noqa: E402 + EXPORT_FUNC_NAME, + LOCK_TIMEOUT, + CacheInfo, + FileLock, + get_cache_path, + jit_cache, +) +from .async_compile import ( # noqa: E402 + CompilePending, + pool_scope, +) + +__all__ = [ + # Persistent .o cache. + "jit_cache", + "CacheInfo", + "EXPORT_FUNC_NAME", + "LOCK_TIMEOUT", + "FileLock", + "get_cache_path", + # Async compilation. + "CompilePending", + "pool_scope", +] diff --git a/build/torch-cuda/quack/cache/_pool_preload.py b/build/torch-cuda/quack/cache/_pool_preload.py new file mode 100644 index 0000000000000000000000000000000000000000..ff0d0454fdecad01e129028fe76b72ce5f620882 --- /dev/null +++ b/build/torch-cuda/quack/cache/_pool_preload.py @@ -0,0 +1,50 @@ +# Copyright (c) 2026, Tri Dao. +"""Forkserver preload module for the async compile pool. + +Imported once inside the multiprocessing *forkserver* process (see +``multiprocessing.set_forkserver_preload``). Every pool worker is then +``fork()``-ed from that warm process and inherits the imported interpreter +state via copy-on-write: worker startup drops from ~13 s (torch 4 s + +cutlass/cute/tvm_ffi 9 s per spawn) to ~0.1 s per fork. + +This is the same architecture as PyTorch Inductor's compile-worker +``SubprocPool``: one sidecar pays the import, workers fork from it. + +Fork-safety: nothing here may initialize CUDA (a forked child of a +CUDA-initialized process is undefined behavior). Importing torch and +cutlass does not create a CUDA context; workers additionally run with +``CUDA_VISIBLE_DEVICES=""`` + ``QUACK_ARCH``/``CUTE_DSL_ARCH`` overrides so +the compile path never touches the driver (the same mechanism the CPU-only +compile workflow uses). +""" + +import os +import subprocess + +# Pin the target arch BEFORE importing quack: import-time code paths (e.g. +# rmsnorm_config._detect_arch_major) consult QUACK_ARCH via +# get_device_capacity and would otherwise initialize CUDA — which both makes +# the forkserver's context leak into children and trips torch's forked-child +# guard. nvidia-smi queries the capability without creating a CUDA context. +if "QUACK_ARCH" not in os.environ: + try: + out = subprocess.run( + ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], + capture_output=True, + text=True, + timeout=10, + ) + cap = out.stdout.strip().splitlines()[0].strip() # e.g. "9.0" + major, minor = cap.split(".") + os.environ["QUACK_ARCH"] = f"{major}{minor}" + os.environ.setdefault( + "CUTE_DSL_ARCH", f"sm_{major}{minor}a" if int(major) >= 9 else f"sm_{major}{minor}" + ) + except Exception: + pass # CPU-only box: rely on user-provided env, as before + +# Belt and suspenders: even if some import still tries to touch CUDA, make +# it see no devices rather than creating a context in the forkserver. +os.environ["CUDA_VISIBLE_DEVICES"] = "" + +from .. import cache # noqa: F401, E402 (pulls torch, cutlass.cute, tvm_ffi) diff --git a/build/torch-cuda/quack/cache/async_compile.py b/build/torch-cuda/quack/cache/async_compile.py new file mode 100644 index 0000000000000000000000000000000000000000..67149f74744daca7eb311a299465c31aafa4c832 --- /dev/null +++ b/build/torch-cuda/quack/cache/async_compile.py @@ -0,0 +1,413 @@ +# Copyright (c) 2026, Tri Dao. +"""Async kernel compilation: defer-and-retry via a pool of CPU subprocesses. + +When a pool is active, ``jit_cache`` handles a ``.o``-cache miss by +submitting the pickled ``(module, qualname, args, kwargs)`` of the +``_compile_*`` function to the pool and raising :class:`CompilePending` +instead of compiling in-process. The caller defers the work item, runs +something else, and retries once the ``.o`` lands (a ~1 ms load). Two +callers implement this loop: + +* the pytest plugin's ``--async-compile`` defer loop (tests are the work + items; see :mod:`quack.testing.pytest_plugin`); +* the autotuner's bench loop under :func:`pool_scope` (candidate configs + are the work items; see :class:`quack.autotuner.Autotuner`). + +Design notes: + +* **The ``.o`` file is the only rendezvous** between workers and consumers — + compiled kernels aren't picklable, so the persistent cache doubles as the + IPC channel, and the per-key ``flock`` in ``jit_cache`` doubles as + cross-process dedupe (multiple pools / xdist workers coexist safely; + :func:`_flock_held_exclusively` lets a consumer defer on a key some other + process is already compiling). +* **Workers never launch kernels by construction**: they call the + tensor-free ``_compile_*`` functions directly, GPU-blind (arch pinned via + ``QUACK_ARCH``/``CUTE_DSL_ARCH``, ``CUDA_VISIBLE_DEVICES=""``). +* **Worker startup is an Inductor-style sidecar**: a ``forkserver`` preloads + torch/cutlass once (:mod:`quack.cache._pool_preload`, ~13 s) and workers + fork from it copy-on-write (~0.1 s each). :func:`_neutral_main` keeps + multiprocessing child prep from re-executing the user's script. +* **Failure semantics**: a failed pool compile is never trusted — the + consumer falls through to an in-process compile so the real exception + surfaces with a local traceback. + +Env knobs: ``QUACK_ASYNC_COMPILE_START=spawn`` (disable the fork sidecar), +``QUACK_COMPILE_WORKERS`` (shared-executor size, default 8). +""" + +from __future__ import annotations + +import base64 +import contextlib +import fcntl +import importlib +import os +import pickle +from concurrent.futures import Future, ProcessPoolExecutor +from multiprocessing import get_context +from typing import Optional + + +def _flock_held_exclusively(lock_path: str) -> bool: + """True if some process currently holds the flock exclusively. + + Used to detect "another process is compiling this key right now" so the + consumer defers instead of submitting a duplicate compile to its own + pool (a duplicate would occupy a pool slot blocked on the same flock). + """ + try: + fd = os.open(lock_path, os.O_RDONLY | os.O_CREAT) + except OSError: + return False + try: + try: + fcntl.flock(fd, fcntl.LOCK_SH | fcntl.LOCK_NB) + fcntl.flock(fd, fcntl.LOCK_UN) + return False + except OSError: + return True + finally: + os.close(fd) + + +class CompilePending(BaseException): + """A jit_cache miss was submitted to the async compile pool. + + The caller (test) cannot proceed until the ``.o`` exists; the test runner + should defer the test and retry it later. Carries the cache ``sha`` so + the runner can poll for completion without re-running the test. + + Derives from :class:`BaseException` (like ``KeyboardInterrupt``) so that + test-body ``except Exception`` / ``pytest.raises(Exception)`` blocks + cannot swallow it and turn a not-yet-run test into a false pass. Only + the plugin's phase hooks are supposed to catch it. + """ + + def __init__(self, sha: str, qualname: str): + super().__init__(f"kernel compile pending in pool: {qualname} [{sha[:12]}]") + self.sha = sha + self.qualname = qualname + + +def _detect_arch_env() -> tuple[Optional[str], Optional[str]]: + """Return (QUACK_ARCH, CUTE_DSL_ARCH) for GPU-blind pool workers. + + An explicit ``QUACK_ARCH`` env override wins — CI cross-compiles for a + different arch than the runner's GPU (e.g. ``QUACK_ARCH=120`` on an + H100), and workers must compile for the *target* arch, not the physical + one. Otherwise detect from the parent's GPU. Either way the workers + themselves never touch the CUDA driver (no context per worker, + fork-safe). + """ + quack_arch = os.environ.get("QUACK_ARCH") + if quack_arch is not None: + cute_arch = os.environ.get("CUTE_DSL_ARCH") + if cute_arch is None: + from ..cute_dsl_utils import _parse_arch_str + + major, minor = _parse_arch_str(quack_arch) + cc = f"{major}{minor}" + cute_arch = f"sm_{cc}a" if major >= 9 else f"sm_{cc}" + return quack_arch, cute_arch + try: + import torch + + if torch.cuda.is_available(): + major, minor = torch.cuda.get_device_capability() + cc = f"{major}{minor}" + return cc, f"sm_{cc}a" if major >= 9 else f"sm_{cc}" + except Exception: + pass + return None, os.environ.get("CUTE_DSL_ARCH") + + +def _pool_initializer(quack_arch: Optional[str], cute_dsl_arch: Optional[str]): + # GPU-blind compilation: hide devices and pin the target arch via the + # same overrides the CPU-only compile workflow uses. Forked workers must + # never initialize CUDA (fork-safety), and spawned workers save the + # ~1-2 s + ~300 MB of a per-worker CUDA context. + if quack_arch is not None: + os.environ["QUACK_ARCH"] = quack_arch + os.environ["CUDA_VISIBLE_DEVICES"] = "" + if cute_dsl_arch is not None: + os.environ["CUTE_DSL_ARCH"] = cute_dsl_arch + # Pay the heavy torch/cutlass import at worker start (no-op under + # forkserver: the preload already imported it before the fork). + from .. import cache # noqa: F401 + + +def _pool_worker(mod_name: str, qualname: str, key_b64: str, o_path: str) -> Optional[str]: + """Compile one key. Returns None on success, error string on failure.""" + try: + obj = importlib.import_module(mod_name) + for part in qualname.split("."): + obj = getattr(obj, part) + args, kwargs = pickle.loads(base64.b64decode(key_b64)) + obj(*args, **kwargs) # jit_cache wrapper: compiles + exports .o + if not os.path.exists(o_path): + return "compile succeeded but .o was not exported" + return None + except Exception as e: + return f"{type(e).__name__}: {e}" + + +def _make_executor(jobs: int) -> ProcessPoolExecutor: + """Build a compile-worker executor (Inductor-style forkserver sidecar). + + Forkserver + preload: one sidecar process pays the ~13 s torch/cutlass + import once, workers fork from it in ~0.1 s each (copy-on-write). The + forkserver singleton is shared per-process, so multiple executors (the + test pool, the autotuner's) fork from the same warm sidecar. Opt out + with QUACK_ASYNC_COMPILE_START=spawn. + """ + start_method = os.environ.get("QUACK_ASYNC_COMPILE_START", "forkserver") + ctx = get_context(start_method) + if start_method == "forkserver": + # vendored under sonic_moe.quack: preload the sibling module by package + # (upstream hardcodes the top-level "quack.cache._pool_preload"). + ctx.set_forkserver_preload([__package__ + "._pool_preload"]) + return ProcessPoolExecutor( + max_workers=jobs, + mp_context=ctx, + initializer=_pool_initializer, + initargs=_detect_arch_env(), + ) + + +_shared_executor: Optional[ProcessPoolExecutor] = None + + +def get_shared_executor() -> ProcessPoolExecutor: + """Executor for ad-hoc compile tasks (e.g. autotuner precompile sweeps). + + Reuses the active :class:`CompilePool`'s executor when one exists (the + pytest ``--async-compile`` session pool); otherwise lazily creates a + process-wide executor sized by ``QUACK_COMPILE_WORKERS`` (default 8). + Deliberately ignores :class:`suppress_pool` — suppression turns off the + *defer-on-miss* behavior of jit_cache, not access to compile workers. + """ + global _shared_executor + if _active_pool is not None: + return _active_pool._executor + if _shared_executor is None: + import atexit + + _shared_executor = _make_executor(int(os.environ.get("QUACK_COMPILE_WORKERS", "8"))) + # Explicit teardown: without this, the executor is GC'd during + # interpreter shutdown after its weakref machinery is already gone, + # printing a spurious "Exception ignored in weakref_cb". + atexit.register(_shared_executor.shutdown, wait=False, cancel_futures=True) + return _shared_executor + + +@contextlib.contextmanager +def _neutral_main(): + """Stop multiprocessing child prep from re-executing the user's script. + + ``Process.start()`` captures preparation data from ``sys.modules['__main__']``: + for a path-based script the *child* re-runs the whole file via + ``runpy.run_path`` (so pickles referencing ``__main__`` resolve). Our + tasks never reference ``__main__`` — they resolve everything by module + name — and a user script that, say, builds CUDA tensors at import time + would kill every worker at spawn with "Cannot re-initialize CUDA in + forked subprocess". Executor workers are spawned synchronously inside + ``executor.submit`` (``_adjust_process_count``), so masking ``__main__`` + with an empty stub for the duration of the submit is sufficient and + scoped. Single-threaded callers only (pytest defer loop, autotune bench + loop). + """ + import sys + import types + + real_main = sys.modules.get("__main__") + sys.modules["__main__"] = types.ModuleType("__main__") # no __file__/__spec__ + try: + yield + finally: + if real_main is not None: + sys.modules["__main__"] = real_main + + +class CompilePool: + """Process pool + in-flight bookkeeping, keyed by jit_cache sha. + + Owns its executor by default; pass ``executor=`` to share one (e.g. + :func:`pool_scope` wraps the session-long shared executor so scoped + pools don't respawn workers per autotune sweep). A shared executor is + not shut down by :meth:`shutdown` — only this pool's futures are + cancelled. + """ + + def __init__(self, jobs: Optional[int] = None, executor: Optional[ProcessPoolExecutor] = None): + self._own_executor = executor is None + self._executor = executor if executor is not None else _make_executor(jobs) + self._futures: dict[str, Future] = {} + # Keys being compiled by *another process* (e.g. a different xdist + # worker's pool), detected via the per-key flock. We defer on them + # without spending one of our own pool slots on a duplicate compile. + # sha -> (o_path, lock_path) + self._external: dict[str, tuple[str, str]] = {} + self.n_submitted = 0 + + def mark_external(self, sha: str, o_path: str, lock_path: str) -> None: + """Record that some other process is compiling ``sha`` right now.""" + if sha not in self._futures: + self._external[sha] = (str(o_path), str(lock_path)) + + def prewarm(self) -> None: + """Start the sidecar + first worker now, off the critical path. + + Same idea as Inductor's ``warm_pool()``: the forkserver's ~13 s + torch/cutlass preload import starts at the first ``Process`` spawn, + which is lazy (first submit). Submitting a no-op at session setup + overlaps that import with pytest collection and the leading warm + tests instead of the first cold compile. + """ + with _neutral_main(): + self._executor.submit(os.getpid) + + def submit_raw(self, sha: str, mod: str, qualname: str, key_b64: str, o_path: str) -> None: + if sha in self._futures: + return + with _neutral_main(): + self._futures[sha] = self._executor.submit(_pool_worker, mod, qualname, key_b64, o_path) + self.n_submitted += 1 + + def submit(self, sha: str, fn, args: tuple, kwargs: dict, o_path) -> bool: + """Submit a live jit_cache miss. Returns False if the key can't be + shipped to a subprocess (unpicklable args, ```` qualname, + fn defined in ``__main__``) — the caller should compile in-process + instead.""" + if sha in self._futures: + return True + if "" in fn.__qualname__ or fn.__module__ == "__main__": + # Not resolvable by module+qualname in a worker; compile in-process. + return False + try: + key_b64 = base64.b64encode(pickle.dumps((args, kwargs))).decode("ascii") + except Exception: + return False + self.submit_raw(sha, fn.__module__, fn.__qualname__, key_b64, str(o_path)) + return True + + def poll(self, sha: str) -> tuple[str, Optional[str]]: + """Return (state, error): state in {"new", "pending", "done", "failed"}.""" + fut = self._futures.get(sha) + if fut is None: + ext = self._external.get(sha) + if ext is not None: + o_path, lock_path = ext + if os.path.exists(o_path): + del self._external[sha] + return "done", None + if _flock_held_exclusively(lock_path): + return "pending", None + # External compiler released the lock without producing a .o + # (crashed / failed): forget it so the next attempt submits + # to our own pool. + del self._external[sha] + return "new", None + if not fut.done(): + return "pending", None + try: + err = fut.result() + except Exception as e: # BrokenProcessPool etc. + err = f"pool worker died: {type(e).__name__}: {e}" + return ("done", None) if err is None else ("failed", err) + + def stats(self) -> dict: + done = sum(1 for f in self._futures.values() if f.done()) + errors = [] + for sha, f in self._futures.items(): + if not f.done() or f.cancelled(): + continue + exc = f.exception() + err = f"{type(exc).__name__}: {exc}" if exc is not None else f.result() + if err: + errors.append((sha, err)) + return { + "submitted": self.n_submitted, + "done": done, + "failed": len(errors), + "errors": errors, + } + + def shutdown(self) -> None: + if self._own_executor: + self._executor.shutdown(wait=False, cancel_futures=True) + else: + for fut in self._futures.values(): + fut.cancel() + + +# --- module-level active pool ----------------------------------------------- + +_active_pool: Optional[CompilePool] = None +_suppress_depth = 0 + + +class suppress_pool: + """Context manager: make :func:`get_active_pool` return None inside. + + Used by the test runner for a deferred test's final attempt: compile + in-process (blocking) so a key that never completes in the pool still + produces a real result or a real traceback instead of deferring forever. + """ + + def __enter__(self): + global _suppress_depth + _suppress_depth += 1 + return self + + def __exit__(self, *exc): + global _suppress_depth + _suppress_depth -= 1 + + +def activate(jobs: int) -> CompilePool: + """Activate the session-wide pool (idempotent). Used by the pytest plugin; + scoped callers should prefer :func:`pool_scope`.""" + global _active_pool + if _active_pool is None: + _active_pool = CompilePool(jobs) + return _active_pool + + +def deactivate() -> None: + global _active_pool + if _active_pool is not None: + _active_pool.shutdown() + _active_pool = None + + +def get_active_pool() -> Optional[CompilePool]: + return None if _suppress_depth > 0 else _active_pool + + +@contextlib.contextmanager +def pool_scope(): + """Activate a compile pool for the duration of the block. + + Reuses the globally active pool when one exists (e.g. the pytest + ``--async-compile`` session pool); otherwise activates a temporary pool + backed by the shared executor and deactivates it on exit — so + ``CompilePending`` can only escape into code inside the block, never + into unrelated user code paths. + + This is how the autotuner overlaps candidate-config compilation with + benchmarking: the bench loop runs inside ``pool_scope()``, catches + ``CompilePending`` per config, and retries a config once its ``.o`` + lands (see ``Autotuner.__call__``). + """ + global _active_pool + if _active_pool is not None: + yield _active_pool + return + pool = CompilePool(executor=get_shared_executor()) + _active_pool = pool + try: + yield pool + finally: + _active_pool = None + pool.shutdown() diff --git a/build/torch-cuda/quack/cache/jit.py b/build/torch-cuda/quack/cache/jit.py new file mode 100644 index 0000000000000000000000000000000000000000..0a29479f0fe896feb1052a7a545cd4f752838c36 --- /dev/null +++ b/build/torch-cuda/quack/cache/jit.py @@ -0,0 +1,332 @@ +# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. +"""Persistent ``.o`` cache for CuTe DSL compiled kernels. + +Compiled kernels are exported as object files (``.o``) via ``export_to_c``. On +subsequent runs the ``.o`` is loaded via tvm_ffi (~1 ms) instead of +re-generating IR + re-JIT'ing (~500 ms per kernel). + +Runtime config (``CACHE_ENABLED``, ``CACHE_DIR``, ``EXTRA_SOURCE_DIRS``) +lives in :mod:`quack.cache` (the package init). + +When an async compile pool is active (see :mod:`quack.cache.async_compile`), +a cold miss is shipped to a CPU worker and :class:`CompilePending` is raised +instead of compiling in-process; the caller (pytest defer loop, autotune +bench loop) retries once the ``.o`` lands. +""" + +from __future__ import annotations + +import fcntl +import functools +import hashlib +import os +import pickle +import sys +import tempfile +import time +from collections import namedtuple +from getpass import getuser +from pathlib import Path + +import cutlass +import cutlass.cute as cute +import tvm_ffi + +# `quack.cache` (the package itself) holds the mutable runtime flags as a +# single source of truth; reads happen via attribute access on `_state` so we +# always see the live value, not a snapshot taken at module import. +from .. import cache as _state # noqa: E402 (intentional partial-import; see __init__.py) + + +EXPORT_FUNC_NAME = "func" +LOCK_TIMEOUT = 60 +CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "maxsize", "currsize"]) + + +def get_cache_path() -> Path: + if _state.CACHE_DIR is not None: + cache_dir = Path(_state.CACHE_DIR) + else: + cache_dir = Path(tempfile.gettempdir()) / getuser() / "quack_cache" + cache_dir.mkdir(parents=True, exist_ok=True) + return cache_dir + + +def _hash_source_dir(h, root: Path) -> None: + """Hash all Python sources under *root* into *h*.""" + for src in sorted(root.rglob("*.py")): + if not src.is_file(): + continue + h.update(src.relative_to(root).as_posix().encode()) + content = src.read_bytes() + h.update(len(content).to_bytes(8, "little")) + h.update(content) + + +@functools.lru_cache(maxsize=1) +def _compute_source_fingerprint() -> str: + """Hash quack + extra source dirs plus runtime ABI stamps into a fingerprint.""" + h = hashlib.sha256() + h.update(f"py{sys.version_info.major}.{sys.version_info.minor}".encode()) + h.update(f"cutlass={cutlass.__version__}".encode()) + h.update(f"tvm_ffi={tvm_ffi.__version__}".encode()) + # Hash the entire `quack` package, not just `quack/cache/`. Resolving via + # the top-level package import keeps the fingerprint stable regardless of + # where inside the package this file lives. + import importlib as _importlib; _quack = _importlib.import_module(__package__.rsplit(".", 1)[0]) + + _hash_source_dir(h, Path(_quack.__file__).resolve().parent) + for extra_dir in _state.EXTRA_SOURCE_DIRS: + _hash_source_dir(h, Path(extra_dir).resolve()) + return h.hexdigest() + + +def _key_to_hash(key: tuple) -> str: + return hashlib.sha256(pickle.dumps(key)).hexdigest() + + +# --------------------------------------------------------------------------- +# File locking +# --------------------------------------------------------------------------- + + +class FileLock: + """Advisory file lock using fcntl.flock with timeout.""" + + def __init__(self, lock_path: Path, exclusive: bool, timeout: float = 15): + self.lock_path = lock_path + self.exclusive = exclusive + self.timeout = timeout + self._fd: int = -1 + + def __enter__(self) -> "FileLock": + flags = os.O_WRONLY | os.O_CREAT if self.exclusive else os.O_RDONLY | os.O_CREAT + lock_type = fcntl.LOCK_EX if self.exclusive else fcntl.LOCK_SH + self._fd = os.open(str(self.lock_path), flags) + deadline = time.monotonic() + self.timeout + while time.monotonic() < deadline: + try: + fcntl.flock(self._fd, lock_type | fcntl.LOCK_NB) + return self + except OSError: + time.sleep(0.1) + os.close(self._fd) + self._fd = -1 + raise RuntimeError(f"Timed out waiting for lock: {self.lock_path}") + + def __exit__(self, *exc) -> None: + if self._fd >= 0: + fcntl.flock(self._fd, fcntl.LOCK_UN) + os.close(self._fd) + self._fd = -1 + + +# --------------------------------------------------------------------------- +# JIT cache decorator +# --------------------------------------------------------------------------- + + +def jit_cache(fn): + """Decorator that caches compiled CuTe DSL kernels in-memory and on disk. + + The decorated function should return a compiled kernel (i.e. call cute.compile). + The disk cache key is (fn.__qualname__, *args, **sorted_kwargs). + + Concurrency model + ----------------- + The disk side uses a per-key ``{sha}.lock`` file (advisory ``flock``): + + * **Fast path (warm cache).** If the ``.o`` file already exists, we take a + shared lock just long enough to ``load_module`` it. Many readers can + proceed concurrently. + * **Slow path (cold cache).** The actual ``fn(*args, **kwargs)`` compile + runs *under* the exclusive lock. This serializes redundant compilations + of the same key across xdist workers / processes: if N processes race + on a cold key, only one calls ``cute.compile``; the rest wait for the + lock, see the ``.o`` appear, and load it. (Previously the compile ran + *between* the shared-lock check and the exclusive-lock export, so all + N processes wasted CPU compiling the same key in parallel — wall time + was unchanged but compile-CPU pressure scaled with concurrency, which + starved other compiles when many keys were cold at once.) + + The lock is per-key, so distinct keys never contend with each other. + """ + cache = {} + hits = 0 + misses = 0 + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + nonlocal hits, misses + cache_key = args + tuple(sorted(kwargs.items())) if kwargs else args + + # Snapshot once per call so a concurrent flip of ``_state.CACHE_ENABLED`` + # mid-call can't desync the disk-path branch. + enabled = _state.CACHE_ENABLED + + # 1. In-memory hit. Same process already compiled or loaded this key. + if cache_key in cache: + hits += 1 + return cache[cache_key] + + # 2. Cache disabled: pure in-process compile, no disk side effects. + if not enabled: + misses += 1 + compiled_fn = fn(*args, **kwargs) + cache[cache_key] = compiled_fn + return compiled_fn + + sha = _key_to_hash((fn.__qualname__,) + cache_key) + cache_path = get_cache_path() / _compute_source_fingerprint() + cache_path.mkdir(parents=True, exist_ok=True) + o_path = cache_path / f"{sha}.o" + lock_path = cache_path / f"{sha}.lock" + + def _load_cached() -> object: + """Load the .o into a callable; caller guarantees existence.""" + m = cute.runtime.load_module(str(o_path), enable_tvm_ffi=True) + return m[EXPORT_FUNC_NAME] + + def _quarantine_corrupt(exc: Exception) -> None: + """A cached .o that fails to load (truncated write from a killed + worker, missing __tvm_ffi_func, ...) is a cache miss, not an error: + delete it so this and future processes recompile instead of failing + forever (the CI cache persists across runs).""" + print( + f"quack cache: corrupt cached object for key {sha} " + f"({type(exc).__name__}: {exc}); deleting and recompiling" + ) + try: + o_path.unlink() + except OSError: + pass + + # 3. Fast path: optimistic existence check, then shared-lock load. + # The unlocked ``.exists()`` is a no-cost short-circuit for warm + # caches; the shared lock guards against reading a partial file + # while a concurrent writer holds the exclusive lock. + if o_path.exists(): + try: + with FileLock(lock_path, exclusive=False, timeout=LOCK_TIMEOUT): + if o_path.exists(): + try: + loaded = _load_cached() + except Exception as e: + # Corrupt entry: recover under the exclusive lock in + # the slow path (shared lock can't safely delete). + _quarantine_corrupt(e) + else: + cache[cache_key] = loaded + hits += 1 + return loaded + except RuntimeError: + pass # lock timeout; fall through to slow path + + # 3b. Async-compile pool: on a cold miss with a pool + # active, ship the key to a CPU subprocess and raise + # CompilePending instead of compiling in-process. The test runner + # defers the test and retries once the worker has exported the + # .o. Pool failures fall through to the in-process compile below + # so the real exception surfaces with a local traceback. + from . import async_compile as _async + + pool = _async.get_active_pool() + if pool is not None: + state, err = pool.poll(sha) + if state == "new": + # If another process (e.g. a different xdist worker's pool) + # holds the exclusive per-key flock, it is compiling this key + # right now: defer on it instead of submitting a duplicate. + if _async._flock_held_exclusively(str(lock_path)): + pool.mark_external(sha, str(o_path), str(lock_path)) + raise _async.CompilePending(sha, fn.__qualname__) + if pool.submit(sha, fn, args, kwargs, o_path): + raise _async.CompilePending(sha, fn.__qualname__) + # unpicklable key / qualname: compile in-process + elif state == "pending": + raise _async.CompilePending(sha, fn.__qualname__) + elif state == "done": + try: + with FileLock(lock_path, exclusive=False, timeout=LOCK_TIMEOUT): + if o_path.exists(): + try: + loaded = _load_cached() + except Exception as e: + _quarantine_corrupt(e) + else: + cache[cache_key] = loaded + hits += 1 + return loaded + except RuntimeError: + pass # lock timeout; fall through to slow path + else: # "failed" + print( + f"quack cache: async compile failed for {fn.__qualname__} " + f"[{sha[:12]}]: {err}; recompiling in-process for a real traceback" + ) + + # 4. Slow path: take EXCLUSIVE lock and compile under it. The recheck + # inside the lock catches the race where another process compiled + # while we were waiting; in that case we just load and return + # without duplicating the compile. + try: + with FileLock(lock_path, exclusive=True, timeout=LOCK_TIMEOUT): + if o_path.exists(): + try: + loaded = _load_cached() + except Exception as e: + _quarantine_corrupt(e) # holds the exclusive lock: safe + else: + cache[cache_key] = loaded + hits += 1 + return loaded + + misses += 1 + compiled_fn = fn(*args, **kwargs) + # Export to a private temp file, then atomically rename into + # place: a process killed mid-export (xdist worker OOM-kill, + # timeout) must never leave a truncated .o at the final path — + # the advisory flock dies with the process, and a persistent + # cache (CI keeps one in $HOME) would then fail every future + # run on this key with "Symbols not found: __tvm_ffi_func". + tmp_path = o_path.with_suffix(f".o.tmp.{os.getpid()}") + try: + compiled_fn.export_to_c( + object_file_path=str(tmp_path), + function_name=EXPORT_FUNC_NAME, + ) + os.replace(tmp_path, o_path) + except Exception as e: + print(f"quack cache: export failed for key {sha}: {e}") + try: + tmp_path.unlink() + except OSError: + pass + cache[cache_key] = compiled_fn + return compiled_fn + except RuntimeError as e: + # Lock acquisition timed out (heavy contention or stuck holder). + # Fall back to in-process compile, no disk write. Better to do + # the work twice than to fail the test. + print( + f"quack cache: lock timeout for key {sha}: {e}; " + f"falling back to in-process compile without disk cache" + ) + misses += 1 + compiled_fn = fn(*args, **kwargs) + cache[cache_key] = compiled_fn + return compiled_fn + + def cache_clear(): + nonlocal hits, misses + cache.clear() + hits = 0 + misses = 0 + + def cache_info(): + return CacheInfo(hits=hits, misses=misses, maxsize=None, currsize=len(cache)) + + wrapper.cache = cache + wrapper.cache_clear = cache_clear + wrapper.cache_info = cache_info + return wrapper diff --git a/build/torch-cuda/quack/cache_utils.py b/build/torch-cuda/quack/cache_utils.py deleted file mode 100644 index b596d7d5916dffb510e549a362f271528344fac7..0000000000000000000000000000000000000000 --- a/build/torch-cuda/quack/cache_utils.py +++ /dev/null @@ -1,195 +0,0 @@ -# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. -"""Persistent .o cache for CuTe DSL compiled kernels. - -Compiled kernels are exported as object files (.o) via export_to_c. -On subsequent runs the .o is loaded via tvm_ffi (~1ms) instead of -re-generating IR + re-JIT'ing (~100ms per kernel). - -Controls: - QUACK_CACHE_ENABLED=0 — disable persistent .o cache (default: enabled) - QUACK_CACHE_DIR=path — override default cache directory -""" - -import fcntl -import functools -import hashlib -import os -import pickle -import sys -import tempfile -import time -from collections import namedtuple -from getpass import getuser -from pathlib import Path - -import cutlass -import cutlass.cute as cute -import tvm_ffi - -CACHE_ENABLED: bool = os.getenv("QUACK_CACHE_ENABLED", "1") == "1" -CACHE_DIR: str | None = os.getenv("QUACK_CACHE_DIR", None) -COMPILE_ONLY: bool = False - -# Downstream projects can append directories here to include their sources -# in the cache fingerprint. Must be set before the first jit_cache call. -EXTRA_SOURCE_DIRS: list[Path] = [] - -EXPORT_FUNC_NAME = "func" -LOCK_TIMEOUT = 60 -CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "maxsize", "currsize"]) - - -def _noop_kernel(*args, **kwargs): - pass - - -def get_cache_path() -> Path: - if CACHE_DIR is not None: - cache_dir = Path(CACHE_DIR) - else: - cache_dir = Path(tempfile.gettempdir()) / getuser() / "quack_cache" - cache_dir.mkdir(parents=True, exist_ok=True) - return cache_dir - - -def _hash_source_dir(h, root: Path) -> None: - """Hash all Python sources under *root* into *h*.""" - for src in sorted(root.rglob("*.py")): - if not src.is_file(): - continue - h.update(src.relative_to(root).as_posix().encode()) - content = src.read_bytes() - h.update(len(content).to_bytes(8, "little")) - h.update(content) - - -@functools.lru_cache(maxsize=1) -def _compute_source_fingerprint() -> str: - """Hash quack + extra source dirs plus runtime ABI stamps into a fingerprint.""" - h = hashlib.sha256() - h.update(f"py{sys.version_info.major}.{sys.version_info.minor}".encode()) - h.update(f"cutlass={cutlass.__version__}".encode()) - h.update(f"tvm_ffi={tvm_ffi.__version__}".encode()) - _hash_source_dir(h, Path(__file__).resolve().parent) - for extra_dir in EXTRA_SOURCE_DIRS: - _hash_source_dir(h, Path(extra_dir).resolve()) - return h.hexdigest() - - -def _key_to_hash(key: tuple) -> str: - return hashlib.sha256(pickle.dumps(key)).hexdigest() - - -# --------------------------------------------------------------------------- -# File locking -# --------------------------------------------------------------------------- - - -class FileLock: - """Advisory file lock using fcntl.flock with timeout.""" - - def __init__(self, lock_path: Path, exclusive: bool, timeout: float = 15): - self.lock_path = lock_path - self.exclusive = exclusive - self.timeout = timeout - self._fd: int = -1 - - def __enter__(self) -> "FileLock": - flags = os.O_WRONLY | os.O_CREAT if self.exclusive else os.O_RDONLY | os.O_CREAT - lock_type = fcntl.LOCK_EX if self.exclusive else fcntl.LOCK_SH - self._fd = os.open(str(self.lock_path), flags) - deadline = time.monotonic() + self.timeout - while time.monotonic() < deadline: - try: - fcntl.flock(self._fd, lock_type | fcntl.LOCK_NB) - return self - except OSError: - time.sleep(0.1) - os.close(self._fd) - self._fd = -1 - raise RuntimeError(f"Timed out waiting for lock: {self.lock_path}") - - def __exit__(self, *exc) -> None: - if self._fd >= 0: - fcntl.flock(self._fd, fcntl.LOCK_UN) - os.close(self._fd) - self._fd = -1 - - -# --------------------------------------------------------------------------- -# JIT cache decorator -# --------------------------------------------------------------------------- - - -def jit_cache(fn): - """Decorator that caches compiled CuTe DSL kernels in-memory and on disk. - - The decorated function should return a compiled kernel (i.e. call cute.compile). - The disk cache key is (fn.__qualname__, *args, **sorted_kwargs). - """ - cache = {} - hits = 0 - misses = 0 - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - nonlocal hits, misses - cache_key = args + tuple(sorted(kwargs.items())) if kwargs else args - - # 1. In-memory hit - if cache_key in cache: - hits += 1 - return _noop_kernel if COMPILE_ONLY else cache[cache_key] - - # 2. Disk hit - disk_key = (fn.__qualname__,) + cache_key - if CACHE_ENABLED: - sha = _key_to_hash(disk_key) - cache_path = get_cache_path() / _compute_source_fingerprint() - cache_path.mkdir(parents=True, exist_ok=True) - o_path = cache_path / f"{sha}.o" - lock_path = cache_path / f"{sha}.lock" - try: - with FileLock(lock_path, exclusive=False, timeout=LOCK_TIMEOUT): - if o_path.exists(): - m = cute.runtime.load_module(str(o_path), enable_tvm_ffi=True) - loaded = m[EXPORT_FUNC_NAME] - cache[cache_key] = loaded - hits += 1 - return _noop_kernel if COMPILE_ONLY else loaded - except RuntimeError: - pass - - # 3. Compile - misses += 1 - compiled_fn = fn(*args, **kwargs) - - # 4. Store - cache[cache_key] = compiled_fn - if CACHE_ENABLED: - try: - with FileLock(lock_path, exclusive=True, timeout=LOCK_TIMEOUT): - if not o_path.exists(): - o_path.parent.mkdir(parents=True, exist_ok=True) - compiled_fn.export_to_c( - object_file_path=str(o_path), - function_name=EXPORT_FUNC_NAME, - ) - except Exception as e: - print(f"quack cache: export failed for key {sha}: {e}") - - return _noop_kernel if COMPILE_ONLY else compiled_fn - - def cache_clear(): - nonlocal hits, misses - cache.clear() - hits = 0 - misses = 0 - - def cache_info(): - return CacheInfo(hits=hits, misses=misses, maxsize=None, currsize=len(cache)) - - wrapper.cache = cache - wrapper.cache_clear = cache_clear - wrapper.cache_info = cache_info - return wrapper diff --git a/build/torch-cuda/quack/compile_utils.py b/build/torch-cuda/quack/compile_utils.py index 4375594669c8f12d6a79d8878316271cb819568a..174031aa86d23c906651c62ae06eb49cc1ace827 100644 --- a/build/torch-cuda/quack/compile_utils.py +++ b/build/torch-cuda/quack/compile_utils.py @@ -6,10 +6,20 @@ import cutlass.cute as cute def make_fake_tensor(dtype, shape, divisibility=1, leading_dim=-1) -> Optional[cute.Tensor]: - if leading_dim < 0: - leading_dim = len(shape) + leading_dim + """Build a fake CuTe tensor with dynamic (sym) strides for tensor-free compilation. + + ``leading_dim`` selects the dim whose stride is statically 1 (matching + ``from_dlpack(...).mark_layout_dynamic(leading_dim=...)``). Pass + ``leading_dim=None`` for a fully-dynamic layout with no static stride-1 dim + (matching ``mark_layout_dynamic()`` on a tensor without a contiguous dim). + + ``divisibility`` is in elements; ``assumed_align`` (bytes) is + ``divisibility * dtype.width // 8``. + """ if dtype is None: return None + if leading_dim is not None and leading_dim < 0: + leading_dim = len(shape) + leading_dim stride = tuple( cute.sym_int64(divisibility=divisibility) if i != leading_dim else 1 for i in range(len(shape)) @@ -17,3 +27,8 @@ def make_fake_tensor(dtype, shape, divisibility=1, leading_dim=-1) -> Optional[c return cute.runtime.make_fake_tensor( dtype, shape, stride=stride, assumed_align=divisibility * dtype.width // 8 ) + + +def make_fake_stream(): + """Fake CUDA stream for tensor-free compilation (real stream comes from the TVM FFI env).""" + return cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) diff --git a/build/torch-cuda/quack/complex.py b/build/torch-cuda/quack/complex.py new file mode 100644 index 0000000000000000000000000000000000000000..e2e2b40940346253c812e58c5f62716b5fb16621 --- /dev/null +++ b/build/torch-cuda/quack/complex.py @@ -0,0 +1,292 @@ +"""Complex64 element type for CuTe-DSL kernels. + +Single-precision complex (re + imj) carried as f64-packed bits (re in the low +32 bits, im in the high 32 bits). f64 is on `cute.MemRefType`'s element-type +allowlist; the natural `complex` MLIR type is not. Arithmetic methods +unpack each f64 into two Float32 lanes, compute, and repack -- the bitcasts +are folded out by ptxas. + +Inherits from `Float32` (with `width=64, mlir_type=T.f64` overrides) so that +Python's subclass-precedence rule routes `Float32 OP Complex64` to our +reflected `__r*__` operators before Numeric's promotion logic sees the +operands. Without this, Float32-on-the-LEFT would silently promote through +Float64 conversion and corrupt the packed bits. + +Boundary convention (tvm-ffi): the compiled kernel's ABI sees f64 storage. +At the call site, pass `torch.complex64` tensors as `t.view(torch.float64)` +-- use `complex_storage(t)` for the conversion. + +See `AI/complex64_design_notes.md` for the why and what's been validated. +""" + +from __future__ import annotations + +import ctypes + +import numpy as np +import torch + +import cutlass.cute as cute +from cutlass import Float32, Numeric +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith +from cutlass._mlir.extras import types as T +from cutlass._mlir_helpers.arith import bitcast as _bitcast +from cutlass.base_dsl.typing import FloatMeta + + +class Complex64(Float32, metaclass=FloatMeta, width=64, mlir_type=T.f64): + """Complex64 carried as f64-packed bits (re in low 32, im in high 32). + + `tensor.element_type is Complex64` inside the kernel; indexing returns + Complex64 instances; `+`, `-`, `*`, `__neg__`, and `conj()` work natively. + """ + + def __init__(self, x, im=None, *, loc=None, ip=None): + # Two-arg lane form: Complex64(re, im) packs (re, im) into f64 bits. + # Coerce both through Float32 so int / float / Float32 / ir.Value(f32) + # all work as inputs. + if im is not None: + # Static fast path: both args are Python int/float -- no MLIR + # context needed (lets host-side code call Complex64(2.5, -1.5)). + if isinstance(x, (int, float)) and isinstance(im, (int, float)): + Complex64.__init__(self, complex(x, im), loc=loc, ip=ip) + return + re_ssa = Float32(x).ir_value() + im_ssa = Float32(im).ir_value() + Numeric.__init__(self, Complex64._pack_ssa(re_ssa, im_ssa)) + return + + # Same-type copy MUST be checked first. `_cvt_to_dest` (cute/tensor.py) + # calls `data.to(element_type)` on every tensor write, which becomes + # `Complex64(complex_instance)`; falling through to the generic-Numeric + # branch below would re-pack as (real_view_of_packed_bits, 0) and + # silently corrupt the data. + if isinstance(x, Complex64): + Numeric.__init__(self, x.value) + return + + if isinstance(x, complex): + f64_val = _pack_python_complex(x) + Numeric.__init__(self, f64_val) + return + + if isinstance(x, ir.Value): + if x.type == T.f64(): + # Already in our storage form (loaded from a Complex64 tensor, + # or output of _pack_ssa). + Numeric.__init__(self, x) + return + if x.type == T.f32(): + packed = Complex64._pack_ssa(x, arith.constant(T.f32(), 0.0)) + Numeric.__init__(self, packed) + return + raise TypeError(f"Complex64: ir.Value of unsupported type {x.type}") + + if isinstance(x, Numeric): + # Float32, Int32, Float64, etc. -> coerce real lane through Float32, + # imag lane = 0. Float32(Float32) is a no-op, so this also handles + # the Float32 case cleanly. + re_ssa = Float32(x).ir_value() + packed = Complex64._pack_ssa(re_ssa, arith.constant(T.f32(), 0.0)) + Numeric.__init__(self, packed) + return + + if isinstance(x, (int, float)): + Complex64.__init__(self, complex(x, 0.0), loc=loc, ip=ip) + return + + raise TypeError(f"Complex64: unsupported source type {type(x)}") + + # ---- packing / unpacking primitives -------------------------------- + + @staticmethod + def _pack_ssa(re_f32, im_f32): + """Pack two f32 SSA lanes into one f64 SSA value (re lo, im hi).""" + re_i32 = _bitcast(re_f32, T.i32()) + im_i32 = _bitcast(im_f32, T.i32()) + re_i64 = arith.extui(T.i64(), re_i32) + im_i64 = arith.extui(T.i64(), im_i32) + hi = arith.shli(im_i64, arith.constant(T.i64(), 32)) + return _bitcast(arith.ori(re_i64, hi), T.f64()) + + def _unpack(self): + """Split self -> (re_f32, im_f32) as Float32 SSA values.""" + i64_ssa = _bitcast(self.ir_value(), T.i64()) + lo32 = arith.trunci(T.i32(), i64_ssa) + hi32 = arith.trunci(T.i32(), arith.shrui(i64_ssa, arith.constant(T.i64(), 32))) + return Float32(_bitcast(lo32, T.f32())), Float32(_bitcast(hi32, T.f32())) + + @staticmethod + def from_re_im(re: Float32, im: Float32) -> "Complex64": + """Build a Complex64 from two Float32 SSA lanes. + + Equivalent to `Complex64(re, im)`; kept as an explicit name for the + hot-path call sites that want to skip the Float32 coercion in __init__. + """ + return Complex64(Complex64._pack_ssa(re.ir_value(), im.ir_value())) + + # Internal alias used by arithmetic methods. + _from_re_im = from_re_im + + # ---- accessors ------------------------------------------------------ + + def real(self) -> Float32: + re, _ = self._unpack() + return re + + def imag(self) -> Float32: + _, im = self._unpack() + return im + + def conj(self) -> "Complex64": + re, im = self._unpack() + return Complex64._from_re_im(re, -im) + + # ---- arithmetic ----------------------------------------------------- + + def __add__(self, other, *, loc=None, ip=None): + a_re, a_im = self._unpack() + b_re, b_im = _other_lanes(other) + return Complex64._from_re_im(a_re + b_re, a_im + b_im) + + def __radd__(self, other, *, loc=None, ip=None): + return self.__add__(other, loc=loc, ip=ip) + + def __sub__(self, other, *, loc=None, ip=None): + a_re, a_im = self._unpack() + b_re, b_im = _other_lanes(other) + return Complex64._from_re_im(a_re - b_re, a_im - b_im) + + def __rsub__(self, other, *, loc=None, ip=None): + a_re, a_im = self._unpack() + b_re, b_im = _other_lanes(other) + return Complex64._from_re_im(b_re - a_re, b_im - a_im) + + def __mul__(self, other, *, loc=None, ip=None): + a_re, a_im = self._unpack() + if isinstance(other, Complex64): + b_re, b_im = other._unpack() + return Complex64._from_re_im(a_re * b_re - a_im * b_im, a_re * b_im + a_im * b_re) + # Real scalar: (re, im) * s = (re*s, im*s) + s = Float32(other) + return Complex64._from_re_im(a_re * s, a_im * s) + + def __rmul__(self, other, *, loc=None, ip=None): + return self.__mul__(other, loc=loc, ip=ip) + + def __neg__(self, *, loc=None, ip=None): + re, im = self._unpack() + return Complex64._from_re_im(-re, -im) + + # ---- runtime arg passing ------------------------------------------- + + def __c_pointers__(self): + # Scalar Complex64 args travel as 8 bytes (the packed-as-f64 value). + if not isinstance(self.value, float): + raise ValueError( + "Complex64 with a dynamic SSA value cannot be passed as a " + "kernel argument; only static values are supported" + ) + return [ctypes.cast(ctypes.pointer(ctypes.c_double(self.value)), ctypes.c_void_p)] + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _pack_python_complex(c: complex) -> float: + """Compute the f64 representation of a complex bit-packed (re, im).""" + re_b = int(np.float32(c.real).view(np.uint32)) + im_b = int(np.float32(c.imag).view(np.uint32)) + return float(np.uint64((im_b << 32) | re_b).view(np.float64)) + + +def _other_lanes(other): + """Unpack the RHS of a binary op into (re_f32, im_f32) Float32 lanes.""" + if isinstance(other, Complex64): + return other._unpack() + return Float32(other), Float32(0.0) + + +def _retag_as_complex64(t): + """Restore `t.element_type is Complex64` after a code path that derived it + from MLIR (where complex64 collapses to Float64 / Int64 because + `Numeric.from_mlir_type` is a many-to-one lookup).""" + t._dtype = Complex64 + return t + + +# --------------------------------------------------------------------------- +# Public helpers +# --------------------------------------------------------------------------- + + +def allocate_smem_complex( + allocator, + layout_or_shape, + byte_alignment: int = 16, + swizzle=None, +): + """Allocate a `Complex64` smem tensor. + + Wraps `cutlass.utils.SmemAllocator.allocate_tensor(Complex64, ...)` and + re-tags the result so `tensor.element_type is Complex64`. Without the + re-tag, the JIT-side tensor's element_type is `Float64` (derived from the + f64 memref) and writes go through `Complex64.to(Float64)` and corrupt the + packed bits. + """ + t = allocator.allocate_tensor( + Complex64, layout_or_shape, byte_alignment=byte_alignment, swizzle=swizzle + ) + return _retag_as_complex64(t) + + +def recast_to_complex64(src: cute.Tensor) -> cute.Tensor: + """Recast any tensor (e.g. Float32, Int64) to a `Complex64` tensor. + + Wraps `cute.recast_tensor(src, Complex64)` and re-tags the result. Same + dtype-loss bug as `allocate_smem_complex` -- the underlying recast goes + through `make_tensor`, which derives element_type from the MLIR memref + (here f64) and gets back Float64. + """ + return _retag_as_complex64(cute.recast_tensor(src, Complex64)) + + +def complex_storage(t: torch.Tensor) -> torch.Tensor: + """View a `torch.complex64` tensor as `torch.float64` with the same memory. + + Compiled kernels declared with `Complex64` element type have an f64 ABI; + use this at the boundary to satisfy tvm-ffi's dtype check without copying. + """ + if t.dtype == torch.float64: + return t + if t.dtype != torch.complex64: + raise TypeError( + f"complex_storage expects torch.complex64 (or torch.float64 for " + f"already-converted storage), got {t.dtype}" + ) + return t.view(torch.float64) + + +# --------------------------------------------------------------------------- +# tvm-ffi registration +# --------------------------------------------------------------------------- + + +def _register_with_tvm_ffi() -> None: + """Teach tvm-ffi that Complex64 has an f64 ABI. + + Both `NumericToTVMFFIDtype` (the type->dtype-string lookup) and + `AcceptableNumericTypesForScalar` (the allowlist for scalar kernel args) + are plain Python collections, so we extend them at import time. + """ + from cutlass.cute import _tvm_ffi_args_spec_converter as _cv + + _cv.NumericToTVMFFIDtype.setdefault(Complex64, "float64") + if Complex64 not in _cv.AcceptableNumericTypesForScalar: + _cv.AcceptableNumericTypesForScalar.append(Complex64) + + +_register_with_tvm_ffi() diff --git a/build/torch-cuda/quack/copy_utils.py b/build/torch-cuda/quack/copy_utils.py index 4966d0edd20d8ea2936952c20498e2216ec5b212..0e01211ea1c4a57057c77866a062ecb8250383c2 100644 --- a/build/torch-cuda/quack/copy_utils.py +++ b/build/torch-cuda/quack/copy_utils.py @@ -1,17 +1,19 @@ -# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. +# Copyright (c) 2025-2026, QuACK team. -from typing import Optional, Type, Tuple, Callable, Sequence +from typing import Any, Optional, Type, Tuple, Callable, Sequence from functools import partial import cutlass import cutlass.cute as cute +import cutlass.utils.blackwell_helpers as sm100_utils from cutlass import Int32, Int16, Boolean, const_expr -from cutlass.cute.nvgpu import cpasync, warp, warpgroup +from cutlass.base_dsl.arch import Arch +from cutlass.cute.nvgpu import cpasync, tcgen05, warp from cutlass.cute.nvgpu.tcgen05.mma import CtaGroup # noqa from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils import LayoutEnum, block_copy import cutlass.pipeline -from cutlass._mlir.dialects import llvm from cutlass._mlir import ir from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir @@ -20,6 +22,97 @@ from .utils import make_vector Sm100MmaPeerBitMask = 0xFEFFFFFF +_TCGEN05_TMEM_OPS = ( + tcgen05.Ld16x128bOp, + tcgen05.Ld16x256bOp, + tcgen05.Ld16x32bx2Op, + tcgen05.Ld16x64bOp, + tcgen05.Ld32x32bOp, + tcgen05.LdRed16x32bx2Op, + tcgen05.LdRed32x32bOp, + tcgen05.St16x128bOp, + tcgen05.St16x256bOp, + tcgen05.St16x32bx2Op, + tcgen05.St16x64bOp, + tcgen05.St32x32bOp, +) +_TCGEN05_TMEM_STORE_OPS = ( + tcgen05.St16x128bOp, + tcgen05.St16x256bOp, + tcgen05.St16x32bx2Op, + tcgen05.St16x64bOp, + tcgen05.St32x32bOp, +) + + +def tmem_store_atom_from_load_atom( + copy_atom_t2r: Any, + src_dtype: Type[cutlass.Numeric], + dst_dtype: Type[cutlass.Numeric], +) -> cute.CopyAtom: + """Return the matching tcgen05 R2T store atom for a selected T2R load atom. + + `src_dtype` is the register fragment dtype loaded by T2R; `dst_dtype` is + the TMEM element dtype to store. Ratio 1 uses CUTLASS's operation-family + mapping directly. Ratio 2 is intentionally narrow: we allow the current + Ld32x32b path by halving repeat, and the widest 16dp path by halving the + vector width. Narrower 16dp cross-family mappings are not mirrored by + CUTLASS's same-family helper, so they assert until validated. + + C++ CuTe's operation-family mapping is `cute::TMEM::tmem_load_to_store`: + https://github.com/NVIDIA/cutlass/blob/main/include/cute/atom/copy_traits_sm100.hpp#L3274 + """ + load_op = copy_atom_t2r.op if const_expr(hasattr(copy_atom_t2r, "op")) else copy_atom_t2r + if const_expr(hasattr(load_op, "op")): + load_op = load_op.op + assert src_dtype.width >= dst_dtype.width, "TMEM R2T helper only supports narrowing stores" + assert src_dtype.width % dst_dtype.width == 0, "TMEM source/destination widths must divide" + ratio = src_dtype.width // dst_dtype.width + assert ratio in (1, 2), "TMEM R2T helper only supports src/dst width ratio 1 or 2" + repeat = load_op.repeat + unpack = tcgen05.Unpack.NONE + if const_expr(getattr(load_op, "pack", None) == tcgen05.Pack.PACK_16b_IN_32b): + unpack = tcgen05.Unpack.UNPACK_32b_IN_16b + if const_expr(isinstance(load_op, tcgen05.Ld16x64bOp)): + assert ratio == 1, "No validated ratio-2 store mapping for Ld16x64bOp" + store_op = tcgen05.St16x64bOp(repeat, unpack) + elif const_expr(isinstance(load_op, tcgen05.Ld16x128bOp)): + assert ratio == 1, "No validated ratio-2 store mapping for Ld16x128bOp" + store_op = tcgen05.St16x128bOp(repeat, unpack) + elif const_expr(isinstance(load_op, tcgen05.Ld16x256bOp)): + store_op = ( + tcgen05.St16x256bOp(repeat, unpack) + if const_expr(ratio == 1) + else tcgen05.St16x128bOp(repeat, unpack) + ) + elif const_expr(isinstance(load_op, tcgen05.Ld16x32bx2Op)): + assert ratio == 1, "No validated ratio-2 store mapping for Ld16x32bx2Op" + store_op = tcgen05.St16x32bx2Op(repeat, unpack) + elif const_expr(isinstance(load_op, tcgen05.Ld32x32bOp)): + if const_expr(ratio == 2): + assert repeat.value % 2 == 0, "Ld32x32b ratio-2 store needs even repeat" + repeat = tcgen05.Repetition(repeat.value // 2) + store_op = tcgen05.St32x32bOp(repeat, unpack) + else: + raise TypeError(f"Unsupported TMEM load op for store conversion: {type(load_op)}") + return cute.make_copy_atom(store_op, dst_dtype) + + +def _tmem_copy_reg_tv_layout(tiled_copy: cute.TiledCopy): + """Return the register-side TV layout for a tcgen05 tmem copy.""" + op = tiled_copy.op + if const_expr(hasattr(op, "op")): + op = op.op + if const_expr(isinstance(op, _TCGEN05_TMEM_OPS)): + # TMEM stores read from registers; all other TMEM copy ops here write + # registers, including LdRed* reductions that upstream is_tmem_load + # intentionally does not classify as plain loads. + return ( + tiled_copy.layout_src_tv_tiled + if const_expr(isinstance(op, _TCGEN05_TMEM_STORE_OPS)) + else tiled_copy.layout_dst_tv_tiled + ) + raise TypeError(f"Cannot infer tmem copy direction from tiled_copy.op={op}") @dsl_user_op @@ -36,9 +129,7 @@ def cvt_copy( ) -> None: assert isinstance(src.iterator, cute.Pointer) and src.memspace == cute.AddressSpace.rmem if const_expr(src.element_type != dst.element_type): - src_cvt = cute.make_rmem_tensor_like(src, dst.element_type) - src_cvt.store(src.load().to(dst.element_type)) - src = src_cvt + src = src.to(dst.element_type, loc=loc, ip=ip) if const_expr(retile): src = tiled_copy.retile(src) cute.copy(tiled_copy, src, dst, pred=pred, loc=loc, ip=ip, **kwargs) @@ -75,6 +166,13 @@ def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: return dst +@dsl_user_op +def contiguous(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: + dst = cute.make_rmem_tensor(src.shape, src.element_type, loc=loc, ip=ip) + cute.autovec_copy(src, dst, loc=loc, ip=ip) + return dst + + @dsl_user_op def load_s2r_retile( tiled_copy: cute.TiledCopy, @@ -95,11 +193,22 @@ def load_s2r_retile( @dsl_user_op def load_t2r( - thr_copy: cute.ThrCopy, shape: cute.Shape, src: cute.Tensor, *, loc=None, ip=None + tiled_copy: cute.TiledCopy, + src: cute.Tensor, + *, + fence: bool = False, + loc=None, + ip=None, ) -> cute.Tensor: - cDst = cute.make_identity_tensor(shape) - dst = cute.make_rmem_tensor(thr_copy.partition_D(cDst).shape, src.element_type, loc=loc, ip=ip) - cute.copy(thr_copy, src, dst, loc=loc, ip=ip) + """Load one tmem tile partition into rmem, deriving the rmem shape from `src`. + + `src` should already be indexed to the tile being copied, with any + stage/subtile modes removed. + """ + dst = tmem_reg_frag(tiled_copy, src, loc=loc, ip=ip) + cute.copy(tiled_copy, src, dst, loc=loc, ip=ip) + if const_expr(fence): + cute.arch.fence_view_async_tmem_load() return dst @@ -321,18 +430,14 @@ def as_position_independent_swizzle_tensor(tensor: cute.Tensor) -> cute.Tensor: return cute.make_tensor(cute.recast_ptr(tensor.iterator, dtype=tensor.element_type), new_layout) -def partition_D_position_independent( - thr_copy: cute.core.ThrCopy, tensor: cute.Tensor -) -> cute.Tensor: +def partition_D_position_independent(thr_copy: cute.ThrCopy, tensor: cute.Tensor) -> cute.Tensor: return cute.make_tensor( swizzle_ptr(thr_copy.partition_D(tensor).iterator), thr_copy.partition_D(as_position_independent_swizzle_tensor(tensor)).layout, ) -def partition_S_position_independent( - thr_copy: cute.core.ThrCopy, tensor: cute.Tensor -) -> cute.Tensor: +def partition_S_position_independent(thr_copy: cute.ThrCopy, tensor: cute.Tensor) -> cute.Tensor: return cute.make_tensor( swizzle_ptr(thr_copy.partition_S(tensor).iterator), thr_copy.partition_S(as_position_independent_swizzle_tensor(tensor)).layout, @@ -373,12 +478,12 @@ def sm90_get_smem_load_op( def get_smem_store_atom( - arch: cutlass.Constexpr[int], element_type: Type[cute.Numeric], transpose: bool = False, major_mode_size: Optional[int] = None, ) -> cute.CopyAtom: - if const_expr(arch < 90 or element_type.width != 16): + arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum() + if const_expr(arch < Arch.sm_90 or element_type.width != 16): return cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), element_type, @@ -397,12 +502,12 @@ def get_smem_store_atom( def get_smem_load_atom( - arch: cutlass.Constexpr[int], element_type: Type[cute.Numeric], transpose: bool = False, major_mode_size: Optional[int] = None, ) -> cute.CopyAtom: - if const_expr(arch < 90 or element_type.width != 16): + arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum() + if const_expr(arch < Arch.sm_90 or element_type.width != 16): return cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), element_type, @@ -421,26 +526,35 @@ def get_smem_load_atom( def get_smem_store_C( - tiled_mma: cute.TiledMma, + tiled_mma: cute.TiledMma | cute.TiledCopy, sC: cute.Tensor, tidx: Int32, - arch: int, transpose: bool = False, position_independent=False, major_mode_size: Optional[int] = None, ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]: dtype = sC.element_type - copy_atom = get_smem_store_atom(arch, dtype, transpose, major_mode_size=major_mode_size) - tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma) + if const_expr(isinstance(tiled_mma, cute.TiledCopy)): + tiled_copy_t2r = tiled_mma + layout = LayoutEnum.COL_MAJOR if const_expr(transpose) else LayoutEnum.ROW_MAJOR + copy_atom = sm100_utils.get_smem_store_op( + layout, dtype, tiled_copy_t2r.value_type, tiled_copy_t2r + ) + tiled_copy = cute.make_tiled_copy_D(copy_atom, tiled_copy_t2r) + else: + copy_atom = get_smem_store_atom(dtype, transpose, major_mode_size=major_mode_size) + tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma) thr_copy = tiled_copy.get_slice(tidx) if const_expr(not position_independent): tRS_sC = thr_copy.partition_D(sC) else: tRS_sC = partition_D_position_independent(thr_copy, sC) - def copy_fn(src: cute.Tensor, dst_idx: Optional[Int32] = None, **new_kwargs): - dst_tensor = tRS_sC if const_expr(dst_idx is None) else tRS_sC[None, None, None, dst_idx] + def copy_fn(src: cute.Tensor, dst_idx: Optional[Int32] = None, fence=False, **new_kwargs): + dst_tensor = tRS_sC if const_expr(dst_idx is None) else tRS_sC[..., dst_idx] cvt_copy(tiled_copy, src, dst_tensor, retile=True, **new_kwargs) + if const_expr(fence): + cute.arch.fence_view_async_shared() return copy_fn, thr_copy, tRS_sC @@ -449,19 +563,18 @@ def get_smem_load_C( tiled_mma: cute.TiledMma, sC: cute.Tensor, tidx: Int32, - arch: int, transpose: bool = False, position_independent=False, ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]: dtype = sC.element_type - copy_atom = get_smem_load_atom(arch, dtype, transpose) + copy_atom = get_smem_load_atom(dtype, transpose) tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma) thr_copy = tiled_copy.get_slice(tidx) if const_expr(not position_independent): tSR_sC = thr_copy.partition_S(sC) else: tSR_sC = partition_S_position_independent(thr_copy, sC) - copy_atom_RS = get_smem_store_atom(arch, dtype, transpose) + copy_atom_RS = get_smem_store_atom(dtype, transpose) thr_copy_RS = cute.make_tiled_copy_C(copy_atom_RS, tiled_mma).get_slice(tidx) tRS_shape = thr_copy_RS.partition_S(cute.make_identity_tensor(sC.shape[:2])).shape @@ -475,10 +588,18 @@ def get_smem_load_C( def epilog_smem_copy_atom( tiled_mma: cute.TiledMma, epi_tile: cute.Shape, transpose: bool = False ) -> cute.TiledCopy: - copy_atom_C = cute.make_copy_atom( - warp.StMatrix8x8x16bOp(transpose, num_matrices=4 if epi_tile[1] % 16 == 0 else 2), - cutlass.Float16, # this is just to get the right source layout - ) + arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum() + if const_expr(arch < Arch.sm_90): + copy_atom_C = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Float16, # this is just to get the right source layout + num_bits_per_copy=(2 if not transpose else 1) * cutlass.Float16.width, + ) + else: + copy_atom_C = cute.make_copy_atom( + warp.StMatrix8x8x16bOp(transpose, num_matrices=4 if epi_tile[1] % 16 == 0 else 2), + cutlass.Float16, # this is just to get the right source layout + ) tiled_copy_C_atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma) return tiled_copy_C_atom @@ -488,13 +609,12 @@ def get_smem_store_epi( epi_tile: cute.Shape, sC: Optional[cute.Tensor], tidx: Int32, - arch: int, transpose: bool = False, position_independent=False, ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor, cute.Tensor]: dtype = sC.element_type if const_expr(sC is not None) else cutlass.Float16 + copy_atom = get_smem_store_atom(dtype, transpose) tiled_copy_C_atom = epilog_smem_copy_atom(tiled_mma, epi_tile) - copy_atom = get_smem_store_atom(arch, dtype, transpose) tiled_copy = cute.make_tiled_copy_S(copy_atom, tiled_copy_C_atom) thr_copy = tiled_copy.get_slice(tidx) tRS_sC = None @@ -515,11 +635,11 @@ def get_smem_store_epi( def get_smem_store_A( - tiled_mma: cute.TiledMma, sA: cute.Tensor, tidx: Int32, arch: int, position_independent=False + tiled_mma: cute.TiledMma, sA: cute.Tensor, tidx: Int32, position_independent=False ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]: dtype = sA.element_type - transpose = tiled_mma.op.a_major_mode == warpgroup.OperandMajorMode.MN - copy_atom = get_smem_store_atom(arch, dtype, transpose) + transpose = tiled_mma.op.a_major_mode == cute.nvgpu.OperandMajorMode.MN + copy_atom = get_smem_store_atom(dtype, transpose) tiled_copy = cute.make_tiled_copy_A(copy_atom, tiled_mma) thr_copy = tiled_copy.get_slice(tidx) if const_expr(not position_independent): @@ -537,13 +657,12 @@ def get_smem_load_A( tiled_mma: cute.TiledMma, sA: cute.Tensor, tidx: Int32, - arch: int, with_dst_tensor: bool = False, position_independent=False, ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]: dtype = sA.element_type - transpose = tiled_mma.op.a_major_mode == warpgroup.OperandMajorMode.MN - copy_atom = get_smem_load_atom(arch, dtype, transpose) + transpose = tiled_mma.op.a_major_mode == cute.nvgpu.OperandMajorMode.MN + copy_atom = get_smem_load_atom(dtype, transpose) tiled_copy = cute.make_tiled_copy_A(copy_atom, tiled_mma) thr_copy = tiled_copy.get_slice(tidx) if const_expr(not position_independent): @@ -563,6 +682,79 @@ def get_smem_load_A( return copy_fn if not with_dst_tensor else copy_fn_w_dst_tensor, thr_copy, tSR_sA +def _cpasync_reduction_kind_name(reduction_kind: Any) -> str: + name = ( + reduction_kind.lower() if isinstance(reduction_kind, str) else reduction_kind.name.lower() + ) + assert name in {"add", "min", "max", "inc", "dec", "and", "or", "xor"}, ( + f"Unsupported cp.reduce.async.bulk reduction kind: {reduction_kind}" + ) + return name + + +def _cpasync_bulk_reduce_suffix( + reduction_kind: Any, + dtype: Type[cutlass.Numeric], +) -> str: + op = _cpasync_reduction_kind_name(reduction_kind) + if dtype is cutlass.Float16: + assert op in {"add", "min", "max"}, f"{op} is not supported for f16 bulk reduce" + return f"{op}.noftz.f16" if op == "add" else f"{op}.f16" + if dtype is cutlass.BFloat16: + assert op in {"add", "min", "max"}, f"{op} is not supported for bf16 bulk reduce" + return f"{op}.noftz.bf16" if op == "add" else f"{op}.bf16" + if dtype is cutlass.Float32: + assert op == "add", f"{op} is not supported for f32 bulk reduce" + return "add.f32" + if dtype is cutlass.Float64: + assert op == "add", f"{op} is not supported for f64 bulk reduce" + return "add.f64" + + signed = getattr(dtype, "signed", None) + width = getattr(dtype, "width", None) + if signed is not None: + assert width in (32, 64), f"Unsupported integer bulk-reduce width: {width}" + if op in {"and", "or", "xor"}: + return f"{op}.b{width}" + if op in {"min", "max", "add"}: + return f"{op}.{'s' if signed else 'u'}{width}" + assert op in {"inc", "dec"} and dtype is cutlass.Uint32, ( + f"{op} bulk reduce is only supported for u32" + ) + return f"{op}.u32" + + raise TypeError(f"Unsupported cp.reduce.async.bulk dtype: {dtype}") + + +@dsl_user_op +def cpasync_bulk_s2g( + smem_ptr: cute.Pointer, + gmem_ptr: cute.Pointer, + store_bytes: int | Int32, + *, + reduction_kind: Optional[Any] = None, + dtype: Optional[Type[cutlass.Numeric]] = None, + loc=None, + ip=None, +): + smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip) + if reduction_kind is None: + ptx = "cp.async.bulk.global.shared::cta.bulk_group [{$r0}], [{$r1}], {$r2};" + else: + assert dtype is not None, "dtype is required for cp.reduce.async.bulk" + ptx = ( + "cp.reduce.async.bulk.global.shared::cta.bulk_group." + f"{_cpasync_bulk_reduce_suffix(reduction_kind, dtype)} " + "[{$r0}], [{$r1}], {$r2};" + ) + cute.arch.inline_ptx( + ptx, + read_only_args=[gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes)], + loc=loc, + ip=ip, + ) + + @dsl_user_op def cpasync_reduce_bulk_add_f32( smem_ptr: cute.Pointer, @@ -572,18 +764,14 @@ def cpasync_reduce_bulk_add_f32( loc=None, ip=None, ): - smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() - # cache_hint = cutlass.Int64(0x14F0000000000000) # EVICT_LAST - llvm.inline_asm( - None, - [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value()], - "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;", - "l,r,r", - # [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value(), cache_hint.ir_value()], - # "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.f32 [$0], [$1], $2, $3;", - # "l,r,r,l", - has_side_effects=True, - is_align_stack=False, + cpasync_bulk_s2g( + smem_ptr, + gmem_ptr, + store_bytes, + reduction_kind=cpasync.ReductionOp.ADD, + dtype=cutlass.Float32, + loc=loc, + ip=ip, ) @@ -674,33 +862,34 @@ def tma_gather4_load( """ if len(row_indices) != 4: raise ValueError(f"gather4 requires exactly 4 row indices, got {len(row_indices)}") - col_val = Int32(col_idx).ir_value() - row_vals = [Int32(row_idx).ir_value() for row_idx in row_indices] + col_val = Int32(col_idx) + row_vals = [Int32(row_idx) for row_idx in row_indices] # Convert pointers to integer addresses - desc_addr = tma_desc_ptr.toint(loc=loc, ip=ip).ir_value() - dst_addr = dst_smem_ptr.toint(loc=loc, ip=ip).ir_value() + desc_addr = tma_desc_ptr.toint(loc=loc, ip=ip) + dst_addr = dst_smem_ptr.toint(loc=loc, ip=ip) mbar_addr = mbarrier_ptr.toint(loc=loc, ip=ip) if num_cta > 1: # Executed by both CTAs. Set peer bit to 0 so that the # transaction bytes will update CTA0's barrier. mbar_addr = mbar_addr & Sm100MmaPeerBitMask - mbar_addr = mbar_addr.ir_value() + mbar_addr = Int32(mbar_addr) # Handle multicast_mask - may already be ir.Value or Python int multicast_mask_val = None if multicast_mask is not None: - multicast_mask_val = Int16(multicast_mask).ir_value() + multicast_mask_val = Int16(multicast_mask) assert multicast_mask_val is None, "multicast is not supported yet" # Emit inline PTX for TMA gather4 # PTX: cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes # [dstMem], [tensorMap, {col, row0, row1, row2, row3}], [smem_bar]; ptx = ( - f"cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes.cta_group::{num_cta} " - "[$0], [$1, {$2, $3, $4, $5, $6}], [$7];" + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes." + f"cta_group::{num_cta} " + "[{$r0}], [{$r1}, {{$r2}, {$r3}, {$r4}, {$r5}, {$r6}}], [{$r7}];" ) - llvm.inline_asm( - None, - [ + cute.arch.inline_ptx( + ptx, + read_only_args=[ dst_addr, desc_addr, col_val, @@ -710,10 +899,6 @@ def tma_gather4_load( row_vals[3], mbar_addr, ], - ptx, - "r,l,r,r,r,r,r,r", # constraints: register, long, 6x register - has_side_effects=True, - is_align_stack=False, loc=loc, ip=ip, ) @@ -723,34 +908,137 @@ def cpasync_bulk_get_copy_fn( src_tensor: cute.Tensor, dst_tensor: cute.Tensor, single_stage: bool = False, + reduction_kind: Optional[cute.nvgpu.cpasync.ReductionKind] = None, **kwargs, ) -> Callable: + src_is_smem = const_expr( + isinstance(src_tensor.iterator, cute.Pointer) + and src_tensor.memspace == cute.AddressSpace.smem + ) + dst_is_smem = const_expr( + isinstance(dst_tensor.iterator, cute.Pointer) + and dst_tensor.memspace == cute.AddressSpace.smem + ) + if const_expr(reduction_kind is not None): + assert src_is_smem and not dst_is_smem, "cp.reduce.async.bulk only supports SMEM -> GMEM" group_rank_src = const_expr(cute.rank(src_tensor) - (1 if not single_stage else 0)) group_rank_dst = const_expr(cute.rank(dst_tensor) - (1 if not single_stage else 0)) # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK) src = cute.group_modes(src_tensor, 0, group_rank_src) dst = cute.group_modes(dst_tensor, 0, group_rank_dst) + if const_expr(src_is_smem and not dst_is_smem): + + def copy_bulk_s2g(src_idx, dst_idx, **new_kwargs): + store_bytes = const_expr(cute.size(src.shape[:-1]) * src.element_type.width // 8) + with cute.arch.elect_one(): + cpasync_bulk_s2g( + src[None, src_idx].iterator, + dst[None, dst_idx].iterator, + store_bytes, + reduction_kind=reduction_kind, + dtype=src.element_type, + **new_kwargs, + **kwargs, + ) + + def copy_bulk_s2g_single_stage(**new_kwargs): + store_bytes = const_expr(cute.size(src.shape) * src.element_type.width // 8) + with cute.arch.elect_one(): + cpasync_bulk_s2g( + src.iterator, + dst.iterator, + store_bytes, + reduction_kind=reduction_kind, + dtype=src.element_type, + **new_kwargs, + **kwargs, + ) + + return copy_bulk_s2g if const_expr(not single_stage) else copy_bulk_s2g_single_stage + def copy_bulk(src_idx, dst_idx, tma_bar_ptr: cute.Pointer, **new_kwargs): + assert dst_is_smem and not src_is_smem, "cp.async.bulk G2S expects GMEM -> SMEM" atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), src.element_type) - with cute.arch.elect_one(): - cute.copy( - atom, - src[None, src_idx], - dst[None, dst_idx], - mbar_ptr=tma_bar_ptr, - **new_kwargs, - **kwargs, - ) + cute.copy( + atom, + src[None, src_idx], + dst[None, dst_idx], + mbar_ptr=tma_bar_ptr, + **new_kwargs, + **kwargs, + ) def copy_bulk_single_stage(tma_bar_ptr: cute.Pointer, **new_kwargs): + assert dst_is_smem and not src_is_smem, "cp.async.bulk G2S expects GMEM -> SMEM" atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), src.element_type) - with cute.arch.elect_one(): - cute.copy(atom, src, dst, mbar_ptr=tma_bar_ptr, **new_kwargs, **kwargs) + cute.copy(atom, src, dst, mbar_ptr=tma_bar_ptr, **new_kwargs, **kwargs) return copy_bulk if const_expr(not single_stage) else copy_bulk_single_stage +def cpasync_bulk_get_store_or_add_fn( + src_tensor: cute.Tensor, + dst_tensor: cute.Tensor, + store_first_contribution: bool, + single_stage: bool = False, + **kwargs, +) -> Callable: + assert not single_stage, "store-or-add helper only supports staged SMEM -> GMEM tensors" + src_is_smem = const_expr( + isinstance(src_tensor.iterator, cute.Pointer) + and src_tensor.memspace == cute.AddressSpace.smem + ) + dst_is_smem = const_expr( + isinstance(dst_tensor.iterator, cute.Pointer) + and dst_tensor.memspace == cute.AddressSpace.smem + ) + assert src_is_smem and not dst_is_smem, "store-or-add helper only supports SMEM -> GMEM" + group_rank_src = const_expr(cute.rank(src_tensor)) + group_rank_dst = const_expr(cute.rank(dst_tensor)) + src = cute.group_modes(src_tensor, 0, group_rank_src - 1) + dst = cute.group_modes(dst_tensor, 0, group_rank_dst - 1) + + @cute.jit + def copy_bulk_s2g_store_or_add(src_idx, dst_idx, idx, **new_kwargs): + store_bytes = const_expr(cute.size(src.shape[:-1]) * src.element_type.width // 8) + src_ptr = src[None, src_idx].iterator + dst_ptr = dst[None, dst_idx].iterator + with cute.arch.elect_one(): + if const_expr(store_first_contribution): + if idx == 0: + cpasync_bulk_s2g( + src_ptr, + dst_ptr, + store_bytes, + reduction_kind=None, + **new_kwargs, + **kwargs, + ) + else: + cpasync_bulk_s2g( + src_ptr, + dst_ptr, + store_bytes, + reduction_kind=cpasync.ReductionOp.ADD, + dtype=src.element_type, + **new_kwargs, + **kwargs, + ) + else: + cpasync_bulk_s2g( + src_ptr, + dst_ptr, + store_bytes, + reduction_kind=cpasync.ReductionOp.ADD, + dtype=src.element_type, + **new_kwargs, + **kwargs, + ) + + return copy_bulk_s2g_store_or_add + + @dsl_user_op def tma_get_copy_fn( atom: cute.CopyAtom, @@ -800,6 +1088,236 @@ def tma_get_copy_fn( return (copy_tma if const_expr(not single_stage) else copy_tma_single_stage), s, g +@dsl_user_op +def tma_get_block_copy_fn( + atom: cute.CopyAtom, + src_tensor: cute.Tensor, + dst_tensor: cute.Tensor, + tma_multicast: Optional[dict] = None, + single_stage: bool = False, + *, + loc=None, + ip=None, + **kwargs, +) -> Callable: + src_is_smem = const_expr( + isinstance(src_tensor.iterator, cute.Pointer) + and src_tensor.memspace == cute.AddressSpace.smem + ) + if const_expr(tma_multicast is not None and "use_2cta_mma_inst" not in tma_multicast): + op = atom.op if const_expr(hasattr(atom, "op")) else atom + tma_multicast = { + **tma_multicast, + "use_2cta_mma_inst": getattr(op, "cta_group", None) == tcgen05.CtaGroup.TWO, + } + smem_tensor, gmem_tensor = (src_tensor, dst_tensor) if src_is_smem else (dst_tensor, src_tensor) + group_rank_smem = const_expr(cute.rank(smem_tensor) - (1 if not single_stage else 0)) + group_rank_gmem = const_expr(cute.rank(gmem_tensor) - (1 if not single_stage else 0)) + s = cute.group_modes(smem_tensor, 0, group_rank_smem) + g = cute.group_modes(gmem_tensor, 0, group_rank_gmem) + src, dst = (s, g) if src_is_smem else (g, s) + + @dsl_user_op + def copy_tma(src_idx, dst_idx, *, loc=None, ip=None, **new_kwargs): + src_cur = src[None, src_idx] + dst_cur = dst[None, dst_idx] + if const_expr(tma_multicast is None): + block_copy(atom, src_cur, dst_cur, **new_kwargs, **kwargs, loc=loc, ip=ip) + else: + block_copy( + atom, + src_cur, + dst_cur, + tma_multicast=tma_multicast, + **new_kwargs, + **kwargs, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def copy_tma_single_stage(*, loc=None, ip=None, **new_kwargs): + if const_expr(tma_multicast is None): + block_copy(atom, src, dst, **new_kwargs, **kwargs, loc=loc, ip=ip) + else: + block_copy( + atom, + src, + dst, + tma_multicast=tma_multicast, + **new_kwargs, + **kwargs, + loc=loc, + ip=ip, + ) + + return copy_tma if const_expr(not single_stage) else copy_tma_single_stage + + +def s2t_get_copy_fn( + src_tensor: cute.Tensor, + dst_tensor: cute.Tensor, + cta_group: tcgen05.CtaGroup, +) -> Callable: + """ + Make tiledCopy for smem to tmem load, then return a copy function over stages. + + :param src_tensor: The source tensor in smem + :param dst_tensor: The destination tensor in tmem + """ + assert src_tensor.element_type == dst_tensor.element_type + # (MMA, MMA_MN, MMA_K, STAGE) + src_compact = cute.filter_zeros(src_tensor) + # (MMA, MMA_MN, MMA_K) + dst_compact = cute.filter_zeros(dst_tensor) + # Make S2T CopyAtom and tiledCopy. + copy_atom = cute.make_copy_atom(tcgen05.Cp4x32x128bOp(cta_group), dst_tensor.element_type) + tiled_copy = tcgen05.make_s2t_copy(copy_atom, dst_compact) + thr_copy = tiled_copy.get_slice(0) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + src_partition = tcgen05.get_s2t_smem_desc_tensor(tiled_copy, thr_copy.partition_S(src_compact)) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + dst_partition = thr_copy.partition_D(dst_compact) + + @dsl_user_op + def copy_s2t(stage_idx, *, loc=None, ip=None, **new_kwargs): + # Stage slice of partitioned source tensor: ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + stage_coord = (None, None, None, None, stage_idx) + cute.copy( + tiled_copy, src_partition[stage_coord], dst_partition, loc=loc, ip=ip, **new_kwargs + ) + + return copy_s2t + + +# tcgen05 TMEM <-> RMEM helpers (t2r loads / r2t stores). +# +# The register-side fragment of a tmem copy is derivable by layout algebra, +# with no reference to the original (pre-partition) tile tensor: +# - the per-thread register VALUE shape is mode 1 of the tiled copy's +# register-side TV layout (`layout_dst_tv_tiled` for loads, +# `layout_src_tv_tiled` for stores). The tmem-side partition can't supply +# it: tmem partitioning is warp-collective, so its value mode counts tmem +# cells across the whole warp, not per-thread register elements. +# - the tile-iteration modes are shared between partition_S and partition_D +# (same tiler over the same tile extent), so they can be read off whichever +# side was already partitioned. +# This kills the make-a-fake/identity-tensor-and-partition_D dance previously +# needed at every t2r site. + + +def tmem_reg_frag( + tiled_copy: cute.TiledCopy, + partitioned: cute.Tensor, + num_extra_modes: int = 0, + dtype: Optional[Type[cutlass.Numeric]] = None, + *, + loc=None, + ip=None, +) -> cute.Tensor: + """Allocate the per-thread register fragment for ONE tile of a tcgen05 + tmem copy, given any partitioned view of it (tmem or otherwise). + + `partitioned` is (V, iter..., extra...) as produced by partition_S/_D; + the trailing `num_extra_modes` modes (stage, epi-subtile, ...) are + excluded from the fragment and indexed at copy time instead. `dtype` + defaults to `partitioned.element_type`. + The register side is inferred from the tcgen05 load/store op: destination + for t2r loads, source for r2t stores.""" + tv = _tmem_copy_reg_tv_layout(tiled_copy) + val_shape = tv.shape[1] + rank = cute.rank(partitioned.shape) + iters = tuple(partitioned.shape[i] for i in range(1, rank - num_extra_modes)) + frag_dtype = partitioned.element_type if const_expr(dtype is None) else dtype + return cute.make_rmem_tensor((val_shape, *iters), frag_dtype, loc=loc, ip=ip) + + +def coord_frag(tiled_copy: cute.TiledCopy, tidx: Int32, shape) -> cute.Tensor: + """Per-thread (row, col) coordinates aligned with a tiled copy's register + fragments (`tmem_reg_frag` / `load_t2r`): the register-side partition of + an identity tensor over `shape`. Deliberately partition_D — partition_S of + a TMEM tiled copy keeps whole warp-addressed atom tiles instead of + distributing elements over lanes.""" + return tiled_copy.get_slice(tidx).partition_D(cute.make_identity_tensor(shape)) + + +def r2s_partition_from_t2r( + tiled_copy_t2r: cute.TiledCopy, + s: cute.Tensor, + tidx: Int32, + transpose: bool = False, + position_independent=False, +) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """SMEM-store (r2s) side chained off a tmem-load tiled copy: the r2s copy + inherits the t2r copy's per-thread value ownership via make_tiled_copy_D, + so the loaded fragment can be stored (post-conversion) without a shuffle. + By default the store atom is selected like SM100 GEMM epilogues: + `get_smem_store_op(layout, dst_dtype, tiled_copy_t2r.value_type, tiled_copy_t2r)`, + so the stmatrix shape follows the tmem-load atom. `transpose=True` maps to + COL_MAJOR, otherwise ROW_MAJOR. + + `s` is the staged SMEM tile; its trailing stage mode is excluded from the + register fragment. `position_independent=True` partitions through a + position-independent swizzle view, matching `get_smem_store_C`. + Returns `(tiled_copy, tRS_r, tRS_s)`; store via + `cute.copy(tiled_copy, tRS_r, tRS_s[..., idx])`.""" + dtype = s.element_type + layout = LayoutEnum.COL_MAJOR if const_expr(transpose) else LayoutEnum.ROW_MAJOR + copy_atom = sm100_utils.get_smem_store_op( + layout, dtype, tiled_copy_t2r.value_type, tiled_copy_t2r + ) + tiled_copy = cute.make_tiled_copy_D(copy_atom, tiled_copy_t2r) + thr_copy = tiled_copy.get_slice(tidx) + if const_expr(not position_independent): + tRS_s = thr_copy.partition_D(s) + else: + tRS_s = partition_D_position_independent(thr_copy, s) + rank = cute.rank(tRS_s.shape) + frag_shape = tuple(tRS_s.shape[i] for i in range(rank - 1)) + tRS_r = cute.make_rmem_tensor(frag_shape, dtype) + return tiled_copy, tRS_r, tRS_s + + +def s2r_partition_from_t2r( + tiled_copy_t2r: cute.TiledCopy, + s: cute.Tensor, + tidx: Int32, + r_layout: cute.Layout, + copy_atom: Optional[cute.CopyAtom] = None, + transpose: bool = False, + position_independent=False, +) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]: + """SMEM-load (s2r) counterpart of `r2s_partition_from_t2r` (ldmatrix vs + stmatrix), for reading an epilogue input that was TMA-staged into an + epi-tile SMEM buffer (e.g. C in gemm, z in ssd) into registers + element-aligned with the t2r fragments. The register fragment is + allocated with `r_layout` (pass the r2s fragment's layout) so its linear + element order matches the t2r/r2s fragments; `tSR_r` is its retiled view + for the collective copy. Per-warp SMEM footprints of this load and the + chained r2s store coincide, so reusing one buffer for input then output + is warp-local (no inter-warp hazard). + + `position_independent=True` partitions through a position-independent + swizzle view, matching `get_smem_load_C`. + + Returns `(tiled_copy, tRS_r, tSR_r, tSR_s)`; load via + `cute.copy(tiled_copy, tSR_s[..., idx], tSR_r)` then read `tRS_r`.""" + dtype = s.element_type + if const_expr(copy_atom is None): + copy_atom = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=transpose, num_matrices=4), dtype + ) + tiled_copy = cute.make_tiled_copy_D(copy_atom, tiled_copy_t2r) + thr_copy = tiled_copy.get_slice(tidx) + if const_expr(not position_independent): + tSR_s = thr_copy.partition_S(s) + else: + tSR_s = partition_S_position_independent(thr_copy, s) + tRS_r = cute.make_rmem_tensor(r_layout, dtype) + tSR_r = tiled_copy.retile(tRS_r) + return tiled_copy, tRS_r, tSR_r, tSR_s + + def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsync): def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs): copy( @@ -812,6 +1330,18 @@ def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsyn return copy_fn +def chain_tma_producer_copy_fns(copy_fns: Sequence[Optional[Callable]]): + if not any(fn is not None for fn in copy_fns): + return None + + def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs): + for fn in copy_fns: + if const_expr(fn is not None): + fn(src_idx=src_idx, producer_state=producer_state, **new_kwargs) + + return copy_fn + + @cute.jit def gather_m_get_copy_fn( thr_copy_A: cute.ThrCopy, @@ -855,7 +1385,7 @@ def gather_m_get_copy_fn( mA_k = cute.logical_divide(mA, (None, tile_K)) - def copy_fn(src_idx, dst_idx, pred: bool = False): + def copy_fn(src_idx, dst_idx, pred: cutlass.Constexpr[bool] = False): tApA_k = None if const_expr(pred): tApA_k = cute.make_rmem_tensor(cols_per_thread, Boolean) @@ -964,9 +1494,7 @@ def gather_k_get_copy_fn( for k in cutlass.range(cols_per_thread): col_idx = tAcA[0, 0, k][1] k_idx[k] = sAIdx_cur[col_idx] - cute.arch.sync_warp() - with cute.arch.elect_one(): - a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state) + a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state) return k_idx, tApA_k def copy_fn( @@ -1078,9 +1606,7 @@ def gather_k_get_tma_copy_fn( ) -> cute.Tensor: a_prefetch_pipeline.consumer_wait(a_prefetch_consumer_state) tSR_rAIdx = load_s2r(tSR_sAIdx[None, None, dst_idx]) - cute.arch.sync_warp() - with cute.arch.elect_one(): - a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state) + a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state) return tSR_rAIdx def copy_fn(src_idx, dst_idx, tSR_rAIdx, tma_bar_ptr: cute.Pointer): diff --git a/build/torch-cuda/quack/cross_entropy.py b/build/torch-cuda/quack/cross_entropy.py index d3057bc44f3c266d9d8a915eb5554696c7e18aab..1985c9fef9b76edd972971c426603d897e641b9f 100644 --- a/build/torch-cuda/quack/cross_entropy.py +++ b/build/torch-cuda/quack/cross_entropy.py @@ -5,7 +5,7 @@ from functools import partial from typing import Optional, Type, Literal import torch -from ._ops_compat import add_quack_op_namespace_prefix +from ._ops_compat import add_op_namespace_prefix from torch import Tensor import cuda.bindings.driver as cuda @@ -18,11 +18,12 @@ from . import utils as utils from . import copy_utils as copy_utils from . import layout_utils as layout_utils from .compile_utils import make_fake_tensor as fake_tensor +from .dsl import cute_op from .reduce import row_reduce, online_softmax_reduce from .reduction_base import ReductionBase -from .cache_utils import jit_cache +from .cache import jit_cache from .cute_dsl_utils import torch2cute_dtype_map -from cutlass.base_dsl import Arch +from cutlass.base_dsl.arch import Arch class CrossEntropy(ReductionBase): @@ -75,6 +76,7 @@ class CrossEntropy(ReductionBase): mLoss: cute.Tensor, # (M,) mLSE: Optional[cute.Tensor], # (M,) mdX: Optional[cute.Tensor], # (M, N) - if provided, compute gradient + mWeight: Optional[cute.Tensor], ignore_index: Int32, # Index to ignore in loss computation stream: cuda.CUstream, ): @@ -97,6 +99,7 @@ class CrossEntropy(ReductionBase): mLoss, mLSE, mdX, + mWeight, ignore_index, tiler_mn, tiled_copy, @@ -117,6 +120,7 @@ class CrossEntropy(ReductionBase): mLoss: cute.Tensor, # (M,) mLSE: Optional[cute.Tensor], # (M,) mdX: Optional[cute.Tensor], # (M, N) - if provided, compute gradient + mWeight: Optional[cute.Tensor], ignore_index: Int32, # Index to ignore in loss computation tiler_mn: cute.Shape, tiled_copy: cute.TiledCopy, @@ -156,8 +160,16 @@ class CrossEntropy(ReductionBase): row = tXcX[0][0] target = Int32.zero + target_weight = Float32.zero if row < shape[0]: target = Int32(mTarget[row]) + if const_expr(mWeight is not None): + # Gate on target != ignore_index: ignore_index may be negative + # (PyTorch default -100), and indexing mWeight at that offset is OOB. + if target != ignore_index: + target_weight = Float32(mWeight[target]) + else: + target_weight = 1.0 if row < shape[0]: copy(tXgX, tXsX, is_async=True) @@ -221,7 +233,7 @@ class CrossEntropy(ReductionBase): ): lse = max_x + cute.math.log(denom, fastmath=True) # Set loss to 0 if this index should be ignored, otherwise compute normally - loss_val = (lse - target_logit) if not should_ignore else Float32.zero + loss_val = target_weight * (lse - target_logit) if not should_ignore else Float32.zero mLoss[row] = mLoss.element_type(loss_val) if const_expr(mLSE is not None): mLSE[row] = lse @@ -239,7 +251,6 @@ class CrossEntropy(ReductionBase): probs = exp_x * denom_inv gdX = cute.local_tile(mdX, tiler_mn, (bidx, cluster_y)) tXgdX = thr_copy.partition_D(gdX) - tXrdX = cute.make_rmem_tensor_like(tXgdX) tXcFull = thr_copy.partition_S(cX) # Compute gradient: probs for all classes, (probs - 1) for target class # If ignored, gradient is already zero @@ -248,46 +259,58 @@ class CrossEntropy(ReductionBase): if not should_ignore: for i in cutlass.range(cute.size(tXrX), unroll_full=True): tXrdX_f32[i] = tXrdX_f32[i] if tXcFull[i][1] != target else tXrdX_f32[i] - 1.0 - tXrdX.store(tXrdX_f32.load().to(tXrdX.element_type)) + if const_expr(mWeight is not None): + tXrdX_f32.store(tXrdX_f32.load() * target_weight) + tXrdX = tXrdX_f32.to(tXgdX.element_type) if row < shape[0]: copy(tXrdX, tXgdX) - -@jit_cache -def _compile_cross_entropy_fwd( - dtype, target_dtype, target_logit_dtype, N, has_lse, has_dx, target_logit_ndim -): - batch_sym = cute.sym_int() - div = math.gcd(128 // dtype.width, N) - x_cute = fake_tensor(dtype, (batch_sym, N), div) - dx_cute = fake_tensor(dtype, (batch_sym, N), div) if has_dx else None - target_cute = fake_tensor(target_dtype, (batch_sym,)) - if target_logit_dtype is not None: - if target_logit_ndim == 2: - target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym, cute.sym_int()), div) + @staticmethod + @jit_cache + def compile( + dtype, + target_dtype, + target_logit_dtype, + N, + has_lse, + has_dx, + weight_dtype, + target_logit_ndim, + ): + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + x_cute = fake_tensor(dtype, (batch_sym, N), div) + dx_cute = fake_tensor(dtype, (batch_sym, N), div) if has_dx else None + target_cute = fake_tensor(target_dtype, (batch_sym,)) + if target_logit_dtype is not None: + if target_logit_ndim == 2: + target_logit_cute = fake_tensor( + target_logit_dtype, (batch_sym, cute.sym_int()), div + ) + else: + target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym,)) else: - target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym,)) - else: - target_logit_cute = None - loss_cute = fake_tensor(Float32, (batch_sym,)) - lse_cute = fake_tensor(Float32, (batch_sym,)) if has_lse else None - # If there's dx, it's faster to not use online softmax since we want the exp(x - max) - cross_entropy_op = CrossEntropy(dtype, N, online_softmax=not has_dx) - return cute.compile( - cross_entropy_op, - x_cute, - target_cute, - target_logit_cute, - loss_cute, - lse_cute, - dx_cute, - Int32(0), # ignore_index, just for compilation - cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), - options="--enable-tvm-ffi", - ) + target_logit_cute = None + loss_cute = fake_tensor(Float32, (batch_sym,)) + lse_cute = fake_tensor(Float32, (batch_sym,)) if has_lse else None + weight_cute = fake_tensor(weight_dtype, (N,)) if weight_dtype is not None else None + # If there's dx, it's faster to not use online softmax since we want the exp(x - max) + return cute.compile( + CrossEntropy(dtype, N, online_softmax=not has_dx), + x_cute, + target_cute, + target_logit_cute, + loss_cute, + lse_cute, + dx_cute, + weight_cute, + Int32(0), # ignore_index, just for compilation + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) -@torch.library.custom_op(add_quack_op_namespace_prefix("cross_entropy_fwd_out"), mutates_args={"loss", "lse", "dx"}) +@cute_op(add_op_namespace_prefix("cross_entropy_fwd_out"), mutates_args={"loss", "lse", "dx"}) def cross_entropy_fwd_out( x: Tensor, target: Tensor, @@ -295,6 +318,7 @@ def cross_entropy_fwd_out( loss: Tensor, lse: Optional[Tensor], dx: Optional[Tensor], + weight: Optional[Tensor], ignore_index: int = -100, ) -> None: """Cross entropy forward pass. @@ -307,6 +331,7 @@ def cross_entropy_fwd_out( loss: Output loss tensor of shape (M,) lse: Optional output log-sum-exp tensor of shape (M,) dx: Optional output gradient tensor of shape (M, N) + weight: Optional weight vector of shape (N,) ignore_index: Index to ignore in loss computation Returns: @@ -314,14 +339,12 @@ def cross_entropy_fwd_out( """ assert x.dim() == 2, "Input must be 2D" assert target.dim() == 1, "Target must be 1D" - assert x.is_cuda and target.is_cuda, "Tensors must be on CUDA device" assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype" assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64" if target_logit is not None: - assert target_logit.is_cuda, "Target logits must be on CUDA device" assert target_logit.dtype in [torch.float16, torch.bfloat16, torch.float32] - if dx is not None: - assert dx.is_cuda, "dx must be on CUDA device" + if x.size(0) == 0: + return N = x.size(1) dtype = torch2cute_dtype_map[x.dtype] target_dtype = torch2cute_dtype_map[target.dtype] @@ -329,54 +352,24 @@ def cross_entropy_fwd_out( torch2cute_dtype_map[target_logit.dtype] if target_logit is not None else None ) target_logit_ndim = target_logit.ndim if target_logit is not None else None - _compile_cross_entropy_fwd( + weight_dtype = torch2cute_dtype_map[weight.dtype] if weight is not None else None + CrossEntropy.compile( dtype, target_dtype, target_logit_dtype, N, lse is not None, dx is not None, + weight_dtype, target_logit_ndim, - )(x, target, target_logit, loss, lse, dx, Int32(ignore_index)) - - -@cross_entropy_fwd_out.register_fake -def _cross_entropy_fwd_out_fake( - x: Tensor, - target: Tensor, - target_logit: Optional[Tensor], - loss: Tensor, - lse: Optional[Tensor], - dx: Optional[Tensor], - ignore_index: int = -100, -) -> None: - # See softmax.py _softmax_fwd_fake for why register_fake is needed. - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY and not isinstance(x.size(1), torch.SymInt): - N = x.size(1) - dtype = torch2cute_dtype_map[x.dtype] - target_dtype = torch2cute_dtype_map[target.dtype] - target_logit_dtype = ( - torch2cute_dtype_map[target_logit.dtype] if target_logit is not None else None - ) - target_logit_ndim = target_logit.ndim if target_logit is not None else None - _compile_cross_entropy_fwd( - dtype, - target_dtype, - target_logit_dtype, - N, - lse is not None, - dx is not None, - target_logit_ndim, - ) - _compile_cross_entropy_backward(dtype, target_dtype, N) + )(x, target, target_logit, loss, lse, dx, weight, Int32(ignore_index)) def cross_entropy_fwd( x: torch.Tensor, target: torch.Tensor, target_logit: Optional[torch.Tensor] = None, + weight: Optional[torch.Tensor] = None, ignore_index: int = -100, return_lse: bool = False, return_dx: bool = False, @@ -387,7 +380,7 @@ def cross_entropy_fwd( loss = torch.empty(M, device=device, dtype=torch.float32) lse = torch.empty(M, device=device, dtype=torch.float32) if return_lse else None dx = (torch.empty_like(x) if not inplace_backward else x) if return_dx else None - cross_entropy_fwd_out(x, target, target_logit, loss, lse, dx, ignore_index) + cross_entropy_fwd_out(x, target, target_logit, loss, lse, dx, weight, ignore_index) if return_lse and return_dx: return loss, lse, dx elif return_lse: @@ -432,6 +425,7 @@ class CrossEntropyBackward: mDLoss: cute.Tensor, mdX: cute.Tensor, mLSE: cute.Tensor, + mWeight: Optional[cute.Tensor], ignore_index: Int32, # Index to ignore in gradient computation stream: cuda.CUstream, ): @@ -451,6 +445,7 @@ class CrossEntropyBackward: mDLoss, mdX, mLSE, + mWeight, ignore_index, mX.shape, tiler_mn, @@ -474,6 +469,7 @@ class CrossEntropyBackward: mDLoss: cute.Tensor, # (M,) mdX: cute.Tensor, # (M, N) mLSE: cute.Tensor, # (M,) + mWeight: Optional[cute.Tensor], ignore_index: Int32, # Index to ignore in gradient computation shape: cute.Shape, tiler_mn: cute.Shape, @@ -507,6 +503,18 @@ class CrossEntropyBackward: copy = partial(copy_utils.copy, pred=tXpX) row = tXcX[0][0] + target = Int32.zero + target_weight = Float32.zero + if row < shape[0]: + target = Int32(mTarget[row]) + if const_expr(mWeight is not None): + # Gate on target != ignore_index: ignore_index may be negative + # (PyTorch default -100), and indexing mWeight at that offset is OOB. + if target != ignore_index: + target_weight = Float32(mWeight[target]) + else: + target_weight = 1.0 + if row < shape[0]: copy(tXgX, tXsX, is_async=True) cute.arch.cp_async_commit_group() @@ -516,13 +524,11 @@ class CrossEntropyBackward: cute.autovec_copy(tXsX, tXrX) x = tXrX.load().to(Float32) - target = Int32.zero dloss = Float32.zero lse = Float32.zero if row < shape[0]: - target = Int32(mTarget[row]) should_ignore = Boolean(target == ignore_index) - # Set dloss to 0 if this index should be ignored + # dloss is set to 0 if this index should be ignored if not should_ignore: dloss = Float32(mDLoss[row]) lse = Float32(mLSE[row]) @@ -534,32 +540,36 @@ class CrossEntropyBackward: for i in cutlass.range(cute.size(tXcFull), unroll_full=True): mask[i] = tXcFull[i][1] == target grad = cute.where(mask.load(), prob_shifted, probs) - grad = grad * dloss + grad = grad * dloss * target_weight tXrdX.store(grad.to(tXrdX.element_type)) if row < shape[0]: copy(tXrdX, tXgdX) - -@jit_cache -def _compile_cross_entropy_backward(dtype, target_dtype, N): - batch_sym = cute.sym_int() - div = math.gcd(128 // dtype.width, N) - x_cute, dx_cute = [fake_tensor(dtype, (batch_sym, N), div)] * 2 - target_cute = fake_tensor(target_dtype, (batch_sym,)) - dloss_cute, lse_cute = [fake_tensor(Float32, (batch_sym,))] * 2 - cross_entropy_backward_op = CrossEntropyBackward(dtype, N) - return cute.compile( - cross_entropy_backward_op, - x_cute, - target_cute, - dloss_cute, - dx_cute, - lse_cute, - Int32(0), # ignore_index, just for compilation - cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), - options="--enable-tvm-ffi", - ) + @staticmethod + @jit_cache + def compile(dtype, target_dtype, N, weight_dtype): + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + x_cute, dx_cute = [fake_tensor(dtype, (batch_sym, N), div)] * 2 + target_cute = fake_tensor(target_dtype, (batch_sym,)) + dloss_cute = cute.runtime.make_fake_tensor( + Float32, (batch_sym,), stride=(cute.sym_int64(),) + ) + lse_cute = fake_tensor(Float32, (batch_sym,)) + weight_cute = fake_tensor(weight_dtype, (N,)) if weight_dtype is not None else None + return cute.compile( + CrossEntropyBackward(dtype, N), + x_cute, + target_cute, + dloss_cute, + dx_cute, + lse_cute, + weight_cute, + Int32(0), # ignore_index, just for compilation + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) def _cross_entropy_backward( @@ -568,6 +578,7 @@ def _cross_entropy_backward( dloss: torch.Tensor, lse: torch.Tensor, dx: torch.Tensor, + weight: Optional[torch.Tensor] = None, ignore_index=-100, ) -> None: """Cross entropy backward pass. @@ -576,8 +587,11 @@ def _cross_entropy_backward( target: Target class indices tensor of shape (M,) dloss: Upstream gradients tensor of shape (M,) lse: Log-sum-exp values tensor of shape (M,) + dx: Output gradient tensor of shape (M, N) + weight: Optional per-class weight tensor of shape (N,) + ignore_index: Index to ignore in gradient computation Returns: - Input gradients tensor of shape (M, N) + None (mutates dx in-place) """ assert x.dim() == 2, "Input must be 2D" assert target.dim() == 1, "Target must be 1D" @@ -586,48 +600,33 @@ def _cross_entropy_backward( assert x.shape[0] == target.shape[0], "Batch dimensions must match" assert x.shape[0] == dloss.shape[0], "Batch dimensions must match" assert x.shape[0] == lse.shape[0], "Batch dimensions must match" - assert x.is_cuda and target.is_cuda and dloss.is_cuda and lse.is_cuda, ( - "Tensors must be on CUDA device" - ) assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype" assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64" + if weight is not None: + assert weight.is_cuda, "weight must be on CUDA device" + assert weight.is_floating_point(), "weight must be a floating-point tensor" + if x.size(0) == 0: + return N = x.size(1) dtype = torch2cute_dtype_map[x.dtype] target_dtype = torch2cute_dtype_map[target.dtype] - _compile_cross_entropy_backward(dtype, target_dtype, N)( - x, target, dloss, dx, lse, Int32(ignore_index) + weight_dtype = torch2cute_dtype_map[weight.dtype] if weight is not None else None + CrossEntropyBackward.compile(dtype, target_dtype, N, weight_dtype)( + x, target, dloss, dx, lse, weight, Int32(ignore_index) ) -@torch.library.custom_op(add_quack_op_namespace_prefix("cross_entropy_bwd_out"), mutates_args={"dx"}) +@cute_op(add_op_namespace_prefix("cross_entropy_bwd_out"), mutates_args={"dx"}) def cross_entropy_bwd_out( x: torch.Tensor, target: torch.Tensor, dloss: torch.Tensor, lse: torch.Tensor, dx: torch.Tensor, + weight: Optional[torch.Tensor] = None, ignore_index: int = -100, ) -> None: - _cross_entropy_backward(x, target, dloss, lse, dx, ignore_index) - - -@cross_entropy_bwd_out.register_fake -def _cross_entropy_bwd_out_fake( - x: torch.Tensor, - target: torch.Tensor, - dloss: torch.Tensor, - lse: torch.Tensor, - dx: torch.Tensor, - ignore_index: int = -100, -) -> None: - # See softmax.py _softmax_fwd_fake for why register_fake is needed. - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY and not isinstance(x.size(1), torch.SymInt): - N = x.size(1) - dtype = torch2cute_dtype_map[x.dtype] - target_dtype = torch2cute_dtype_map[target.dtype] - _compile_cross_entropy_backward(dtype, target_dtype, N) + _cross_entropy_backward(x, target, dloss, lse, dx, weight, ignore_index) def cross_entropy_bwd( @@ -635,51 +634,90 @@ def cross_entropy_bwd( target: torch.Tensor, dloss: torch.Tensor, lse: torch.Tensor, + weight: Optional[torch.Tensor] = None, ignore_index: int = -100, inplace_backward: bool = False, ) -> None: if inplace_backward and not torch.compiler.is_compiling(): dx = x _cross_entropy_backward( - x=x, target=target, dloss=dloss, lse=lse, dx=x, ignore_index=ignore_index + x=x, + target=target, + dloss=dloss, + lse=lse, + dx=x, + weight=weight, + ignore_index=ignore_index, ) else: dx = torch.empty_like(x) cross_entropy_bwd_out( - x=x, target=target, dloss=dloss, lse=lse, dx=dx, ignore_index=ignore_index + x=x, + target=target, + dloss=dloss, + lse=lse, + dx=dx, + weight=weight, + ignore_index=ignore_index, ) return dx class CrossEntropyFunction(torch.autograd.Function): @staticmethod - def forward(ctx, x, target, lse_partial=None, ignore_index=-100, inplace_backward=False): + def forward( + ctx, + x, + target, + lse_partial=None, + weight=None, + ignore_index=-100, + inplace_backward=False, + ): if lse_partial is None: - loss, lse = cross_entropy_fwd(x, target, ignore_index=ignore_index, return_lse=True) + loss, lse = cross_entropy_fwd( + x, + target, + weight=weight, + ignore_index=ignore_index, + return_lse=True, + ) else: # if we already compute partial lse, then to compute the final lse we treat # @lse_partial as @x and @x as @target_logit loss, lse = cross_entropy_fwd( - lse_partial, target, target_logit=x, ignore_index=ignore_index, return_lse=True + lse_partial, + target, + target_logit=x, + weight=weight, + ignore_index=ignore_index, + return_lse=True, ) - ctx.save_for_backward(x, target, lse) + ctx.save_for_backward(x, target, lse, weight) ctx.ignore_index = ignore_index ctx.inplace_backward = inplace_backward return loss @staticmethod def backward(ctx, dloss): - x, target, lse = ctx.saved_tensors + x, target, lse, weight = ctx.saved_tensors dx = cross_entropy_bwd( - x, target, dloss, lse, ctx.ignore_index, inplace_backward=ctx.inplace_backward + x, + target, + dloss, + lse, + weight=weight, + ignore_index=ctx.ignore_index, + inplace_backward=ctx.inplace_backward, ) - return dx, None, None, None, None + return dx, None, None, None, None, None def cross_entropy( x: torch.Tensor, target: torch.Tensor, lse_partial: Optional[torch.Tensor] = None, + weight: Optional[torch.Tensor] = None, ignore_index: int = -100, reduction: Literal["none", "mean", "sum"] = "mean", inplace_backward: bool = False, @@ -690,12 +728,13 @@ def cross_entropy( x: Input logits tensor of shape (M, N) target: Target class indices tensor of shape (M,) lse_partial: Optional precomputed log-sum-exp partial results + weight: Optional per-class weight tensor of shape (N,) + ignore_index: Index to ignore in loss computation (loss will be 0 for these indices) reduction: Specifies the reduction to apply to the output: 'none': no reduction will be applied (default) 'mean': the sum of the output will be divided by the number of elements 'sum': the output will be summed inplace_backward: Whether to perform backward pass in-place - ignore_index: Index to ignore in loss computation (loss will be 0 for these indices) Returns: Cross entropy loss tensor: @@ -703,8 +742,19 @@ def cross_entropy( - If reduction='mean': scalar tensor with mean loss - If reduction='sum': scalar tensor with sum of losses """ - loss = CrossEntropyFunction.apply(x, target, lse_partial, ignore_index, inplace_backward) + loss = CrossEntropyFunction.apply( + x, + target, + lse_partial, + weight, + ignore_index, + inplace_backward, + ) if reduction == "mean": + if weight is not None: + valid = target != ignore_index + denom = (weight[target.clamp(min=0)] * valid).sum() + return loss.sum() / denom return loss.sum() / (target != ignore_index).sum().float() elif reduction == "sum": return loss.sum() diff --git a/build/torch-cuda/quack/cute_dsl_utils.py b/build/torch-cuda/quack/cute_dsl_utils.py index 0988b6e72b8b84c14677de71eb0e8939aeec1955..af3c30d0543e71afb5d5049dac903492d047a311 100644 --- a/build/torch-cuda/quack/cute_dsl_utils.py +++ b/build/torch-cuda/quack/cute_dsl_utils.py @@ -59,12 +59,32 @@ torch2cute_dtype_map = { torch.float32: Float32, torch.int32: Int32, torch.int64: Int64, + torch.float8_e4m3fn: cutlass.Float8E4M3FN, + torch.float8_e5m2: cutlass.Float8E5M2, + torch.float8_e8m0fnu: cutlass.Float8E8M0FNU, + # Packed fp4: dlpack presents the logical (doubled) K extent to the DSL. + torch.float4_e2m1fn_x2: cutlass.Float4E2M1FN, } @lru_cache -def get_max_active_clusters(cluster_size): - return cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size=cluster_size) +def get_device_multiprocessor_count(device_id: int = 0) -> int: + return cutlass.utils.HardwareInfo(device_id).get_device_multiprocessor_count() + + +@lru_cache +def get_max_active_clusters( + cluster_size: int, + device_capacity: Tuple[int, int] | None = None, + device_id: int = 0, +) -> int: + if device_capacity is None: + device_capacity = get_device_capacity() + if device_capacity[0] < 9: + if cluster_size != 1: + raise ValueError("SM8x kernels do not support CTA clusters; cluster_size must be 1") + return get_device_multiprocessor_count(device_id) + return cutlass.utils.HardwareInfo(device_id).get_max_active_clusters(cluster_size=cluster_size) def _parse_arch_str(arch_str: str) -> Tuple[int, int]: @@ -148,8 +168,46 @@ def _namedtuple_new_from_mlir_values(self, values): return self.__class__(*new_fields) +def _namedtuple_dynamic_fields(self): + """Yield fields that are represented by MLIR/runtime values. + + Keep this in sync with ``_namedtuple_new_from_mlir_values``: fields that + are ``None`` or plain Python compile-time constants are preserved from the + compile-time template and do not consume MLIR block arguments. + """ + for field_val in self: + if field_val is None or isinstance(field_val, StaticTypes): + continue + yield field_val + + +def _namedtuple_c_pointers(self): + """Generic ``JitArgument.__c_pointers__`` for ``@mlir_namedtuple`` classes.""" + from cutlass.base_dsl.typing import get_c_pointers + + ptrs = [] + for field_val in _namedtuple_dynamic_fields(self): + ptrs.extend(get_c_pointers(field_val)) + return ptrs + + +def _namedtuple_get_mlir_types(self): + """Generic ``JitArgument.__get_mlir_types__`` for ``@mlir_namedtuple`` classes.""" + from cutlass.base_dsl.typing import get_mlir_types + + types = [] + for field_val in _namedtuple_dynamic_fields(self): + types.extend(get_mlir_types(field_val)) + return types + + def mlir_namedtuple(cls): - """Decorator that adds MLIR value reconstruction to a NamedTuple class. + """Decorator that makes a NamedTuple usable as a CuTe JIT argument. + + Adds the full ``JitArgument`` protocol. This matters even when a concrete + instance has no dynamic fields: newer CuTe DSL versions warn for non- + constexpr compile-only arguments that flatten to zero MLIR values unless + they explicitly implement the protocol. Usage:: @@ -158,6 +216,8 @@ def mlir_namedtuple(cls): tensor_arg: cute.Tensor const_arg: cutlass.Constexpr[int] = 0 """ + cls.__c_pointers__ = _namedtuple_c_pointers + cls.__get_mlir_types__ = _namedtuple_get_mlir_types cls.__new_from_mlir_values__ = _namedtuple_new_from_mlir_values return cls diff --git a/build/torch-cuda/quack/dsl/__init__.py b/build/torch-cuda/quack/dsl/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..88dfc2a14c50847c8d68d4e8a555496df710121c --- /dev/null +++ b/build/torch-cuda/quack/dsl/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. + +"""CuTe DSL helpers and integration hooks.""" + +from . import cute_tensor_indexing # noqa: F401 +from . import cute_tensor # noqa: F401 +from .torch_library_op import cute_op + +__all__ = ["cute_op"] + + +def __getattr__(name: str): + if name == "cute_op": + from .torch_library_op import cute_op + + return cute_op + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/build/torch-cuda/quack/cute_dsl_ptxas.py b/build/torch-cuda/quack/dsl/cute_dsl_ptxas.py similarity index 80% rename from build/torch-cuda/quack/cute_dsl_ptxas.py rename to build/torch-cuda/quack/dsl/cute_dsl_ptxas.py index ed4e78701d5962e0bb3723a5ae6cc86f56650679..ad8bb3b25327eaacf4f0216fe164f5d45286fd02 100644 --- a/build/torch-cuda/quack/cute_dsl_ptxas.py +++ b/build/torch-cuda/quack/dsl/cute_dsl_ptxas.py @@ -3,35 +3,61 @@ System ptxas replacement for CUTLASS DSL. Usage:: - CUTE_DSL_KEEP_PTX=1 CUTE_DSL_PTXAS_PATH=/usr/local/cuda/bin/ptxas pytest tests/ + CUTE_DSL_PTXAS_PATH=/usr/local/cuda/bin/ptxas pytest tests/ Environment variables: CUTE_DSL_PTXAS_PATH - Path to ptxas (e.g., /usr/local/cuda/bin/ptxas) - CUTE_DSL_KEEP_PTX - Must be set to 1 before cutlass is imported + CUTE_DSL_KEEP=ptx - Optional; keep PTX files instead of deleting them CUTE_DSL_PTXAS_VERBOSE - Set to 1 for verbose output CUTE_DSL_DUMP_DIR - Directory for dumped PTX files (default: cwd) - CUTE_DSL_KEEP_CUBIN - Set to 1 to save compiled cubin files + CUTE_DSL_KEEP_CUBIN - Set to 1 to save system-ptxas cubin files """ +import ctypes import os -import sys import re -import ctypes import subprocess +import sys from pathlib import Path -import cutlass - CUTE_DSL_PTXAS_PATH = os.environ.get("CUTE_DSL_PTXAS_PATH", None) + +def _keep_tokens() -> set[str]: + keep = os.environ.get("CUTE_DSL_KEEP", "") + return {token.strip().lower() for token in keep.split(",") if token.strip()} + + +def _env_requests_ptx() -> bool: + tokens = _keep_tokens() + return "all" in tokens or "ptx" in tokens or os.environ.get("CUTE_DSL_KEEP_PTX") == "1" + + +_USER_WANTED_PTX = _env_requests_ptx() + + +def _force_keep_ptx_env() -> None: + tokens = _keep_tokens() + if "all" not in tokens and "ptx" not in tokens: + tokens.add("ptx") + os.environ["CUTE_DSL_KEEP"] = ",".join(sorted(tokens)) + # Keep the deprecated switch too so older CUTLASS DSL builds that do not + # understand CUTE_DSL_KEEP still dump PTX. + os.environ.setdefault("CUTE_DSL_KEEP_PTX", "1") + + if CUTE_DSL_PTXAS_PATH: - os.environ["CUTE_DSL_KEEP_PTX"] = "1" + # Must happen before CuTeDSL's EnvironmentVarManager is instantiated. + _force_keep_ptx_env() + +import cutlass # noqa: E402 + VERBOSE = os.environ.get("CUTE_DSL_PTXAS_VERBOSE", "0") == "1" _original_load_cuda_library = None _original_create_tvm_ffi_function = None -_user_wanted_ptx = False # True if user originally set CUTE_DSL_KEEP_PTX=1 +_user_wanted_ptx = False # True if user originally requested KEEP=ptx / KEEP_PTX=1 def _log(msg: str): @@ -201,18 +227,36 @@ def _patched_create_tvm_ffi_function(self): return _original_create_tvm_ffi_function(self) +def _force_live_keep_ptx() -> None: + """Update already-created CuTe DSL environment managers, if any. + + quack imports this module before importing its kernels, but keeping this + fallback makes direct/late calls to ``patch()`` less fragile. + """ + for cls_name in ("CuTeDSL", "CuteExperimentalDSL"): + dsl_cls = getattr(cutlass.cutlass_dsl, cls_name, None) + if dsl_cls is None: + continue + try: + envar = dsl_cls._get_dsl().envar + except Exception: + continue + envar.keep_ptx = True + if hasattr(envar, "keep_tokens"): + envar.keep_tokens = frozenset(set(envar.keep_tokens) | {"ptx"}) + + def patch(): - """Install system ptxas hook. Call before importing cutlass.""" + """Install system ptxas hook.""" global _original_load_cuda_library, _original_create_tvm_ffi_function, _user_wanted_ptx assert CUTE_DSL_PTXAS_PATH is not None if not os.path.isfile(CUTE_DSL_PTXAS_PATH) or not os.access(CUTE_DSL_PTXAS_PATH, os.X_OK): raise RuntimeError(f"ptxas not found: {CUTE_DSL_PTXAS_PATH}") - _user_wanted_ptx = os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1" - assert os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1", ( - "Require CUTE_DSL_KEEP_PTX=1 to use system's ptxas" - ) + _user_wanted_ptx = _USER_WANTED_PTX + _force_keep_ptx_env() + _force_live_keep_ptx() patched = False cuda_jit_function_cls = cutlass.cutlass_dsl.cuda_jit_executor.CudaDialectJitCompiledFunction diff --git a/build/torch-cuda/quack/dsl/cute_tensor.py b/build/torch-cuda/quack/dsl/cute_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..fcae5860730ff62c5d8bdba0e17bc6f1802cabfc --- /dev/null +++ b/build/torch-cuda/quack/dsl/cute_tensor.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026, Tri Dao. +"""Small CuTe tensor convenience helpers. + +Importing this module intentionally mutates CuTe's tensor class process-wide so +fragments can be written in a more PyTorch-like style:: + + rmem_f32 = rmem_f16.to(Float32) + rmem_copy = rmem_view.clone() + rmem_contig = rmem_view.contiguous() + +``tensor.to(dtype)`` is exactly the explicit CuTe sequence:: + + dst = cute.make_rmem_tensor_like(src, dtype) + dst.store(src.load().to(dtype)) + +``tensor.to(dtype, force_materialize=True)`` uses the same value conversion, then inserts an +opaque no-op SSA boundary on packed f16/bf16 lanes. This is useful when downstream codegen +would otherwise rematerialize a vector truncation for multiple consumers. + +``tensor.clone()`` materializes into ``cute.make_rmem_tensor_like(src)``. +``tensor.contiguous()`` mirrors ``quack.copy_utils.contiguous``. +""" + +from __future__ import annotations + +from typing import Any + +import cutlass +import cutlass.cute as cute +import cutlass.cute.tensor as _cute_tensor +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass._mlir.dialects import llvm + + +_ORIGINAL_TO_ATTR = "_quack_original_to" +_PATCHED_TO_ATTR = "_quack_rmem_tensor_to" +_ORIGINAL_CLONE_ATTR = "_quack_original_clone" +_PATCHED_CLONE_ATTR = "_quack_tensor_clone" +_ORIGINAL_CONTIGUOUS_ATTR = "_quack_original_contiguous" +_PATCHED_CONTIGUOUS_ATTR = "_quack_tensor_contiguous" + + +@dsl_user_op +def _black_box_b32(x: cutlass.Int32, *, loc: Any = None, ip: Any = None) -> cutlass.Int32: + """Opaque identity for packed registers. + + The empty asm emits no PTX instruction. The tied input constraint (``0``) makes the + output use the same register class/value as the input, while still creating an SSA boundary. + """ + + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int32(x).ir_value(loc=loc, ip=ip)], + "", + "=r,0", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@cute.jit +def _to_f16_materialized(src: Any, dtype: Any) -> Any: + assert src.element_type is cutlass.Float32, "src must be Float32" + assert dtype in (cutlass.BFloat16, cutlass.Float16), "dtype must be BFloat16 or Float16" + assert cute.size(src.shape) % 2 == 0, "src must have an even number of elements" + + # Why this exists: plain Tensor.to lowers f32 -> f16/bf16 as a vector truncation. + # In some kernels the converted fragment has multiple consumers (e.g. STSM and RS WGMMA + # in SM90 FlashAttention backward). Leaving the truncation as a high-level vector value + # lets later codegen rematerialize it for each consumer, producing extra packed converts + # and a worse WGMMA schedule. Storing the .to result, viewing the f16/bf16 pairs as i32, + # then passing each packed lane through an empty tied-operand asm creates an opaque SSA + # boundary. The asm emits no PTX instruction, but it forces one materialized packed value + # that downstream users share. + tmp = cute.make_rmem_tensor_like(src, dtype) + tmp.store(src.load().to(dtype)) + dst = cute.make_rmem_tensor_like(src, dtype) + tmp_i32 = cute.recast_tensor(tmp, cutlass.Int32) + dst_i32 = cute.recast_tensor(dst, cutlass.Int32) + assert cute.size(dst_i32.shape) * 2 == cute.size(src.shape) + for i in cutlass.range(cute.size(dst_i32), unroll_full=True): + dst_i32[i] = _black_box_b32(tmp_i32[i]) + return dst + + +def _make_to() -> Any: + @dsl_user_op + def _to( + self: Any, + dtype: Any, + *, + force_materialize: bool = False, + loc: Any = None, + ip: Any = None, + ) -> Any: + if self.memspace != cute.AddressSpace.rmem: + raise ValueError("Tensor.to(dtype) is only supported for rmem tensors") + + if force_materialize: + return _to_f16_materialized(self, dtype) + + dst = cute.make_rmem_tensor_like(self, dtype, loc=loc, ip=ip) + dst.store(self.load(loc=loc, ip=ip).to(dtype, loc=loc, ip=ip), loc=loc, ip=ip) + return dst + + return _to + + +def _make_clone() -> Any: + @dsl_user_op + def _clone(self: Any, *, loc: Any = None, ip: Any = None) -> Any: + dst = cute.make_rmem_tensor_like(self, loc=loc, ip=ip) + cute.autovec_copy(self, dst, loc=loc, ip=ip) + return dst + + return _clone + + +def _make_contiguous() -> Any: + @dsl_user_op + def _contiguous(self: Any, *, loc: Any = None, ip: Any = None) -> Any: + dst = cute.make_rmem_tensor(self.shape, self.element_type, loc=loc, ip=ip) + cute.autovec_copy(self, dst, loc=loc, ip=ip) + return dst + + return _contiguous + + +def patch_cute_tensor() -> None: + """Monkey patch CuTe tensors with QuACK convenience methods. + + The patch is idempotent. CuTe's immutable ``TensorSSA.to`` already handles + value conversion; this installs the analogous materializing conversion on + mutable register-backed ``_Tensor`` fragments, plus a ``contiguous`` method + equivalent to :func:`quack.copy_utils.contiguous`, and a ``clone`` method + that copies into a matching compact rmem tensor. + """ + tensor_cls = _cute_tensor._Tensor + if _PATCHED_TO_ATTR not in tensor_cls.__dict__: + original_to = getattr(tensor_cls, "to", None) + if original_to is not None: + setattr(tensor_cls, _ORIGINAL_TO_ATTR, original_to) + tensor_cls.to = _make_to() # type: ignore[method-assign] + setattr(tensor_cls, _PATCHED_TO_ATTR, True) + + if _PATCHED_CLONE_ATTR not in tensor_cls.__dict__: + original_clone = getattr(tensor_cls, "clone", None) + if original_clone is not None: + setattr(tensor_cls, _ORIGINAL_CLONE_ATTR, original_clone) + tensor_cls.clone = _make_clone() # type: ignore[method-assign] + setattr(tensor_cls, _PATCHED_CLONE_ATTR, True) + + if _PATCHED_CONTIGUOUS_ATTR not in tensor_cls.__dict__: + original_contiguous = getattr(tensor_cls, "contiguous", None) + if original_contiguous is not None: + setattr(tensor_cls, _ORIGINAL_CONTIGUOUS_ATTR, original_contiguous) + tensor_cls.contiguous = _make_contiguous() # type: ignore[method-assign] + setattr(tensor_cls, _PATCHED_CONTIGUOUS_ATTR, True) + + +patch_cute_tensor() + + +__all__ = ["patch_cute_tensor"] diff --git a/build/torch-cuda/quack/dsl/cute_tensor_indexing.py b/build/torch-cuda/quack/dsl/cute_tensor_indexing.py new file mode 100644 index 0000000000000000000000000000000000000000..d6342665d680f71a5f3242005471fa1856df6040 --- /dev/null +++ b/build/torch-cuda/quack/dsl/cute_tensor_indexing.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026, Tri Dao. +"""Small compatibility layer for more Pythonic CuTe tensor indexing. + +CuTe uses ``None`` as its underscore/full-mode slice marker, e.g. ``A[i, None]``. +This module teaches CuTe tensors the equivalent Python spelling ``:`` and expands +``...`` to the right number of full-mode slices. + +Importing this module intentionally mutates CuTe's tensor classes process-wide. +The original methods are retained as ``_quack_original_getitem`` and +``_quack_original_setitem`` for debugging or manual rollback. +""" + +from __future__ import annotations + +from typing import Any + +from cutlass.cutlass_dsl import dsl_user_op +import cutlass.cute.tensor as _cute_tensor + + +_ORIGINAL_GETITEM_ATTR = "_quack_original_getitem" +_ORIGINAL_SETITEM_ATTR = "_quack_original_setitem" +_PATCHED_ATTR = "_quack_extended_indexing" +_PATCHED_SETITEM_ATTR = f"{_PATCHED_ATTR}_setitem" + + +def _is_full_slice(idx: Any) -> bool: + return isinstance(idx, slice) and idx.start is None and idx.stop is None and idx.step is None + + +def _index_uses_ellipsis(idx: Any) -> bool: + if idx is Ellipsis: + return True + if isinstance(idx, tuple): + return any(_index_uses_ellipsis(item) for item in idx) + return False + + +def _shape_rank(shape: Any, idx: Any = None) -> int: + if shape is None: + suffix = f" in {idx!r}" if idx is not None else "" + raise ValueError(f"tensor shape is required to expand ellipsis{suffix}") + return _cute_tensor.rank(shape) + + +def _shape_mode(shape: Any, mode: int) -> Any: + if isinstance(shape, tuple) and mode < len(shape): + return shape[mode] + return None + + +def _canonicalize_cute_tensor_index(idx: Any, tensor_shape: Any = None) -> Any: + """Convert Python indexing sugar to CuTe's coordinate convention. + + ``:`` becomes ``None`` (CuTe's full-mode/underscore marker) and ``...`` expands + within the current hierarchy level using ``tensor_shape``. Other slices like + ``1:4`` are intentionally rejected because CuTe tensor slicing only supports + keeping an entire mode or selecting a single coordinate. + """ + if idx is Ellipsis: + return (None,) * _shape_rank(tensor_shape, idx) + if _is_full_slice(idx): + return None + if isinstance(idx, slice): + raise ValueError(f"CuTe Tensor indexing only supports full slices ':', got {idx!r}") + if not isinstance(idx, tuple): + return idx + + ellipsis_count = sum(item is Ellipsis for item in idx) + if ellipsis_count > 1: + raise ValueError("CuTe Tensor indexing supports at most one ellipsis per tuple level") + + explicit_modes = len(idx) - ellipsis_count + fill_modes = 0 + if ellipsis_count: + tensor_rank = _shape_rank(tensor_shape, idx) + fill_modes = tensor_rank - explicit_modes + if fill_modes < 0: + raise ValueError( + f"ellipsis cannot expand index {idx!r} for rank-{tensor_rank} CuTe Tensor mode" + ) + + result: list[Any] = [] + mode = 0 + for item in idx: + if item is Ellipsis: + result.extend([None] * fill_modes) + mode += fill_modes + else: + result.append(_canonicalize_cute_tensor_index(item, _shape_mode(tensor_shape, mode))) + mode += 1 + return tuple(result) + + +def _make_getitem(original_getitem: Any) -> Any: + @dsl_user_op + def _getitem(self: Any, idx: Any, *, loc: Any = None, ip: Any = None) -> Any: + tensor_shape = self.shape if _index_uses_ellipsis(idx) else None + idx = _canonicalize_cute_tensor_index(idx, tensor_shape) + return original_getitem(self, idx, loc=loc, ip=ip) + + return _getitem + + +def _make_setitem(original_setitem: Any) -> Any: + @dsl_user_op + def _setitem(self: Any, idx: Any, data: Any, *, loc: Any = None, ip: Any = None) -> Any: + tensor_shape = self.shape if _index_uses_ellipsis(idx) else None + idx = _canonicalize_cute_tensor_index(idx, tensor_shape) + return original_setitem(self, idx, data, loc=loc, ip=ip) + + return _setitem + + +def patch_cute_tensor_indexing() -> None: + """Monkey patch CuTe Tensor indexing with ``:``, ``...`` sugar. + + The patch is idempotent and keeps the original CuTe implementation for all + canonical coordinates, so existing ``A[i, j, None]`` code continues to behave + exactly as before. It is a process-wide mutation of CuTe's tensor classes. + """ + for cls in (_cute_tensor._Tensor, _cute_tensor.TensorSSA): + if _PATCHED_ATTR not in cls.__dict__: + setattr(cls, _ORIGINAL_GETITEM_ATTR, cls.__getitem__) + cls.__getitem__ = _make_getitem(cls.__getitem__) # type: ignore[method-assign] + setattr(cls, _PATCHED_ATTR, True) + + # TensorSSA has no upstream __setitem__, so only _Tensor needs the store path patched. + tensor_cls = _cute_tensor._Tensor + if _PATCHED_SETITEM_ATTR not in tensor_cls.__dict__: + setattr(tensor_cls, _ORIGINAL_SETITEM_ATTR, tensor_cls.__setitem__) + tensor_cls.__setitem__ = _make_setitem(tensor_cls.__setitem__) # type: ignore[method-assign] + setattr(tensor_cls, _PATCHED_SETITEM_ATTR, True) + + +patch_cute_tensor_indexing() + + +__all__ = ["patch_cute_tensor_indexing"] diff --git a/build/torch-cuda/quack/dsl/smem_struct.py b/build/torch-cuda/quack/dsl/smem_struct.py new file mode 100644 index 0000000000000000000000000000000000000000..d4cae236f7bfe0cd50a497c8cba5bd4099646bd6 --- /dev/null +++ b/build/torch-cuda/quack/dsl/smem_struct.py @@ -0,0 +1,99 @@ +# Copyright (c) 2026, Tri Dao. +# SPDX-License-Identifier: BSD-3-Clause + +"""Per-field smem-partition annotations for SharedStorage declarations. + +A `cute.struct` lowers to ONE smem_alloca op carrying ONE `smem.partition_id` +attribute, so a single struct cannot mix RESERVED and USER fields at the IR +level. `Reserved[...]` + `@partitioned_struct` provide the single-declaration +sugar instead: at trace time the decorator splits the annotated class into two +plain cute.structs — fields wrapped in `Reserved[...]` go to a struct +allocated with `partition=SmemPartition.RESERVED` (low addresses, packing with +the pipeline mbarriers and the TMEM holding buf under the 1KB that is +otherwise alignment pad ahead of the 1024-aligned USER buffers), everything +else stays USER — and `.allocate(smem)` returns one namespace exposing all +fields uniformly. + + @partitioned_struct + class SharedStorage: + sdt: Reserved[spec.smem_struct(128)] # RESERVED partition + sX: X.smem_struct(1024) # USER partition + ... + + storage = SharedStorage.allocate(smem) + storage.sdt, storage.sX # fields, regardless of partition + +This is trace-time-only machinery (no monkey-patching of the DSL). +""" + +from types import SimpleNamespace + +import cutlass.cute as cute +from cutlass.utils import SmemPartition + + +class Reserved: + """Annotation marker: allocate this field in the RESERVED smem partition.""" + + def __init__(self, inner): + self.inner = inner + + def __class_getitem__(cls, inner): + return cls(inner) + + +class PartitionedStruct: + """A SharedStorage declaration split by partition. Not a cute.struct itself: + holds one cute.struct per partition and allocates/merges them.""" + + def __init__(self, cls): + annotations = dict(cls.__annotations__) + reserved_ann = {k: v.inner for k, v in annotations.items() if isinstance(v, Reserved)} + user_ann = {k: v for k, v in annotations.items() if not isinstance(v, Reserved)} + self._user_struct = ( + cute.struct(type(cls.__name__, (), {"__annotations__": user_ann})) if user_ann else None + ) + self._reserved_struct = ( + cute.struct(type(cls.__name__ + "Reserved", (), {"__annotations__": reserved_ann})) + if reserved_ann + else None + ) + self._user_fields = list(user_ann) + self._reserved_fields = list(reserved_ann) + + def size_in_bytes(self) -> int: + """USER-partition footprint (what counts against smem_capacity - 1KB).""" + return self._user_struct.size_in_bytes() if self._user_struct is not None else 0 + + def reserved_size_in_bytes(self) -> int: + """RESERVED-partition footprint of the declared fields (the pipeline + mbarriers / TMEM holding buf allocate there separately).""" + return self._reserved_struct.size_in_bytes() if self._reserved_struct is not None else 0 + + def allocate(self, smem) -> SimpleNamespace: + """Allocate both partitions (RESERVED first, at the partition base) and + return a namespace exposing every declared field. A partition whose + struct is empty for this config (every field zero-sized) is skipped — + smem_alloca rejects 0-byte layouts — and its fields come back as None; + callers only touch such fields under the same has_* guards that made + them zero-sized.""" + fields = {} + if self._reserved_struct is not None: + if self._reserved_struct.size_in_bytes() > 0: + inst = smem.allocate(self._reserved_struct, partition=SmemPartition.RESERVED) + for name in self._reserved_fields: + fields[name] = getattr(inst, name) + else: + fields.update(dict.fromkeys(self._reserved_fields)) + if self._user_struct is not None: + if self._user_struct.size_in_bytes() > 0: + inst = smem.allocate(self._user_struct) + for name in self._user_fields: + fields[name] = getattr(inst, name) + else: + fields.update(dict.fromkeys(self._user_fields)) + return SimpleNamespace(**fields) + + +def partitioned_struct(cls) -> PartitionedStruct: + return PartitionedStruct(cls) diff --git a/build/torch-cuda/quack/dsl/torch_library_op.py b/build/torch-cuda/quack/dsl/torch_library_op.py new file mode 100644 index 0000000000000000000000000000000000000000..7042799b9d9e6f9047dd2723366d7311d8a659e4 --- /dev/null +++ b/build/torch-cuda/quack/dsl/torch_library_op.py @@ -0,0 +1,65 @@ +# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. +"""``cute_op``: ``torch.library.custom_op`` for CuTe DSL kernels. + +Same trick as ``torch.library.triton_op`` (register the impl as the fake/meta +kernel too), specialized for our setup: the fake is a pure no-op. Our ops +only mutate their inputs, so Dynamo / AOT autograd need no shape effect from +the fake, and kernel compilation is owned entirely by ``jit_cache`` (plus +the async compile pool) at real execution time. + +This removes the need for hand-written ``_*_fake`` twins on each op. + +Note: we deliberately do NOT gate on ``torch.compiler.is_compiling()`` — +that flag's underlying ``_is_compiling_flag`` is only set during +``torch.export``, never during ``torch.compile``. Dynamo's +``_get_fake_value_impl`` would otherwise run the body and surface +any ``_compile_*`` ``ValueError`` as a ``TorchRuntimeError`` graph break. +""" + +from __future__ import annotations + +from typing import Any, Callable, Iterable, Optional, Union + +import torch + +__all__ = ["cute_op"] + + +def cute_op( + name: str, + *, + mutates_args: Union[str, Iterable[str]], + schema: Optional[str] = None, + device_types: Optional[Union[str, Iterable[str]]] = None, +) -> Callable: + """Like ``torch.library.triton_op``, but for CuTe DSL kernels. + + Args: + name: ``"namespace::op_name"``. + mutates_args: Names of mutated tensor args. + schema: Optional explicit schema. Required when mutating an + ``Optional[Tensor]`` arg (PyTorch can't infer those). + device_types: Optional device-type restriction. + """ + + def dec(fn: Callable) -> Any: + kwargs: dict[str, Any] = {"mutates_args": mutates_args} + if schema is not None: + kwargs["schema"] = schema + if device_types is not None: + kwargs["device_types"] = device_types + op = torch.library.custom_op(name, fn, **kwargs) + + @op.register_fake + def _fake(*args, **kw): + # Pure no-op: our ops only mutate their input tensors, so under + # torch.compile / AOT autograd tracing there is no fake output to + # produce, and running the body would pay compile latency at + # dynamo trace time (or crash for shape/dtype combos the kernel + # intentionally rejects). Kernel compilation is handled by + # jit_cache + the async compile pool at real execution time. + return + + return op + + return dec diff --git a/build/torch-cuda/quack/epi_composable.py b/build/torch-cuda/quack/epi_composable.py index 8185e2526d5d3e6ff12a25244741732586437a76..7821cf6395816e723acddace41923c3652b5302d 100644 --- a/build/torch-cuda/quack/epi_composable.py +++ b/build/torch-cuda/quack/epi_composable.py @@ -1,35 +1,45 @@ # Copyright (c) 2025, Tri Dao. """ComposableEpiMixin: composes EpiOps into epilogue hook methods. -Subclasses declare _epi_ops as a tuple of EpiOp instances. The mixin auto-generates -epi_smem_bytes_per_stage, epi_get_smem_struct, epi_get_smem_tensors, epi_begin, -epi_begin_loop, epi_end, and EpilogueParams by querying each op. +Subclasses declare _epi_ops as a class-level tuple of EpiOp instances — the +static *schema* for the epilogue. The mixin auto-generates epi_smem_bytes, +epi_get_smem_struct, epi_get_smem_tensors, epi_begin, epi_begin_loop, +epi_end_loop, epi_end, and EpilogueParams by querying each op. -epi_begin and epi_begin_loop return dicts keyed by op name, so epi_visit_subtile -can access values by name (e.g. epi_loop_tensors["alpha"]). +Host-side, `_epi_ops_to_params_dict` (called from each subclass's +`epi_to_underlying_arguments`) filters `_epi_ops` automatically: it shadows +the class-level tuple with an instance-level tuple containing only the ops +whose argument tensor is non-None. All later iteration (host- and +device-side) walks the filtered tuple, so each EpiOp's hook methods can +assume their `param`/`arg_tensor` is non-None. -EpilogueParams is auto-generated from _epi_ops (via param_fields()) plus any -_extra_param_fields declared on the subclass. Subclasses still define -EpilogueArguments and epi_to_underlying_arguments manually. +The two host-side hooks that run *before* `epi_to_underlying_arguments` +(`resolve_epi_m_major` and the classmethod `epi_smem_bytes`) filter inline +from `args`, preserving the same non-None invariant for `op.epi_m_major_score` +and `op.smem_bytes`. They have to run first because `epi_to_underlying_arguments` +itself depends on static attributes set up before it (e.g. `gemm.epi_tile`, +`gemm.epi_c_stage`), and those attributes are themselves derived from +`epi_m_major` and the epi smem budget — a chicken-and-egg ordering we resolve by +letting these two hooks see the raw `args` and filter inline. + +epi_get_smem_tensors, epi_begin, and epi_begin_loop all return dicts keyed by +op name, so consumers access values by name (e.g. epi_smem_tensors["mAuxOut"], +epi_loop_tensors["alpha"]). Because inactive ops are filtered out, consumers +must use `.get(name)` (returns None for inactive ops) rather than `[name]`. + +EpilogueParams is auto-generated from the full class-level _epi_ops (via +param_fields()) plus any _extra_param_fields declared on the subclass. +Subclasses still define EpilogueArguments and epi_to_underlying_arguments +manually. """ from dataclasses import make_dataclass, MISSING import cutlass.cute as cute -from cutlass import const_expr - -from .epi_ops import EpiContext, Scalar - +from cutlass import Int32, const_expr -def _compute_smem_map(ops): - """Pre-compute name → smem tensor index for each non-Scalar op.""" - smem_map = {} - idx = 0 - for op in ops: - if not isinstance(op, Scalar): - smem_map[op.name] = idx - idx += 1 - return smem_map +from .cute_dsl_utils import ParamsBase +from .epi_ops import EpiContext, EpiSmemBytes, Scalar def _make_epi_params(epi_ops, extra_fields, bases): @@ -52,15 +62,11 @@ class ComposableEpiMixin: _epi_ops = () _extra_param_fields = () # [(name, type, default), ...] for non-op params (e.g. act_fn) - _epi_param_bases = () # Base classes for EpilogueParams (e.g. (ParamsBase,)) - _epi_smem_map = {} - _epi_has_async_ops = False + _epi_param_bases = (ParamsBase,) # Base classes for the auto-generated EpilogueParams def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) if cls._epi_ops: - cls._epi_smem_map = _compute_smem_map(cls._epi_ops) - cls._epi_has_async_ops = any(op.needs_async_fence() for op in cls._epi_ops) # Auto-generate EpilogueParams if not explicitly defined on this class if "EpilogueParams" not in cls.__dict__: cls.EpilogueParams = _make_epi_params( @@ -69,39 +75,89 @@ class ComposableEpiMixin: # --- Host-side: args → params --- + def _filter_epi_ops(self, args): + """Shadow `_epi_ops` with an instance-level tuple of only the ops whose + arg is non-None. Called automatically by `_epi_ops_to_params_dict`, so + subclass `epi_to_underlying_arguments` methods don't need to invoke it + directly. After this runs, op hook methods can assume their + `param`/`arg_tensor` is non-None.""" + self._epi_ops = tuple( + op for op in type(self)._epi_ops if getattr(args, op.name, None) is not None + ) + def _epi_ops_to_params_dict(self, args): - """Merge each op's to_params into a single dict. Subclasses call this, - add custom fields, then construct self.EpilogueParams(**d).""" + """Filter `_epi_ops` to active ops, then merge each op's to_params into + a single dict. Subclasses call this from epi_to_underlying_arguments, + add custom fields, then construct self.EpilogueParams(**d). Filtering + here means every later iteration of self._epi_ops (host- and + device-side) walks only active ops, and each op hook can assume its + arg is non-None.""" + self._filter_epi_ops(args) d = {} for op in self._epi_ops: d.update(op.to_params(self, args)) return d + def resolve_epi_m_major(self, args): + # Runs inside _setup_attributes, before epi_to_underlying_arguments, + # because epi_m_major drives epi_tile / smem layout choices that + # epi_to_underlying_arguments later consumes. self._epi_ops is still + # the class-level schema at this point, so we filter inline from args + # to keep op.epi_m_major_score's non-None invariant. + score = 0 + for op in type(self)._epi_ops: + arg = getattr(args, op.name, None) + if arg is not None: + score += op.epi_m_major_score(arg, self) + return score >= 0 + # --- Host-side: smem allocation (queried from ops) --- @classmethod - def epi_smem_bytes_per_stage(cls, args, cta_tile_shape_mnk, epi_tile): - return sum( - op.smem_bytes(getattr(args, op.name, None), cta_tile_shape_mnk, epi_tile) - for op in cls._epi_ops - ) + def epi_smem_bytes(cls, args, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None): + # Runs inside _compute_stages, before epi_to_underlying_arguments, + # because the AB/epi stage counts (and therefore epi_c_stage) depend + # on the epi smem budget that this returns; epi_to_underlying_arguments + # then consumes epi_c_stage to build TileLoad's staged smem layout. + # Stays a classmethod because _compute_stages is a classmethod and may + # be invoked without an instance, so we filter inline from args. + result = EpiSmemBytes() + for op in cls._epi_ops: + arg = getattr(args, op.name, None) + if arg is not None: + result += op.smem_bytes(arg, cta_tile_shape_mnk, epi_tile, warp_shape_mnk) + return result def epi_get_smem_struct(self, params): - fields = {} + fields = [] for op in self._epi_ops: result = op.smem_struct_field(self, params) if result is not None: - name, ftype = result - fields[name] = ftype - EpiSharedStorage = type("EpiSharedStorage", (), {"__annotations__": fields}) + fields.append(result) + + # cute.struct rejects empty annotations. When every active op contributes + # no smem (e.g. only Scalar ops, or no active ops at all), return a + # zero-byte placeholder matching gemm_base's default epi struct. + if not fields: + return cute.struct.MemRange[Int32, 0] + + # Sort smallest-to-largest so smaller fields pack ahead of larger + # higher-aligned fields, reducing smem wasted to alignment padding. + def _field_bytes(name_ftype): + wrapper = type("_F", (), {"__annotations__": {name_ftype[0]: name_ftype[1]}}) + return cute.struct(wrapper).size_in_bytes() + + fields.sort(key=_field_bytes) + annotations = {name: ftype for name, ftype in fields} + EpiSharedStorage = type("EpiSharedStorage", (), {"__annotations__": annotations}) return cute.struct(EpiSharedStorage) def epi_get_smem_tensors(self, params, storage): - return tuple( - op.get_smem_tensor(self, params, storage.epi) + return { + op.name: op.get_smem_tensor(self, params, storage.epi) for op in self._epi_ops if not isinstance(op, Scalar) - ) + } def epi_get_tma_atoms(self, params, *, loc=None, ip=None): atoms = [] @@ -123,6 +179,7 @@ class ComposableEpiMixin: varlen_manager, epilogue_barrier, tidx, + tRS_rD_layout=None, ): ctx = EpiContext( self, @@ -133,27 +190,24 @@ class ComposableEpiMixin: varlen_manager, epilogue_barrier, tidx, + tRS_rD_layout, ) - smem_map = self._epi_smem_map results = { op.name: op.begin( self, - getattr(params, op.name, None), - epi_smem_tensors[smem_map[op.name]] if op.name in smem_map else None, + getattr(params, op.name), + epi_smem_tensors.get(op.name), ctx, ) for op in self._epi_ops } - if const_expr(self._epi_has_async_ops): - has_async_data = any( - getattr(params, op.name, None) is not None - for op in self._epi_ops - if op.needs_async_fence() - ) - if const_expr(has_async_data): - cute.arch.cp_async_commit_group() - cute.arch.cp_async_wait_group(0) - epilogue_barrier.arrive_and_wait() + # self._epi_ops is filtered to active ops, so any op needing a fence + # has a non-None tensor; no inner None check required. + has_async_data = any(op.needs_async_fence() for op in self._epi_ops) + if const_expr(has_async_data): + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + epilogue_barrier.arrive_and_wait() return results def epi_begin_loop(self, params, epi_tensors, epi_coord): @@ -161,6 +215,59 @@ class ComposableEpiMixin: op.name: op.begin_loop(self, epi_tensors[op.name], epi_coord) for op in self._epi_ops } + def epi_tile_load_g2s_copy_fns( + self, + params, + epi_smem_tensors, + tile_coord_mnkl, + varlen_manager, + epi_pipeline, + ): + return tuple( + op.load_g2s_copy_fn( + self, + params, + epi_smem_tensors.get(op.name), + tile_coord_mnkl, + varlen_manager, + epi_pipeline, + ) + for op in self._epi_ops + if op.is_tile_load() + ) + + @cute.jit + def epi_tile_load_s2r(self, params, epi_tensors, stage_idx): + for op in self._epi_ops: + op.load_s2r(self, getattr(params, op.name), epi_tensors[op.name], stage_idx) + + @cute.jit + def epi_end_loop( + self, + params, + epi_tensors, + epi_coord, + epi_tile, + tiled_copy_t2r, + tiled_copy_r2s, + tile_coord_mnkl, + varlen_manager, + tidx, + ): + for op in self._epi_ops: + op.end_loop( + self, + getattr(params, op.name), + epi_tensors[op.name], + epi_coord, + epi_tile, + tiled_copy_t2r, + tiled_copy_r2s, + tile_coord_mnkl, + varlen_manager, + tidx, + ) + @cute.jit def epi_end( self, @@ -176,7 +283,7 @@ class ComposableEpiMixin: for op in self._epi_ops: op.end( self, - getattr(params, op.name, None), + getattr(params, op.name), epi_tensors[op.name], epi_tile, tiled_copy_t2r, diff --git a/build/torch-cuda/quack/epi_ops.py b/build/torch-cuda/quack/epi_ops.py index 19f873b5c120798b32b336146c033532829905f2..7dbf9f67c63717e5c20bd7abd410e6541c305b8f 100644 --- a/build/torch-cuda/quack/epi_ops.py +++ b/build/torch-cuda/quack/epi_ops.py @@ -5,14 +5,18 @@ Each EpiOp encapsulates a single tensor kind's behavior across the epilogue life smem allocation, begin (one-time per-tile setup), begin_loop (per-subtile extraction), end (cleanup). -The ops are composed via ComposableEpiMixin which iterates over a static _epi_ops tuple -to generate epi_smem_bytes_per_stage, epi_get_smem_struct, epi_get_smem_tensors, -epi_begin, and epi_begin_loop automatically. +The ops are composed via ComposableEpiMixin. Class-level `_epi_ops` is the +static schema; `_epi_ops_to_params_dict` (called from each subclass's +`epi_to_underlying_arguments`) shadows it with an instance-level tuple of only +the active ops (those whose arg tensor is non-None). All EpiOp hook methods +below therefore assume their `param` / `arg_tensor` is non-None — the +framework guarantees inactive ops are never iterated. """ import math import operator from functools import partial +from typing import NamedTuple import cutlass import cutlass.cute as cute @@ -26,7 +30,12 @@ from . import layout_utils as layout_utils class EpiContext: - """Shared context passed to EpiOp.begin methods. Bundles common arguments.""" + """Shared context passed to EpiOp.begin methods. Bundles common arguments. + + `tRS_rD_layout` is only populated by callers that need TileLoad — it's the + register layout of the matmul output tile, which TileLoad uses to shape its + own register tile so it lines up element-wise with tRS_rD in epi_visit_subtile. + """ __slots__ = ( "epi_tile", @@ -36,6 +45,7 @@ class EpiContext: "varlen_manager", "epilogue_barrier", "tidx", + "tRS_rD_layout", "partition_for_epilogue_fn", "num_epi_threads", "batch_idx", @@ -53,6 +63,7 @@ class EpiContext: varlen_manager, epilogue_barrier, tidx, + tRS_rD_layout=None, ): self.epi_tile = epi_tile self.tiled_copy_t2r = tiled_copy_t2r @@ -61,6 +72,7 @@ class EpiContext: self.varlen_manager = varlen_manager self.epilogue_barrier = epilogue_barrier self.tidx = tidx + self.tRS_rD_layout = tRS_rD_layout self.tile_M = gemm.cta_tile_shape_mnk[0] self.tile_N = gemm.cta_tile_shape_mnk[1] self.batch_idx = tile_coord_mnkl[3] @@ -120,6 +132,31 @@ def _get_lane_warp_layouts(tiled_copy, reference_src=True): return lane_layout_MN, warp_layout_MN +class EpiSmemBytes(NamedTuple): + """Shared-memory accounting for one epilogue op. + + unstaged: allocated once per CTA tile. + d_stage: allocated per D/store epilogue stage. + c_stage: allocated per C/load epilogue stage. + """ + + unstaged: int = 0 + d_stage: int = 0 + c_stage: int = 0 + + def __add__(self, other): + return EpiSmemBytes( + self.unstaged + other.unstaged, + self.d_stage + other.d_stage, + self.c_stage + other.c_stage, + ) + + def __radd__(self, other): + if other == 0: + return self + return self.__add__(other) + + class EpiOp: """Base class for composable epilogue operations.""" @@ -137,11 +174,15 @@ class EpiOp: Returns dict of {param_name: value}. Like EVT's to_underlying_arguments.""" return {} - # --- Host-side: smem allocation --- - def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile): - """Bytes of smem needed per stage. arg_tensor is the EpilogueArguments field.""" + def epi_m_major_score(self, arg_tensor, gemm): + """Preference for epilogue subtile order. Positive prefers M-major, negative N-major.""" return 0 + # --- Host-side: smem allocation --- + def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None): + """Bytes of smem needed by unstaged / D-stage / C-stage storage.""" + return EpiSmemBytes() + def smem_struct_field(self, gemm, params): """Return (field_name, field_type) for @cute.struct, or None if no smem needed. params is the full EpilogueParams object.""" @@ -156,6 +197,22 @@ class EpiOp: """Return list of TMA atoms for this op.""" return [] + def is_tile_load(self): + """Whether this op is a tile-sized epilogue input loaded through the C pipeline.""" + return False + + def load_g2s_copy_fn( + self, + gemm, + params, + smem_tensor, + tile_coord_mnkl, + varlen_manager, + epi_pipeline, + ): + """Return a per-subtile gmem->smem copy function, or None.""" + return None + # --- Device-side: kernel execution --- @cute.jit def begin(self, gemm, param, smem_tensor, ctx): @@ -166,6 +223,27 @@ class EpiOp: """Per-subtile extraction. Returns value for epi_visit_subtile.""" return state + @cute.jit + def load_s2r(self, gemm, param, state, stage_idx): + """Issue this op's tile-load smem->register copy for one epilogue stage.""" + pass + + def end_loop( + self, + gemm, + param, + state, + epi_coord, + epi_tile, + tiled_copy_t2r, + tiled_copy_r2s, + tile_coord_mnkl, + varlen_manager, + tidx, + ): + """Per-subtile cleanup after epi_visit_subtile.""" + pass + def needs_async_fence(self): """Whether this op issues async copies that need a fence.""" return False @@ -201,14 +279,9 @@ class Scalar(EpiOp): @cute.jit def begin(self, gemm, param, smem_tensor, ctx): - result = None - if const_expr(param is not None): - result = ( - utils.load_scalar_or_pointer(param, dtype=self.dtype) - if const_expr(self.dtype is not None) - else utils.load_scalar_or_pointer(param) - ) - return result + if const_expr(self.dtype is not None): + return utils.load_scalar_or_pointer(param, dtype=self.dtype) + return utils.load_scalar_or_pointer(param) class VecLoad(EpiOp): @@ -239,23 +312,20 @@ class VecLoad(EpiOp): def _coord_idx(self): return 1 if self.dim == 1 else 0 - def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile): - if arg_tensor is None: - return 0 - return self._tile_size(cta_tile_shape_mnk) * (arg_tensor.element_type.width // 8) + def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None): + return EpiSmemBytes( + unstaged=self._tile_size(cta_tile_shape_mnk) * (arg_tensor.element_type.width // 8) + ) def smem_struct_field(self, gemm, params): - tensor = getattr(params, self.name, None) - if tensor is None: - size, dtype = 0, Float32 - else: - size = self._tile_size(gemm.cta_tile_shape_mnk) - dtype = tensor.element_type - return (f"s_{self.name}", cute.struct.Align[cute.struct.MemRange[dtype, size], 16]) + tensor = getattr(params, self.name) + size = self._tile_size(gemm.cta_tile_shape_mnk) + return ( + f"s_{self.name}", + cute.struct.Align[cute.struct.MemRange[tensor.element_type, size], 16], + ) def get_smem_tensor(self, gemm, params, storage_epi): - if getattr(params, self.name, None) is None: - return None return getattr(storage_epi, f"s_{self.name}").get_tensor( cute.make_layout(self._tile_size(gemm.cta_tile_shape_mnk)) ) @@ -263,49 +333,61 @@ class VecLoad(EpiOp): def needs_async_fence(self): return True + def epi_m_major_score(self, arg_tensor, gemm): + # It costs more registers (say 4x) to keep rowvec in register vs keeping colvec in register + return 4 if self.dim == 1 else -1 + def _get_gmem_vec(self, param, ctx): """Get the global memory vector for this tile. Override for varlen.""" return param[ctx.batch_idx, None] @cute.jit def begin(self, gemm, param, smem_tensor, ctx): - tDsV = None - if const_expr(param is not None): - dtype = param.element_type - num_copy_elems = const_expr(max(32, dtype.width)) // dtype.width - thr_copy = copy_utils.tiled_copy_1d( - dtype, ctx.num_epi_threads, num_copy_elems, is_async=True - ).get_slice(ctx.tidx) - mVec = self._get_gmem_vec(param, ctx) - tile_dim = self._tile_dim(ctx) - coord_idx = ctx.tile_coord_mnkl[self._coord_idx()] - gVec = cute.local_tile(mVec, (tile_dim,), (coord_idx,)) - tVgV = thr_copy.partition_S(gVec) - tVsV = thr_copy.partition_D(smem_tensor) - tVcV = thr_copy.partition_S(cute.make_identity_tensor(tile_dim)) - limit = min(cute.size(mVec, mode=[0]) - coord_idx * tile_dim, tile_dim) - pred = cute.make_rmem_tensor((1, cute.size(tVsV.shape[1])), Boolean) - for m in cutlass.range(cute.size(tVsV.shape[1]), unroll_full=True): - pred[0, m] = tVcV[0, m] < limit - cute.copy(thr_copy, tVgV, tVsV, pred=pred) - tDsV = ctx.partition_for_epilogue_fn( - cute.make_tensor( - smem_tensor.iterator, - cute.make_layout((ctx.tile_M, ctx.tile_N), stride=self._broadcast_stride()), - ) + dtype = param.element_type + num_copy_elems = const_expr(max(32, dtype.width)) // dtype.width + thr_copy = copy_utils.tiled_copy_1d( + dtype, ctx.num_epi_threads, num_copy_elems, is_async=True + ).get_slice(ctx.tidx) + mVec = self._get_gmem_vec(param, ctx) + tile_dim = self._tile_dim(ctx) + coord_idx = ctx.tile_coord_mnkl[self._coord_idx()] + gVec = cute.local_tile(mVec, (tile_dim,), (coord_idx,)) + tVgV = thr_copy.partition_S(gVec) + tVsV = thr_copy.partition_D(smem_tensor) + tVcV = thr_copy.partition_S(cute.make_identity_tensor(tile_dim)) + limit = min(cute.size(mVec, mode=[0]) - coord_idx * tile_dim, tile_dim) + for m in cutlass.range(cute.size(tVsV.shape[1]), unroll_full=True): + if tVcV[0, m] < tile_dim: # Guard to avoid writing beyond the smem we've allocated + pred = cute.make_rmem_tensor(1, Boolean) + pred[0] = tVcV[0, m] < limit + cute.copy(thr_copy, tVgV[None, m], tVsV[None, m], pred=pred) + tDsV = ctx.partition_for_epilogue_fn( + cute.make_tensor( + smem_tensor.iterator, + cute.make_layout((ctx.tile_M, ctx.tile_N), stride=self._broadcast_stride()), ) - if const_expr(ctx.tiled_copy_t2r is not None): - tDsV = ctx.tiled_copy_r2s.retile(tDsV) - return tDsV + ) + if const_expr(ctx.tiled_copy_t2r is not None): + tDsV = ctx.tiled_copy_r2s.retile(tDsV) + # Pre-allocate register tensor reused across begin_loop calls + tDsV_sub = cute.group_modes(tDsV, 3, cute.rank(tDsV))[None, None, None, 0] + tDrV_cvt = cute.make_rmem_tensor(tDsV_sub.layout, gemm.acc_dtype) + return [tDsV, tDrV_cvt] @cute.jit def begin_loop(self, gemm, state, epi_coord): - tDrV_cvt = None - if const_expr(state is not None): - tDsV_cur = cute.group_modes(state, 3, cute.rank(state))[None, None, None, epi_coord] + tDsV, tDrV_cvt = state[0], state[1] + should_load = Boolean(True) + if const_expr(self.dim == 1): + if const_expr(gemm.epi_m_major): + should_load = epi_coord[0] == 0 + else: + if const_expr(not gemm.epi_m_major): + should_load = epi_coord[1] == 0 + if should_load: + tDsV_cur = cute.group_modes(tDsV, 3, cute.rank(tDsV))[None, None, None, epi_coord] tDrV = cute.make_rmem_tensor(tDsV_cur.layout, tDsV_cur.element_type) cute.autovec_copy(cute.filter_zeros(tDsV_cur), cute.filter_zeros(tDrV)) - tDrV_cvt = cute.make_rmem_tensor_like(tDrV, gemm.acc_dtype) tDrV_cvt.store(tDrV.load().to(gemm.acc_dtype)) return tDrV_cvt @@ -338,63 +420,47 @@ class ColVecLoad(VecLoad): @cute.jit def begin(self, gemm, param, smem_tensor, ctx): - tDsV = None - tDrV_cvt = None - if const_expr(param is not None): - dtype = param.element_type - num_copy_elems = const_expr(max(32, dtype.width)) // dtype.width - thr_copy = copy_utils.tiled_copy_1d( - dtype, ctx.num_epi_threads, num_copy_elems, is_async=True - ).get_slice(ctx.tidx) - mVec = self._get_gmem_vec(param, ctx) - tile_dim = self._tile_dim(ctx) - coord_idx = ctx.tile_coord_mnkl[self._coord_idx()] - gVec = cute.local_tile(mVec, (tile_dim,), (coord_idx,)) - tVgV = thr_copy.partition_S(gVec) - tVsV = thr_copy.partition_D(smem_tensor) - tVcV = thr_copy.partition_S(cute.make_identity_tensor(tile_dim)) - # ColVec uses varlen-aware limit - limit = min( - ctx.varlen_manager.len_m(ctx.batch_idx) - coord_idx * tile_dim, - tile_dim, - ) - pred = cute.make_rmem_tensor((1, cute.size(tVsV.shape[1])), Boolean) - for m in cutlass.range(cute.size(tVsV.shape[1]), unroll_full=True): - pred[0, m] = tVcV[0, m] < limit - cute.copy(thr_copy, tVgV, tVsV, pred=pred) - tDsV = ctx.partition_for_epilogue_fn( - cute.make_tensor( - smem_tensor.iterator, - cute.make_layout((ctx.tile_M, ctx.tile_N), stride=self._broadcast_stride()), - ) + dtype = param.element_type + num_copy_elems = const_expr(max(32, dtype.width)) // dtype.width + thr_copy = copy_utils.tiled_copy_1d( + dtype, ctx.num_epi_threads, num_copy_elems, is_async=True + ).get_slice(ctx.tidx) + mVec = self._get_gmem_vec(param, ctx) + tile_dim = self._tile_dim(ctx) + coord_idx = ctx.tile_coord_mnkl[self._coord_idx()] + gVec = cute.local_tile(mVec, (tile_dim,), (coord_idx,)) + tVgV = thr_copy.partition_S(gVec) + tVsV = thr_copy.partition_D(smem_tensor) + tVcV = thr_copy.partition_S(cute.make_identity_tensor(tile_dim)) + # ColVec uses varlen-aware limit + limit = min( + ctx.varlen_manager.len_m(ctx.batch_idx) - coord_idx * tile_dim, + tile_dim, + ) + for m in cutlass.range(cute.size(tVsV.shape[1]), unroll_full=True): + if tVcV[0, m] < tile_dim: # Guard to avoid writing beyond the smem we've allocated + pred = cute.make_rmem_tensor(1, Boolean) + pred[0] = tVcV[0, m] < limit + cute.copy(thr_copy, tVgV[None, m], tVsV[None, m], pred=pred) + tDsV = ctx.partition_for_epilogue_fn( + cute.make_tensor( + smem_tensor.iterator, + cute.make_layout((ctx.tile_M, ctx.tile_N), stride=self._broadcast_stride()), ) - if const_expr(ctx.tiled_copy_t2r is not None): - tDsV = ctx.tiled_copy_r2s.retile(tDsV) - # Pre-allocate register tensor reused across begin_loop calls - tDsV_sub = cute.group_modes(tDsV, 3, cute.rank(tDsV))[None, None, None, 0] - tDrV_cvt = cute.make_rmem_tensor(tDsV_sub.layout, gemm.acc_dtype) + ) + if const_expr(ctx.tiled_copy_t2r is not None): + tDsV = ctx.tiled_copy_r2s.retile(tDsV) + # Pre-allocate register tensor reused across begin_loop calls + tDsV_sub = cute.group_modes(tDsV, 3, cute.rank(tDsV))[None, None, None, 0] + tDrV_cvt = cute.make_rmem_tensor(tDsV_sub.layout, gemm.acc_dtype) return [tDsV, tDrV_cvt] - @cute.jit - def begin_loop(self, gemm, state, epi_coord): - tDsV, tDrV_cvt = state[0], state[1] - if const_expr(tDsV is not None): - # Col vector is constant across N subtiles — only copy on first N subtile. - # Assumes N-major epi subtile order: epi_tile_layout = ordered_layout(..., order=(1,0)) - epi_n = epi_coord[1] - if epi_n == 0: - tDsV_cur = cute.group_modes(tDsV, 3, cute.rank(tDsV))[None, None, None, epi_coord] - tDrV = cute.make_rmem_tensor(tDsV_cur.layout, tDsV_cur.element_type) - cute.autovec_copy(cute.filter_zeros(tDsV_cur), cute.filter_zeros(tDrV)) - tDrV_cvt.store(tDrV.load().to(gemm.acc_dtype)) - return tDrV_cvt - class TileStore(EpiOp): """Tile-sized output tensor stored via TMA (e.g. postact). Args: - name: field name in EpilogueArguments/Params (e.g. "mPostAct") + name: field name in EpilogueArguments/Params (e.g. "mAuxOut") epi_tile_fn: optional (gemm, epi_tile) -> epi_tile for half-tile (GemmGated) """ @@ -412,13 +478,13 @@ class TileStore(EpiOp): return f"epi_tile_{self.name}" def param_fields(self): - from dataclasses import MISSING - + # Defaults are None so EpilogueParams can be constructed when this op is + # filtered out (inactive). Active calls always set all four via to_params. return [ - (self._tma_atom_key(), object, MISSING), - (self.name, object, MISSING), - (self._smem_layout_key(), object, MISSING), - (self._epi_tile_key(), object, MISSING), + (self._tma_atom_key(), object, None), + (self.name, object, None), + (self._smem_layout_key(), object, None), + (self._epi_tile_key(), object, None), ] def to_params(self, gemm, args): @@ -434,50 +500,207 @@ class TileStore(EpiOp): self._epi_tile_key(): epi_tile_out, } - def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile): - if arg_tensor is None: - return 0 + def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None): if self.epi_tile_fn is not None: epi_tile = self.epi_tile_fn(None, epi_tile) - return cute.size(cute.shape(epi_tile)) * (arg_tensor.element_type.width // 8) + # epi_tile may contain Layout entries (from SM100's compute_epilogue_tile_shape + # fixup path), so extract the int shape first. + return EpiSmemBytes( + d_stage=cute.size(cute.shape(epi_tile)) * (arg_tensor.element_type.width // 8) + ) def smem_struct_field(self, gemm, params): - smem_layout_key = self._smem_layout_key() - if not hasattr(params, smem_layout_key): - return (f"s_{self.name}", cute.struct.MemRange[Float32, 0]) + smem_layout = getattr(params, self._smem_layout_key()) return ( f"s_{self.name}", cute.struct.Align[ cute.struct.MemRange[ - gemm.postact_dtype, - cute.cosize(getattr(params, smem_layout_key)), + gemm.aux_out_dtype, + cute.cosize(smem_layout), ], gemm.buffer_align_bytes, ], ) def get_smem_tensor(self, gemm, params, storage_epi): - smem_layout_key = self._smem_layout_key() - if not hasattr(params, smem_layout_key): - return None - smem_layout = getattr(params, smem_layout_key) + smem_layout = getattr(params, self._smem_layout_key()) return getattr(storage_epi, f"s_{self.name}").get_tensor( smem_layout.outer, swizzle=smem_layout.inner, ) def tma_atoms(self, gemm, params): - tma_key = self._tma_atom_key() - if hasattr(params, tma_key): - return [getattr(params, tma_key)] - return [] + return [getattr(params, self._tma_atom_key())] + + +class _TileLoadState(NamedTuple): + """Per-tile register state produced by TileLoad.begin and consumed by load_s2r / + begin_loop. tRS_rTile is the register tile partitioned to match tRS_rD's layout; + tSR_sTile / tSR_rTile drive the per-stage smem→register copy.""" + + tiled_copy_s2r: object + tRS_rTile: object + tSR_rTile: object + tSR_sTile: object + + +class TileLoad(EpiOp): + """Tile-sized auxiliary input loaded through the epilogue load pipeline. + + TileLoad uses the same staged gmem->smem->register pipeline as GEMM's C operand, + but it is exposed to the epilogue as ``epi_loop_tensors[name]`` instead of as + ``tRS_rC``. That lets custom epilogues consume extra MxN tensors without using + the GEMM C argument. + + Its shared memory is accounted as ``EpiSmemBytes.c_stage``, so it is allocated + per epilogue load stage. Multiple TileLoads are supported: each has its own TMA + descriptor and smem buffer, and the pipeline transaction count includes C plus + all enabled TileLoad buffers. Supported on SM90, SM100, and SM120. + """ + + def __init__(self, name, epi_tile_fn=None): + super().__init__(name) + self.epi_tile_fn = epi_tile_fn + + def _tma_atom_key(self): + return f"tma_atom_{self.name}" + + def _smem_layout_key(self): + return f"epi_{self.name}_smem_layout_staged" + + def _epi_tile_key(self): + return f"epi_tile_{self.name}" + + # The original LayoutEnum and element_type can't be recovered from the + # TMA-prepared tensor that ends up in params (`from_tensor` returns a typing + # annotation post-TMA, not a Numeric class). We stash both on the gemm at + # to_params time and read them back in begin(). The dtype is also exposed on + # the params dataclass for smem_struct_field. + def _layout_gemm_attr(self): + return f"_tile_load_layout_{self.name}" + + def _dtype_gemm_attr(self): + return f"_tile_load_dtype_{self.name}" + + def _dtype_field(self): + return f"{self.name}_dtype" + + def param_fields(self): + # Defaults are None so EpilogueParams can be constructed when this op is + # filtered out (inactive). Active calls always set all five via to_params. + return [ + (self._tma_atom_key(), object, None), + (self.name, object, None), + (self._smem_layout_key(), object, None), + (self._epi_tile_key(), object, None), + (self._dtype_field(), object, None), + ] + + def to_params(self, gemm, args): + tensor = getattr(args, self.name) + setattr(gemm, self._layout_gemm_attr(), cutlass.utils.LayoutEnum.from_tensor(tensor)) + setattr(gemm, self._dtype_gemm_attr(), tensor.element_type) + epi_tile = self.epi_tile_fn(gemm, gemm.epi_tile) if self.epi_tile_fn else None + tma_atom, tma_tensor, smem_layout, epi_tile_out = setup_epi_tensor( + gemm, tensor, epi_tile=epi_tile, op_type="load", stage=gemm.epi_c_stage + ) + return { + self._tma_atom_key(): tma_atom, + self.name: tma_tensor, + self._smem_layout_key(): smem_layout, + self._epi_tile_key(): epi_tile_out, + self._dtype_field(): tensor.element_type, + } + + def is_tile_load(self): + return True + + def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None): + if self.epi_tile_fn is not None: + epi_tile = self.epi_tile_fn(None, epi_tile) + # epi_tile may contain Layout entries from SM100's compute_epilogue_tile_shape + # fixup; extract the int shape first. + return EpiSmemBytes( + c_stage=cute.size(cute.shape(epi_tile)) * (arg_tensor.element_type.width // 8) + ) + + def smem_struct_field(self, gemm, params): + smem_layout = getattr(params, self._smem_layout_key()) + dtype = getattr(params, self._dtype_field()) + return ( + f"s_{self.name}", + cute.struct.Align[ + cute.struct.MemRange[dtype, cute.cosize(smem_layout)], + gemm.buffer_align_bytes, + ], + ) + + def get_smem_tensor(self, gemm, params, storage_epi): + smem_layout = getattr(params, self._smem_layout_key()) + return getattr(storage_epi, f"s_{self.name}").get_tensor( + smem_layout.outer, + swizzle=smem_layout.inner, + ) + + def tma_atoms(self, gemm, params): + return [getattr(params, self._tma_atom_key())] + + def load_g2s_copy_fn( + self, + gemm, + params, + smem_tensor, + tile_coord_mnkl, + varlen_manager, + epi_pipeline, + ): + tensor = getattr(params, self.name) + batch_idx = tile_coord_mnkl[3] + copy_tile_fn, _, _ = gemm.epilog_gmem_copy_and_partition( + getattr(params, self._tma_atom_key()), + varlen_manager.offset_batch_epi(tensor, batch_idx), + gemm.cta_tile_shape_mnk[:2], + getattr(params, self._epi_tile_key()), + smem_tensor, + tile_coord_mnkl, + ) + return copy_utils.tma_producer_copy_fn(copy_tile_fn, epi_pipeline) + + @cute.jit + def begin(self, gemm, param, smem_tensor, ctx): + assert gemm.arch in (90, 100, 120), "TileLoad requires the SM90/SM100/SM120 epilogue path" + assert ctx.tRS_rD_layout is not None + smem_load_ref = ctx.tiled_copy_t2r if const_expr(gemm.arch == 100) else gemm.tiled_mma + tiled_copy_s2r, tRS_rTile, tSR_rTile, tSR_sTile = gemm.epilog_smem_load_and_partition( + smem_load_ref, + getattr(gemm, self._layout_gemm_attr()), + getattr(gemm, self._dtype_gemm_attr()), + smem_tensor, + ctx.tRS_rD_layout, + ctx.tidx, + ) + # Shape: (s2r-copy-handle, register-tile-as-rD-layout, smem→r retile target, + # smem→r staged source). begin_loop returns tRS_rTile; load_s2r uses the rest. + return _TileLoadState(tiled_copy_s2r, tRS_rTile, tSR_rTile, tSR_sTile) + + @cute.jit + def load_s2r(self, gemm, param, state, stage_idx): + cute.copy( + state.tiled_copy_s2r, + state.tSR_sTile[None, None, None, stage_idx], + state.tSR_rTile, + ) + + @cute.jit + def begin_loop(self, gemm, state, epi_coord): + return state.tRS_rTile @cute.jit def vec_multiply(gemm, tRS_rD, tDrColVec, tDrRowVec): - """Multiply tRS_rD by colvec and/or rowvec in-place. Uses packed f32x2 on SM100+.""" + """Multiply tRS_rD by colvec and/or rowvec in-place. Uses packed f32x2 on SM100.""" if const_expr(tDrColVec is not None): - if const_expr(gemm.arch < 100): + if const_expr(gemm.arch != 100): for i in cutlass.range(cute.size(tDrColVec), unroll_full=True): tRS_rD[i] *= tDrColVec[i] else: @@ -487,7 +710,7 @@ def vec_multiply(gemm, tRS_rD, tDrColVec, tDrRowVec): (tDrColVec[2 * i], tDrColVec[2 * i + 1]), ) if const_expr(tDrRowVec is not None): - if const_expr(gemm.arch < 100): + if const_expr(gemm.arch != 100): for i in cutlass.range(cute.size(tDrRowVec), unroll_full=True): tRS_rD[i] *= tDrRowVec[i] else: @@ -503,13 +726,13 @@ def colvec_reduce_accumulate(gemm, tDrReduce, tRS_rInput, transform_fn=None, rSc """Accumulate transform_fn(input) or input * rScale into a ColVecReduce buffer. If transform_fn is provided, accumulates transform_fn(input[i]). - If rScale is provided, accumulates input[i] * rScale[i] (uses mul/fma for SM100). + If rScale is provided, accumulates input[i] * rScale[i] (uses packed mul/fma for SM100). If neither, accumulates input directly (identity). """ if const_expr(tDrReduce is not None): if const_expr(transform_fn is None): transform_fn = lambda x: x - if const_expr(gemm.arch < 100): + if const_expr(gemm.arch != 100): for i in cutlass.range(cute.size(tDrReduce), unroll_full=True): val = transform_fn(tRS_rInput[i]) tDrReduce[i] += val * rScale[i] if const_expr(rScale is not None) else val @@ -521,6 +744,7 @@ def colvec_reduce_accumulate(gemm, tDrReduce, tRS_rInput, transform_fn=None, rSc for m in cutlass.range(cute.size(tDrReduce_mn, mode=[0]), unroll_full=True): inp = lambda n: (tRS_rInput_mn[m, 2 * n], tRS_rInput_mn[m, 2 * n + 1]) val0 = transform_fn(inp(0)) + assert cute.size(tDrReduce_mn, mode=[1]) % 2 == 0 if const_expr(rScale is not None): row_sum = cute.arch.mul_packed_f32x2(val0, (rScale_mn[m, 0], rScale_mn[m, 1])) else: @@ -536,14 +760,48 @@ def colvec_reduce_accumulate(gemm, tDrReduce, tRS_rInput, transform_fn=None, rSc tDrReduce_mn[m, 0] += row_sum[0] + row_sum[1] -class ColVecReduce(EpiOp): - """Column vector reduction: accumulates across N subtiles in registers, - then warp-reduces and writes to gmem in epi_end. +@cute.jit +def rowvec_reduce_accumulate(gemm, tDrReduce, tRS_rInput, transform_fn=None, rScale=None): + """Accumulate transform_fn(input) or input * rScale into a RowVecReduce buffer. - No smem. The accumulation itself happens in epi_visit_subtile (user code). - This op handles the register allocation (begin), per-subtile slicing (begin_loop), - and final warp reduction + gmem write (end). + Reduces along M dimension, keeping N. The zero-stride layout on M ensures + elements at different M positions but same N column accumulate correctly. """ + if const_expr(tDrReduce is not None): + if const_expr(transform_fn is None): + transform_fn = lambda x: x + if const_expr(gemm.arch != 100): + for i in cutlass.range(cute.size(tDrReduce), unroll_full=True): + val = transform_fn(tRS_rInput[i]) + tDrReduce[i] += val * rScale[i] if const_expr(rScale is not None) else val + else: + # Keep CUTLASS's linear fragment indexing, but use packed f32x2 arithmetic + # for any transform that accepts and returns an f32x2 tuple. + # We have to be careful to avoid tDrReduce[2 * i] and tDrReduce[2 * i + 1] aliasing + # each other. For SM100, tDrReduce has layout ((32,1),1,1):((1,0),0,0) or + # (((2,2,4),1),2,1):(((1,0,8),0),0,0), so this works. But it's error-prone. + for i in cutlass.range(cute.size(tRS_rInput) // 2, unroll_full=True): + acc = (tDrReduce[2 * i], tDrReduce[2 * i + 1]) + val = (tRS_rInput[2 * i], tRS_rInput[2 * i + 1]) + val = transform_fn(val) + if const_expr(rScale is not None): + scale = (rScale[2 * i], rScale[2 * i + 1]) + tDrReduce[2 * i], tDrReduce[2 * i + 1] = cute.arch.fma_packed_f32x2( + val, scale, acc + ) + else: + tDrReduce[2 * i], tDrReduce[2 * i + 1] = cute.arch.add_packed_f32x2(val, acc) + if const_expr(cute.size(tRS_rInput) % 2 != 0): + i = cute.size(tRS_rInput) - 1 + val = transform_fn(tRS_rInput[i]) + tDrReduce[i] += val * rScale[i] if const_expr(rScale is not None) else val + + +class VecReduce(EpiOp): + """Base class for row/column vector reductions.""" + + dim = 0 # 0 for colvec output along M, 1 for rowvec output along N + epi_m_major_preference = 0 def param_fields(self): return [(self.name, object, None)] @@ -551,31 +809,83 @@ class ColVecReduce(EpiOp): def to_params(self, gemm, args): return {self.name: assume_stride_divisibility(getattr(args, self.name))} + def epi_m_major_score(self, arg_tensor, gemm): + return self.epi_m_major_preference + + def _tile_size(self, cta_tile_shape_mnk): + return cta_tile_shape_mnk[self.dim] + + def _broadcast_stride(self): + # Col: stride (1,0) broadcasts along N. Row: stride (0,1) broadcasts along M. + return (1, 0) if self.dim == 0 else (0, 1) + + def _reduce_dim(self): + return 1 - self.dim + + def _smem_warps(self, warp_shape_mnk): + warps = warp_shape_mnk[self._reduce_dim()] if warp_shape_mnk is not None else 1 + return max(warps - 1, 0) + + def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None): + smem_warps = self._smem_warps(warp_shape_mnk) + if smem_warps == 0: + return EpiSmemBytes() + return EpiSmemBytes( + unstaged=self._tile_size(cta_tile_shape_mnk) * smem_warps * (Float32.width // 8) + ) + + def smem_struct_field(self, gemm, params): + smem_warps = self._smem_warps(gemm.epi_smem_warp_shape_mnk()) + if smem_warps == 0: + return None + size = self._tile_size(gemm.cta_tile_shape_mnk) * smem_warps + return (f"s_{self.name}", cute.struct.Align[cute.struct.MemRange[Float32, size], 16]) + + def get_smem_tensor(self, gemm, params, storage_epi): + smem_warps = self._smem_warps(gemm.epi_smem_warp_shape_mnk()) + if smem_warps == 0: + return None + return getattr(storage_epi, f"s_{self.name}").get_tensor( + cute.make_layout((self._tile_size(gemm.cta_tile_shape_mnk), smem_warps)) + ) + @cute.jit def begin(self, gemm, param, smem_tensor, ctx): - tDrReduce = None - if const_expr(param is not None): - colvec_mma_layout = cute.make_layout((ctx.tile_M, ctx.tile_N), stride=(1, 0)) - tDrReduce_layout = ctx.partition_for_epilogue_fn( - cute.make_rmem_tensor(colvec_mma_layout, Float32) - ).layout - tDrReduce = cute.make_rmem_tensor(tDrReduce_layout, Float32) - cute.filter_zeros(tDrReduce).fill(0.0) - return tDrReduce + vec_mma_layout = cute.make_layout((ctx.tile_M, ctx.tile_N), stride=self._broadcast_stride()) + tDrReduce_layout = ctx.partition_for_epilogue_fn( + cute.make_rmem_tensor(vec_mma_layout, Float32) + ).layout + tDrReduce = cute.make_rmem_tensor(tDrReduce_layout, Float32) + return (tDrReduce, smem_tensor) @cute.jit def begin_loop(self, gemm, state, epi_coord): - result = None - if const_expr(state is not None): - result = cute.group_modes(state, 3, cute.rank(state))[None, None, None, epi_coord] + tDrReduce = state[0] + result = tDrReduce[None, None, None, epi_coord[0], epi_coord[1]] + if const_expr(epi_coord[self._reduce_dim()] == 0): + cute.filter_zeros(result).fill(0.0) return result + +class ColVecReduce(VecReduce): + """Column vector reduction: accumulates across N subtiles in registers, + then reduces across N lanes/warps and writes to gmem per completed M stripe. + + The accumulation itself happens in epi_visit_subtile (user code). + This op handles the register allocation (begin), per-subtile slicing (begin_loop), + and reduction + gmem write (end_loop). + """ + + dim = 0 + epi_m_major_preference = -1 + @cute.jit - def end( + def end_loop( self, gemm, param, state, + epi_coord, epi_tile, tiled_copy_t2r, tiled_copy_r2s, @@ -583,9 +893,13 @@ class ColVecReduce(EpiOp): varlen_manager, tidx, ): - """Intra-warp shuffle reduction across N lanes, then direct gmem write.""" - if const_expr(param is not None): - tDrReduce = state + """Flush the current M stripe when the last N subtile has accumulated.""" + epi_tile_shape = cute.zipped_divide( + cute.make_layout(gemm.cta_tile_shape_mnk[:2]), epi_tile + ).shape[1] + if const_expr(epi_coord[1] == epi_tile_shape[1] - 1): + tDrReduce, sDrReduce = state[0], state[1] + tDrReduce_cur = tDrReduce[None, None, None, epi_coord[0], epi_coord[1]] tiled_copy = tiled_copy_t2r if tiled_copy_t2r is not None else tiled_copy_r2s reference_src = tiled_copy_t2r is None @@ -593,26 +907,147 @@ class ColVecReduce(EpiOp): lane_layout_MN, warp_layout_MN = _get_lane_warp_layouts(tiled_copy, reference_src) # For ColVecReduce: reduce across N lanes (lanes_in_N threads share same M row) lanes_in_N = cute.size(lane_layout_MN, mode=[1]) + is_lane_n_leader = cute.arch.lane_idx() % lanes_in_N == 0 # Typically lanes_in_N is 4 for Sm90 assert lanes_in_N == 1 << int(math.log2(lanes_in_N)), ( "lanes_in_N must be a power of 2 for butterfly reduction" ) - # ── Intra-warp shuffle reduction across N lanes ── + # Intra-warp shuffle reduction across N lanes if const_expr(lanes_in_N > 1): + # Assumes threads for each M row are contiguous along N, so + # warp_reduction over groups of lanes_in_N matches lane_layout_MN. assert lane_layout_MN.stride[1] == 1 - tDrReduce_flt = cute.filter_zeros(tDrReduce) + tDrReduce_flt = cute.filter_zeros(tDrReduce_cur) for i in cutlass.range(cute.size(tDrReduce_flt), unroll_full=True): tDrReduce_flt[i] = cute.arch.warp_reduction( tDrReduce_flt[i], operator.add, threads_in_group=lanes_in_N ) warp_N = warp_layout_MN[1] - assert cute.size(warp_N) == 1, ( - "ColVecReduce assumes all reduction cols are within the same warp" + warps_in_N = const_expr(cute.size(warp_N)) + partition_for_epilogue_fn = partial( + partition_for_epilogue, + epi_tile=epi_tile, + tiled_copy=tiled_copy, + tidx=tidx, + reference_src=tiled_copy_t2r is None, + ) + tile_M, tile_N = gemm.cta_tile_shape_mnk[:2] + tDcD = partition_for_epilogue_fn(cute.make_identity_tensor((tile_M, tile_N))) + tDcD_cur = tDcD[None, None, None, epi_coord[0], epi_coord[1]] + tDrReduce_m = layout_utils.convert_layout_zero_stride( + tDrReduce_cur, tDrReduce_cur.layout + )[None, 0] + tDcD_m = layout_utils.convert_layout_zero_stride(tDcD_cur, tDrReduce_cur.layout)[ + None, 0 + ] + + # Inter-warp reduction through smem + warp_idx = cute.arch.make_warp_uniform(tidx // cute.arch.WARP_SIZE) + warp_n_idx = warp_layout_MN.get_hier_coord(warp_idx)[1] + if const_expr(warps_in_N > 1): + if warp_n_idx > 0 and is_lane_n_leader: + for m in cutlass.range(cute.size(tDcD_m, mode=[0])): + row_idx = tDcD_m[m][0] + sDrReduce[row_idx, warp_n_idx - 1] = tDrReduce_m[m] + gemm.epilogue_barrier.arrive_and_wait() + if warp_n_idx == 0 and is_lane_n_leader: + for m in cutlass.range(cute.size(tDcD_m, mode=[0])): + row_idx = tDcD_m[m][0] + for warp_n in cutlass.range_constexpr(1, warps_in_N): + tDrReduce_m[m] += sDrReduce[row_idx, warp_n - 1] + + # Write to gmem + batch_idx = tile_coord_mnkl[3] + limit_m = min(varlen_manager.len_m(batch_idx) - tile_coord_mnkl[0] * tile_M, tile_M) + limit_n_tiles = param.shape[2] if not varlen_manager.varlen_m else param.shape[1] + if const_expr(not varlen_manager.varlen_m): + mColVec = param[batch_idx, None, tile_coord_mnkl[1]] + else: + mColVec = cute.domain_offset( + (varlen_manager.params.cu_seqlens_m[batch_idx],), + param[None, tile_coord_mnkl[1]], + ) + gColVec = cute.local_tile(mColVec, (tile_M,), (tile_coord_mnkl[0],)) + should_write_gmem = ( + is_lane_n_leader + if const_expr(warps_in_N == 1) + else warp_n_idx == 0 and is_lane_n_leader ) + if tile_coord_mnkl[1] < limit_n_tiles and should_write_gmem: + for m in cutlass.range(cute.size(tDcD_m, mode=[0])): + row_idx = tDcD_m[m][0] + if row_idx < limit_m: + gColVec[row_idx] = tDrReduce_m[m] + - # ── Direct gmem write (no inter-warp reduction needed: warps_in_N == 1) ── +class RowVecReduce(VecReduce): + """Row vector reduction: accumulates across M subtiles in registers, + then reduces across M lanes/warps and writes to gmem per completed N stripe. + + Output shape is (L, ceildiv(M, tile_M), N): one partial sum per CTA-M tile per + N column. This mirrors ColVecReduce with M/N swapped. + """ + + dim = 1 + epi_m_major_preference = 4 + + @cute.jit + def end_loop( + self, + gemm, + param, + state, + epi_coord, + epi_tile, + tiled_copy_t2r, + tiled_copy_r2s, + tile_coord_mnkl, + varlen_manager, + tidx, + ): + """Flush the current N stripe when the last M subtile has accumulated.""" + epi_tile_shape = cute.zipped_divide( + cute.make_layout(gemm.cta_tile_shape_mnk[:2]), epi_tile + ).shape[1] + if const_expr(epi_coord[0] == epi_tile_shape[0] - 1): + tDrReduce, sDrReduce = state[0], state[1] + tDrReduce_cur = tDrReduce[None, None, None, epi_coord[0], epi_coord[1]] + tiled_copy = tiled_copy_t2r if tiled_copy_t2r is not None else tiled_copy_r2s + reference_src = tiled_copy_t2r is None + + # ── Derive lane layout from tiled_copy ── + lane_layout_MN, warp_layout_MN = _get_lane_warp_layouts(tiled_copy, reference_src) + # For RowVecReduce: reduce across M lanes (lanes_in_M threads share same N col) + lanes_in_M = cute.size(lane_layout_MN, mode=[0]) + lanes_in_N = cute.size(lane_layout_MN, mode=[1]) + is_lane_m_leader = cute.arch.lane_idx() < lanes_in_N + assert lanes_in_M == 1 << int(math.log2(lanes_in_M)), ( + "lanes_in_M must be a power of 2 for butterfly reduction" + ) + if const_expr(lanes_in_N > 1): + assert lane_layout_MN.stride[1] == 1, ( + "RowVecReduce assumes contiguous N lanes when lanes_in_N > 1" + ) + + # Intra-warp shuffle reduction across M lanes. M lanes may be either contiguous + # (SM100 N-major output) or strided by N lanes (SM100 M-major output). + tDrReduce_n = layout_utils.convert_layout_zero_stride( + tDrReduce_cur, tDrReduce_cur.layout + )[None, 0] + if const_expr(lanes_in_M > 1): + for n in cutlass.range(cute.size(tDrReduce_n), unroll_full=True): + reduction_rows = lanes_in_M // 2 + while reduction_rows > 0: + tDrReduce_n[n] += cute.arch.shuffle_sync_bfly( + tDrReduce_n[n], + offset=cute.crd2idx((reduction_rows, 0), lane_layout_MN), + ) + reduction_rows = reduction_rows // 2 + + warp_M = warp_layout_MN[0] + warps_in_M = const_expr(cute.size(warp_M)) partition_for_epilogue_fn = partial( partition_for_epilogue, epi_tile=epi_tile, @@ -621,28 +1056,46 @@ class ColVecReduce(EpiOp): reference_src=tiled_copy_t2r is None, ) tile_M, tile_N = gemm.cta_tile_shape_mnk[:2] + tDcD = partition_for_epilogue_fn(cute.make_identity_tensor((tile_M, tile_N))) + tDcD_cur = tDcD[None, None, None, epi_coord[0], epi_coord[1]] + tDcD_n = layout_utils.convert_layout_zero_stride(tDcD_cur, tDrReduce_cur.layout)[ + None, 0 + ] + + # Inter-warp reduction through smem + warp_idx = cute.arch.make_warp_uniform(tidx // cute.arch.WARP_SIZE) + warp_m_idx = warp_layout_MN.get_hier_coord(warp_idx)[0] + if const_expr(warps_in_M > 1): + if warp_m_idx > 0 and is_lane_m_leader: + for n in cutlass.range(cute.size(tDcD_n, mode=[0])): + col_idx = tDcD_n[n][1] + sDrReduce[col_idx, warp_m_idx - 1] = tDrReduce_n[n] + gemm.epilogue_barrier.arrive_and_wait() + if warp_m_idx == 0 and is_lane_m_leader: + for n in cutlass.range(cute.size(tDcD_n, mode=[0])): + col_idx = tDcD_n[n][1] + for warp_m in cutlass.range_constexpr(1, warps_in_M): + tDrReduce_n[n] += sDrReduce[col_idx, warp_m - 1] + + # Write to gmem batch_idx = tile_coord_mnkl[3] - limit_n = param.shape[2] if not varlen_manager.varlen_m else param.shape[1] - if tile_coord_mnkl[1] < limit_n: - if const_expr(not varlen_manager.varlen_m): - mColVec = param[batch_idx, None, tile_coord_mnkl[1]] - else: - mColVec = cute.domain_offset( - (varlen_manager.params.cu_seqlens_m[batch_idx],), - param[None, tile_coord_mnkl[1]], - ) - gColVec = cute.local_tile(mColVec, (tile_M,), (tile_coord_mnkl[0],)) - limit_m = min( - varlen_manager.len_m(batch_idx) - tile_coord_mnkl[0] * tile_M, - tile_M, - ) - tDcD = partition_for_epilogue_fn(cute.make_identity_tensor((tile_M, tile_N))) - tDrReduce_m = layout_utils.convert_layout_zero_stride(tDrReduce, tDrReduce.layout)[ - None, 0 - ] - tDcD_m = layout_utils.convert_layout_zero_stride(tDcD, tDrReduce.layout)[None, 0] - if tDcD_m[0][1] == 0: - for m in cutlass.range(cute.size(tDcD_m, mode=[0])): - row_idx = tDcD_m[m][0] - if row_idx < limit_m: - gColVec[row_idx] = tDrReduce_m[m] + limit_m_tiles = param.shape[1] if not varlen_manager.varlen_m else param.shape[0] + if const_expr(not varlen_manager.varlen_m): + mRowVec = param[batch_idx, tile_coord_mnkl[0], None] + else: + mRowVec = param[tile_coord_mnkl[0], None] + gRowVec = cute.local_tile(mRowVec, (tile_N,), (tile_coord_mnkl[1],)) + limit_n = min( + cute.size(mRowVec, mode=[0]) - tile_coord_mnkl[1] * tile_N, + tile_N, + ) + should_write_gmem = ( + is_lane_m_leader + if const_expr(warps_in_M == 1) + else warp_m_idx == 0 and is_lane_m_leader + ) + if tile_coord_mnkl[0] < limit_m_tiles and should_write_gmem: + for n in cutlass.range(cute.size(tDcD_n, mode=[0])): + col_idx = tDcD_n[n][1] + if col_idx < limit_n: + gRowVec[col_idx] = tDrReduce_n[n] diff --git a/build/torch-cuda/quack/epi_utils.py b/build/torch-cuda/quack/epi_utils.py index 500d380cc7a10ced3d8d65ceed471e1734721d32..2973642fd259b0820f07349100c7e47b76f74a0d 100644 --- a/build/torch-cuda/quack/epi_utils.py +++ b/build/torch-cuda/quack/epi_utils.py @@ -32,27 +32,32 @@ def assume_broadcast_strides(*tensors): return [assume_stride_divisibility(t) for t in tensors] -def setup_epi_tensor(gemm, tensor, epi_tile=None, op_type="store"): - """Create TMA atom + smem layout for a supplemental epilogue tensor. +def setup_epi_tensor(gemm, tensor, epi_tile=None, op_type="store", stage=None): + """Create copy metadata + smem layout for a supplemental epilogue tensor. Args: - gemm: The GEMM object (provides arch, epi_stage, _make_tma_epi_atoms_and_tensors). - tensor: The global memory tensor to set up TMA for. + gemm: The GEMM object (provides arch, epi_stage, and epilogue layout helpers). + tensor: The global memory tensor to set up for the epilogue. epi_tile: Epilogue tile shape. Defaults to gemm.epi_tile. op_type: "store" or "load". Returns: - (tma_atom, tma_tensor, smem_layout_staged, epi_tile) + (copy_atom, tensor, smem_layout_staged, epi_tile). copy_atom is None for pre-TMA archs. """ if epi_tile is None: epi_tile = gemm.epi_tile + if stage is None: + stage = gemm.epi_stage dtype = tensor.element_type layout = cutlass.utils.LayoutEnum.from_tensor(tensor) utils_cls = sm100_utils if gemm.arch >= 100 else sm90_utils - smem_layout_staged = utils_cls.make_smem_layout_epi(dtype, layout, epi_tile, gemm.epi_stage) + smem_layout_staged = utils_cls.make_smem_layout_epi(dtype, layout, epi_tile, stage) + # Ragging-for-TMA is for varlen_m stores that need a per-batch row offset baked + # into the TMA descriptor. Loads don't currently support varlen_m, so skip the + # ragging conversion. tma_input = ( copy_utils.create_ragged_tensor_for_tma(tensor, ragged_dim=0, ptr_shift=True) - if cute.rank(tensor) == 2 + if op_type != "load" and cute.rank(tensor) == 2 else tensor ) tma_atom, tma_tensor = gemm._make_tma_epi_atoms_and_tensors( diff --git a/build/torch-cuda/quack/fast_math.py b/build/torch-cuda/quack/fast_math.py index 73bbd2ecb04686d17a621fdd46bc67a78b1e2331..6ce2c925693420c85447c08f4bdd25f746023031 100644 --- a/build/torch-cuda/quack/fast_math.py +++ b/build/torch-cuda/quack/fast_math.py @@ -2,32 +2,78 @@ import cutlass import cutlass.cute as cute +from cutlass import Int32, Uint32, Uint64 +from cutlass._mlir.dialects import llvm from cutlass.base_dsl.typing import Integer -from cutlass.cutlass_dsl import dsl_user_op - - -class FastDivmod(cute.FastDivmodDivisor): - """We store the divisor along with the FastDivmodDivisor.""" - - @dsl_user_op - def __init__( - self, - divisor: Integer, - is_power_of_2: bool = None, - *, - loc=None, - ip=None, - ): - super().__init__(divisor, is_power_of_2=is_power_of_2, loc=loc, ip=ip) + + +def ceil_log2(x: Integer) -> Int32: + """ceil(log2(x)) for 1 <= x < 2^31, as 32 - ctlz(x - 1). The llvm.intr.ctlz + lowers to lzcnt on the host (launch-prep path) and clz on device; + is_zero_poison=False makes ctlz(0) == 32, so x == 1 correctly yields 0.""" + xm1 = (Int32(x) - 1).ir_value() + return Int32(32) - Int32(llvm.intr_ctlz(xm1, False)) + + +class FastDivmod: + """Magic-number unsigned divmod: q = umulhi(n, magic) >> shift, r = n - q * divisor. + + The multiplier and shift are precomputed on the host (kernel params), so the + device-side divmod is 3 uniform-datapath ops (IMAD.WIDE.U32 + USHF + IMAD) plus + a select handling divisor == 1 (magic == 0 sentinel, like nvjet). This is the + lean form without the add-back correction the stock cute FastDivmodDivisor + emits: with shift = max(ceil_log2(d) - 1, 0) the multiplier ceil(2^(32+s)/d) + fits 32 bits (worst case d = 2^(c-1)+1 gives m <= 2^32 - 1) and the result is + exact for all dividends < 2^31, i.e. any non-negative Int32. Same contract and + algorithm as C++ cutlass::FastDivmod. Negative dividends are OUT of contract + (they reinterpret through Uint32 to values >= 2^31 and silently divide wrong); + use cute.FastDivmodDivisor if signed or full-u32 dividends are ever needed. + """ + + def __init__(self, divisor: Integer): + if isinstance(divisor, int): + assert 0 < divisor < 1 << 31 + divisor = Int32(divisor) # constants fold through the arithmetic below self.divisor = divisor + # Runs on the CPU at launch prep (host-side jit); the udiv is once per launch. + s = cutlass.max(ceil_log2(divisor) - 1, Int32(0)) + pow_s = Uint64(Uint32(1) << Uint32(s)) + numer = Uint64(0x100000000) * pow_s + magic = (numer + Uint64(Uint32(divisor)) - 1) // Uint64(Uint32(divisor)) + self.magic = Uint32(magic & 0xFFFFFFFF) # 0 when divisor == 1 + self.shift = Uint32(s) + + def __rdivmod__(self, dividend: Integer): + q = Uint32(cute.arch.mul_hi(Uint32(dividend), self.magic)) >> self.shift + # divisor == 1 sentinel: magic wrapped to 0, so q is 0; select the dividend + # instead. The remainder then self-corrects: r = n - n * 1 = 0. + q = Int32(cutlass.select_(self.magic == Uint32(0), Int32(dividend), Int32(q))) + r = Int32(dividend) - q * Int32(self.divisor) + return q, r + + def __rfloordiv__(self, dividend: Integer) -> Int32: + q, _ = self.__rdivmod__(dividend) + return q + + def __rmod__(self, dividend: Integer) -> Int32: + _, r = self.__rdivmod__(dividend) + return r def __extract_mlir_values__(self): - """Extract MLIR values for Host->Device transfer.""" - return [self._divisor] + cutlass.extract_mlir_values(self.divisor) + values = [] + self._values_pos = [] + for obj in [self.magic, self.shift, self.divisor]: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values def __new_from_mlir_values__(self, values): - """Reconstruct FastDivmodDivisor from MLIR values.""" new_obj = object.__new__(FastDivmod) - new_obj._divisor = values[0] - new_obj.divisor = cutlass.new_from_mlir_values(self.divisor, values[1:]) + for name, n_items in zip(["magic", "shift", "divisor"], self._values_pos): + setattr( + new_obj, name, cutlass.new_from_mlir_values(getattr(self, name), values[:n_items]) + ) + values = values[n_items:] + new_obj._values_pos = self._values_pos return new_obj diff --git a/build/torch-cuda/quack/gemm.py b/build/torch-cuda/quack/gemm.py index c4b41e6969ce370ffde9946edc67fe8d760010f8..0a32c2f1a69cc2d06692eb817760f3bcd3c51602 100644 --- a/build/torch-cuda/quack/gemm.py +++ b/build/torch-cuda/quack/gemm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025-2026, Tri Dao. +# Copyright (c) 2025-2026, QuACK team. # GEMM compilation via TVM-FFI with fake tensors and NamedTuple args. from typing import Optional @@ -9,11 +9,12 @@ import cutlass.cute as cute from cutlass import Int32, Float32 from cutlass.cute.runtime import make_ptr -from .cache_utils import jit_cache +from .cache import jit_cache from .compile_utils import make_fake_tensor as fake_tensor from .cute_dsl_utils import get_device_capacity, get_max_active_clusters, torch2cute_dtype_map from .gemm_default_epi import ( GemmDefaultEpiMixin, + GemmDefaultSm80, GemmDefaultSm90, GemmDefaultSm100, GemmDefaultSm120, @@ -28,7 +29,9 @@ from .gemm_tvm_ffi_utils import ( make_fake_scheduler_args, make_fake_varlen_args, make_fake_gemm_tensors, + make_fake_sf_tensor, compile_gemm_kernel, + validate_blockscaled_sf, ) @@ -62,9 +65,12 @@ def _compile_gemm( device_capacity, rounding_mode, sr_seed_mode, - has_trace_ptr, + num_warps, + sf_dtype=None, + sf_vec_size=None, ): sm_to_cls = { + 8: GemmDefaultSm80, 9: GemmDefaultSm90, 10: GemmDefaultSm100, 11: GemmDefaultSm100, @@ -111,10 +117,18 @@ def _compile_gemm( sr_seed=fake_scalar(sr_seed_mode, dtype=Int32), ) scheduler_args = make_fake_scheduler_args( - (is_dynamic_persistent and device_capacity[0] == 9), has_batch_idx_permute, l + (is_dynamic_persistent and device_capacity[0] <= 9), has_batch_idx_permute, l ) aidx_len = m if varlen_m else (k if varlen_k else None) varlen_args = make_fake_varlen_args(varlen_m, varlen_k, gather_A, aidx_len) + if sf_dtype is not None: + # Padded SF buffers have a static batch dim of exactly 1 (not l): SFA for + # varlen_m (M-padded) and varlen_k (K-padded); SFB is K-padded too for + # varlen_k but stays per-batch (l, rn, rk, ...) for varlen_m. + mSFA = make_fake_sf_tensor(sf_dtype, 1 if (varlen_m or varlen_k) else l) + mSFB = make_fake_sf_tensor(sf_dtype, 1 if varlen_k else l) + else: + mSFA, mSFB = None, None return compile_gemm_kernel( GemmCls, a_dtype, @@ -132,9 +146,12 @@ def _compile_gemm( epi_args, scheduler_args, varlen_args, - has_trace_ptr=has_trace_ptr, + mSFA=mSFA, + mSFB=mSFB, use_tma_gather=use_tma_gather, concat_layout=concat_layout or None, + num_warps=num_warps, + sf_vec_size=sf_vec_size, ) @@ -149,6 +166,8 @@ def gemm( tile_N: int, cluster_M: int, cluster_N: int, + cluster_K: int = 1, + tile_K: int | None = None, pingpong: bool = False, persistent: bool = True, is_dynamic_persistent: bool = False, @@ -166,18 +185,24 @@ def gemm( sr_seed: int | Tensor = 0, use_tma_gather: bool = False, concat_layout: dict | None = None, - trace_ptr=None, # Optional Int64 from TraceSession.ptr + num_warps: Optional[int] = None, + # SFA/SFB: (l, rm/rn, rk, 32, 4, 4) blocked scale factors. For varlen_m, SFA is + # M-padded (1, total_padded_rm, rk, 32, 4, 4) while SFB stays per-batch. For + # varlen_k, BOTH are K-padded (1, rm/rn, total_padded_rk, 32, 4, 4); pad bytes + # may be arbitrary (the kernel skips the MMA instructions covering them). + # See AI/varlen_blockscaled_sf_layout.md. + SFA: Optional[Tensor] = None, + SFB: Optional[Tensor] = None, ) -> None: varlen_m = cu_seqlens_m is not None varlen_k = cu_seqlens_k is not None varlen = varlen_m or varlen_k gather_A = A_idx is not None + blockscaled = SFA is not None assert not (varlen_m and varlen_k), "Only one of cu_seqlens_m and cu_seqlens_k" if gather_A: assert varlen, "gather_A requires varlen" assert cluster_N == 1, "gather_A requires cluster_N=1" - if varlen: - assert persistent, "varlen requires persistent=True" if add_to_output: assert not varlen_m, "Add to output not supported with varlen_m" if varlen_m: @@ -188,15 +213,35 @@ def gemm( assert B.stride(-2) == 1, "varlen_k requires B to be n-major" device_capacity = get_device_capacity(A.device) - assert device_capacity[0] in [9, 10, 11, 12], "Only SM90, SM100, SM110, and SM120 are supported" + assert device_capacity[0] in [8, 9, 10, 11, 12], ( + "Only SM8x, SM90, SM100, SM110, and SM120 are supported" + ) + sf_dtype, sf_vec_size = None, None + if blockscaled: + assert not gather_A, "Blockscaled GEMM does not support gather_A yet" + assert not concat_layout, "Blockscaled GEMM does not support concat_layout" + assert tile_K is None, "Blockscaled GEMM derives tile_K from the MMA instruction" + if varlen_m: + num_batches = cu_seqlens_m.shape[0] - 1 + elif varlen_k: + num_batches = cu_seqlens_k.shape[0] - 1 + else: + num_batches = None + sf_dtype, sf_vec_size = validate_blockscaled_sf( + A, B, SFA, SFB, device_capacity, num_batches=num_batches, varlen_k=varlen_k + ) if use_tma_gather: assert device_capacity[0] in [10, 11], "TMA gather currently requires SM100/SM110" if rounding_mode == RoundingMode.RS: assert device_capacity[0] == 10, "Stochastic rounding (RoundingMode.RS) requires SM100" - if is_dynamic_persistent and device_capacity[0] == 9: + if is_dynamic_persistent and device_capacity[0] <= 9: assert tile_count_semaphore is not None, ( - "Dynamic persistent tile scheduler in SM90 requires a semaphore in GMEM" + "Dynamic persistent tile scheduler for SM8x and SM90 requires a semaphore in GMEM" ) + if device_capacity[0] == 8: + if add_to_output: + C = D + add_to_output = False A_p, B_p, D_p, C_p = perm3d(A, B, D, C, varlen_m=varlen_m, varlen_k=varlen_k) a_major, b_major, d_major, c_major = get_majors(A_p, B_p, D_p, C_p) @@ -210,6 +255,7 @@ def gemm( sr_seed_mode = ( 2 if isinstance(sr_seed, Tensor) else (1 if rounding_mode == RoundingMode.RS else 0) ) + tile_shape_mnk = (tile_M, tile_N) if tile_K is None else (tile_M, tile_N, tile_K) compiled_fn = _compile_gemm( a_dtype, b_dtype, @@ -219,8 +265,8 @@ def gemm( b_major, d_major, c_major, - (tile_M, tile_N), - (cluster_M, cluster_N, 1), + tile_shape_mnk, + (cluster_M, cluster_N, cluster_K), pingpong, persistent, is_dynamic_persistent, @@ -239,14 +285,11 @@ def gemm( device_capacity, rounding_mode, sr_seed_mode, - trace_ptr is not None, + num_warps, + sf_dtype, + sf_vec_size, ) - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY: - return - def scalar_arg(scalar, mode, dtype=Float32): if mode == 0: return None @@ -255,7 +298,10 @@ def gemm( else: return scalar.data_ptr() - max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0 + cluster_size = cluster_M * cluster_N * cluster_K + max_active_clusters = ( + get_max_active_clusters(cluster_size, device_capacity=device_capacity) if persistent else 0 + ) epi_args = GemmDefaultEpiMixin.EpilogueArguments( alpha=scalar_arg(alpha, alpha_mode), @@ -269,14 +315,15 @@ def gemm( scheduler_args = make_scheduler_args( max_active_clusters, max_swizzle_size, - tile_count_semaphore, + # Must mirror make_fake_scheduler_args in _compile_gemm: only the SM8x/SM90 + # dynamic scheduler consumes the semaphore; SM100 uses CLC instead, and the + # compiled signature has None there. + tile_count_semaphore if (is_dynamic_persistent and device_capacity[0] <= 9) else None, batch_idx_permute, ) varlen_args = make_varlen_args(cu_seqlens_m, cu_seqlens_k, A_idx) if device_capacity[0] in [10, 11]: - compiled_fn( - A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None, trace_ptr - ) + compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, SFA, SFB) else: - compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, trace_ptr) + compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args) diff --git a/build/torch-cuda/quack/gemm_act.py b/build/torch-cuda/quack/gemm_act.py index ea0cdbd98b8b5d50032fdc013b50a4e9d6763410..9fd19e998c0a217340c34c6a0a532fec732215c4 100644 --- a/build/torch-cuda/quack/gemm_act.py +++ b/build/torch-cuda/quack/gemm_act.py @@ -1,26 +1,27 @@ # Copyright (c) 2025, Wentao Guo, Tri Dao. from __future__ import annotations -from typing import NamedTuple, Tuple, Optional, Callable -from functools import partial +import math +from typing import NamedTuple, Tuple, Optional, Callable, Type from torch import Tensor import cutlass import cutlass.cute as cute -import cutlass.utils.hopper_helpers as sm90_utils_og import cutlass.utils.blackwell_helpers as sm100_utils from cutlass import Int32, Float32, const_expr from cutlass.cute.runtime import make_ptr +from cutlass.cute.nvgpu import warp from .compile_utils import make_fake_tensor as fake_tensor from .cute_dsl_utils import ( - ParamsBase, mlir_namedtuple, get_device_capacity, get_max_active_clusters, torch2cute_dtype_map, ) -from .epi_ops import TileStore +from .epi_composable import ComposableEpiMixin +from .epi_ops import ColVecLoad, RowVecLoad, Scalar, TileStore +from .gemm_sm80 import GemmSm80 from .gemm_sm90 import GemmSm90 from .gemm_sm100 import GemmSm100 from .gemm_sm120 import GemmSm120 @@ -34,23 +35,32 @@ from .gemm_tvm_ffi_utils import ( make_fake_varlen_args, div_for_dtype, make_fake_gemm_tensors, + make_fake_sf_tensor, compile_gemm_kernel, + validate_blockscaled_sf, ) -from .cache_utils import jit_cache +from .cache import jit_cache from . import layout_utils as layout_utils +from . import copy_utils as copy_utils from .layout_utils import permute_gated_Cregs_b16 from .activation import act_fn_map, gate_fn_map -from .rounding import RoundingMode +from .rounding import RoundingMode, convert_f32_to_bf16_sr, epilogue_aux_out_sr_seed -class GemmActMixin(GemmDefaultEpiMixin): - _epi_ops = (*GemmDefaultEpiMixin._epi_ops, TileStore("mPostAct")) +class GemmActMixin(ComposableEpiMixin): + _epi_ops = ( + Scalar("alpha"), + Scalar("beta"), + Scalar("sr_seed", dtype=Int32), + RowVecLoad("mRowVecBroadcast"), + ColVecLoad("mColVecBroadcast"), + TileStore("mAuxOut"), + ) _extra_param_fields = (("act_fn", cutlass.Constexpr, None),) - _epi_param_bases = (ParamsBase,) @mlir_namedtuple class EpilogueArguments(NamedTuple): - mPostAct: cute.Tensor + mAuxOut: cute.Tensor act_fn: cutlass.Constexpr[Optional[Callable]] = None alpha: Optional[Float32 | cute.Tensor] = None beta: Optional[Float32 | cute.Tensor] = None @@ -63,20 +73,39 @@ class GemmActMixin(GemmDefaultEpiMixin): def epi_to_underlying_arguments(self, args: EpilogueArguments, *, loc=None, ip=None): self.rounding_mode = args.rounding_mode - self.postact_dtype = args.mPostAct.element_type - self.postact_layout = cutlass.utils.LayoutEnum.from_tensor(args.mPostAct) - self.cta_tile_shape_postact_mn = self.cta_tile_shape_mnk[:2] + self.aux_out_dtype = args.mAuxOut.element_type + self.aux_out_layout = cutlass.utils.LayoutEnum.from_tensor(args.mAuxOut) + self.cta_tile_shape_aux_out_mn = self.cta_tile_shape_mnk[:2] d = self._epi_ops_to_params_dict(args) d["act_fn"] = args.act_fn for key in ("mRowVecBroadcast", "mColVecBroadcast"): - if key in self.concat_layout and key in d and d[key] is not None: + if key in self.concat_layout and key in d: d[key] = layout_utils.concat_to_interleave(d[key], 1) return self.EpilogueParams(**d) - # epi_get_tma_atoms, epi_smem_bytes_per_stage, epi_get_smem_struct, + # epi_get_tma_atoms, epi_smem_bytes, epi_get_smem_struct, # epi_get_smem_tensors are all inherited from ComposableEpiMixin via _epi_ops. - def epi_setup_postact( + def epi_make_aux_out_copy_atom_r2s(self, params, tiled_copy_t2r): + """Build the register-to-shared copy atom used by aux outputs.""" + if self.arch == 100: + return sm100_utils.get_smem_store_op( + self.aux_out_layout, self.aux_out_dtype, self.acc_dtype, tiled_copy_t2r + ) + else: + return copy_utils.get_smem_store_atom( + self.aux_out_dtype, + transpose=self.aux_out_layout != cutlass.utils.LayoutEnum.ROW_MAJOR, + major_mode_size=cute.size(params.epi_tile_mAuxOut, mode=[1]) + // self.atom_layout_mnk[1], + ) + + def epi_make_aux_out_tiled_copy_r2s(self, params, tiled_copy_r2s, tiled_copy_t2r): + """Build the register-to-shared tiled copy used by aux outputs.""" + copy_atom_aux_out_r2s = self.epi_make_aux_out_copy_atom_r2s(params, tiled_copy_t2r) + return cute.make_tiled_copy_S(copy_atom_aux_out_r2s, tiled_copy_r2s) + + def epi_setup_aux_out( self, params, epi_smem_tensors, @@ -86,61 +115,56 @@ class GemmActMixin(GemmDefaultEpiMixin): varlen_manager, tidx, ): - """Setup postact TMA copies and partitions before the epilogue loop.""" - sPostAct = epi_smem_tensors[self._epi_smem_map["mPostAct"]] - get_smem_store_op = ( - partial(sm100_utils.get_smem_store_op, tiled_tmem_load=tiled_copy_t2r) - if self.arch == 100 - else sm90_utils_og.sm90_get_smem_store_op - ) - copy_atom_postact_r2s = get_smem_store_op( - self.postact_layout, self.postact_dtype, self.acc_dtype + """Setup aux output TMA copies and partitions before the epilogue loop. + + Returns an empty tuple when mAuxOut wasn't supplied so the framework + skips the aux-out path. + """ + if getattr(params, "mAuxOut", None) is None: + return () + sAuxOut = epi_smem_tensors["mAuxOut"] + tiled_copy_aux_out_r2s = self.epi_make_aux_out_tiled_copy_r2s( + params, tiled_copy_r2s, tiled_copy_t2r ) - tiled_copy_postact_r2s = cute.make_tiled_copy_S(copy_atom_postact_r2s, tiled_copy_r2s) - tRS_sPostAct = tiled_copy_postact_r2s.get_slice(tidx).partition_D(sPostAct) + tRS_sAuxOut = tiled_copy_aux_out_r2s.get_slice(tidx).partition_D(sAuxOut) batch_idx = tile_coord_mnkl[3] - copy_postact, _, _ = self.epilog_gmem_copy_and_partition( - params.tma_atom_mPostAct, - varlen_manager.offset_batch_epi(params.mPostAct, batch_idx), - self.cta_tile_shape_postact_mn, - params.epi_tile_mPostAct, - sPostAct, + copy_aux_out, _, _ = self.epilog_gmem_copy_and_partition( + params.tma_atom_mAuxOut, + varlen_manager.offset_batch_epi(params.mAuxOut, batch_idx), + self.cta_tile_shape_aux_out_mn, + params.epi_tile_mAuxOut, + sAuxOut, tile_coord_mnkl, ) - return tiled_copy_postact_r2s, tRS_sPostAct, copy_postact + return ((tiled_copy_aux_out_r2s, tRS_sAuxOut, copy_aux_out),) @cute.jit - def epi_convert_postact( - self, tRS_rPostAct, sr_seed, tidx, tile_coord_mnkl, num_prev_subtiles, epi_idx + def epi_convert_aux_out( + self, + output_idx: cutlass.Constexpr[int], + tRS_rAuxOut, + sr_seed, + tidx, + tile_coord_mnkl, + num_prev_subtiles, + epi_idx, ): - """Convert postact from acc_dtype to postact_dtype. Override for custom postprocessing.""" + """Convert aux output from acc_dtype to aux_out_dtype. Override for custom postprocessing.""" if const_expr( self.rounding_mode == RoundingMode.RS - and tRS_rPostAct.element_type == cutlass.Float32 - and self.postact_dtype == cutlass.BFloat16 + and tRS_rAuxOut.element_type == cutlass.Float32 + and self.aux_out_dtype == cutlass.BFloat16 ): - from .rounding import convert_f32_to_bf16_sr from cutlass.cute.tensor import TensorSSA - # Salt with 0x9E3779B1 to avoid sharing entropy with the D output seed - seed = ( - sr_seed - + 0x9E3779B1 - + ( - tile_coord_mnkl[0] * 65537 - + tile_coord_mnkl[1] * 257 - + tile_coord_mnkl[3] * 17 - + (num_prev_subtiles + epi_idx) * 7 - ) - ) - tRS_rPostAct_out = cute.make_rmem_tensor_like(tRS_rPostAct, self.postact_dtype) - src_vec = tRS_rPostAct.load() + seed = epilogue_aux_out_sr_seed(sr_seed, tile_coord_mnkl, num_prev_subtiles + epi_idx) + tRS_rAuxOut_out = cute.make_rmem_tensor_like(tRS_rAuxOut, self.aux_out_dtype) + src_vec = tRS_rAuxOut.load() raw_vec = convert_f32_to_bf16_sr(src_vec, seed, tidx) - tRS_rPostAct_out.store(TensorSSA(raw_vec, src_vec.shape, self.postact_dtype)) + tRS_rAuxOut_out.store(TensorSSA(raw_vec, src_vec.shape, self.aux_out_dtype)) else: - tRS_rPostAct_out = cute.make_rmem_tensor_like(tRS_rPostAct, self.postact_dtype) - tRS_rPostAct_out.store(tRS_rPostAct.load().to(self.postact_dtype)) - return tRS_rPostAct_out + tRS_rAuxOut_out = tRS_rAuxOut.to(self.aux_out_dtype) + return tRS_rAuxOut_out @cute.jit def epi_visit_subtile( @@ -149,29 +173,28 @@ class GemmActMixin(GemmDefaultEpiMixin): epi_loop_tensors: Tuple[cute.Tensor, ...], tRS_rD: cute.Tensor, tRS_rC: Optional[cute.Tensor] = None, - ) -> Optional[cute.Tensor]: + ) -> Tuple[cute.Tensor, ...]: GemmDefaultEpiMixin.epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC) # Apply activation function if provided # If we don't have .shape here, the compiler generates local stores and loads if const_expr(params.act_fn is not None): - tRS_rPostAct = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype) - if const_expr(self.arch < 100): - for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True): - tRS_rPostAct[i] = params.act_fn(tRS_rD[i]) - else: - for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True): - tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1] = params.act_fn( - (tRS_rD[2 * i], tRS_rD[2 * i + 1]) - ) + tRS_rAuxOut = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype) + vectorize = const_expr(self.arch == 100) + for i in cutlass.range(cute.size(tRS_rAuxOut), unroll_full=True, vectorize=vectorize): + tRS_rAuxOut[i] = params.act_fn(tRS_rD[i]) else: - tRS_rPostAct = tRS_rD - return tRS_rPostAct + tRS_rAuxOut = tRS_rD + return (tRS_rAuxOut,) class GemmActSm90(GemmActMixin, GemmSm90): pass +class GemmActSm80(GemmActMixin, GemmSm80): + pass + + class GemmActSm100(GemmActMixin, GemmSm100): pass @@ -189,33 +212,37 @@ def _gated_epi_tile_fn(gemm, epi_tile): class GemmGatedMixin(GemmActMixin): _epi_ops = ( - *GemmDefaultEpiMixin._epi_ops, - TileStore("mPostAct", epi_tile_fn=_gated_epi_tile_fn), + Scalar("alpha"), + Scalar("beta"), + Scalar("sr_seed", dtype=Int32), + RowVecLoad("mRowVecBroadcast"), + ColVecLoad("mColVecBroadcast"), + TileStore("mAuxOut", epi_tile_fn=_gated_epi_tile_fn), ) def epi_to_underlying_arguments( self, args: GemmActMixin.EpilogueArguments, *, loc=None, ip=None ) -> GemmActMixin.EpilogueParams: - assert args.mPostAct.element_type.width == 16, ( + assert args.mAuxOut.element_type.width == 16, ( "GemmGated only supports 16bit postact for now" ) assert self.d_layout is None or self.d_layout.is_n_major_c() - assert cutlass.utils.LayoutEnum.from_tensor(args.mPostAct).is_n_major_c() + assert cutlass.utils.LayoutEnum.from_tensor(args.mAuxOut).is_n_major_c() if self.arch == 90: assert self.cta_tile_shape_mnk[1] % 32 == 0, ( "GemmGatedSm90 requires tileN to be divisible by 32" ) self.rounding_mode = args.rounding_mode - self.postact_dtype = args.mPostAct.element_type - self.postact_layout = cutlass.utils.LayoutEnum.from_tensor(args.mPostAct) - self.cta_tile_shape_postact_mn = ( + self.aux_out_dtype = args.mAuxOut.element_type + self.aux_out_layout = cutlass.utils.LayoutEnum.from_tensor(args.mAuxOut) + self.cta_tile_shape_aux_out_mn = ( self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1] // 2, ) d = self._epi_ops_to_params_dict(args) d["act_fn"] = args.act_fn for key in ("mRowVecBroadcast", "mColVecBroadcast"): - if key in self.concat_layout and key in d and d[key] is not None: + if key in self.concat_layout and key in d: d[key] = layout_utils.concat_to_interleave(d[key], 1) return self.EpilogueParams(**d) @@ -226,43 +253,96 @@ class GemmGatedMixin(GemmActMixin): epi_loop_tensors: Tuple[cute.Tensor, ...], tRS_rD: cute.Tensor, tRS_rC: Optional[cute.Tensor] = None, - ) -> Optional[cute.Tensor]: + ) -> Tuple[cute.Tensor, ...]: GemmDefaultEpiMixin.epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC) - tRS_rPostAct_layout = cute.recast_layout(2, 1, tRS_rD.layout) + tRS_rAuxOut_layout = cute.recast_layout(2, 1, tRS_rD.layout) # If we don't have .shape here, the compiler generates local stores and loads - tRS_rPostAct = cute.make_rmem_tensor(tRS_rPostAct_layout.shape, self.acc_dtype) - if const_expr(self.arch < 100): - for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True): - tRS_rPostAct[i] = params.act_fn(tRS_rD[2 * i], tRS_rD[2 * i + 1]) - else: - for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True): - tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1] = params.act_fn( - (tRS_rD[4 * i], tRS_rD[4 * i + 2]), (tRS_rD[4 * i + 1], tRS_rD[4 * i + 3]) - ) - return tRS_rPostAct + tRS_rAuxOut = cute.make_rmem_tensor(tRS_rAuxOut_layout.shape, self.acc_dtype) + tRS_rD_pair = cute.flat_divide(tRS_rD, cute.make_layout(2)) + tRS_rGate = tRS_rD_pair[0, ...] + tRS_rUp = tRS_rD_pair[1, ...] + vectorize = const_expr(self.arch == 100) + for i in cutlass.range(cute.size(tRS_rAuxOut), unroll_full=True, vectorize=vectorize): + tRS_rAuxOut[i] = params.act_fn(tRS_rGate[i], tRS_rUp[i]) + return (tRS_rAuxOut,) @cute.jit - def epi_convert_postact( - self, tRS_rPostAct, sr_seed, tidx, tile_coord_mnkl, num_prev_subtiles, epi_idx + def epi_convert_aux_out( + self, + output_idx: cutlass.Constexpr[int], + tRS_rAuxOut, + sr_seed, + tidx, + tile_coord_mnkl, + num_prev_subtiles, + epi_idx, ): - tRS_rPostAct_out = GemmActMixin.epi_convert_postact( - self, tRS_rPostAct, sr_seed, tidx, tile_coord_mnkl, num_prev_subtiles, epi_idx + tRS_rAuxOut_out = GemmActMixin.epi_convert_aux_out( + self, + output_idx, + tRS_rAuxOut, + sr_seed, + tidx, + tile_coord_mnkl, + num_prev_subtiles, + epi_idx, ) - if const_expr(self.arch == 90): + if const_expr(self.arch in (90, 120)): # Only need this if we're using STSM - permute_gated_Cregs_b16(tRS_rPostAct_out) - return tRS_rPostAct_out + permute_gated_Cregs_b16(tRS_rAuxOut_out) + return tRS_rAuxOut_out class GemmGatedSm90(GemmGatedMixin, GemmSm90): pass +class GemmGatedSm80(GemmGatedMixin, GemmSm80): + pass + + class GemmGatedSm100(GemmGatedMixin, GemmSm100): pass -class GemmGatedSm120(GemmGatedMixin, GemmSm120): +class GemmGatedSm120Mixin: + @staticmethod + def _compute_tile_shape_or_override( + cta_tile_shape_mnk: Tuple[int, int, int], + atom_layout_mnk: Tuple[int, int, int], + element_type: Optional[Type[cutlass.Numeric]] = None, + epi_tile_override: Tuple[int, int] | None = None, + ) -> Tuple[int, int]: + if epi_tile_override is not None: + return epi_tile_override + # Typically epi_tile is (64, 32) but since we want tile_n = 64 (see below), we might set + # tile_m = 32 if there's only 2 warps along the M direction. + tile_m = math.gcd(atom_layout_mnk[0] * 16, cute.size(cta_tile_shape_mnk, mode=[0])) + atom_n = atom_layout_mnk[1] + # E.g. if we have 2 warps along N direction, we want each warp to have 32 elems so that + # postact has 16 elements, which means tile_n should be 64. + tile_n = math.gcd(atom_n * 8 * 4, cute.size(cta_tile_shape_mnk, mode=[1])) + return (tile_m, tile_n) + + def epi_make_aux_out_tiled_copy_r2s(self, params, tiled_copy_r2s, tiled_copy_t2r): + copy_atom_aux_out_r2s = self.epi_make_aux_out_copy_atom_r2s(params, tiled_copy_t2r) + copy_atom_postact_c = self.epi_make_aux_out_copy_atom_r2s(params, cutlass.Float16) + op = warp.MmaF16BF16Op(self.a_dtype, self.acc_dtype, self.mma_inst_mnk) + tC = cute.make_layout(self.atom_layout_mnk) + atom_m, atom_n, atom_k = self.atom_layout_mnk + permutation_mnk = ( + self.mma_inst_mnk[0] * atom_m, + self.mma_inst_mnk[1] * atom_n * 2, + self.mma_inst_mnk[2] * atom_k, + ) + tiled_mma_gated_postact = cute.make_tiled_mma(op, tC, permutation_mnk=permutation_mnk) + tiled_copy_aux_out_c_atom = cute.make_tiled_copy_C_atom( + copy_atom_postact_c, tiled_mma_gated_postact + ) + return cute.make_tiled_copy_S(copy_atom_aux_out_r2s, tiled_copy_aux_out_c_atom) + + +class GemmGatedSm120(GemmGatedSm120Mixin, GemmGatedMixin, GemmSm120): pass @@ -295,13 +375,25 @@ def _compile_gemm_act( rounding_mode=RoundingMode.RN, sr_seed_mode=0, use_tma_gather=False, + sf_dtype=None, + sf_vec_size=None, ): sm_to_cls = { - "act": {9: GemmActSm90, 10: GemmActSm100, 11: GemmActSm100, 12: GemmActSm120}, - "gated": {9: GemmGatedSm90, 10: GemmGatedSm100, 11: GemmGatedSm100, 12: GemmGatedSm120}, + "act": { + 8: GemmActSm80, + 9: GemmActSm90, + 10: GemmActSm100, + 11: GemmActSm100, + 12: GemmActSm120, + }, + "gated": { + 8: GemmGatedSm80, + 9: GemmGatedSm90, + 10: GemmGatedSm100, + 11: GemmGatedSm100, + 12: GemmGatedSm120, + }, } - if device_capacity[0] == 12 and gemm_cls_name == "act": - raise NotImplementedError("SM120 non-gated activation GEMM epilogue is not yet supported") GemmCls = sm_to_cls[gemm_cls_name][device_capacity[0]] pa_leading = 1 if postact_major == "n" else 0 mA, mB, mD, mC, m, n, k, l = make_fake_gemm_tensors( @@ -320,7 +412,7 @@ def _compile_gemm_act( div_pa = div_for_dtype(postact_dtype) pa_leading_dim = 1 if gemm_cls_name == "gated" else pa_leading pa_shape = (m, pa_n) if varlen_m else (m, pa_n, l) - mPostAct = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading_dim, divisibility=div_pa) + mAuxOut = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading_dim, divisibility=div_pa) mRowVec = fake_tensor(rowvec_dtype, (l, n), leading_dim=1, divisibility=4) if colvec_ndim == 2: @@ -341,7 +433,7 @@ def _compile_gemm_act( return make_ptr(dtype, 0, cute.AddressSpace.gmem, assumed_align=4) epi_args = GemmCls.EpilogueArguments( - mPostAct, + mAuxOut, act_fn, mRowVecBroadcast=mRowVec, mColVecBroadcast=mColVec, @@ -352,6 +444,11 @@ def _compile_gemm_act( (is_dynamic_persistent and device_capacity[0] == 9), False, l ) varlen_args = make_fake_varlen_args(varlen_m, False, gather_A, m if varlen_m else None) + if sf_dtype is not None: + mSFA = make_fake_sf_tensor(sf_dtype, l) + mSFB = make_fake_sf_tensor(sf_dtype, l) + else: + mSFA, mSFB = None, None return compile_gemm_kernel( GemmCls, a_dtype, @@ -369,8 +466,11 @@ def _compile_gemm_act( epi_args, scheduler_args, varlen_args, + mSFA=mSFA, + mSFB=mSFB, use_tma_gather=use_tma_gather, concat_layout=concat_layout or None, + sf_vec_size=sf_vec_size, ) @@ -386,6 +486,7 @@ def gemm_act( tile_N: int, cluster_M: int, cluster_N: int, + tile_K: int | None = None, pingpong: bool = False, persistent: bool = True, is_dynamic_persistent: bool = False, @@ -398,6 +499,8 @@ def gemm_act( sr_seed: int | Tensor = 0, use_tma_gather: bool = False, concat_layout: tuple | None = None, + SFA: Optional[Tensor] = None, # (l, rm, rk, 32, 4, 4) blocked scale factors + SFB: Optional[Tensor] = None, # (l, rn, rk, 32, 4, 4) ) -> None: if activation in gate_fn_map: gemm_cls_name = "gated" @@ -407,6 +510,7 @@ def gemm_act( varlen_m = cu_seqlens_m is not None gather_A = A_idx is not None + blockscaled = SFA is not None if varlen_m: assert persistent, "varlen_m requires persistent=True" assert A.stride(-1) == 1, "varlen_m requires A to be k-major" @@ -437,7 +541,16 @@ def gemm_act( colvec_ndim = colvec_bias.ndim if colvec_bias is not None else 0 device_capacity = get_device_capacity(A.device) - assert device_capacity[0] in [9, 10, 11, 12], "Only SM90, SM100, SM110, and SM120 are supported" + assert device_capacity[0] in [8, 9, 10, 11, 12], ( + "Only SM8x, SM90, SM100, SM110, and SM120 are supported" + ) + sf_dtype, sf_vec_size = None, None + if blockscaled: + assert not varlen_m and not gather_A, "Blockscaled GEMM does not support varlen/gather yet" + assert not concat_layout, "Blockscaled GEMM does not support concat_layout" + assert tile_K is None, "Blockscaled GEMM derives tile_K from the MMA instruction" + # A / B are still (l, m, k) / (l, n, k) here (perm3d only made views). + sf_dtype, sf_vec_size = validate_blockscaled_sf(A, B, SFA, SFB, device_capacity) if rounding_mode == RoundingMode.RS: assert device_capacity[0] == 10, "Stochastic rounding (RoundingMode.RS) requires SM100" @@ -461,7 +574,7 @@ def gemm_act( d_major, c_major, postact_major, - (tile_M, tile_N), + (tile_M, tile_N, tile_K) if tile_K is not None else (tile_M, tile_N), (cluster_M, cluster_N, 1), pingpong, persistent, @@ -478,13 +591,10 @@ def gemm_act( rounding_mode=rounding_mode, sr_seed_mode=sr_seed_mode, use_tma_gather=use_tma_gather, + sf_dtype=sf_dtype, + sf_vec_size=sf_vec_size, ) - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY: - return - max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0 def scalar_arg(scalar, mode, dtype=Int32): @@ -511,9 +621,9 @@ def gemm_act( varlen_args = make_varlen_args(cu_seqlens_m, None, A_idx) if device_capacity[0] in [10, 11]: - compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None, None) + compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, SFA, SFB) else: - compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None) + compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args) gemm_gated = gemm_act diff --git a/build/torch-cuda/quack/gemm_base.py b/build/torch-cuda/quack/gemm_base.py new file mode 100644 index 0000000000000000000000000000000000000000..3df0ea8ea3fa02adca6befe32922fe23dd00b546 --- /dev/null +++ b/build/torch-cuda/quack/gemm_base.py @@ -0,0 +1,731 @@ +# Copyright (c) 2026, Tri Dao. + +import enum +import math +from dataclasses import dataclass +from typing import Callable, Dict, Literal, Optional, Sequence, Tuple + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass import Boolean, Int32, const_expr +from cutlass.cute.nvgpu import cpasync +from cutlass.utils import LayoutEnum + +from . import copy_utils as copy_utils +from .cute_dsl_utils import ParamsBase +from .epi_ops import EpiSmemBytes +from .pipeline import PipelineTmaAsync, PipelineTmaCpAsync +from .rounding import RoundingMode, epilogue_sr_seed +from .tile_scheduler import ( + PersistenceMode, + TileScheduler, + TileSchedulerArguments, + VarlenMTileScheduler, + VarlenMTileSchedulerArguments, +) +from .varlen_utils import VarlenManager + + +class NamedBarrierGemm(enum.IntEnum): + Epilogue = enum.auto() # starts from 1 as barrier 0 is reserved for sync_threads() + # For mainloop load warps to signal that the epilogue load warp can start. + # This is to avoid loading C too early, interfering with loading A and B. + EpilogueLoad = enum.auto() + MmaWG0 = enum.auto() + MmaWG1 = enum.auto() + EpiWG0 = enum.auto() + EpiWG1 = enum.auto() + TmemPtr = enum.auto() + # CLC-multicast throttle: CTA0 load warp arrives once per tile started, + # CTA0 scheduler warp syncs once per CLC query (2 warps, 64 threads). + ClcThrottle = enum.auto() + + +class GemmBase: + """Common non-mainloop pieces shared by GEMM architectures.""" + + arch = 0 + + @dataclass + class EpilogueArguments: + pass + + EpilogueParams = ParamsBase + + def epi_smem_warp_shape_mnk(self): + return (self.num_epi_warps, 1, 1) + + @cute.jit + def epilogue( + self, + params: EpilogueParams, + epi_smem_tensors: Dict[str, cute.Tensor], + epi_pipeline: Optional[cutlass.pipeline.PipelineAsync], + epi_store_pipeline: Optional[cutlass.pipeline.PipelineAsync], + epi_read_state: Optional[cutlass.pipeline.PipelineState], + epi_producer_state: Optional[cutlass.pipeline.PipelineState], + epi_tile: cute.Tile, + load_acc_subtile: Callable, + tRS_rD: cute.Tensor, + tRS_rC: Optional[cute.Tensor], + tiled_copy_t2r: Optional[cute.TiledCopy], # Only for Sm100 + tiled_copy_r2s: cute.TiledCopy, + tRS_sD: cute.Tensor, + tiled_copy_s2r: Optional[cute.ThrCopy], + tSR_rC: Optional[cute.Tensor], + tSR_sC: Optional[cute.Tensor], + copy_D: Optional[Callable], + copy_C: Optional[Callable], + tile_coord_mnkl: cute.Coord, + varlen_manager: VarlenManager, + epilogue_barrier: cutlass.pipeline.NamedBarrier, + tile_scheduler, + tidx: Int32, + is_tma_warp: cutlass.Boolean, + ) -> Tuple[cutlass.pipeline.PipelineState, cutlass.pipeline.PipelineState]: + has_C = const_expr(tRS_rC is not None) + has_epi_load = const_expr(self.epi_c_stage > 0) + has_D = const_expr(copy_D is not None) + use_tma_epi = const_expr(epi_store_pipeline is not None) + use_tma_c = const_expr(epi_pipeline is not None) + inline_epi_load = const_expr(copy_C is not None) + use_stochastic_rounding = const_expr( + self.rounding_mode == RoundingMode.RS + and self.acc_dtype == cutlass.Float32 + and self.d_dtype == cutlass.BFloat16 + ) + + # Setup aux outputs. Returns a tuple of ``(tiled_copy_r2s, + # tRS_sAuxOut, copy_aux_out)`` triples — empty for the default + # epilogue, one entry for the standard ``GemmAct``/``GemmGated`` + # single-output mixins, multiple entries for multi-output mixins + # (e.g. ``T*tanh`` + ``1-tanh^2`` from one GEMM). + aux_out_ctxs = self.epi_setup_aux_out( + params, + epi_smem_tensors, + tiled_copy_r2s, + tiled_copy_t2r, + tile_coord_mnkl, + varlen_manager, + tidx, + ) + + epi_tile_shape = cute.zipped_divide( + cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile + ).shape[1] + epi_tile_layout = cute.make_ordered_layout( + epi_tile_shape, order=(0, 1) if const_expr(self.epi_m_major) else (1, 0) + ) + epi_tile_num = cute.size(epi_tile_shape) + num_prev_subtiles = tile_scheduler.num_tiles_executed * epi_tile_num + + epi_tensors = self.epi_begin( + params, + epi_smem_tensors, + epi_tile, + tiled_copy_t2r, + tiled_copy_r2s, + tile_coord_mnkl, + varlen_manager, + epilogue_barrier, + tidx, + tRS_rD.layout, + ) + + if const_expr(inline_epi_load): + for epi_idx in cutlass.range(min(epi_tile_num, self.epi_c_stage), unroll=1): + epi_coord_C = epi_tile_layout.get_hier_coord(epi_idx) + if const_expr(use_tma_c): + if is_tma_warp: + epi_pipeline.producer_acquire(epi_producer_state) + copy_C(src_idx=epi_coord_C, producer_state=epi_producer_state) + epi_pipeline.producer_commit(epi_producer_state) + epi_producer_state.advance() + else: + # TODO: turn this to cp.async instead of direct G2R copy + copy_C(src_idx=epi_coord_C, dst_idx=epi_idx % self.epi_c_stage) + if const_expr(use_tma_c): + epilogue_barrier.arrive_and_wait() + + for epi_idx in cutlass.range_constexpr(epi_tile_num): + epi_coord = epi_tile_layout.get_hier_coord(epi_idx) # (epi_m, epi_n) + # Copy from acc to D registers + load_acc_subtile(tRS_rD, epi_coord) + if const_expr(has_epi_load): + if const_expr(use_tma_c): + epi_pipeline.consumer_wait(epi_read_state) + if const_expr(has_C): + cute.copy( + tiled_copy_s2r, tSR_sC[None, None, None, epi_read_state.index], tSR_rC + ) + self.epi_tile_load_s2r(params, epi_tensors, epi_read_state.index) + cute.arch.fence_view_async_shared() + epi_pipeline.consumer_release(epi_read_state) + epi_read_state.advance() + else: + c_buffer = epi_idx % self.epi_c_stage + cute.copy(tiled_copy_s2r, tSR_sC[None, None, None, c_buffer], tSR_rC) + # TODO: cp.async wait once we switch to cp.async + epilogue_barrier.arrive_and_wait() + epi_loop_tensors = self.epi_begin_loop(params, epi_tensors, epi_coord) + if const_expr(inline_epi_load and epi_idx + self.epi_c_stage < epi_tile_num): + epi_coord_C = epi_tile_layout.get_hier_coord(epi_idx + self.epi_c_stage) + if const_expr(use_tma_c): + if is_tma_warp: + epi_pipeline.producer_acquire(epi_producer_state) + copy_C(src_idx=epi_coord_C, producer_state=epi_producer_state) + epi_pipeline.producer_commit(epi_producer_state) + epi_producer_state.advance() + else: + epilogue_barrier.arrive_and_wait() + copy_C( + src_idx=epi_coord_C, + dst_idx=(epi_idx + self.epi_c_stage) % self.epi_c_stage, + ) + # Returns a tuple of register tensors — one per aux output. + # Length matches ``aux_out_ctxs``. ``()`` for the default + # epilogue (no aux output). + tRS_rAuxOuts = self.epi_visit_subtile(params, epi_loop_tensors, tRS_rD, tRS_rC) + self.epi_end_loop( + params, + epi_tensors, + epi_coord, + epi_tile, + tiled_copy_t2r, + tiled_copy_r2s, + tile_coord_mnkl, + varlen_manager, + tidx, + ) + # Convert each output to its storage dtype. + tRS_rAuxOuts_out = tuple( + self.epi_convert_aux_out( + i, + tRS_rAuxOuts[i], + epi_loop_tensors.get("sr_seed"), + tidx, + tile_coord_mnkl, + num_prev_subtiles, + epi_idx, + ) + for i in range(len(aux_out_ctxs)) + ) + if const_expr(use_tma_epi): + if is_tma_warp: + epi_store_pipeline.producer_acquire() + else: + epilogue_barrier.arrive_and_wait() + if const_expr(use_tma_epi): + epilogue_barrier.arrive_and_wait() + epi_buffer = (num_prev_subtiles + epi_idx) % self.epi_stage + if const_expr(has_D): + tRS_sD_cur = tRS_sD[None, None, None, epi_buffer] + if const_expr(use_stochastic_rounding): + seed = epilogue_sr_seed( + epi_loop_tensors.get("sr_seed"), + tile_coord_mnkl, + num_prev_subtiles + epi_idx, + ) + copy_utils.sr_cvt_copy(tiled_copy_r2s, tRS_rD, tRS_sD_cur, seed, tidx) + else: + copy_utils.cvt_copy(tiled_copy_r2s, tRS_rD, tRS_sD_cur) + # Copy each aux output from registers to shared memory. All share + # the same ``epi_buffer`` index so the s2g TMA stores below happen + # in lockstep after the fence. + for i in cutlass.range_constexpr(len(aux_out_ctxs)): + tiled_copy_aux_out_r2s, tRS_sAuxOut, _ = aux_out_ctxs[i] + cute.copy( + tiled_copy_aux_out_r2s, + # Need contiguous for Sm80 and Sm120 where acc layout is ((2, 2), MMA_M, MMA_N) + tiled_copy_aux_out_r2s.retile(tRS_rAuxOuts_out[i]).contiguous(), + tRS_sAuxOut[None, None, None, epi_buffer], + ) + if const_expr(use_tma_epi): + cute.arch.fence_view_async_shared() + epilogue_barrier.arrive_and_wait() + if is_tma_warp: + if const_expr(has_D): + copy_D(src_idx=epi_buffer, dst_idx=epi_coord) + for i in cutlass.range_constexpr(len(aux_out_ctxs)): + _, _, copy_aux_out = aux_out_ctxs[i] + copy_aux_out(src_idx=epi_buffer, dst_idx=epi_coord) + epi_store_pipeline.producer_commit() + else: + epilogue_barrier.arrive_and_wait() + if const_expr(has_D): + copy_D(src_idx=epi_buffer, dst_idx=epi_coord) + for i in cutlass.range_constexpr(len(aux_out_ctxs)): + _, _, copy_aux_out = aux_out_ctxs[i] + copy_aux_out(src_idx=epi_buffer, dst_idx=epi_coord) + epilogue_barrier.arrive_and_wait() + + self.epi_end( + params, + epi_tensors, + epi_tile, + tiled_copy_t2r, + tiled_copy_r2s, + tile_coord_mnkl, + varlen_manager, + tidx, + ) + + return epi_read_state, epi_producer_state + + def get_scheduler_class(self, varlen_m: bool = False): + """Return the scheduler class to use. Override in subclasses for custom schedulers.""" + return TileScheduler if not varlen_m else VarlenMTileScheduler + + def resolve_epi_m_major(self, epilogue_args: EpilogueArguments): + return True + + def get_scheduler_arguments( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mD: Optional[cute.Tensor], + scheduler_args, + varlen_args, + epilogue_args, + ): + """Create scheduler arguments. Override in subclasses for custom schedulers.""" + if const_expr(not self.is_persistent): + persistence_mode = PersistenceMode.NONE + else: + if const_expr(self.arch >= 100 and self.use_clc_persistence): + persistence_mode = PersistenceMode.CLC + elif const_expr(scheduler_args.tile_count_semaphore is not None): + persistence_mode = PersistenceMode.DYNAMIC + else: + persistence_mode = PersistenceMode.STATIC + if const_expr(varlen_args.mCuSeqlensM is None): + num_problems = ( + mD.shape[2] + if mD is not None + else ( + mB.shape[2] + if varlen_args.mCuSeqlensK is None + else varlen_args.mCuSeqlensK.shape[0] - 1 + ) + ) + problem_shape_ntile_mnl = ( + cute.ceil_div(cute.size(mA, mode=[0]), self.cta_tile_shape_mnk[0]), + cute.ceil_div(cute.size(mB, mode=[0]), self.cta_tile_shape_mnk[1]), + num_problems, + ) + tile_sched_args = TileSchedulerArguments( + problem_shape_ntile_mnl=problem_shape_ntile_mnl, + raster_order=scheduler_args.raster_order, + group_size=scheduler_args.max_swizzle_size, + cluster_shape_mnk=self.cluster_shape_mnk, + tile_count_semaphore=scheduler_args.tile_count_semaphore, + batch_idx_permute=scheduler_args.batch_idx_permute, + persistence_mode=persistence_mode, + ) + else: + assert (mD is not None) or (epilogue_args.mAuxOut is not None) or (not self.gather_A) + problem_shape_ntile_mnl = ( + None, + cute.ceil_div(cute.size(mB, mode=[0]), self.cta_tile_shape_mnk[1]), + varlen_args.mCuSeqlensM.shape[0] - 1, + ) + tile_sched_args = VarlenMTileSchedulerArguments( + problem_shape_ntile_mnl=problem_shape_ntile_mnl, + total_m=( + mD.shape[0] + if mD is not None + else ( + varlen_args.mAIdx.shape[0] + if varlen_args.mAIdx is not None + else cute.size(mA, mode=[0]) + ) + ), + cu_seqlens_m=varlen_args.mCuSeqlensM, + max_active_clusters=scheduler_args.max_active_clusters, + raster_order=scheduler_args.raster_order, + group_size=scheduler_args.max_swizzle_size, + tile_shape_mn=self.cta_tile_shape_mnk[:2], + cluster_shape_mnk=self.cluster_shape_mnk, + tile_count_semaphore=scheduler_args.tile_count_semaphore, + persistence_mode=persistence_mode, + ) + return tile_sched_args + + @cute.jit + def epi_load_acc_subtile( + self, + tRS_rAcc: cute.Tensor, + tRS_rD: cute.Tensor, + epi_coord, # (int, int) + ): + cute.autovec_copy(tRS_rAcc[None, None, None, epi_coord], tRS_rD) + + @cute.jit + def epi_begin( + self, + params: EpilogueParams, + epi_smem_tensors: Dict[str, cute.Tensor], + epi_tile: cute.Tile, + tiled_copy_t2r: Optional[cute.TiledCopy], + tiled_copy_r2s: cute.TiledCopy, + tile_coord_mnkl: cute.Coord, + varlen_manager: VarlenManager, + epilogue_barrier: cutlass.pipeline.NamedBarrier, + tidx: Int32, + tRS_rD_layout=None, + ) -> Tuple[cute.Tensor, ...]: + return () + + def epi_begin_loop( + self, params: EpilogueParams, epi_tensors: Tuple[cute.Tensor, ...], epi_coord: cute.Coord + ) -> Tuple[cute.Tensor, ...]: + return () + + def epi_visit_subtile( + self, + params: EpilogueParams, + epi_loop_tensors: Tuple[cute.Tensor, ...], + tRS_rD: cute.Tensor, + tRS_rC: Optional[cute.Tensor] = None, + ) -> Tuple[cute.Tensor, ...]: + return () + + def epi_visit_acc( + self, + params: EpilogueParams, + acc: cute.Tensor, + tiled_mma: cute.TiledMma, + tile_coord_mnkl: cute.Coord, + tidx: Int32, + ) -> None: + pass + + @cute.jit + def epi_end_loop( + self, + params: EpilogueParams, + epi_tensors: Tuple[cute.Tensor, ...], + epi_coord: cute.Coord, + epi_tile: cute.Tile, + tiled_copy_t2r: Optional[cute.TiledCopy], + tiled_copy_r2s: cute.TiledCopy, + tile_coord_mnkl: cute.Coord, + varlen_manager, + tidx, + ) -> None: + pass + + @cute.jit + def epi_end( + self, + params: EpilogueParams, + epi_tensors: Tuple[cute.Tensor, ...], + epi_tile: cute.Tile, + tiled_copy_t2r: Optional[cute.TiledCopy], + tiled_copy_r2s: cute.TiledCopy, + tile_coord_mnkl: cute.Coord, + varlen_manager, + tidx, + ) -> None: + pass + + def epi_to_underlying_arguments( + self, args: EpilogueArguments, *, loc=None, ip=None + ) -> EpilogueParams: + return self.EpilogueParams() + + def epi_get_tma_atoms( + self, params: EpilogueParams, *, loc=None, ip=None + ) -> list[cute.CopyAtom]: + """Subclasses can override this.""" + return [] + + def epi_tile_load_g2s_copy_fns( + self, + params, + epi_smem_tensors, + tile_coord_mnkl, + varlen_manager, + epi_pipeline, + ): + return () + + @cute.jit + def epi_tile_load_s2r(self, params, epi_tensors, stage_idx): + pass + + @staticmethod + def epi_smem_bytes( + args: Optional[EpilogueArguments], + cta_tile_shape_mnk: Tuple[int, int, int], + epi_tile: cute.Tile, + warp_shape_mnk: Tuple[int, int, int] | None = None, + ) -> EpiSmemBytes: + return EpiSmemBytes() + + def epi_get_smem_struct(self, params: EpilogueParams): + return cute.struct.MemRange[Int32, 0] # Dummy struct + + def epi_get_smem_tensors(self, params: EpilogueParams, storage) -> Dict[str, cute.Tensor]: + return {} + + def epi_setup_aux_out( + self, + params, + epi_smem_tensors, + tiled_copy_r2s, + tiled_copy_t2r, + tile_coord_mnkl, + varlen_manager, + tidx, + ): + """Return a tuple of ``(tiled_copy_r2s, tRS_sAuxOut, copy_aux_out)`` + triples — one per aux output. The default epilogue has no aux output, + so the tuple is empty. + """ + return () + + @cute.jit + def epi_convert_aux_out( + self, + output_idx: cutlass.Constexpr[int], + tRS_rAuxOut, + sr_seed, + tidx, + tile_coord_mnkl, + num_prev_subtiles, + epi_idx, + ): + """Convert one aux output register tensor from acc_dtype to its storage + dtype. ``output_idx`` selects which aux output this call is for + (single-output mixins can ignore it). + """ + return tRS_rAuxOut + + +class GemmTmaBase(GemmBase): + """Common TMA descriptor and pipeline helpers for SM90+ GEMM paths.""" + + @cute.jit + def load_tma( + self, + pipeline: cutlass.pipeline.PipelineAsync, + producer_state: cutlass.pipeline.PipelineState, + copy_fns: Sequence[Optional[Callable]], + k_tile_cnt: Int32, + ) -> cutlass.pipeline.PipelineState: + # Peek (try_wait) AB buffer empty for k_block = prefetch_k_tile_cnt. + peek_empty_status = Boolean(True) + if 0 < k_tile_cnt: + peek_empty_status = pipeline.producer_try_acquire(producer_state) + # TMA load + for k_tile in cutlass.range(k_tile_cnt, unroll=1): + # Wait for A/B buffers to be empty before loading into them. + # Also sets the transaction barrier for the A/B buffers. + pipeline.producer_acquire(producer_state, peek_empty_status) + tma_bar_ptr = pipeline.producer_get_barrier(producer_state) + smem_idx = producer_state.index + for copy_fn in copy_fns: + if const_expr(copy_fn is not None): + copy_fn(k_tile, smem_idx, tma_bar_ptr=tma_bar_ptr) + # Mainloop pipeline's producer commit is a NOP for TMA pipelines. + pipeline.producer_commit(producer_state) + producer_state.advance() + peek_empty_status = Boolean(True) + if k_tile + 1 < k_tile_cnt: + peek_empty_status = pipeline.producer_try_acquire(producer_state) + return producer_state + + def _make_gmem_tiled_copy_A(self, dtype, major_mode, num_threads, copy_bits=128): + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + dtype, + num_bits_per_copy=copy_bits, + ) + copy_elems = copy_bits // dtype.width + loads_per_cache_line = 128 * 8 // copy_bits # 128 bytes per cache line + shape_dim_1 = cute.size(self.cta_tile_shape_mnk[2]) // copy_elems + if shape_dim_1 > loads_per_cache_line: + shape_dim_1 = math.gcd(shape_dim_1, loads_per_cache_line) + # thread layout for copy + thread_layout = cute.make_layout( + (num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1) + ) + if major_mode != LayoutEnum.ROW_MAJOR: + shape_dim_0 = cute.size(self.cta_tile_shape_mnk[0]) // copy_elems + if shape_dim_0 > loads_per_cache_line: + shape_dim_0 = math.gcd(shape_dim_0, loads_per_cache_line) + thread_layout = cute.make_layout( + (shape_dim_0, num_threads // shape_dim_0), stride=(1, shape_dim_0) + ) + # Value layout for copy + value_layout = ( + cute.make_layout((1, copy_elems)) + if major_mode == LayoutEnum.ROW_MAJOR + else cute.make_layout((copy_elems, 1)) + ) + return cute.make_tiled_copy_tv(atom_async_copy, thread_layout, value_layout) + + def make_tma_load_atoms_and_tensors( + self, + mA: cute.Tensor, + mB: cute.Tensor, + a_smem_layout: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + varlen_k: bool, + ): + tma_atom_a, tma_tensor_a = None, None + if const_expr(not self.gather_A): + tma_atom_a, tma_tensor_a = self._make_tma_atoms_and_tensors( + copy_utils.create_ragged_tensor_for_tma(mA, ragged_dim=1) + if varlen_k and not self.gather_A + else mA, + a_smem_layout, + (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[2]), + self.cluster_shape_mnk[1], + ) + tma_atom_b, tma_tensor_b = self._make_tma_atoms_and_tensors( + copy_utils.create_ragged_tensor_for_tma(mB, ragged_dim=1) if varlen_k else mB, + b_smem_layout, + (self.cta_tile_shape_mnk[1], self.cta_tile_shape_mnk[2]), + self.cluster_shape_mnk[0], + ) + return tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b + + def make_tma_epilogue_atoms_and_tensors( + self, + mD: Optional[cute.Tensor], + mC: Optional[cute.Tensor], + epilogue_args, + varlen_m: bool, + ): + tma_atom_d, tma_tensor_d = None, None + if const_expr(mD is not None): + tma_atom_d, tma_tensor_d = self._make_tma_epi_atoms_and_tensors( + copy_utils.create_ragged_tensor_for_tma(mD, ragged_dim=0, ptr_shift=True) + if varlen_m + else mD, + self.epi_smem_layout_staged, + self.epi_tile, + op_type="store" + if not (hasattr(epilogue_args, "add_to_output") and epilogue_args.add_to_output) + else "add", + ) + tma_atom_c, tma_tensor_c = None, None + if const_expr(mC is not None): + tma_atom_c, tma_tensor_c = self._make_tma_epi_atoms_and_tensors( + mC, self.epi_c_smem_layout_staged, self.epi_tile, op_type="load" + ) + return tma_atom_d, tma_tensor_d, tma_atom_c, tma_tensor_c + + def epilog_gmem_copy_and_partition( + self, + atom: cute.CopyAtom | cute.TiledCopy, + mD_mn: cute.Tensor, + tile_shape_mn: cute.Tile, + epi_tile: cute.Tile, + sD: cute.Tensor, + tile_coord_mnkl: cute.Coord, + ) -> Tuple[cute.Tensor, cute.Tensor]: + gD = cute.local_tile(mD_mn, tile_shape_mn, tile_coord_mnkl[:2]) # (bM, bN) + tDgD_for_tma_partition = cute.zipped_divide(gD, epi_tile) + is_s2g = isinstance( + atom.op, (cpasync.CopyBulkTensorTileS2GOp, cpasync.CopyReduceBulkTensorTileS2GOp) + ) + src_tensor, dst_tensor = ( + (sD, tDgD_for_tma_partition) if is_s2g else (tDgD_for_tma_partition, sD) + ) + return copy_utils.tma_get_copy_fn( + atom, + cta_coord=0, + cta_layout=cute.make_layout(1), + src_tensor=src_tensor, + dst_tensor=dst_tensor, + ) + + def make_ab_pipeline( + self, + tiled_mma: cute.TiledMma, + cluster_layout_vmnk: cute.Layout, + ): + # Threads/warps participating in this pipeline + producer_cnt = 1 if const_expr(not self.gather_A) else 1 + self.num_ab_load_warps * 32 + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt) + # Each warp will contribute to the arrive count with the number of mcast size + mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + consumer_arrive_cnt = mcast_size * tiled_mma.size // cute.arch.WARP_SIZE + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, consumer_arrive_cnt + ) + pipeline_cls = pipeline.PipelineTmaAsync if not self.gather_A else PipelineTmaCpAsync + return pipeline_cls.create( + num_stages=self.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, + ) + + def make_epi_pipeline( + self, + tx_count: int, + ): + epi_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + # Each warp will contribute 1 to the arrive count + consumer_arrive_cnt = self.num_epi_warps + epi_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, consumer_arrive_cnt + ) + return PipelineTmaAsync.create( + num_stages=self.epi_c_stage, + producer_group=epi_pipeline_producer_group, + consumer_group=epi_pipeline_consumer_group, + tx_count=tx_count, + defer_sync=True, + elect_one_release=True, + syncwarp_before_release=True, + ) + + def make_epi_store_pipeline(self): + num_epi_threads = self.num_epi_warps * cute.arch.WARP_SIZE + epi_store_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_epi_threads) + return pipeline.PipelineTmaStore.create( + num_stages=self.epi_stage, producer_group=epi_store_producer_group + ) + + @staticmethod + def _make_tma_epi_atoms_and_tensors( + tensor_d: cute.Tensor, + epi_smem_layout_staged: cute.ComposedLayout, + epi_tile: Tuple[int, int], + op_type: Literal["store", "load", "add"], + ) -> Tuple[cute.CopyAtom, cute.Tensor]: + """Create TMA atoms and tensors for storing D or loading C.""" + assert op_type in ["load", "store", "add"] + epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0)) + d_cta_v_layout = cute.composition(cute.make_identity_layout(tensor_d.shape), epi_tile) + op = { + "load": cpasync.CopyBulkTensorTileG2SOp(), + "store": cpasync.CopyBulkTensorTileS2GOp(), + "add": cpasync.CopyReduceBulkTensorTileS2GOp(cpasync.ReductionOp.ADD), + }[op_type] + tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom( + op, tensor_d, epi_smem_layout, d_cta_v_layout + ) + return tma_atom_d, tma_tensor_d + + @staticmethod + def _make_tma_atoms_and_tensors( + tensor: cute.Tensor, + smem_layout: cute.ComposedLayout, + smem_tile: Tuple[int, int], + mcast_dim: int, + ) -> Tuple[cute.CopyAtom, cute.Tensor]: + """Create TMA atoms and tensors for input tensors.""" + # block_copy takes compiler-driven multicast metadata at the copy site, + # so the TMA atom itself must stay the non-multicast variant here. + op = cpasync.CopyBulkTensorTileG2SOp() + tma_atom, tma_tensor = cpasync.make_tiled_tma_atom(op, tensor, smem_layout, smem_tile) + return tma_atom, tma_tensor diff --git a/build/torch-cuda/quack/gemm_blockscaled_interface.py b/build/torch-cuda/quack/gemm_blockscaled_interface.py deleted file mode 100644 index 0bc39749e5dee049cc137e6329a93223ac8b3c81..0000000000000000000000000000000000000000 --- a/build/torch-cuda/quack/gemm_blockscaled_interface.py +++ /dev/null @@ -1,326 +0,0 @@ -# Copyright (c) 2026, Tri Dao. -"""PyTorch-friendly interface for the SM100 MXFP8 blockscaled GEMM. - -Shape / layout conventions (matches torch.matmul, torch._scaled_mm, cuBLAS): - A: (M, K) or (L, M, K) dtype float8_e4m3fn, K-contiguous (row-major) - B: (K, N) or (L, K, N) dtype float8_e4m3fn, K-contiguous (col-major) - A_scale: (M, K/32) or (L, M, K/32) dtype float8_e8m0fnu, K-contiguous - B_scale: (K/32, N) or (L, K/32, N) dtype float8_e8m0fnu, K-contiguous - out: (M, N) or (L, M, N) dtype bfloat16/float16, contiguous - -"K-contiguous" means stride 1 on the K axis. This matches how torchao/cuBLAS -use `torch._scaled_mm(a, b.t(), ...)`: - - you store a weight as nn.Linear-style `W` of shape `(N, K)` row-major - - you pass `W.mT` (a zero-copy view of shape (K, N) with K-contig) as B -The interface applies `.mT` internally to reach the `(N, K) K-major` layout -the quack kernel consumes. No data is copied. -""" - -from functools import lru_cache -from typing import Optional, Tuple - -import torch -from torch import Tensor - -import cutlass - -from .blockscaled_gemm_utils import ( - ceil_div, - compile_blockscaled_gemm_tvm_ffi, - pack_scale_2d_to_blocked_contig, - scale_blocked_for_cublas, - scale_view_for_kernel, -) -from .gemm_default_epi import GemmDefaultSm100 -from .mx_utils import to_mx - -_SF_VEC_SIZE = 32 -_TORCH_TO_CUTLASS_D = { - torch.bfloat16: cutlass.BFloat16, - torch.float16: cutlass.Float16, - torch.float32: cutlass.Float32, -} - - -def _default_tiler_cluster(m: int, n: int) -> Tuple[Tuple[int, int], Tuple[int, int]]: - """Pick a reasonable default (mma_tiler_mn, cluster_shape_mn).""" - if m >= 512 and n >= 128: - return (256, 128), (2, 1) - return (128, 128), (1, 1) - - -@lru_cache(maxsize=64) -def _compile_cached( - m: int, - n: int, - k: int, - l: int, - mma_tiler_mn: Tuple[int, int], - cluster_shape_mn: Tuple[int, int], - out_torch_dtype, - ab_dtype_cutlass, - sf_dtype_cutlass, -): - """Compile kernel for a given (shape, dtype, tiler, cluster) and cache it.""" - dev = torch.device("cuda") - rm = ceil_div(m, 128) - rn = ceil_div(n, 128) - rk = ceil_div(k // _SF_VEC_SIZE, 4) - # K-major: (l, m, k) contiguous, viewed as (m, k, l) strides (k, 1, m*k) - fake_mA = torch.empty(l, m, k, dtype=torch.float8_e4m3fn, device=dev).permute(1, 2, 0) - fake_mB = torch.empty(l, n, k, dtype=torch.float8_e4m3fn, device=dev).permute(1, 2, 0) - # N-major: (l, m, n) contiguous, viewed as (m, n, l) strides (n, 1, m*n) - fake_mD = torch.empty(l, m, n, dtype=out_torch_dtype, device=dev).permute(1, 2, 0) - fake_sc_A = torch.empty(l, rm, rk, 512, dtype=torch.float8_e8m0fnu, device=dev) - fake_sc_B = torch.empty(l, rn, rk, 512, dtype=torch.float8_e8m0fnu, device=dev) - fake_mSFA = scale_view_for_kernel(fake_sc_A, m, k // _SF_VEC_SIZE, l) - fake_mSFB = scale_view_for_kernel(fake_sc_B, n, k // _SF_VEC_SIZE, l) - return compile_blockscaled_gemm_tvm_ffi( - ab_dtype_cutlass, - sf_dtype_cutlass, - _SF_VEC_SIZE, - _TORCH_TO_CUTLASS_D[out_torch_dtype], - mma_tiler_mn, - cluster_shape_mn, - fake_mA, - fake_mB, - fake_mD, - fake_mSFA, - fake_mSFB, - ) - - -def _as_3d(x: Tensor, ndim_in: int) -> Tensor: - """Add a leading batch dim if input is 2D. Returns a view.""" - if ndim_in == 2: - return x.unsqueeze(0) - return x - - -def _to_kernel_layout( - A: Tensor, - B: Tensor, - A_scale: Tensor, - B_scale: Tensor, -) -> Tuple[int, int, int, int, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, bool]: - """Normalize shapes/strides, validate, and repack scales. Returns - (m, n, k, l, mA_mkl, mB_nkl, sc_contig_A, sc_contig_B, sfa_view, sfb_view, was_2d). - - A: (M,K) or (L,M,K) K-contig. B: (K,N) or (L,K,N) K-contig. - A_scale: (M,K/32) or (L,M,K/32) K-contig. B_scale: (K/32,N) or (L,K/32,N) K-contig. - """ - assert A.dtype == torch.float8_e4m3fn, f"A dtype must be float8_e4m3fn, got {A.dtype}" - assert B.dtype == torch.float8_e4m3fn, f"B dtype must be float8_e4m3fn, got {B.dtype}" - assert A_scale.dtype == torch.float8_e8m0fnu - assert B_scale.dtype == torch.float8_e8m0fnu - was_2d = A.dim() == 2 - # Flip B from (K,N) to (N,K) via .mT (zero-copy). User's B K-contig → .mT K-contig. - A3 = _as_3d(A, A.dim()) # (l, m, k) K-contig row-major expected - B3 = _as_3d(B, B.dim()).mT # (l, n, k) K-contig (view) from (l, k, n) - l, m, k = A3.shape - l2, n, k2 = B3.shape - assert l == l2, f"batch mismatch: A={l}, B={l2}" - assert k == k2, f"K mismatch: A K={k}, B K={k2}" - assert k % _SF_VEC_SIZE == 0, f"K ({k}) must be divisible by {_SF_VEC_SIZE}" - assert A3.stride(-1) == 1, "A must be K-contiguous (stride 1 on K)" - assert B3.stride(-1) == 1, ( - "B must be K-contiguous on its K axis (pass .mT of an (N,K) row-major tensor)" - ) - sf_k = k // _SF_VEC_SIZE - as3 = _as_3d(A_scale, A_scale.dim()) # expected (l, m, sf_k) K-contig row-major - bs3 = _as_3d(B_scale, B_scale.dim()).mT # (l, n, sf_k) K-contig (view) from (l, sf_k, n) - assert as3.stride(-1) == 1, "A_scale must be K-contiguous" - assert bs3.stride(-1) == 1, ( - "B_scale must be K-contiguous on its K axis (pass .mT of an (N, K/32) row-major tensor)" - ) - assert as3.shape == (l, m, sf_k), ( - f"A_scale shape: expected (l={l},m={m},sf_k={sf_k}) K-contig, got {tuple(as3.shape)}" - ) - assert bs3.shape == (l, n, sf_k), ( - f"B_scale shape: expected .mT of (l={l},sf_k={sf_k},n={n}) -> ({l},{n},{sf_k}), got {tuple(bs3.shape)}" - ) - # Force row-major contiguous for packer/kernel consumption. - # A3 / B3 are views — .contiguous() materializes (l,m,k) / (l,n,k) row-major. - A3_c = A3.contiguous() - B3_c = B3.contiguous() - # (l, m, k) -> (m, k, l) K-major view (no copy; strides (k, 1, m*k)) - mA_mkl = A3_c.permute(1, 2, 0) - mB_nkl = B3_c.permute(1, 2, 0) - sc_contig_A = pack_scale_2d_to_blocked_contig(as3.contiguous()) - sc_contig_B = pack_scale_2d_to_blocked_contig(bs3.contiguous()) - sfa_view = scale_view_for_kernel(sc_contig_A, m, sf_k, l) - sfb_view = scale_view_for_kernel(sc_contig_B, n, sf_k, l) - return m, n, k, l, mA_mkl, mB_nkl, sc_contig_A, sc_contig_B, sfa_view, sfb_view, was_2d - - -def mxfp8_gemm_out( - A: Tensor, - B: Tensor, - A_scale: Tensor, - B_scale: Tensor, - out: Tensor, - *, - mma_tiler_mn: Optional[Tuple[int, int]] = None, - cluster_shape_mn: Optional[Tuple[int, int]] = None, -) -> None: - """MXFP8 blockscaled GEMM with pre-allocated output. See module doc for shape conventions.""" - m, n, k, l, mA, mB, _scA, _scB, sfa, sfb, was_2d = _to_kernel_layout(A, B, A_scale, B_scale) - out_dtype = out.dtype - assert out_dtype in _TORCH_TO_CUTLASS_D, f"unsupported out dtype: {out_dtype}" - expected_out_shape = (m, n) if was_2d else (l, m, n) - assert tuple(out.shape) == expected_out_shape, ( - f"out shape {tuple(out.shape)} != expected {expected_out_shape}" - ) - assert out.is_contiguous(), "out must be contiguous" - # View caller's contiguous (M,N) or (L,M,N) as (M,N,L) N-major strided view, no copy. - out_3d = out.unsqueeze(0) if was_2d else out # (l, m, n) - mD = out_3d.permute(1, 2, 0) # (m, n, l), strides (n, 1, m*n) - if mma_tiler_mn is None or cluster_shape_mn is None: - tlr, clu = _default_tiler_cluster(m, n) - mma_tiler_mn = mma_tiler_mn or tlr - cluster_shape_mn = cluster_shape_mn or clu - if not GemmDefaultSm100.can_implement_blockscaled( - cutlass.Float8E4M3FN, - cutlass.Float8E8M0FNU, - _SF_VEC_SIZE, - _TORCH_TO_CUTLASS_D[out_dtype], - mma_tiler_mn, - cluster_shape_mn, - m, - n, - k, - l, - "k", - "k", - "n", - ): - raise ValueError( - f"unsupported config: m={m}, n={n}, k={k}, l={l}, " - f"tiler={mma_tiler_mn}, cluster={cluster_shape_mn}" - ) - runner = _compile_cached( - m, - n, - k, - l, - mma_tiler_mn, - cluster_shape_mn, - out_dtype, - cutlass.Float8E4M3FN, - cutlass.Float8E8M0FNU, - ) - runner(mA, mB, mD, sfa, sfb) - - -def mxfp8_gemm( - A: Tensor, - B: Tensor, - A_scale: Tensor, - B_scale: Tensor, - out: Optional[Tensor] = None, - out_dtype: torch.dtype = torch.bfloat16, - *, - mma_tiler_mn: Optional[Tuple[int, int]] = None, - cluster_shape_mn: Optional[Tuple[int, int]] = None, -) -> Tensor: - """MXFP8 blockscaled GEMM. Allocates output if not provided.""" - if out is None: - # A: (M,K) or (L,M,K); B: (K,N) or (L,K,N); out: (M,N) or (L,M,N) - if A.dim() == 2: - out_shape = (A.shape[0], B.shape[1]) - else: - out_shape = (A.shape[0], A.shape[1], B.shape[2]) - out = torch.empty(out_shape, dtype=out_dtype, device=A.device) - mxfp8_gemm_out( - A, - B, - A_scale, - B_scale, - out, - mma_tiler_mn=mma_tiler_mn, - cluster_shape_mn=cluster_shape_mn, - ) - return out - - -def mxfp8_quantize(x: Tensor) -> Tuple[Tensor, Tensor]: - """Quantize a (..., K) bf16/fp32 tensor to MXFP8. Returns (qdata, scale_2d) - in torchao-convention layout. Last dim (K) must be divisible by 32.""" - assert x.shape[-1] % _SF_VEC_SIZE == 0, ( - f"last dim ({x.shape[-1]}) must be divisible by {_SF_VEC_SIZE}" - ) - return to_mx(x.contiguous(), _SF_VEC_SIZE) - - -def mxfp8_gemm_quantize( - A: Tensor, - B: Tensor, - out: Optional[Tensor] = None, - out_dtype: torch.dtype = torch.bfloat16, - *, - mma_tiler_mn: Optional[Tuple[int, int]] = None, - cluster_shape_mn: Optional[Tuple[int, int]] = None, -) -> Tensor: - """High-level: quantize bf16 A, B_as_NK to MXFP8, then run C = A @ B_as_NK.mT. - Inputs: A=(M,K)/(L,M,K), B_as_NK=(N,K)/(L,N,K) bf16/fp32. Quantization - scales along the last (K) dim. Returned output has shape (M,N)/(L,M,N).""" - A_q, A_sc = mxfp8_quantize(A) - B_q, B_sc = mxfp8_quantize(B) - # B_q, B_sc are (..., N, K) / (..., N, K/32). Flip to (..., K, N) / (..., K/32, N) - # K-contig zero-copy views to match the interface convention. - return mxfp8_gemm( - A_q, - B_q.mT, - A_sc, - B_sc.mT, - out=out, - out_dtype=out_dtype, - mma_tiler_mn=mma_tiler_mn, - cluster_shape_mn=cluster_shape_mn, - ) - - -def mxfp8_gemm_cublas( - A: Tensor, - B: Tensor, - A_scale: Tensor, - B_scale: Tensor, - out_dtype: torch.dtype = torch.bfloat16, -) -> Tensor: - """Reference path via torch._scaled_mm. Requires l=1 (or 2D inputs).""" - m, n, k, l, _mA, _mB, sc_A, sc_B, _sfa, _sfb, was_2d = _to_kernel_layout(A, B, A_scale, B_scale) - assert l == 1, "torch._scaled_mm MXFP8 path is 2D only; pass 2D inputs or l=1" - # torch._scaled_mm: A=(M,K) row-major, B=(K,N) col-major (both K-contig) -- same layout user gave us. - a2d = A if A.dim() == 2 else A.squeeze(0) - b2d = B if B.dim() == 2 else B.squeeze(0) - sca = scale_blocked_for_cublas(sc_A, m, k // _SF_VEC_SIZE, 0) - scb = scale_blocked_for_cublas(sc_B, n, k // _SF_VEC_SIZE, 0) - out = torch._scaled_mm( - a2d, - b2d, - scale_a=sca, - scale_b=scb, - out_dtype=out_dtype, - ) - return out if was_2d else out.unsqueeze(0) - - -def mxfp8_gemm_ref( - A: Tensor, - B: Tensor, - A_scale: Tensor, - B_scale: Tensor, - out_dtype: torch.dtype = torch.bfloat16, -) -> Tensor: - """Dequantize + plain matmul reference. A=(M,K), B=(K,N).""" - was_2d = A.dim() == 2 - # (l, m, k) - A3 = _as_3d(A, A.dim()).float() - # B is (K, N)/(L, K, N); flip to (l, n, k) for dequant by last-dim - B3 = _as_3d(B, B.dim()).mT.contiguous().float() - as3 = _as_3d(A_scale, A_scale.dim()).float() - bs3 = _as_3d(B_scale, B_scale.dim()).mT.contiguous().float() - a_dq = A3 * as3.repeat_interleave(_SF_VEC_SIZE, dim=-1) - b_dq = B3 * bs3.repeat_interleave(_SF_VEC_SIZE, dim=-1) - out3 = torch.einsum("lmk,lnk->lmn", a_dq, b_dq).to(out_dtype) - return out3.squeeze(0) if was_2d else out3 diff --git a/build/torch-cuda/quack/gemm_config.py b/build/torch-cuda/quack/gemm_config.py index d989a7dbe0a901343f83b65d5d2ca1a1cd749f1a..6718d1eabbe76f07408e67b8ebaa201ced4f92dd 100644 --- a/build/torch-cuda/quack/gemm_config.py +++ b/build/torch-cuda/quack/gemm_config.py @@ -1,4 +1,4 @@ -# Copyright (C) 2025, Fri Dao. +# Copyright (C) 2025, Tri Dao. import itertools from typing import Optional, List from functools import partial @@ -9,11 +9,14 @@ from dataclasses import dataclass class GemmConfig: tile_m: int = 128 tile_n: int = 192 + tile_k: int | None = None + num_warps: int | None = None pingpong: bool = True # by default, we use dynamic persistent tile scheduler on SM100 but not on SM90 is_dynamic_persistent: bool = True cluster_m: int = 2 cluster_n: int = 1 + cluster_k: int = 1 swap_ab: bool = False # raster_order: int = 1 max_swizzle_size: int = 8 @@ -70,6 +73,39 @@ def _get_sm90_configs( ] +def _get_sm80_configs() -> List[GemmConfig]: + tile_mn_warps_vals = [ + (128, 128, 4), + (128, 128, 8), + (128, 160, 4), + # TODO: Make 128x160 work with 8 warps. It currently makes the accumulator + # N layout odd and fails epilogue retile. + (128, 192, 4), + (128, 192, 8), + (128, 256, 8), + (128, 64, 4), + (64, 128, 4), + ] + return [ + GemmConfig( + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + num_warps=num_warps, + pingpong=False, + cluster_m=1, + cluster_n=1, + swap_ab=swap_ab, + device_capacity=8, + is_dynamic_persistent=False, + use_tma_gather=False, + ) + for (tile_m, tile_n, num_warps), tile_k, swap_ab in itertools.product( + tile_mn_warps_vals, [32, 64], [False, True] + ) + ] + + def _get_sm100_configs( epilogue: Optional[str] = None, ) -> List[GemmConfig]: @@ -141,14 +177,15 @@ def get_all_configs( epilogue: Optional[str] = None, tune_coop: bool = True, ) -> List[GemmConfig]: - """Return autotuning configs for all supported device capabilities (sm90 + sm100 + sm120). + """Return autotuning configs for all supported device capabilities. Each GemmConfig is tagged with its target device_capacity, so the caller can filter at runtime based on the actual device. This avoids querying the device (and initializing a CUDA context) at import time. """ return ( - _get_sm90_configs(epilogue, tune_coop) + _get_sm80_configs() + + _get_sm90_configs(epilogue, tune_coop) + _get_sm100_configs(epilogue) + _get_sm120_configs(epilogue, tune_coop) ) diff --git a/build/torch-cuda/quack/gemm_dact.py b/build/torch-cuda/quack/gemm_dact.py index 625ddfaa2b92cefcd6a8d9ba794b6217aa6645a3..5237bbe22be3e3cabfbd23c2b8f30b24f04c737e 100644 --- a/build/torch-cuda/quack/gemm_dact.py +++ b/build/torch-cuda/quack/gemm_dact.py @@ -8,15 +8,15 @@ from torch import Tensor import cutlass import cutlass.cute as cute from cutlass import Int32, Float32, const_expr +from .gemm_sm80 import GemmSm80 from .gemm_sm90 import GemmSm90 from .gemm_sm100 import GemmSm100 from .gemm_sm120 import GemmSm120 from .gemm_default_epi import GemmDefaultEpiMixin from .gemm_act import GemmActMixin -from .epi_ops import ColVecReduce, colvec_reduce_accumulate +from .epi_ops import ColVecLoad, ColVecReduce, Scalar, TileStore, colvec_reduce_accumulate from .compile_utils import make_fake_tensor as fake_tensor from .cute_dsl_utils import ( - ParamsBase, mlir_namedtuple, torch2cute_dtype_map, get_device_capacity, @@ -33,10 +33,10 @@ from .gemm_tvm_ffi_utils import ( make_fake_gemm_tensors, compile_gemm_kernel, ) -from .cache_utils import jit_cache +from .cache import jit_cache from .rounding import RoundingMode -from . import layout_utils as layout_utils from .activation import dact_fn_map, dgate_fn_map +from . import layout_utils class GemmDActMixin(GemmActMixin): @@ -51,36 +51,30 @@ class GemmDActMixin(GemmActMixin): epi_loop_tensors: Tuple[cute.Tensor, ...], tRS_rD: cute.Tensor, tRS_rC: Optional[cute.Tensor] = None, - ) -> Optional[cute.Tensor]: + ) -> Tuple[cute.Tensor, ...]: assert tRS_rC is not None # We don't add C to the accumulator GemmDefaultEpiMixin.epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC=None) - tRS_rC_acc = cute.make_rmem_tensor_like(tRS_rC, self.acc_dtype) - tRS_rC_acc.store(tRS_rC.load().to(self.acc_dtype)) + tRS_rC_acc = tRS_rC.to(self.acc_dtype) # If we don't have .shape here, the compiler generates local stores and loads if const_expr(params.act_fn is not None): - tRS_rPostAct = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype) - if const_expr(self.arch < 100): - for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True): - tRS_rD[i], tRS_rPostAct[i] = params.act_fn(tRS_rC_acc[i], tRS_rD[i]) - else: - for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True): - ( - (tRS_rD[2 * i], tRS_rD[2 * i + 1]), - (tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1]), - ) = params.act_fn( - (tRS_rC_acc[2 * i], tRS_rC_acc[2 * i + 1]), - (tRS_rD[2 * i], tRS_rD[2 * i + 1]), - ) + tRS_rAuxOut = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype) + vectorize = const_expr(self.arch == 100) + for i in cutlass.range(cute.size(tRS_rAuxOut), unroll_full=True, vectorize=vectorize): + tRS_rD[i], tRS_rAuxOut[i] = params.act_fn(tRS_rC_acc[i], tRS_rD[i]) else: - tRS_rPostAct = tRS_rC_acc - return tRS_rPostAct + tRS_rAuxOut = tRS_rC_acc + return (tRS_rAuxOut,) class GemmDActSm90(GemmDActMixin, GemmSm90): pass +class GemmDActSm80(GemmDActMixin, GemmSm80): + pass + + class GemmDActSm100(GemmDActMixin, GemmSm100): pass @@ -92,17 +86,18 @@ class GemmDActSm120(GemmDActMixin, GemmSm120): class GemmDGatedMixin(GemmActMixin): # Different from GemmActMixin, here act_bwd_fn must take in 3 arguments (x, y, dout) # and return 3 arguments (dx, dy, out) - _epi_ops = (*GemmActMixin._epi_ops, ColVecReduce("mColVecReduce")) + _epi_ops = ( + ColVecLoad("mColVecBroadcast"), + Scalar("sr_seed", dtype=Int32), + TileStore("mAuxOut"), + ColVecReduce("mColVecReduce"), + ) _extra_param_fields = (("act_bwd_fn", cutlass.Constexpr, None),) - _epi_param_bases = (ParamsBase,) @mlir_namedtuple class EpilogueArguments(NamedTuple): - mPostAct: cute.Tensor + mAuxOut: cute.Tensor act_bwd_fn: cutlass.Constexpr[Callable] = None - alpha: Optional[Float32 | cute.Tensor] = None - beta: Optional[Float32 | cute.Tensor] = None - mRowVecBroadcast: Optional[cute.Tensor] = None mColVecBroadcast: Optional[cute.Tensor] = None mColVecReduce: Optional[cute.Tensor] = None rounding_mode: cutlass.Constexpr[int] = RoundingMode.RN @@ -117,9 +112,9 @@ class GemmDGatedMixin(GemmActMixin): assert self.d_dtype.width == 32, "D storage type must be 32 bit" assert self.c_dtype.width == 32, "C storage type must be 32 bit" self.rounding_mode = args.rounding_mode - self.postact_dtype = args.mPostAct.element_type - self.postact_layout = cutlass.utils.LayoutEnum.from_tensor(args.mPostAct) - self.cta_tile_shape_postact_mn = self.cta_tile_shape_mnk[:2] + self.aux_out_dtype = args.mAuxOut.element_type + self.aux_out_layout = cutlass.utils.LayoutEnum.from_tensor(args.mAuxOut) + self.cta_tile_shape_aux_out_mn = self.cta_tile_shape_mnk[:2] d = self._epi_ops_to_params_dict(args) d["act_bwd_fn"] = args.act_bwd_fn return self.EpilogueParams(**d) @@ -133,13 +128,9 @@ class GemmDGatedMixin(GemmActMixin): epi_loop_tensors: Tuple[cute.Tensor, ...], tRS_rD: cute.Tensor, tRS_rC: Optional[cute.Tensor] = None, - ) -> Optional[cute.Tensor]: - alpha = epi_loop_tensors["alpha"] - beta = epi_loop_tensors["beta"] - tDrRowVec = epi_loop_tensors["mRowVecBroadcast"] - tDrColVec = epi_loop_tensors["mColVecBroadcast"] - tDrColVecReduce = epi_loop_tensors["mColVecReduce"] - assert alpha is None and beta is None and tDrRowVec is None # We don't use these for now + ) -> Tuple[cute.Tensor, ...]: + tDrColVec = epi_loop_tensors.get("mColVecBroadcast") + tDrColVecReduce = epi_loop_tensors.get("mColVecReduce") assert tRS_rC is not None implicit_dtype = self.implicit_dtype assert implicit_dtype.width == 16, "GemmDGatedMixin only supports 16bit for now" @@ -150,7 +141,7 @@ class GemmDGatedMixin(GemmActMixin): tRS_rOut = cute.make_rmem_tensor_like(tRS_rD, Float32) tRS_rD_scaled = cute.make_rmem_tensor_like(tRS_rD) if const_expr(tDrColVec is not None): # Scale D by colvec - if const_expr(self.arch < 100): + if const_expr(self.arch != 100): tRS_rD_scaled.store(tRS_rD.load() * tDrColVec.load().to(tRS_rD.element_type)) else: tDrColVec_mn = layout_utils.convert_layout_zero_stride(tDrColVec, tDrColVec.layout) @@ -159,63 +150,45 @@ class GemmDGatedMixin(GemmActMixin): tRS_rD_scaled, tDrColVec.layout ) for m in cutlass.range(cute.size(tDrColVec_mn, mode=[0]), unroll_full=True): + scale = tDrColVec_mn[m, 0] for n in cutlass.range( - cute.size(tDrColVec_mn, mode=[1]) // 2, unroll_full=True + cute.size(tDrColVec_mn, mode=[1]), unroll_full=True, vectorize=True ): - ( - tRS_rD_scaled_mn[m, 2 * n], - tRS_rD_scaled_mn[m, 2 * n + 1], - ) = cute.arch.mul_packed_f32x2( - (tRS_rD_mn[m, 2 * n], tRS_rD_mn[m, 2 * n + 1]), - (tDrColVec_mn[m, 0], tDrColVec_mn[m, 0]), - ) + tRS_rD_scaled_mn[m, n] = tRS_rD_mn[m, n] * scale else: tRS_rD_scaled.store(tRS_rD.load()) - if const_expr(self.arch < 100): - for i in cutlass.range(cute.size(tRS_rD)): - ( - tRS_rdXY_f32x2[2 * i], - tRS_rdXY_f32x2[2 * i + 1], - tRS_rOut[i], - ) = params.act_bwd_fn( - tRS_rXY_f32x2[2 * i], tRS_rXY_f32x2[2 * i + 1], tRS_rD_scaled[i] - ) - else: - for i in cutlass.range(cute.size(tRS_rD) // 2): - ( - (tRS_rdXY_f32x2[4 * i], tRS_rdXY_f32x2[4 * i + 2]), - (tRS_rdXY_f32x2[4 * i + 1], tRS_rdXY_f32x2[4 * i + 3]), - (tRS_rOut[2 * i], tRS_rOut[2 * i + 1]), - ) = params.act_bwd_fn( - (tRS_rXY_f32x2[4 * i], tRS_rXY_f32x2[4 * i + 2]), - (tRS_rXY_f32x2[4 * i + 1], tRS_rXY_f32x2[4 * i + 3]), - (tRS_rD_scaled[2 * i], tRS_rD_scaled[2 * i + 1]), - ) + tRS_rXY_pair = cute.flat_divide(tRS_rXY_f32x2, cute.make_layout(2)) + tRS_rX = tRS_rXY_pair[0, ...] + tRS_rY = tRS_rXY_pair[1, ...] + tRS_rdXY_pair = cute.flat_divide(tRS_rdXY_f32x2, cute.make_layout(2)) + tRS_rdX = tRS_rdXY_pair[0, ...] + tRS_rdY = tRS_rdXY_pair[1, ...] + vectorize = const_expr(self.arch == 100) + for i in cutlass.range(cute.size(tRS_rD), vectorize=vectorize): + tRS_rdX[i], tRS_rdY[i], tRS_rOut[i] = params.act_bwd_fn( + tRS_rX[i], tRS_rY[i], tRS_rD_scaled[i] + ) if const_expr(tDrColVecReduce is not None): # Accumulate postact * dout before D is scaled by colvec_scale colvec_reduce_accumulate(self, tDrColVecReduce, tRS_rOut, rScale=tRS_rD) if const_expr(tDrColVec is not None): # Scale Out by colvec - if const_expr(self.arch < 100): + if const_expr(self.arch != 100): tRS_rOut.store(tRS_rOut.load() * tDrColVec.load().to(tRS_rD.element_type)) else: tDrColVec_mn = layout_utils.convert_layout_zero_stride(tDrColVec, tDrColVec.layout) tRS_rOut_mn = layout_utils.convert_layout_zero_stride(tRS_rOut, tDrColVec.layout) for m in cutlass.range(cute.size(tDrColVec_mn, mode=[0]), unroll_full=True): + scale = tDrColVec_mn[m, 0] for n in cutlass.range( - cute.size(tDrColVec_mn, mode=[1]) // 2, unroll_full=True + cute.size(tDrColVec_mn, mode=[1]), unroll_full=True, vectorize=True ): - tRS_rOut_mn[m, 2 * n], tRS_rOut_mn[m, 2 * n + 1] = ( - cute.arch.mul_packed_f32x2( - (tRS_rOut_mn[m, 2 * n], tRS_rOut_mn[m, 2 * n + 1]), - (tDrColVec_mn[m, 0], tDrColVec_mn[m, 0]), - ) - ) + tRS_rOut_mn[m, n] = tRS_rOut_mn[m, n] * scale # Type conversion tRS_rdXY_f16x2 = cute.make_rmem_tensor(tRS_rdXY_f32x2.layout, implicit_dtype) tRS_rdXY_f16x2.store(tRS_rdXY_f32x2.load().to(implicit_dtype)) tRS_rD.store(cute.recast_tensor(tRS_rdXY_f16x2, Float32).load()) - return tRS_rOut + return (tRS_rOut,) # epi_end is inherited from ComposableEpiMixin → delegates to ColVecReduce.end() @@ -224,6 +197,10 @@ class GemmDGatedSm90(GemmDGatedMixin, GemmSm90): pass +class GemmDGatedSm80(GemmDGatedMixin, GemmSm80): + pass + + class GemmDGatedSm100(GemmDGatedMixin, GemmSm100): pass @@ -263,16 +240,21 @@ def _compile_gemm_dact( ): is_dgated = gemm_cls_name == "dgated" sm_to_cls = { - "dact": {9: GemmDActSm90, 10: GemmDActSm100, 11: GemmDActSm100, 12: GemmDActSm120}, + "dact": { + 8: GemmDActSm80, + 9: GemmDActSm90, + 10: GemmDActSm100, + 11: GemmDActSm100, + 12: GemmDActSm120, + }, "dgated": { + 8: GemmDGatedSm80, 9: GemmDGatedSm90, 10: GemmDGatedSm100, 11: GemmDGatedSm100, 12: GemmDGatedSm120, }, } - if device_capacity[0] == 12 and gemm_cls_name == "dact": - raise NotImplementedError("SM120 non-gated dactivation GEMM epilogue is not yet supported") GemmCls = sm_to_cls[gemm_cls_name][device_capacity[0]] mA, mB, mD, mC, m, n, k, l = make_fake_gemm_tensors( a_dtype, @@ -289,7 +271,7 @@ def _compile_gemm_dact( div_pa = div_for_dtype(postact_dtype) pa_leading = 1 if postact_major == "n" else 0 pa_shape = (m, n) if varlen_m else (m, n, l) - mPostAct = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading, divisibility=div_pa) + mAuxOut = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading, divisibility=div_pa) if is_dgated: act_fn = dgate_fn_map[activation] @@ -316,7 +298,7 @@ def _compile_gemm_dact( divisibility=1, ) epi_args = GemmCls.EpilogueArguments( - mPostAct, + mAuxOut, act_fn, mColVecBroadcast=mColVec, mColVecReduce=mColVecReduce, @@ -328,7 +310,7 @@ def _compile_gemm_dact( post_init = _set_implicit_dtype else: act_fn = dact_fn_map[activation] - epi_args = GemmCls.EpilogueArguments(mPostAct, act_fn) + epi_args = GemmCls.EpilogueArguments(mAuxOut, act_fn) post_init = None scheduler_args = make_fake_scheduler_args( @@ -369,6 +351,7 @@ def gemm_dact( tile_N: int, cluster_M: int, cluster_N: int, + tile_K: int | None = None, pingpong: bool = True, persistent: bool = True, is_dynamic_persistent: bool = False, @@ -432,7 +415,9 @@ def gemm_dact( postact_dtype = torch2cute_dtype_map[PostAct.dtype] device_capacity = get_device_capacity(A.device) - assert device_capacity[0] in [9, 10, 11, 12], "Only SM90, SM100, SM110, and SM120 are supported" + assert device_capacity[0] in [8, 9, 10, 11, 12], ( + "Only SM8x, SM90, SM100, SM110, and SM120 are supported" + ) if is_dynamic_persistent and device_capacity[0] == 9: assert tile_count_semaphore is not None, ( @@ -451,7 +436,7 @@ def gemm_dact( d_major, c_major, postact_major, - (tile_M, tile_N), + (tile_M, tile_N, tile_K) if tile_K is not None else (tile_M, tile_N), (cluster_M, cluster_N, 1), pingpong, persistent, @@ -468,11 +453,6 @@ def gemm_dact( use_tma_gather=use_tma_gather, ) - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY: - return - max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0 if is_dgated: epi_args = GemmDGatedMixin.EpilogueArguments( @@ -498,11 +478,9 @@ def gemm_dact( varlen_args = make_varlen_args(cu_seqlens_m, None, A_idx) if device_capacity[0] in [10, 11]: - compiled_fn( - A_p, B_p, Out_p, PreAct_p, epi_args, scheduler_args, varlen_args, None, None, None - ) + compiled_fn(A_p, B_p, Out_p, PreAct_p, epi_args, scheduler_args, varlen_args, None, None) else: - compiled_fn(A_p, B_p, Out_p, PreAct_p, epi_args, scheduler_args, varlen_args, None) + compiled_fn(A_p, B_p, Out_p, PreAct_p, epi_args, scheduler_args, varlen_args) gemm_dgated = gemm_dact diff --git a/build/torch-cuda/quack/gemm_default_epi.py b/build/torch-cuda/quack/gemm_default_epi.py index edb2612a1c709f531c397468649ce6dbf03c51cc..0c115d446fe3d3d31fe8fab86988431253396eb6 100644 --- a/build/torch-cuda/quack/gemm_default_epi.py +++ b/build/torch-cuda/quack/gemm_default_epi.py @@ -1,5 +1,5 @@ # Copyright (c) 2025, Wentao Guo, Tri Dao. -from typing import NamedTuple, Optional +from typing import NamedTuple, Optional, Tuple import cutlass import cutlass.cute as cute @@ -8,6 +8,7 @@ from cutlass import Int32, Float32, const_expr from .cute_dsl_utils import mlir_namedtuple from .epi_composable import ComposableEpiMixin from .epi_ops import Scalar, RowVecLoad, ColVecLoad +from .gemm_sm80 import GemmSm80 from .gemm_sm90 import GemmSm90 from .gemm_sm100 import GemmSm100 from .gemm_sm120 import GemmSm120 @@ -41,7 +42,7 @@ class GemmDefaultEpiMixin(ComposableEpiMixin): self.rounding_mode = args.rounding_mode d = self._epi_ops_to_params_dict(args) for key in ("mRowVecBroadcast", "mColVecBroadcast"): - if key in self.concat_layout and key in d and d[key] is not None: + if key in self.concat_layout and key in d: d[key] = layout_utils.concat_to_interleave(d[key], 1) return self.EpilogueParams(**d) @@ -52,11 +53,18 @@ class GemmDefaultEpiMixin(ComposableEpiMixin): epi_loop_tensors, tRS_rD: cute.Tensor, tRS_rC: Optional[cute.Tensor] = None, - ) -> Optional[cute.Tensor]: - alpha = epi_loop_tensors["alpha"] - beta = epi_loop_tensors["beta"] - tDrRowVec = epi_loop_tensors["mRowVecBroadcast"] - tDrColVec = epi_loop_tensors["mColVecBroadcast"] + ) -> Tuple[cute.Tensor, ...]: + """Return a tuple of register tensors (one per aux output). + + The returned tuple must be the same length as the tuple returned + from :meth:`epi_setup_aux_out`. The default impl returns ``()`` — + no aux outputs. + """ + # Use .get(): inactive ops are filtered out of epi_loop_tensors. + alpha = epi_loop_tensors.get("alpha") + beta = epi_loop_tensors.get("beta") + tDrRowVec = epi_loop_tensors.get("mRowVecBroadcast") + tDrColVec = epi_loop_tensors.get("mColVecBroadcast") rD = tRS_rD.load() # Apply alpha scaling to accumulator if alpha is provided (not None) if const_expr(hasattr(params, "alpha") and params.alpha is not None): @@ -77,27 +85,11 @@ class GemmDefaultEpiMixin(ComposableEpiMixin): if const_expr(tDrColVec is not None): for i in cutlass.range(cute.size(tDrColVec), unroll_full=True): tRS_rD[i] += tDrColVec[i] - return None + return () - def epi_setup_postact( - self, - params, - epi_smem_tensors, - tiled_copy_r2s, - tiled_copy_t2r, - tile_coord_mnkl, - varlen_manager, - tidx, - ): - """Returns None — default epilogue has no postact output.""" - return None - @cute.jit - def epi_convert_postact( - self, tRS_rPostAct, sr_seed, tidx, tile_coord_mnkl, num_prev_subtiles, epi_idx - ): - """Convert postact from acc_dtype to output dtype. Override for custom postprocessing.""" - return tRS_rPostAct +class GemmDefaultSm80(GemmDefaultEpiMixin, GemmSm80): + pass class GemmDefaultSm90(GemmDefaultEpiMixin, GemmSm90): diff --git a/build/torch-cuda/quack/gemm_interface.py b/build/torch-cuda/quack/gemm_interface.py index ce710b47aa6c85c87421c185caf43b783dbe0571..310686152b1576c82cc96d3eb33c6ceb64f6301b 100644 --- a/build/torch-cuda/quack/gemm_interface.py +++ b/build/torch-cuda/quack/gemm_interface.py @@ -3,7 +3,7 @@ from typing import Optional, Tuple, Literal from functools import partial import torch -from ._ops_compat import add_quack_op_namespace_prefix +from ._ops_compat import add_op_namespace_prefix import torch.nn.functional as F from torch import Tensor @@ -21,12 +21,49 @@ from .rms_final_reduce import rms_final_reduce from .rounding import RoundingMode +def _empty_k_matmul_into( + out: Tensor, + *, + bias: Optional[Tensor] = None, + C: Optional[Tensor] = None, + beta: float | Tensor = 1.0, +) -> None: + """K=0 fast path: write `beta * C + bias` (or zero if neither) into `out`. + + Used by every gemm-flavored wrapper to skip a kernel launch when the + contraction dim is empty. The matmul A @ B contributes zero, so the only + remaining terms are the C term and the (broadcast) bias. + """ + if C is not None: + if isinstance(beta, float) and beta == 1.0: + out.copy_(C) + else: + torch.mul(C, beta, out=out) + else: + out.zero_() + if bias is not None: + out += bias + + +def _silu_tanh(x: Tensor) -> Tensor: + x_half = 0.5 * x + return x_half * torch.tanh(x_half) + x_half + + +def _swiglu_oai_tanh(gate: Tensor, up: Tensor, alpha: float = 1.702) -> Tensor: + gate_half = 0.5 * gate + return (gate_half * torch.tanh(alpha * gate_half) + gate_half) * (up + 1) + + # Dictionary mapping activation names to PyTorch functions act_to_pytorch_fn_map = { None: lambda x: x, + "silu": F.silu, + "silu-tanh": _silu_tanh, "relu": F.relu, "relu_sq": lambda x: F.relu(x).square(), "gelu_tanh_approx": partial(F.gelu, approximate="tanh"), + "tanh": torch.tanh, } @@ -34,22 +71,37 @@ act_to_pytorch_fn_map = { # Each function takes (gate, up) and returns postact gated_to_pytorch_fn_map = { "swiglu": lambda gate, up: F.silu(gate) * up, + "swiglu-tanh": lambda gate, up: _silu_tanh(gate) * up, "swiglu_oai": lambda gate, up: gate * torch.sigmoid(1.702 * gate) * (up + 1), + "swiglu_oai-tanh": _swiglu_oai_tanh, "reglu": lambda gate, up: F.relu(gate) * up, "geglu": lambda gate, up: F.gelu(gate, approximate="tanh") * up, "glu": lambda gate, up: torch.sigmoid(gate) * up, } -ActActivation = Literal[None, "relu", "relu_sq", "gelu_tanh_approx"] -GatedActivation = Literal["swiglu", "swiglu_oai", "reglu", "geglu", "glu"] +ActActivation = Literal[None, "silu", "silu-tanh", "relu", "relu_sq", "gelu_tanh_approx", "tanh"] +GatedActivation = Literal[ + "swiglu", + "swiglu-tanh", + "swiglu_oai", + "swiglu_oai-tanh", + "reglu", + "geglu", + "glu", +] Activation = Literal[ None, + "silu", + "silu-tanh", "relu", "relu_sq", "gelu_tanh_approx", + "tanh", "swiglu", + "swiglu-tanh", "swiglu_oai", + "swiglu_oai-tanh", "reglu", "geglu", "glu", @@ -68,9 +120,88 @@ def _concat_interleave_bias(t): return t.unflatten(-1, (2, half)).transpose(-2, -1).flatten(-2, -1) +# ── Blockscaled (MXFP8 / MXFP4 / NVFP4) helpers ───────────────────────────── +# +# A and B may be passed as ``(data, scale_factor)`` tuples. Scale factors use +# the canonical cuBLAS/CUTLASS 128x4 blocked layout: shape +# ``(M // 128, K // VEC // 4, 32, 4, 4)`` (optionally with a leading batch L), +# where the ``(32, 4, 4)`` inner block is ``(m % 32, (m // 32) % 4, k_block % 4)`` +# with strides ``(16, 4, 1)`` — one contiguous 512-byte atom per 128 rows x 4 +# K-blocks, matching torchao's ``to_blocked`` and ``torch._scaled_mm``. +# VEC (the quantization block along K) is implied by the SF dtype: +# float8_e8m0fnu -> 32 (MX formats), float8_e4m3fn -> 16 (NVFP4). +# fp4 operands use ``torch.float4_e2m1fn_x2`` storage: shapes carry packed K +# (two elements per byte); K here always refers to logical K. + + +def _unpack_operand(X) -> Tuple[Tensor, Optional[Tensor]]: + """Split an ``A`` / ``B`` argument into (data, scale_factor or None).""" + if isinstance(X, (tuple, list)): + data, sf = X + return data, sf + return X, None + + +def _sf_normalize(SF: Tensor, name: str) -> Tensor: + """Validate a user scale-factor tensor and add the batch dim if missing. + + Requires ``(rm, rk, 32, 4, 4)`` or ``(L, rm, rk, 32, 4, 4)`` with the inner + ``(32, 4, 4)`` block contiguous (strides ``(16, 4, 1)`` — one 512 B atom); + outer strides are free, so slices of a larger scale buffer are accepted. + Returns a zero-copy ``(L, rm, rk, 32, 4, 4)`` view, which is what the + compiled kernel consumes directly (the kernel reads only the base pointer + and the outer strides; the inner atom layout is hardware-fixed). + """ + assert SF.ndim in (5, 6) and tuple(SF.shape[-3:]) == (32, 4, 4), ( + f"{name}: expected (rm, rk, 32, 4, 4) or (L, rm, rk, 32, 4, 4) blocked scale factors, " + f"got shape {tuple(SF.shape)}" + ) + assert SF.stride()[-3:] == (16, 4, 1), ( + f"{name}: inner (32, 4, 4) block must be contiguous with strides (16, 4, 1), " + f"got {SF.stride()[-3:]}" + ) + return SF.unsqueeze(0) if SF.ndim == 5 else SF + + +def _logical_k(X: Tensor) -> int: + """Logical contraction extent of the last dim (fp4 packs two elements per byte).""" + return X.shape[-1] * (2 if X.dtype == torch.float4_e2m1fn_x2 else 1) + + +def _sf_encode(SF: Tensor) -> Tensor: + """View e8m0 scale factors as uint8 for the custom-op boundary. + + Upstream PyTorch bug: a ``float8_e8m0fnu`` input to any mutable custom op + makes Inductor's ``decompose_auto_functionalized`` pass fail with + "auto_functionalized_v2 was not removed" (e4m3 is unaffected), breaking + torch.compile. The uint8 view is zero-copy and unambiguous (uint8 == e8m0, + e4m3 stays itself); :func:`_sf_decode` restores the dtype inside the op. + """ + return SF.view(torch.uint8) if SF.dtype == torch.float8_e8m0fnu else SF + + +def _sf_decode(SF: Optional[Tensor]) -> Optional[Tensor]: + """Inverse of :func:`_sf_encode` (uint8 -> e8m0), applied inside op bodies.""" + if SF is not None and SF.dtype == torch.uint8: + SF = SF.view(torch.float8_e8m0fnu) + return SF + + def default_config(device): cap = get_device_capacity(device)[0] - if cap in [10, 11]: + if cap == 8: + return GemmConfig( + tile_m=128, + tile_n=128, + tile_k=32, + num_warps=4, + cluster_m=1, + cluster_n=1, + pingpong=False, + is_dynamic_persistent=False, + device_capacity=8, + ) + elif cap in [10, 11]: return GemmConfig( tile_m=256, tile_n=256, @@ -101,6 +232,31 @@ def default_config(device): ) +def blockscaled_default_config(m: int, n: int) -> GemmConfig: + """Default SM100 config for blockscaled GEMM. + + Large shapes use a (256, 256) tile: it makes num_acc_stage == 1, which turns + on ``overlap_accum_sf`` (a second TMEM accumulator stage) so the per-tile + scale-apply + TMEM drain overlaps the next tile's MMA instead of + serializing after it. + """ + if m >= 512 and n >= 256: + tile_m, tile_n, cluster = 256, 256, (2, 1) + elif m >= 512 and n >= 128: + tile_m, tile_n, cluster = 256, 128, (2, 1) + else: + tile_m, tile_n, cluster = 128, 128, (1, 1) + return GemmConfig( + tile_m=tile_m, + tile_n=tile_n, + cluster_m=cluster[0], + cluster_n=cluster[1], + pingpong=False, + is_dynamic_persistent=True, + device_capacity=10, + ) + + def nvmmh_config(A, B, device_capacity): """Use nvMatmulHeuristics to pick a config for pure GEMM (no varlen/gather/epilogue). @@ -130,6 +286,21 @@ def prune_invalid_gemm_configs(configs, named_args: dict, **kwargs): # use_tma_gather only valid when gather_A is active on SM100/SM110 if not gather_A or device_capacity not in [10, 11]: configs = [conf for conf in configs if not conf.kwargs["config"].use_tma_gather] + if kwargs.get("SFA", None) is not None: # blockscaled (SM100 tcgen05 MMA constraints) + + def _blockscaled_ok(c: GemmConfig) -> bool: + return ( + c.device_capacity in (10, 11) + and not c.swap_ab # untested with blockscaled; SFA/SFB would swap too + and c.tile_k is None # tile_k is derived from the MMA instruction + and c.tile_m in (128, 256) + and c.tile_n in (64, 128, 192, 256) + # SF multicast is limited to 4 CTAs per cluster dim + and c.cluster_m <= 4 + and c.cluster_n <= 4 + ) + + configs = [conf for conf in configs if _blockscaled_ok(conf.kwargs["config"])] return configs @@ -157,26 +328,39 @@ def gemm_tuned( rounding_mode: int = RoundingMode.RN, sr_seed: int | Tensor = 0, concat_layout: tuple | None = None, # tensors whose non-contiguous dim is concat [gate; up] + SFA: Optional[Tensor] = None, # (L, rm, rk, 32, 4, 4) blocked scale factors + SFB: Optional[Tensor] = None, # (L, rn, rk, 32, 4, 4) ) -> None: + blockscaled = SFA is not None + if blockscaled: + SFA, SFB = _sf_decode(SFA), _sf_decode(SFB) if config is None: - # Use nvMMH heuristic for pure GEMM (no varlen, no gather, no epilogue) - is_pure_gemm = ( - cu_seqlens_m is None - and cu_seqlens_k is None - and A_idx is None - and C is None - and bias is None - and not add_to_output - ) - if is_pure_gemm: - device_capacity = get_device_capacity(A.device)[0] - config = nvmmh_config(A, B, device_capacity) - if config is None: - config = default_config(A.device) + if blockscaled: + m = A.shape[-2] + config = blockscaled_default_config(m, B.shape[-1]) + else: + # Use nvMMH heuristic for pure GEMM (no varlen, no gather, no epilogue) + is_pure_gemm = ( + cu_seqlens_m is None + and cu_seqlens_k is None + and A_idx is None + and C is None + and bias is None + and not add_to_output + ) + if is_pure_gemm: + device_capacity = get_device_capacity(A.device)[0] + config = nvmmh_config(A, B, device_capacity) + if config is None: + config = default_config(A.device) varlen_m = cu_seqlens_m is not None varlen_k = cu_seqlens_k is not None varlen = varlen_m or varlen_k gather_A = A_idx is not None + if blockscaled: + assert not gather_A, "Blockscaled GEMM does not support gather_A yet" + assert not concat_layout, "Blockscaled GEMM does not support concat_layout" + assert not config.swap_ab, "Blockscaled GEMM does not support swap_ab yet" if gather_A: assert varlen, "gather_A requires either varlen_m or varlen_k" assert config.cluster_n == 1, "gather_A requires cluster_n=1" @@ -235,7 +419,8 @@ def gemm_tuned( config.tile_n, config.cluster_m, config.cluster_n, - config.pingpong, + config.cluster_k, + pingpong=config.pingpong, persistent=True, is_dynamic_persistent=dynamic_scheduler, max_swizzle_size=config.max_swizzle_size, @@ -252,6 +437,10 @@ def gemm_tuned( sr_seed=sr_seed, use_tma_gather=config.use_tma_gather, concat_layout=swapped_concat, + num_warps=config.num_warps, + tile_K=config.tile_k, + SFA=SFA, + SFB=SFB, ) @@ -274,12 +463,23 @@ def gemm_act_tuned( A_idx: Optional[Tensor] = None, # (total_M,) if gather_A with varlen_m dynamic_scheduler: bool = False, config: Optional[GemmConfig] = None, + SFA: Optional[Tensor] = None, # (L, rm, rk, 32, 4, 4) blocked scale factors + SFB: Optional[Tensor] = None, # (L, rn, rk, 32, 4, 4) ) -> None: + blockscaled = SFA is not None + if blockscaled: + SFA, SFB = _sf_decode(SFA), _sf_decode(SFB) if config is None: - config = default_config(A.device) + if blockscaled: + config = blockscaled_default_config(A.shape[-2], B.shape[-1]) + else: + config = default_config(A.device) varlen_m = cu_seqlens_m is not None if varlen_m: assert not config.swap_ab, "Variable-length sequences not supported with swap_ab" + if blockscaled: + assert not varlen_m and A_idx is None, "Blockscaled GEMM does not support varlen/gather yet" + assert not config.swap_ab, "Blockscaled GEMM does not support swap_ab yet" if A.ndim == 2 and not varlen_m: A = A.unsqueeze(0) # (1, M, K) B = B.mT # (N, K) or (L, N, K) @@ -315,7 +515,8 @@ def gemm_act_tuned( config.tile_n, config.cluster_m, config.cluster_n, - config.pingpong, + tile_K=config.tile_k, + pingpong=config.pingpong, persistent=True, is_dynamic_persistent=dynamic_scheduler, max_swizzle_size=config.max_swizzle_size, @@ -324,6 +525,8 @@ def gemm_act_tuned( cu_seqlens_m=cu_seqlens_m, A_idx=A_idx, use_tma_gather=config.use_tma_gather, + SFA=SFA, + SFB=SFB, ) @@ -383,7 +586,8 @@ def gemm_dact_tuned( config.tile_n, config.cluster_m, config.cluster_n, - config.pingpong, + tile_K=config.tile_k, + pingpong=config.pingpong, persistent=True, is_dynamic_persistent=dynamic_scheduler, max_swizzle_size=config.max_swizzle_size, @@ -395,8 +599,10 @@ def gemm_dact_tuned( def gemm( # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (M, total_K) if varlen_k or (whatever, K) if gather_A with varlen_m or (M, whatever) if gather_A with varlen_k - A: Tensor, - B: Tensor, # (K, N) or (L, K, N) or (total_K, N) if varlen_k + # For blockscaled (MXFP8/MXFP4/NVFP4): a tuple (A, SFA) with A fp8/fp4 and SFA the + # blocked scale factors (rm, rk, 32, 4, 4) or (L, rm, rk, 32, 4, 4) — see helpers above. + A: Tensor | Tuple[Tensor, Tensor], + B: Tensor | Tuple[Tensor, Tensor], # (K, N) or (L, K, N) or (total_K, N) if varlen_k out: Optional[Tensor] = None, # (M, N) or (L, M, N) or (total_M, N) if varlen_m bias: Optional[Tensor] = None, # (N,) or (L, N) alpha: float | Tensor = 1.0, @@ -412,8 +618,16 @@ def gemm( concat_layout: tuple | None = None, # tensors whose non-contiguous dim is concat [gate; up] ) -> Tensor: """GEMM with optional output tensor and tuning control.""" + A, SFA = _unpack_operand(A) + B, SFB = _unpack_operand(B) + assert (SFA is None) == (SFB is None), "A and B must both (or neither) carry scale factors" + if SFA is not None: + SFA = _sf_encode(_sf_normalize(SFA, "SFA")) + SFB = _sf_encode(_sf_normalize(SFB, "SFB")) if out is None: - out_dtype = A.dtype if out_dtype is None else out_dtype + if out_dtype is None: + # Blockscaled inputs are fp8/fp4; default to bf16 output. + out_dtype = torch.bfloat16 if SFA is not None else A.dtype varlen_m = cu_seqlens_m is not None varlen_k = cu_seqlens_k is not None if varlen_m: @@ -428,6 +642,15 @@ def gemm( (A.shape[0], B.shape[-1]) if A.ndim == 2 else (A.shape[0], A.shape[-2], B.shape[-1]) ) out = torch.empty(out_shape, dtype=out_dtype, device=A.device) + # Empty-input fast path: skip kernel launch. + # M=0 / N=0 — the tile scheduler's ceil_div over a zero dim divides by zero. + # K=0 — the kernel rejects stride-0 inputs (stride must be divisible by 8); + # semantically the empty contraction yields a zero matrix. + if out.numel() == 0: + return out + if A.numel() == 0: + _empty_k_matmul_into(out, bias=bias) + return out alpha_tensor = alpha if not isinstance(alpha, float) else None alpha = alpha if isinstance(alpha, float) else 1.0 sr_seed_tensor = sr_seed if isinstance(sr_seed, Tensor) else None @@ -450,12 +673,14 @@ def gemm( sr_seed=sr_seed_int, sr_seed_tensor=sr_seed_tensor, concat_layout=concat_str, + SFA=SFA, + SFB=SFB, ) return out @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_out"), + add_op_namespace_prefix("gemm_out"), mutates_args=("out",), device_types="cuda", # We have to split out alpha and alpha_tensor since torch.library requires @@ -480,11 +705,18 @@ def gemm_out( sr_seed: int = 0, sr_seed_tensor: Optional[Tensor] = None, concat_layout: Optional[str] = None, + # Blockscaled scale factors, (L, rm/rn, rk, 32, 4, 4); tuples are unpacked + # to these flat args before the custom-op boundary since torch.library + # schemas have no (Tensor, Tensor) argument type. + SFA: Optional[Tensor] = None, + SFB: Optional[Tensor] = None, ) -> None: """GEMM with pre-allocated output tensor.""" fn = gemm_tuned if tuned else partial(gemm_tuned.fn, config=None) - alpha = alpha_tensor if alpha_tensor is not None else alpha - sr_seed_arg = sr_seed_tensor if sr_seed_tensor is not None else sr_seed + # Shared helpers: drift between this eager body and the register_fake side + # is structurally impossible because both call the same functions. + alpha = _merge_tensor(alpha, alpha_tensor) + sr_seed_arg = _merge_tensor(sr_seed, sr_seed_tensor) fn( A, B, @@ -499,7 +731,9 @@ def gemm_out( dynamic_scheduler=dynamic_scheduler, rounding_mode=rounding_mode, sr_seed=sr_seed_arg, - concat_layout=tuple(concat_layout.split(",")) if concat_layout else None, + concat_layout=_parse_concat_layout(concat_layout), + SFA=SFA, + SFB=SFB, ) @@ -572,10 +806,43 @@ def gemm_ref( return out +def gemm_blockscaled_ref( + A: Tuple[Tensor, Tensor], # ((M, K) or (L, M, K) fp8/fp4x2, blocked SFA) + B: Tuple[Tensor, Tensor], # ((K, N) or (L, K, N) fp8/fp4x2 K-contig, blocked SFB) + alpha: float | Tensor = 1.0, + out_dtype: torch.dtype = torch.bfloat16, +) -> Tensor: + """Dequantize-and-matmul reference for blockscaled GEMM.""" + from .blockscaled.utils import dequant_operand, unpack_scale_blocked_to_2d + + A, SFA = _unpack_operand(A) + B, SFB = _unpack_operand(B) + SFA = _sf_normalize(SFA, "SFA") + SFB = _sf_normalize(SFB, "SFB") + sf_vec = 32 if SFA.dtype == torch.float8_e8m0fnu else 16 + batched = A.ndim == 3 + a3 = A if batched else A.unsqueeze(0) # (l, m, k_packed) + b3 = (B if batched else B.unsqueeze(0)).mT # (l, n, k_packed) + a_val = dequant_operand(a3) # (l, m, k) fp32 + b_val = dequant_operand(b3) + l, m, k = a_val.shape + n = b_val.shape[1] + sfa = unpack_scale_blocked_to_2d(SFA, m, k // sf_vec).float() + sfb = unpack_scale_blocked_to_2d(SFB, n, k // sf_vec).float() + a_dq = a_val * sfa.repeat_interleave(sf_vec, dim=-1) + b_dq = b_val * sfb.repeat_interleave(sf_vec, dim=-1) + out = torch.einsum("lmk,lnk->lmn", a_dq, b_dq) + if not (isinstance(alpha, float) and alpha == 1.0): + out = out * alpha + out = out.to(out_dtype) + return out if batched else out.squeeze(0) + + def gemm_add( # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (M, total_K) if varlen_k or (whatever, K) if gather_A with varlen_m or (M, whatever) if gather_A with varlen_k - A: Tensor, - B: Tensor, # (K, N) or (L, K, N) or (total_K, N) if varlen_k + # For blockscaled: a tuple (A, SFA) — see gemm(). + A: Tensor | Tuple[Tensor, Tensor], + B: Tensor | Tuple[Tensor, Tensor], # (K, N) or (L, K, N) or (total_K, N) if varlen_k C: Tensor, # (M, N) or (L, M, N) or (total_M, N) if varlen_m or (L, M, N) if varlen_k out: Optional[Tensor] = None, # (M, N) or (L, M, N) or (total_M, N) if varlen_m alpha: float | Tensor = 1.0, @@ -590,8 +857,15 @@ def gemm_add( concat_layout: tuple | None = None, # tensors whose non-contiguous dim is concat [gate; up] ) -> Tensor: """GEMM with addition and optional output tensor.""" + A, SFA = _unpack_operand(A) + B, SFB = _unpack_operand(B) + assert (SFA is None) == (SFB is None), "A and B must both (or neither) carry scale factors" + if SFA is not None: + SFA = _sf_encode(_sf_normalize(SFA, "SFA")) + SFB = _sf_encode(_sf_normalize(SFB, "SFB")) if out is None: - out_dtype = A.dtype if out_dtype is None else out_dtype + if out_dtype is None: + out_dtype = torch.bfloat16 if SFA is not None else A.dtype varlen_m = cu_seqlens_m is not None varlen_k = cu_seqlens_k is not None if varlen_m: @@ -608,17 +882,26 @@ def gemm_add( ) out = torch.empty(out_shape, dtype=out_dtype, device=A.device) add_to_output = C is out and isinstance(beta, float) and beta == 1.0 and cu_seqlens_m is None + # Empty-input fast path: skip kernel launch (see gemm() for rationale). + # K=0 reduces D = alpha*A@B + beta*C to D = beta*C. + if out.numel() == 0: + return out + if A.numel() == 0: + if add_to_output: + return out # out IS C, and out += alpha * 0 is a no-op + _empty_k_matmul_into(out, C=C, beta=beta) + return out alpha_tensor = alpha if not isinstance(alpha, float) else None alpha = alpha if isinstance(alpha, float) else 1.0 beta_tensor = beta if not isinstance(beta, float) else None beta = beta if isinstance(beta, float) else 1.0 - alpha_arg = alpha_tensor if alpha_tensor is not None else alpha - beta_arg = beta_tensor if beta_tensor is not None else beta + alpha_arg = _merge_tensor(alpha, alpha_tensor) + beta_arg = _merge_tensor(beta, beta_tensor) concat_str = ",".join(concat_layout) if concat_layout else None if add_to_output: gemm_add_inplace( - A, - B, + A if SFA is None else (A, SFA), + B if SFB is None else (B, SFB), out, alpha=alpha_arg, beta=beta_arg, @@ -648,12 +931,14 @@ def gemm_add( dynamic_scheduler=dynamic_scheduler, tuned=tuned, concat_layout=concat_str, + SFA=SFA, + SFB=SFB, ) return out @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_add_out"), + add_op_namespace_prefix("gemm_add_out"), mutates_args=("out",), device_types="cuda", # We have to split out alpha and alpha_tensor since torch.library requires @@ -678,11 +963,13 @@ def gemm_add_out( dynamic_scheduler: bool = False, tuned: bool = True, concat_layout: Optional[str] = None, + SFA: Optional[Tensor] = None, # blocked scale factors, (L, rm, rk, 32, 4, 4) (see gemm_out) + SFB: Optional[Tensor] = None, ) -> None: """GEMM with addition and pre-allocated output tensor.""" fn = gemm_tuned if tuned else partial(gemm_tuned.fn, config=None) - alpha = alpha_tensor if alpha_tensor is not None else alpha - beta = beta_tensor if beta_tensor is not None else beta + alpha = _merge_tensor(alpha, alpha_tensor) + beta = _merge_tensor(beta, beta_tensor) fn( A, B, @@ -690,13 +977,15 @@ def gemm_add_out( C, alpha=alpha, beta=beta, + SFA=SFA, + SFB=SFB, cu_seqlens_m=cu_seqlens_m, cu_seqlens_k=cu_seqlens_k, A_idx=A_idx, batch_idx_permute=batch_idx_permute, add_to_output=add_to_output, dynamic_scheduler=dynamic_scheduler, - concat_layout=tuple(concat_layout.split(",")) if concat_layout else None, + concat_layout=_parse_concat_layout(concat_layout), ) @@ -783,8 +1072,9 @@ def gemm_add_ref( def gemm_add_inplace( # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (M, total_K) if varlen_k or (whatever, K) if gather_A with varlen_m or (M, whatever) if gather_A with varlen_k - A: Tensor, - B: Tensor, # (K, N) or (L, K, N) or (total_K, N) if varlen_k + # For blockscaled: a tuple (A, SFA) — see gemm(). + A: Tensor | Tuple[Tensor, Tensor], + B: Tensor | Tuple[Tensor, Tensor], # (K, N) or (L, K, N) or (total_K, N) if varlen_k out: Tensor, # (M, N) or (L, M, N) or (total_M, N) if varlen_m or (L, M, N) if varlen_k alpha: float | Tensor = 1.0, beta: float | Tensor = 1.0, @@ -808,10 +1098,24 @@ def gemm_add_inplace( dynamic_scheduler: Whether to use dynamic scheduler tuned: Whether to use autotuned configuration """ + A, SFA = _unpack_operand(A) + B, SFB = _unpack_operand(B) + assert (SFA is None) == (SFB is None), "A and B must both (or neither) carry scale factors" + if SFA is not None: + SFA = _sf_encode(_sf_normalize(SFA, "SFA")) + SFB = _sf_encode(_sf_normalize(SFB, "SFB")) alpha_tensor = alpha if not isinstance(alpha, float) else None alpha = alpha if isinstance(alpha, float) else 1.0 beta_tensor = beta if not isinstance(beta, float) else None beta = beta if isinstance(beta, float) else 1.0 + # Empty-input fast path: out += alpha * A@B with K=0 reduces to out *= beta. + # The matmul contributes zero, so use the helper with C=out. + if out.numel() == 0: + return + if A.numel() == 0: + if beta != 1.0 or beta_tensor is not None: + out.mul_(_merge_tensor(beta, beta_tensor)) + return gemm_add_inplace_op( A, B, @@ -829,11 +1133,13 @@ def gemm_add_inplace( concat_layout=",".join(concat_layout) if isinstance(concat_layout, tuple) else concat_layout, + SFA=SFA, + SFB=SFB, ) @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_add_inplace"), + add_op_namespace_prefix("gemm_add_inplace"), mutates_args=("out",), device_types="cuda", # We have to split out alpha and alpha_tensor since torch.library requires @@ -856,10 +1162,12 @@ def gemm_add_inplace_op( dynamic_scheduler: bool = False, tuned: bool = True, concat_layout: Optional[str] = None, + SFA: Optional[Tensor] = None, # blocked scale factors, (L, rm, rk, 32, 4, 4) (see gemm_out) + SFB: Optional[Tensor] = None, ) -> None: fn = gemm_tuned if tuned else partial(gemm_tuned.fn, config=None) - alpha = alpha_tensor if alpha_tensor is not None else alpha - beta = beta_tensor if beta_tensor is not None else beta + alpha = _merge_tensor(alpha, alpha_tensor) + beta = _merge_tensor(beta, beta_tensor) add_to_output = isinstance(beta, float) and beta == 1.0 and cu_seqlens_m is None # Use out as both input bias and output fn( @@ -875,13 +1183,19 @@ def gemm_add_inplace_op( batch_idx_permute=batch_idx_permute, add_to_output=add_to_output, dynamic_scheduler=dynamic_scheduler, - concat_layout=tuple(concat_layout.split(",")) if concat_layout else None, + concat_layout=_parse_concat_layout(concat_layout), + SFA=SFA, + SFB=SFB, ) def gemm_act( - A: Tensor, # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m - B: Tensor, # (K, N) or (L, K, N) + # For blockscaled: a tuple (A, SFA) — see gemm(). + A: Tensor + | Tuple[ + Tensor, Tensor + ], # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m + B: Tensor | Tuple[Tensor, Tensor], # (K, N) or (L, K, N) C: Optional[Tensor] = None, # (M, N) or (L, M, N) or (total_M, N) if varlen_m bias: Optional[Tensor] = None, # (N,) or (L, N) activation: Activation = None, @@ -897,9 +1211,16 @@ def gemm_act( concat_layout: tuple | None = None, # tensors whose non-contiguous dim is concat [gate; up] ) -> Tuple[Optional[Tensor], Tensor]: """GEMM with activation (or gated activation) and optional output tensors.""" + A, SFA = _unpack_operand(A) + B, SFB = _unpack_operand(B) + assert (SFA is None) == (SFB is None), "A and B must both (or neither) carry scale factors" + if SFA is not None: + SFA = _sf_encode(_sf_normalize(SFA, "SFA")) + SFB = _sf_encode(_sf_normalize(SFB, "SFB")) is_gated = activation in gated_to_pytorch_fn_map - out_dtype = A.dtype if out_dtype is None else out_dtype - postact_dtype = A.dtype if postact_dtype is None else postact_dtype + default_dtype = torch.bfloat16 if SFA is not None else A.dtype + out_dtype = default_dtype if out_dtype is None else out_dtype + postact_dtype = default_dtype if postact_dtype is None else postact_dtype varlen_m = cu_seqlens_m is not None # Determine output shape based on gather_A if varlen_m: @@ -914,6 +1235,14 @@ def gemm_act( preact_out = torch.empty(out_shape, dtype=out_dtype, device=A.device) if postact_out is None: postact_out = torch.empty(postact_shape, dtype=postact_dtype, device=A.device) + # Empty-input fast path. For M=0 or N=0 the outputs are empty; for K=0 + # (A@B == 0) the no-bias / no-C surface yields preact=0 and act(0)=0 for + # every supported activation, so both outputs are zero. + if postact_out.numel() == 0 or A.numel() == 0: + if preact_out is not None: + _empty_k_matmul_into(preact_out) + _empty_k_matmul_into(postact_out) + return preact_out, postact_out concat_str = ",".join(concat_layout) if concat_layout else None if is_gated: gemm_gated_out( @@ -929,6 +1258,8 @@ def gemm_act( dynamic_scheduler, tuned, concat_layout=concat_str, + SFA=SFA, + SFB=SFB, ) else: gemm_act_out( @@ -943,6 +1274,8 @@ def gemm_act( A_idx, dynamic_scheduler, tuned, + SFA=SFA, + SFB=SFB, ) return preact_out, postact_out @@ -951,10 +1284,10 @@ gemm_gated = gemm_act @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_act_out"), + add_op_namespace_prefix("gemm_act_out"), mutates_args=("preact_out", "postact_out"), device_types="cuda", - schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? bias=None, str? activation=None, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=False, bool tuned=True) -> ()", + schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? bias=None, str? activation=None, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=False, bool tuned=True, Tensor? SFA=None, Tensor? SFB=None) -> ()", ) def gemm_act_out( A: Tensor, # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m @@ -968,10 +1301,25 @@ def gemm_act_out( A_idx: Optional[Tensor] = None, # (total_M,) if gather_A with varlen_m dynamic_scheduler: bool = False, tuned: bool = True, + SFA: Optional[Tensor] = None, # blocked scale factors, (L, rm, rk, 32, 4, 4) (see gemm_out) + SFB: Optional[Tensor] = None, ) -> None: """GEMM with activation and pre-allocated output tensors.""" fn = gemm_act_tuned if tuned else partial(gemm_act_tuned.fn, config=None) - fn(A, B, preact_out, postact_out, C, bias, activation, cu_seqlens_m, A_idx, dynamic_scheduler) + fn( + A, + B, + preact_out, + postact_out, + C, + bias, + activation, + cu_seqlens_m, + A_idx, + dynamic_scheduler, + SFA=SFA, + SFB=SFB, + ) def gemm_act_ref( @@ -1048,6 +1396,16 @@ def gemm_dact( dx_out = torch.empty(out_shape, dtype=out_dtype, device=A.device) if postact_out is None: postact_out = torch.empty(postact_shape, dtype=postact_dtype, device=A.device) + # Empty-input fast path: M=0 / N=0 → outputs are empty; K=0 (A.numel()==0) + # makes the upstream GEMM gradient zero, so dx is zero regardless of activation. + if dx_out.numel() == 0 or A.numel() == 0: + _empty_k_matmul_into(dx_out) + _empty_k_matmul_into(postact_out) + results = [dx_out, postact_out] + if colvec_reduce: + colvec_shape = (*out_shape[:-1],) + results.append(torch.zeros(colvec_shape, dtype=torch.float32, device=A.device)) + return tuple(results) if is_dgated: colvec_reduce_final = gemm_dgated_out( A, @@ -1063,10 +1421,10 @@ def gemm_dact( dynamic_scheduler, tuned, ) - if not colvec_reduce: - return dx_out, postact_out - else: - return dx_out, postact_out, colvec_reduce_final + results = [dx_out, postact_out] + if colvec_reduce: + results.append(colvec_reduce_final) + return tuple(results) else: gemm_dact_out( A, @@ -1080,14 +1438,15 @@ def gemm_dact( dynamic_scheduler, tuned, ) - return dx_out, postact_out + results = [dx_out, postact_out] + return tuple(results) gemm_dgated = gemm_dact @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_dact_out"), + add_op_namespace_prefix("gemm_dact_out"), mutates_args=("dx_out", "postact_out"), device_types="cuda", schema="(Tensor A, Tensor B, Tensor PreAct, Tensor(a3!) dx_out, Tensor(a4!) postact_out, str? activation=None, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=True, bool tuned=True) -> ()", @@ -1152,11 +1511,27 @@ def gemm_dact_ref( gemm_dgated_ref = gemm_dact_ref +def _symmetric_gemm_config(sm: int) -> tuple[int, int, int, bool]: + configs = { + 8: (128, 128, 1, False), + 9: (128, 256, 2, False), + 10: (256, 256, 2, False), + 11: (256, 256, 2, False), + 12: (128, 128, 1, True), + } + if sm not in configs: + raise NotImplementedError( + "gemm_symmetric is only supported on SM8x, SM90, SM100, SM110, and SM120" + ) + return configs[sm] + + @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_symmetric_out"), + add_op_namespace_prefix("gemm_symmetric_out"), mutates_args=("out",), device_types="cuda", - schema="(Tensor A, Tensor B, Tensor(a2!) out, Tensor? C=None, bool dynamic_scheduler=False, float alpha=1.0, float beta=1.0) -> ()", + # alpha/beta split into float + Tensor pair because torch.library requires + # each schema arg to have a fixed type. See gemm_add_out for the pattern. ) def gemm_symmetric_out( A: Tensor, # (M, K) or (L, M, K) @@ -1166,8 +1541,12 @@ def gemm_symmetric_out( dynamic_scheduler: bool = False, alpha: float = 1.0, beta: float = 1.0, + alpha_tensor: Optional[Tensor] = None, + beta_tensor: Optional[Tensor] = None, ) -> None: """GEMM with guaranteed symmetric output.""" + alpha = _merge_tensor(alpha, alpha_tensor) + beta = _merge_tensor(beta, beta_tensor) if A.ndim == 2: A = A.unsqueeze(0) # (1, M, K) B = B.mT # (M, K) or (L, M, K) @@ -1184,12 +1563,7 @@ def gemm_symmetric_out( ) sm = get_device_capacity(A.device)[0] # We want square tile per cluster - tile_m, tile_n, cluster_m, pingpong = { - 9: (128, 256, 2, False), - 10: (256, 256, 2, False), - 11: (256, 256, 2, False), - 12: (128, 128, 1, True), - }[sm] + tile_m, tile_n, cluster_m, pingpong = _symmetric_gemm_config(sm) gemm_symmetric_dispatch( A, B, @@ -1229,11 +1603,29 @@ def gemm_symmetric( if out is None: out = torch.empty(out_shape, dtype=out_dtype, device=A.device) + alpha_tensor = alpha if not isinstance(alpha, float) else None alpha_val = alpha if isinstance(alpha, float) else 1.0 + beta_tensor = beta if not isinstance(beta, float) else None beta_val = beta if isinstance(beta, float) else 1.0 + # Empty-input fast path: out = alpha * A@A.T + beta * C reduces to beta * C + # when K=0 (or just zeros / empty for M=0). + if out.numel() == 0: + return out + if A.numel() == 0: + _empty_k_matmul_into(out, C=C, beta=beta) + return out + gemm_symmetric_out( - A, B, out, C, dynamic_scheduler=dynamic_scheduler, alpha=alpha_val, beta=beta_val + A, + B, + out, + C, + dynamic_scheduler=dynamic_scheduler, + alpha=alpha_val, + beta=beta_val, + alpha_tensor=alpha_tensor, + beta_tensor=beta_tensor, ) return out @@ -1258,12 +1650,24 @@ def gemm_gated_tuned( dynamic_scheduler: bool = False, config: Optional[GemmConfig] = None, concat_layout: tuple | None = None, # tensors whose non-contiguous dim is concat [gate; up] + SFA: Optional[Tensor] = None, # (L, rm, rk, 32, 4, 4) blocked scale factors + SFB: Optional[Tensor] = None, # (L, rn, rk, 32, 4, 4) ) -> None: + blockscaled = SFA is not None + if blockscaled: + SFA, SFB = _sf_decode(SFA), _sf_decode(SFB) if config is None: - config = default_config(A.device) + if blockscaled: + config = blockscaled_default_config(A.shape[-2], B.shape[-1]) + else: + config = default_config(A.device) varlen_m = cu_seqlens_m is not None if varlen_m: assert not config.swap_ab, "Variable-length sequences not supported with swap_ab" + if blockscaled: + assert not varlen_m and A_idx is None, "Blockscaled GEMM does not support varlen/gather yet" + assert not concat_layout, "Blockscaled GEMM does not support concat_layout" + assert not config.swap_ab, "Blockscaled GEMM does not support swap_ab yet" if A.ndim == 2 and not varlen_m: A = A.unsqueeze(0) # (1, M, K) B = B.mT # (N, K) or (L, N, K) @@ -1307,7 +1711,8 @@ def gemm_gated_tuned( config.tile_n, config.cluster_m, config.cluster_n, - config.pingpong, + tile_K=config.tile_k, + pingpong=config.pingpong, persistent=True, is_dynamic_persistent=dynamic_scheduler, max_swizzle_size=config.max_swizzle_size, @@ -1317,6 +1722,8 @@ def gemm_gated_tuned( A_idx=A_idx, use_tma_gather=config.use_tma_gather, concat_layout=concat_layout, + SFA=SFA, + SFB=SFB, ) @@ -1403,7 +1810,8 @@ def gemm_dgated_tuned( config.tile_n, config.cluster_m, config.cluster_n, - config.pingpong, + tile_K=config.tile_k, + pingpong=config.pingpong, persistent=True, is_dynamic_persistent=dynamic_scheduler, max_swizzle_size=config.max_swizzle_size, @@ -1423,10 +1831,10 @@ def gemm_dgated_tuned( @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_gated_out"), + add_op_namespace_prefix("gemm_gated_out"), mutates_args=("preact_out", "postact_out"), device_types="cuda", - schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? bias=None, str activation='swiglu', Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=False, bool tuned=True, str? concat_layout=None) -> ()", + schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? bias=None, str activation='swiglu', Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=False, bool tuned=True, str? concat_layout=None, Tensor? SFA=None, Tensor? SFB=None) -> ()", ) def gemm_gated_out( A: Tensor, # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m @@ -1441,6 +1849,8 @@ def gemm_gated_out( dynamic_scheduler: bool = False, tuned: bool = True, concat_layout: Optional[str] = None, + SFA: Optional[Tensor] = None, # blocked scale factors, (L, rm, rk, 32, 4, 4) (see gemm_out) + SFB: Optional[Tensor] = None, ) -> None: """GEMM with gated activation and pre-allocated output tensors.""" fn = gemm_gated_tuned if tuned else partial(gemm_gated_tuned.fn, config=None) @@ -1455,12 +1865,14 @@ def gemm_gated_out( cu_seqlens_m, A_idx, dynamic_scheduler, - concat_layout=tuple(concat_layout.split(",")) if concat_layout else None, + concat_layout=_parse_concat_layout(concat_layout), + SFA=SFA, + SFB=SFB, ) @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_dgated_out"), + add_op_namespace_prefix("gemm_dgated_out"), mutates_args=("dx_out", "postact_out"), device_types="cuda", schema="(Tensor A, Tensor B, Tensor PreAct, Tensor(a!) dx_out, Tensor(b!) postact_out, Tensor? colvec_scale=None, str activation='swiglu', bool colvec_reduce=False, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=True, bool tuned=True) -> Tensor", @@ -1499,7 +1911,7 @@ def gemm_dgated_out( return result -@torch.library.register_fake(add_quack_op_namespace_prefix("gemm_dgated_out")) +@torch.library.register_fake(add_op_namespace_prefix("gemm_dgated_out")) def gemm_dgated_out_fake( A: Tensor, B: Tensor, @@ -1514,20 +1926,6 @@ def gemm_dgated_out_fake( dynamic_scheduler: bool = True, tuned: bool = True, ) -> Tensor: - _precompile_default_config( - gemm_dgated_tuned, - A, - B, - PreAct, - dx_out, - postact_out, - colvec_scale=colvec_scale, - activation=activation, - colvec_reduce=colvec_reduce, - cu_seqlens_m=cu_seqlens_m, - A_idx=A_idx, - dynamic_scheduler=dynamic_scheduler, - ) if not colvec_reduce: return torch.empty(0, dtype=torch.float32, device=A.device) else: @@ -1541,25 +1939,6 @@ def gemm_dgated_out_fake( return torch.empty(out_shape, dtype=torch.float32, device=A.device) -def _precompile_default_config(autotuned_fn, *args, **kwargs): - """Compile the default config in COMPILE_ONLY mode. - - Checks COMPILE_ONLY flag and SymInt guard, then calls the unwrapped function with - config=None (which selects the default config), triggering compilation (exports .o) - without benchmarking or kernel launch. - Tests use tuned=False which also selects the default config, so this is sufficient. - """ - from .cache_utils import COMPILE_ONLY - - A = args[0] if args else kwargs.get("A") - if not COMPILE_ONLY or A is None or isinstance(A.shape[0], torch.SymInt): - return - try: - autotuned_fn.fn(*args, config=None, **kwargs) - except Exception: - pass - - @gemm_add_inplace_op.register_fake def gemm_add_inplace_fake( A: Tensor, @@ -1575,73 +1954,71 @@ def gemm_add_inplace_fake( batch_idx_permute: Optional[Tensor] = None, dynamic_scheduler: bool = False, tuned: bool = True, + concat_layout: Optional[str] = None, + SFA: Optional[Tensor] = None, + SFB: Optional[Tensor] = None, ) -> None: - alpha_val = alpha_tensor if alpha_tensor is not None else alpha - beta_val = beta_tensor if beta_tensor is not None else beta - add_to_output = isinstance(beta_val, float) and beta_val == 1.0 and cu_seqlens_m is None - _precompile_default_config( - gemm_tuned, - A, - B, - out, - out if not add_to_output else None, - alpha=alpha_val, - beta=beta_val, - cu_seqlens_m=cu_seqlens_m, - cu_seqlens_k=cu_seqlens_k, - A_idx=A_idx, - batch_idx_permute=batch_idx_permute, - add_to_output=add_to_output, - dynamic_scheduler=dynamic_scheduler, - ) + # Pure no-op: the op only mutates ``out``; kernel compilation is owned + # by jit_cache + the async compile pool at real execution time. + return + + +# --------------------------------------------------------------------------- +# Shared schema-split helpers. +# +# torch.library.custom_op requires a concrete type per arg, so union-typed +# autotuned args (e.g. ``alpha: Union[float, Tensor]``, ``sr_seed: Union[int, +# Tensor]``) are split into two fixed-typed schema kwargs at the custom_op +# boundary (``alpha: float`` + ``alpha_tensor: Optional[Tensor]``). The eager +# bodies merge them back into the unified form via :func:`_merge_tensor` +# before calling the autotuned fn. +# --------------------------------------------------------------------------- + + +def _merge_tensor(value, tensor_value): + """Return ``tensor_value`` if non-None, else ``value``. + + Single source of truth for the ``Union[scalar, Tensor]`` schema-split + merge. Used both inside eager bodies (where ``value = alpha, + tensor_value = alpha_tensor``) and inside the fake path (which derives + the split pairs from the custom_op signature). + """ + return tensor_value if tensor_value is not None else value -def _register_precompile_fake(custom_op, autotuned_fn, rewrite=None): - """Register a fake that precompiles the default config in COMPILE_ONLY mode. +def _parse_concat_layout(value): + """Coerce ``concat_layout`` from schema form (``Optional[str]``) to + autotuned form (``Optional[tuple[str, ...]]``). - For custom_ops that forward args to their autotuned fn. Binds all args by name, - strips 'tuned', applies optional rewrite(kw), then calls _precompile_default_config. - PyTorch normalizes all custom_op args to positional, so we use inspect.signature - to recover keyword names. + custom_op schemas can't express ``tuple[str, ...]``, so callers pass a + comma-separated string. The autotuned fn keys on a tuple (via + ``tuple(sorted(concat_layout))``); a stray string would be iterated + char-by-char and silently produce a wrong, never-used compile signature. + Single source of truth used by both eager bodies and the fake path. """ - import inspect + if value is None or isinstance(value, tuple): + return value + return tuple(value.split(",")) if value else None - sig = inspect.signature(custom_op._init_fn) - @custom_op.register_fake - def _fake(*args, **kwargs): - bound = sig.bind(*args, **kwargs) - bound.apply_defaults() - kw = dict(bound.arguments) - kw.pop("tuned", None) - if rewrite is not None: - rewrite(kw) - _precompile_default_config(autotuned_fn, **kw) - - -def _rewrite_merge_alpha(kwargs): - """Merge alpha_tensor into alpha for gemm_tuned; add C=None.""" - at = kwargs.pop("alpha_tensor", None) - if at is not None: - kwargs["alpha"] = at - kwargs.setdefault("C", None) +def _register_noop_fake(custom_op): + """Register a pure no-op fake for a mutating custom op. + These ops only mutate their ``out`` argument, so Dynamo / AOT autograd + need no shape effect from the fake; kernel compilation is owned by + jit_cache + the async compile pool at real execution time. + """ -def _rewrite_merge_alpha_beta(kwargs): - """Merge alpha_tensor/beta_tensor into alpha/beta for gemm_tuned.""" - at = kwargs.pop("alpha_tensor", None) - if at is not None: - kwargs["alpha"] = at - bt = kwargs.pop("beta_tensor", None) - if bt is not None: - kwargs["beta"] = bt + @custom_op.register_fake + def _fake(*args, **kwargs): + return -_register_precompile_fake(gemm_out, gemm_tuned, rewrite=_rewrite_merge_alpha) -_register_precompile_fake(gemm_add_out, gemm_tuned, rewrite=_rewrite_merge_alpha_beta) -_register_precompile_fake(gemm_act_out, gemm_act_tuned) -_register_precompile_fake(gemm_dact_out, gemm_dact_tuned) -_register_precompile_fake(gemm_gated_out, gemm_gated_tuned) +_register_noop_fake(gemm_out) +_register_noop_fake(gemm_add_out) +_register_noop_fake(gemm_act_out) +_register_noop_fake(gemm_dact_out) +_register_noop_fake(gemm_gated_out) @gemm_symmetric_out.register_fake @@ -1653,35 +2030,12 @@ def gemm_symmetric_out_fake( dynamic_scheduler: bool = False, alpha: float = 1.0, beta: float = 1.0, + alpha_tensor: Optional[Tensor] = None, + beta_tensor: Optional[Tensor] = None, ) -> None: - from .cache_utils import COMPILE_ONLY - - if not COMPILE_ONLY or isinstance(A.shape[0], torch.SymInt): - return - # gemm_symmetric is not autotuned, compile the single fixed config directly - sm = get_device_capacity(A.device)[0] - tile_m = 256 if sm == 10 else 128 - tile_n = 128 if sm == 12 else 256 - cluster_m = 1 if sm == 12 else 2 - try: - gemm_symmetric_dispatch( - A.unsqueeze(0) if A.ndim == 2 else A, - (B.mT.unsqueeze(0) if B.ndim == 2 else B.mT), - out.unsqueeze(0) if out.ndim == 2 else out, - (C.unsqueeze(0) if C.ndim == 2 else C) if C is not None else None, - torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None, - tile_M=tile_m, - tile_N=tile_n, - cluster_M=cluster_m, - cluster_N=1, - pingpong=False, - persistent=True, - max_swizzle_size=8, - alpha=alpha, - beta=beta, - ) - except Exception: - pass + # Pure no-op: the op only mutates ``out``; kernel compilation is owned + # by jit_cache + the async compile pool at real execution time. + return ## ── gemm_rms ──────────────────────────────────────────────────────────────── @@ -1704,6 +2058,7 @@ def _gemm_rms_tuned( out: Tensor, # (M, N) or (L, M, N) C: Optional[Tensor] = None, # (M, N) or (L, M, N) norm_weight: Optional[Tensor] = None, # (N,) or (L, N) + premult_out: Optional[Tensor] = None, # (M, N) or (L, M, N) — pre-norm_weight snapshot eps: float = 1e-6, dynamic_scheduler: bool = False, config: Optional[GemmConfig] = None, @@ -1723,6 +2078,8 @@ def _gemm_rms_tuned( C = C.unsqueeze(0) if norm_weight is not None and norm_weight.ndim == 1: norm_weight = norm_weight.unsqueeze(0) # (L, N) + if premult_out is not None and premult_out.ndim == 2: + premult_out = premult_out.unsqueeze(0) # Allocate partial reduction buffer tile_n = config.tile_n n_tiles = (N + tile_n - 1) // tile_n @@ -1746,11 +2103,13 @@ def _gemm_rms_tuned( config.tile_n, config.cluster_m, config.cluster_n, - config.pingpong, + tile_K=config.tile_k, + pingpong=config.pingpong, persistent=True, is_dynamic_persistent=dynamic_scheduler, max_swizzle_size=config.max_swizzle_size, rowvec=norm_weight, + aux_out=premult_out, ) # Final reduction: rstd = rsqrt(sum(partials) / N + eps) scale = 1.0 / N @@ -1763,10 +2122,10 @@ def _gemm_rms_tuned( @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_rms_out"), - mutates_args=("out",), + add_op_namespace_prefix("gemm_rms_out"), + mutates_args=("out", "premult_out"), device_types="cuda", - schema="(Tensor A, Tensor B, Tensor(a!) out, Tensor? C=None, Tensor? norm_weight=None, float eps=1e-6, bool dynamic_scheduler=False, bool tuned=True) -> Tensor", + schema="(Tensor A, Tensor B, Tensor(a!) out, Tensor? C=None, Tensor? norm_weight=None, Tensor(a2!)? premult_out=None, float eps=1e-6, bool dynamic_scheduler=False, bool tuned=True) -> Tensor", ) def _gemm_rms_out( A: Tensor, @@ -1774,6 +2133,7 @@ def _gemm_rms_out( out: Tensor, C: Optional[Tensor] = None, norm_weight: Optional[Tensor] = None, + premult_out: Optional[Tensor] = None, eps: float = 1e-6, dynamic_scheduler: bool = False, tuned: bool = True, @@ -1781,6 +2141,7 @@ def _gemm_rms_out( """GEMM + RMS + optional rowvec scaling. D_raw = A @ B (+ C), rstd = rsqrt(mean(D_raw^2) + eps), D_out = D_raw * norm_weight. + If premult_out is provided, D_raw (the pre-norm_weight value) is also written to it. """ fn = _gemm_rms_tuned if tuned else partial(_gemm_rms_tuned.fn, config=None) return fn( @@ -1789,32 +2150,24 @@ def _gemm_rms_out( out, C=C, norm_weight=norm_weight, + premult_out=premult_out, eps=eps, dynamic_scheduler=dynamic_scheduler, ) -@torch.library.register_fake(add_quack_op_namespace_prefix("gemm_rms_out")) +@torch.library.register_fake(add_op_namespace_prefix("gemm_rms_out")) def _gemm_rms_out_fake( A: Tensor, B: Tensor, out: Tensor, C: Optional[Tensor] = None, norm_weight: Optional[Tensor] = None, + premult_out: Optional[Tensor] = None, eps: float = 1e-6, dynamic_scheduler: bool = False, tuned: bool = True, ) -> Tensor: - _precompile_default_config( - _gemm_rms_tuned, - A, - B, - out, - C=C, - norm_weight=norm_weight, - eps=eps, - dynamic_scheduler=dynamic_scheduler, - ) rstd_shape = A.shape[:-1] return torch.empty(rstd_shape, dtype=torch.float32, device=A.device) @@ -1844,6 +2197,7 @@ def gemm_rms( norm_weight: Optional[Tensor] = None, # (N,) or (L, N) out: Optional[Tensor] = None, # (M, N) or (L, M, N) out_dtype: Optional[torch.dtype] = None, + premult_out: Optional[Tensor] = None, # (M, N) or (L, M, N) — pre-norm_weight snapshot eps: float = 1e-6, dynamic_scheduler: bool = False, tuned: bool = True, @@ -1851,6 +2205,7 @@ def gemm_rms( """GEMM + RMS statistics + optional rowvec scaling. D_raw = A @ B (+ C), rstd = rsqrt(mean(D_raw^2) + eps), D_out = D_raw * norm_weight. + If premult_out is provided, D_raw (the pre-norm_weight value) is also written to it. Returns (D_out, rstd). """ out_dtype = A.dtype if out_dtype is None else out_dtype @@ -1858,12 +2213,28 @@ def gemm_rms( if out is None: out_shape = (*A.shape[:-1], N) out = torch.empty(out_shape, dtype=out_dtype, device=A.device) + # Empty-input fast path. Skipping the kernel also avoids a torch.library + # adinplaceorview_impl IndexError that fires on empty inputs because + # premult_out's positional slot isn't materialized in the boxed args tuple. + # K=0 with no C reduces the matmul to zero, so D = 0 and rstd = rsqrt(eps). + if out.numel() == 0 or A.numel() == 0: + _empty_k_matmul_into(out) + if premult_out is not None: + _empty_k_matmul_into(premult_out) + rstd_shape = A.shape[:-1] + if A.numel() == 0 and out.numel() > 0: + # K=0: rstd = rsqrt(0 + eps) for every row. + rstd = torch.full(rstd_shape, eps**-0.5, dtype=torch.float32, device=A.device) + else: + rstd = torch.empty(rstd_shape, dtype=torch.float32, device=A.device) + return out, rstd rstd = _gemm_rms_out( A, B, out, C=C, norm_weight=norm_weight, + premult_out=premult_out, eps=eps, dynamic_scheduler=dynamic_scheduler, tuned=tuned, @@ -1927,7 +2298,8 @@ def gemm_norm_act_tuned( config.tile_n, config.cluster_m, config.cluster_n, - config.pingpong, + tile_K=config.tile_k, + pingpong=config.pingpong, persistent=True, is_dynamic_persistent=dynamic_scheduler, max_swizzle_size=config.max_swizzle_size, @@ -1989,7 +2361,8 @@ def gemm_norm_gated_tuned( config.tile_n, config.cluster_m, config.cluster_n, - config.pingpong, + tile_K=config.tile_k, + pingpong=config.pingpong, persistent=True, is_dynamic_persistent=dynamic_scheduler, max_swizzle_size=config.max_swizzle_size, @@ -1999,7 +2372,7 @@ def gemm_norm_gated_tuned( @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_norm_act_out"), + add_op_namespace_prefix("gemm_norm_act_out"), mutates_args=("preact_out", "postact_out"), device_types="cuda", schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? rstd=None, str? activation=None, bool dynamic_scheduler=False, bool tuned=True) -> ()", @@ -2019,23 +2392,11 @@ def gemm_norm_act_out( fn(A, B, preact_out, postact_out, C, rstd, activation, dynamic_scheduler) -@torch.library.register_fake(add_quack_op_namespace_prefix("gemm_norm_act_out")) -def _gemm_norm_act_out_fake( - A, - B, - preact_out, - postact_out, - C=None, - rstd=None, - activation=None, - dynamic_scheduler=False, - tuned=True, -) -> None: - pass +_register_noop_fake(gemm_norm_act_out) @torch.library.custom_op( - add_quack_op_namespace_prefix("gemm_norm_gated_out"), + add_op_namespace_prefix("gemm_norm_gated_out"), mutates_args=("preact_out", "postact_out"), device_types="cuda", schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? rstd=None, str activation='swiglu', bool dynamic_scheduler=False, bool tuned=True) -> ()", @@ -2055,19 +2416,7 @@ def gemm_norm_gated_out( fn(A, B, preact_out, postact_out, C, rstd, activation, dynamic_scheduler) -@torch.library.register_fake(add_quack_op_namespace_prefix("gemm_norm_gated_out")) -def _gemm_norm_gated_out_fake( - A, - B, - preact_out, - postact_out, - C=None, - rstd=None, - activation="swiglu", - dynamic_scheduler=False, - tuned=True, -) -> None: - pass +_register_noop_fake(gemm_norm_gated_out) def gemm_norm_act( @@ -2101,6 +2450,13 @@ def gemm_norm_act( preact_out = torch.empty(out_shape, dtype=out_dtype, device=A.device) if postact_out is None: postact_out = torch.empty(postact_shape, dtype=postact_dtype, device=A.device) + # Empty-input fast path: skip kernel; zero both outputs (act(0)=0 for all + # supported activations under the no-bias/no-C path of this test surface). + if postact_out.numel() == 0 or A.numel() == 0: + if preact_out is not None: + _empty_k_matmul_into(preact_out) + _empty_k_matmul_into(postact_out) + return preact_out, postact_out if is_gated: gemm_norm_gated_out( A, @@ -2152,13 +2508,12 @@ def gemm_norm_act_ref( if rstd is not None: D = D * rstd.unsqueeze(-1) preact = D.to(out_dtype) if store_preact else None - _act_map = {**act_to_pytorch_fn_map, "silu": F.silu} if is_gated: gate = D[..., ::2] up = D[..., 1::2] postact = gated_to_pytorch_fn_map[activation](gate, up).to(postact_dtype) else: - postact = _act_map[activation](D).to(postact_dtype) + postact = act_to_pytorch_fn_map[activation](D).to(postact_dtype) return preact, postact diff --git a/build/torch-cuda/quack/gemm_norm_act.py b/build/torch-cuda/quack/gemm_norm_act.py index c360b5a361027e9ef9d3fc2a0c334bb04a4eed01..b0b4652e05f7566203577c877e5957b850fddaca 100644 --- a/build/torch-cuda/quack/gemm_norm_act.py +++ b/build/torch-cuda/quack/gemm_norm_act.py @@ -18,13 +18,14 @@ from .cute_dsl_utils import ( get_device_capacity, get_max_active_clusters, ) +from .gemm_sm80 import GemmSm80 from .gemm_sm90 import GemmSm90 from .gemm_sm100 import GemmSm100 from .gemm_sm120 import GemmSm120 -from .gemm_act import GemmActMixin, GemmGatedMixin +from .gemm_act import GemmActMixin, GemmGatedMixin, GemmGatedSm120Mixin from .epi_ops import vec_multiply from .activation import act_fn_map, gate_fn_map -from .cache_utils import jit_cache +from .cache import jit_cache from .rounding import RoundingMode from .gemm_tvm_ffi_utils import ( get_major, @@ -54,9 +55,9 @@ class GemmNormActMixin(GemmActMixin): epi_loop_tensors: Tuple[cute.Tensor, ...], tRS_rD: cute.Tensor, tRS_rC: Optional[cute.Tensor] = None, - ) -> Optional[cute.Tensor]: - tDrRowVec = epi_loop_tensors["mRowVecBroadcast"] - tDrColVec = epi_loop_tensors["mColVecBroadcast"] + ) -> Tuple[cute.Tensor, ...]: + tDrRowVec = epi_loop_tensors.get("mRowVecBroadcast") + tDrColVec = epi_loop_tensors.get("mColVecBroadcast") # Load accumulator and apply alpha/beta/C rD = tRS_rD.load() if const_expr(hasattr(params, "alpha") and params.alpha is not None): @@ -73,24 +74,28 @@ class GemmNormActMixin(GemmActMixin): vec_multiply(self, tRS_rD, tDrColVec, tDrRowVec) # Apply activation if const_expr(params.act_fn is not None): - tRS_rPostAct = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype) - if const_expr(self.arch < 100): - for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True): - tRS_rPostAct[i] = params.act_fn(tRS_rD[i]) + tRS_rAuxOut = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype) + if const_expr(self.arch != 100): + for i in cutlass.range(cute.size(tRS_rAuxOut), unroll_full=True): + tRS_rAuxOut[i] = params.act_fn(tRS_rD[i]) else: - for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True): - tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1] = params.act_fn( + for i in cutlass.range(cute.size(tRS_rAuxOut) // 2, unroll_full=True): + tRS_rAuxOut[2 * i], tRS_rAuxOut[2 * i + 1] = params.act_fn( (tRS_rD[2 * i], tRS_rD[2 * i + 1]) ) else: - tRS_rPostAct = tRS_rD - return tRS_rPostAct + tRS_rAuxOut = tRS_rD + return (tRS_rAuxOut,) class GemmNormActSm90(GemmNormActMixin, GemmSm90): pass +class GemmNormActSm80(GemmNormActMixin, GemmSm80): + pass + + class GemmNormActSm100(GemmNormActMixin, GemmSm100): pass @@ -109,9 +114,9 @@ class GemmNormGatedMixin(GemmGatedMixin): epi_loop_tensors: Tuple[cute.Tensor, ...], tRS_rD: cute.Tensor, tRS_rC: Optional[cute.Tensor] = None, - ) -> Optional[cute.Tensor]: - tDrRowVec = epi_loop_tensors["mRowVecBroadcast"] - tDrColVec = epi_loop_tensors["mColVecBroadcast"] + ) -> Tuple[cute.Tensor, ...]: + tDrRowVec = epi_loop_tensors.get("mRowVecBroadcast") + tDrColVec = epi_loop_tensors.get("mColVecBroadcast") # Load accumulator and apply alpha/beta/C rD = tRS_rD.load() if const_expr(hasattr(params, "alpha") and params.alpha is not None): @@ -127,29 +132,30 @@ class GemmNormGatedMixin(GemmGatedMixin): # Multiply by colvec (rstd) and rowvec (norm_weight) vec_multiply(self, tRS_rD, tDrColVec, tDrRowVec) # Gated activation on normalized D - tRS_rPostAct_layout = cute.recast_layout(2, 1, tRS_rD.layout) - tRS_rPostAct = cute.make_rmem_tensor(tRS_rPostAct_layout.shape, self.acc_dtype) - if const_expr(self.arch < 100): - for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True): - tRS_rPostAct[i] = params.act_fn(tRS_rD[2 * i], tRS_rD[2 * i + 1]) - else: - for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True): - tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1] = params.act_fn( - (tRS_rD[4 * i], tRS_rD[4 * i + 2]), - (tRS_rD[4 * i + 1], tRS_rD[4 * i + 3]), - ) - return tRS_rPostAct + tRS_rAuxOut_layout = cute.recast_layout(2, 1, tRS_rD.layout) + tRS_rAuxOut = cute.make_rmem_tensor(tRS_rAuxOut_layout.shape, self.acc_dtype) + tRS_rD_pair = cute.flat_divide(tRS_rD, cute.make_layout(2)) + tRS_rGate = tRS_rD_pair[0, ...] + tRS_rUp = tRS_rD_pair[1, ...] + vectorize = const_expr(self.arch == 100) + for i in cutlass.range(cute.size(tRS_rAuxOut), unroll_full=True, vectorize=vectorize): + tRS_rAuxOut[i] = params.act_fn(tRS_rGate[i], tRS_rUp[i]) + return (tRS_rAuxOut,) class GemmNormGatedSm90(GemmNormGatedMixin, GemmSm90): pass +class GemmNormGatedSm80(GemmNormGatedMixin, GemmSm80): + pass + + class GemmNormGatedSm100(GemmNormGatedMixin, GemmSm100): pass -class GemmNormGatedSm120(GemmNormGatedMixin, GemmSm120): +class GemmNormGatedSm120(GemmGatedSm120Mixin, GemmNormGatedMixin, GemmSm120): pass @@ -183,12 +189,14 @@ def _compile_gemm_norm_act( ): sm_to_cls = { "norm_act": { + 8: GemmNormActSm80, 9: GemmNormActSm90, 10: GemmNormActSm100, 11: GemmNormActSm100, 12: GemmNormActSm120, }, "norm_gated": { + 8: GemmNormGatedSm80, 9: GemmNormGatedSm90, 10: GemmNormGatedSm100, 11: GemmNormGatedSm100, @@ -213,7 +221,7 @@ def _compile_gemm_norm_act( pa_n = cute.sym_int() if gemm_cls_name == "norm_gated" else n pa_leading_dim = 1 if gemm_cls_name == "norm_gated" else pa_leading pa_shape = (m, pa_n) if varlen_m else (m, pa_n, l) - mPostAct = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading_dim, divisibility=div_pa) + mAuxOut = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading_dim, divisibility=div_pa) mRowVec = fake_tensor(rowvec_dtype, (l, n), leading_dim=1, divisibility=4) if colvec_ndim == 2: @@ -234,7 +242,7 @@ def _compile_gemm_norm_act( return make_ptr(dtype, 0, cute.AddressSpace.gmem, assumed_align=4) epi_args = GemmCls.EpilogueArguments( - mPostAct, + mAuxOut, act_fn, mRowVecBroadcast=mRowVec, mColVecBroadcast=mColVec, @@ -277,6 +285,7 @@ def gemm_norm_act_fn( tile_N: int, cluster_M: int, cluster_N: int, + tile_K: int | None = None, pingpong: bool = False, persistent: bool = True, is_dynamic_persistent: bool = False, @@ -326,7 +335,9 @@ def gemm_norm_act_fn( colvec_ndim = colvec.ndim if colvec is not None else 0 device_capacity = get_device_capacity(A.device) - assert device_capacity[0] in [9, 10, 11, 12], "Only SM90, SM100, SM110, and SM120 are supported" + assert device_capacity[0] in [8, 9, 10, 11, 12], ( + "Only SM8x, SM90, SM100, SM110, and SM120 are supported" + ) if rounding_mode == RoundingMode.RS: assert device_capacity[0] == 10, "Stochastic rounding requires SM100" @@ -349,7 +360,7 @@ def gemm_norm_act_fn( d_major, c_major, postact_major, - (tile_M, tile_N), + (tile_M, tile_N, tile_K) if tile_K is not None else (tile_M, tile_N), (cluster_M, cluster_N, 1), pingpong, persistent, @@ -366,11 +377,6 @@ def gemm_norm_act_fn( sr_seed_mode=sr_seed_mode, ) - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY: - return - max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0 def scalar_arg(scalar, mode, dtype=Int32): @@ -395,6 +401,6 @@ def gemm_norm_act_fn( varlen_args = make_varlen_args(cu_seqlens_m, None, A_idx) if device_capacity[0] in [10, 11]: - compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None, None) + compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None) else: - compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None) + compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args) diff --git a/build/torch-cuda/quack/gemm_sm100.py b/build/torch-cuda/quack/gemm_sm100.py index 7fe29bf161ca39f622ae2a274e5df7facd711b9e..1c58015dfd7d9ca2295dd586bca1f6e6a8a4ae81 100644 --- a/build/torch-cuda/quack/gemm_sm100.py +++ b/build/torch-cuda/quack/gemm_sm100.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025-2026, Tri Dao. +# Copyright (c) 2025-2026, QuACK team. # Based on the cute-dsl example: # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py @@ -23,11 +23,22 @@ from cutlass.cute.nvgpu.warp import ( ) from cutlass import Int32, Float32, Boolean, const_expr from cutlass.utils import LayoutEnum +from cutlass.cute.experimental import iket -from .pipeline import PipelineTmaUmma, PipelineTmaCpAsyncUmma + +from .pipeline import ( + PipelineAsync as QuackPipelineAsync, + PipelineCpAsync, + PipelineTmaUmma, + PipelineTmaCpAsyncUmma, + PipelineUmmaAsync, + mbarrier_arrive_release_cluster, + mbarrier_acquire_cluster, +) +from .dsl.smem_struct import Reserved, partitioned_struct from .tile_scheduler import TileSchedulerOptions from .varlen_utils import VarlenArguments, VarlenManager -from .gemm_sm90 import GemmSm90, NamedBarrierGemm +from .gemm_base import GemmTmaBase, NamedBarrierGemm from . import layout_utils from . import copy_utils as copy_utils from . import sm100_utils as quack_sm100_utils @@ -64,29 +75,7 @@ SM100 tcgen05.mma instructions operate as follows: - 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 same as dense_gemm.py. - -.. code-block:: bash - - python examples/blackwell/dense_gemm_persistent.py \ - --ab_dtype Float16 --d_dtype Float16 --acc_dtype Float32 \ - --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ - --mnkl 8192,8192,8192,1 \ - --use_2cta_instrs - -To collect performance with NCU profiler: - -.. code-block:: bash - - ncu python examples/blackwell/dense_gemm_persistent.py \ - --ab_dtype Float16 --d_dtype Float16 --acc_dtype Float32 \ - --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ - --mnkl 8192,8192,8192,1 \ - --use_2cta_instrs \ - --warmup_iterations 1 --iterations 10 --skip_ref_check - - -Constraints are same as dense_gemm.py: +Constraints: * Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2), see detailed valid dtype combinations in below GemmSm100 class documentation * A/B tensor must have the same data type @@ -99,15 +88,15 @@ Constraints are same as dense_gemm.py: """ -class GemmSm100(GemmSm90): +class GemmSm100(GemmTmaBase): """This class implements batched matrix multiplication (C = A x B) with support for various data types and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. :param acc_dtype: Data type for accumulation during computation :type acc_dtype: type[cutlass.Numeric] - :param mma_tiler_mn: Shape of the MMA tile. Pass (M, N) to default K to + :param mma_tiler_mnk: Shape of the MMA tile. Pass (M, N) to default K to 4 MMA instructions, or (M, N, K) to set the K tile size explicitly. - :type mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]] + :type mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]] :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing :type cluster_shape_mn: Tuple[int, int] @@ -141,7 +130,7 @@ class GemmSm100(GemmSm90): Example: >>> gemm = GemmSm100( ... acc_dtype=Float32, - ... mma_tiler_mn=(128, 128), + ... mma_tiler_mnk=(128, 128), ... cluster_shape_mn=(2, 2) ... ) >>> gemm(mA, mB, mD, max_active_clusters, stream) @@ -149,14 +138,14 @@ class GemmSm100(GemmSm90): arch = 100 - EpilogueArguments = GemmSm90.EpilogueArguments - EpilogueParams = GemmSm90.EpilogueParams + EpilogueArguments = GemmTmaBase.EpilogueArguments + EpilogueParams = GemmTmaBase.EpilogueParams def __init__( self, acc_dtype: Type[cutlass.Numeric], a_dtype: Type[cutlass.Numeric], # ignored for now - mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]], + mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]], cluster_shape_mnk: Tuple[int, int, int], sf_vec_size: Optional[int] = None, gather_A: bool = False, @@ -171,7 +160,7 @@ class GemmSm100(GemmSm90): 1. MMA Instruction Settings (tcgen05): - acc_dtype: Data types for MMA accumulator. - - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + - mma_tiler_mnk: The (M, N) or (M, N, K) shape of the MMA instruction tiler. - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant with cta_group=2 should be used. @@ -180,27 +169,29 @@ class GemmSm100(GemmSm90): :param acc_dtype: Data type of the accumulator. :type acc_dtype: type[cutlass.Numeric] - :param mma_tiler_mn: (M, N) or (M, N, K) shape of the MMA tile. + :param mma_tiler_mnk: (M, N) or (M, N, K) shape of the MMA tile. If only (M, N) is given, K defaults to 4 * instruction K. - :type mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]] + :type mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]] :param cluster_shape_mnk: Tuple (ClusterM, ClusterN) shape of the cluster. :type cluster_shape_mnk: Tuple[int, int] """ self.acc_dtype: Type[cutlass.Numeric] = acc_dtype - self.use_2cta_instrs = mma_tiler_mn[0] in (256,) + self.sf_vec_size = sf_vec_size + self.blockscaled = sf_vec_size is not None + assert len(mma_tiler_mnk) in [2, 3], "MMA tiler must be (M, N) or (M, N, K)" + valid_2cta_m = (128, 256) if not self.blockscaled else (256,) + self.use_2cta_instrs = cluster_shape_mnk[0] % 2 == 0 and mma_tiler_mnk[0] in valid_2cta_m self.cluster_shape_mnk = cluster_shape_mnk assert cluster_shape_mnk[2] == 1, "Cluster shape K must be 1" # K dimension: if user provides 3 values, use their K; otherwise default in _setup_attributes - if len(mma_tiler_mn) == 3: - self.mma_tiler = tuple(mma_tiler_mn) + if len(mma_tiler_mnk) == 3: + self.mma_tiler = tuple(mma_tiler_mnk) else: - self.mma_tiler = (*mma_tiler_mn, 0) - self.sf_vec_size = sf_vec_size - self.blockscaled = sf_vec_size is not None + self.mma_tiler = (*mma_tiler_mnk, 0) self.is_persistent = True - self.pingpong = False # for compatibility with GemmSm90 self.use_clc_persistence = use_clc_persistence + self.epi_m_major = True self.gather_A = gather_A self.concat_layout = concat_layout or () self.use_tma_gather = use_tma_gather @@ -229,6 +220,23 @@ class GemmSm100(GemmSm90): barrier_id=int(NamedBarrierGemm.Epilogue), num_threads=self.num_epi_warps * cute.arch.WARP_SIZE, ) + # CLC throttle: paces query issue to tile consumption so the multi-stage + # lookahead can't over-cancel the pending pool. Producer = CTA0 load warp + # (arrive per tile started), consumer = CTA0 scheduler warp (sync per + # query); 2 warps => 64 threads. Lives here (not the scheduler) because a + # NamedBarrier id is a whole-CTA resource coordinated by NamedBarrierGemm, + # and the participating-thread count is arch-specific (warp layout). A + # single barrier suffices: the dependency chain commit(k+1) <- fetch(k+1) + # <- query(k+1) <- sync(k) forces strict producer/consumer alternation, so + # <= 1 credit is ever outstanding; bar.sync also gives a hardware wakeup vs + # an mbarrier pipeline's PHASECHK + NANOSLEEP polling. + self.clc_throttle_barrier = ( + pipeline.NamedBarrier( + barrier_id=int(NamedBarrierGemm.ClcThrottle), num_threads=2 * cute.arch.WARP_SIZE + ) + if self.use_clc_persistence + else None + ) # Register reallocation for gather_A (3 warp groups, 504 regs total, 168 per WG default). # Heavy epilogues (e.g. colvec_reduce in DGated) override these to avoid register spilling. # Without gather_A there are only 2 WGs (512 total, 256 per WG = max), no reallocation needed. @@ -250,6 +258,15 @@ class GemmSm100(GemmSm90): # Multiple of 4 warps to increase/decrease number of registers assert self.threads_per_cta % 128 == 0 + def epi_smem_warp_shape_mnk(self): + # Mirrors cutlass.utils.blackwell_helpers.compute_epilogue_tile_shape: + # the epilogue tmem layout uses two M warps and two N warps when the + # per-CTA M tile is 64 and the kernel uses 2-CTA instructions. + warp_m, warp_n = ( + (2, 2) if self.cta_tile_shape_mnk[0] == 64 and self.use_2cta_instrs else (4, 1) + ) + return (warp_m, warp_n, 1) + def _setup_attributes(self, epilogue_args: EpilogueArguments, varlen_args: VarlenArguments): """Set up configurations that are dependent on GEMM inputs @@ -264,6 +281,8 @@ class GemmSm100(GemmSm90): - Computing A/B/C shared memory layout - Computing tensor memory allocation columns """ + self.epi_m_major = self.resolve_epi_m_major(epilogue_args) + # Compute mma instruction shapes mma_inst_bits_k = 256 # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K) @@ -284,6 +303,7 @@ class GemmSm100(GemmSm90): if const_expr(not self.blockscaled): self.tiled_mma = sm100_utils.make_trivial_tiled_mma( self.a_dtype, + self.b_dtype, self.a_major_mode, self.b_major_mode, self.acc_dtype, @@ -294,6 +314,7 @@ class GemmSm100(GemmSm90): else: self.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, @@ -303,6 +324,7 @@ class GemmSm100(GemmSm90): ) self.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, @@ -313,6 +335,10 @@ class GemmSm100(GemmSm90): # Compute mma/cluster/tile shapes if self.mma_tiler[2] > 0: + assert self.mma_tiler[2] % self.mma_inst_shape_mnk[2] == 0, ( + f"MMA tiler K ({self.mma_tiler[2]}) must be divisible by " + f"MMA instruction K ({self.mma_inst_shape_mnk[2]})" + ) mma_inst_tile_k = self.mma_tiler[2] // self.mma_inst_shape_mnk[2] else: mma_inst_tile_k = 4 @@ -340,6 +366,25 @@ class GemmSm100(GemmSm90): self.mma_tiler_sfb[1], self.mma_tiler_sfb[2], ) + # The SF atom fixed by the tcgen05 MMA (BlockScaledBasicChunk) is 128 wide + # in N, but cta_tile_n need not be a multiple of 128. Two derived + # quantities localize all the resulting special handling: + # - sfb_tiles_per_atom: adjacent N-tiles that share one 128-wide atom + # (tile_n=64) load the same SFB atom; gmem N-tile coords are divided + # by this. + # - sfb_n_atom_misaligned: tile_n an odd multiple of 64 (64, 192) puts + # odd N-tiles 64 into an atom; the MMA's SFB tmem base shifts by 2 + # columns for odd N-tile coords, and tile_n=192 additionally needs + # the overlapped-window TMA remap at the SFB TMA setup in __call__. + # tile_n=224 is rejected: its tiles start at 32-column offsets within + # the atom ((224*j) % 128 cycles 0/96/64/32), which neither mechanism + # covers. + assert self.cta_tile_shape_mnk[1] in (64, 128, 192, 256), ( + f"blockscaled tile_n must be in (64, 128, 192, 256), " + f"got {self.cta_tile_shape_mnk[1]}" + ) + self.sfb_tiles_per_atom = max(128 // self.cta_tile_shape_mnk[1], 1) + self.sfb_n_atom_misaligned = (self.cta_tile_shape_mnk[1] // 64) % 2 == 1 else: self.cta_tile_shape_mnk_sfb = None @@ -368,13 +413,25 @@ class GemmSm100(GemmSm90): self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 # Compute epilogue subtile + tile_load_layout = None + tile_load_dtype = None + # If TileLoad exists without C, use the first non-None tile-load tensor as + # the C-like input for SM100's epilogue tile shape. Multiple TileLoads + # share the same epi_tile shape. + for op in getattr(self, "_epi_ops", ()): + if op.is_tile_load(): + tile_load_tensor = getattr(epilogue_args, op.name, None) + if tile_load_tensor is not None: + tile_load_layout = LayoutEnum.from_tensor(tile_load_tensor) + tile_load_dtype = tile_load_tensor.element_type + break self.epi_tile = sm100_utils.compute_epilogue_tile_shape( self.cta_tile_shape_mnk, self.use_2cta_instrs, self.d_layout if self.d_layout is not None else LayoutEnum.ROW_MAJOR, self.d_dtype if self.d_dtype is not None else cutlass.BFloat16, - layout_c=self.c_layout, - elem_ty_c=self.c_dtype, + layout_c=self.c_layout if self.c_layout is not None else tile_load_layout, + elem_ty_c=self.c_dtype if self.c_dtype is not None else tile_load_dtype, ) # TMA store tile starts must stay aligned when advancing across CTA-N tiles. # There's a bug w compute_epilogue_tile_shape (as of cutlass-dsl 4.4.2) where if @@ -415,8 +472,16 @@ class GemmSm100(GemmSm90): prefetch_A_idx, cutlass.utils.get_smem_capacity_in_bytes(f"sm_{self.arch}"), # smem_capacity self.occupancy, + self.epi_smem_warp_shape_mnk(), ) - self.sched_stage = 1 + # With CLC the try_cancel response lands directly in the consumer slot, so + # the next query can only be issued once all consumers (cluster-wide) + # release that slot. >=2 stages keep a query in flight while the previous + # tile's info is still being read (cutlass's SchedulerPipelineStageCount + # >= 2); the 3rd stage buys response slack for epilogue-bound tiles (e.g. + # symmetric's double store, ~3% at M=8192 K=512) and costs only 12 smem + # ints + one mbarrier pair. + self.sched_stage = 3 if self.use_clc_persistence else 1 self.a_prefetch_stage = ( 0 if not self.gather_A @@ -512,7 +577,6 @@ class GemmSm100(GemmSm90): stream: cuda.CUstream, mSFA: Optional[cute.Tensor] = None, mSFB: Optional[cute.Tensor] = None, - trace_ptr: Optional[cutlass.Int64] = None, ): """Execute the GEMM operation in steps: - Setup static attributes before smem/grid/tma computation @@ -576,7 +640,7 @@ class GemmSm100(GemmSm90): # so non-packed buffers work (e.g. a slice of a larger scale tensor). # Only the innermost 512-B tile must be contiguous. # For varlen_m, mSFA is sized for per-expert 128-row-padded storage - # (dQaccum format), so use its own M dim (= total_padded_rm * 128) + # (tile-aligned per-batch padding), so use its own M dim (= total_padded_rm * 128) # instead of mA.shape[0] (= total_m, unpadded). if const_expr(cute.rank(mA) == 3): sfa_shape = mA.shape @@ -599,8 +663,12 @@ class GemmSm100(GemmSm90): a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) tma_atom_a, tma_tensor_a = None, None - a_op = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mnk, self.tiled_mma.thr_id + a_op = ( + cpasync.CopyBulkTensorTileG2SOp(self.cta_group) + if const_expr(not self.gather_A) + else sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mnk, self.tiled_mma.thr_id + ) ) if const_expr(not self.gather_A): tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( @@ -628,9 +696,9 @@ class GemmSm100(GemmSm90): tma_smem_layout.shape, internal_type=(cutlass.TFloat32 if mA.element_type is Float32 else None), ) - b_op = sm100_utils.cluster_shape_to_tma_atom_B( - self.cluster_shape_mnk, self.tiled_mma.thr_id - ) + # block_copy takes compiler-driven multicast metadata at the copy site, + # so the TMA atom itself must stay the non-multicast variant here. + b_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group) tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( b_op, copy_utils.create_ragged_tensor_for_tma(mB, ragged_dim=1) if varlen_k else mB, @@ -645,9 +713,7 @@ class GemmSm100(GemmSm90): tma_atom_sfb, tma_tensor_sfb = None, None if const_expr(self.blockscaled): # Setup TMA load for SFA - sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mnk, self.tiled_mma.thr_id - ) + sfa_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group) 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, @@ -659,9 +725,7 @@ class GemmSm100(GemmSm90): internal_type=cutlass.Int16, ) # Setup TMA load for SFB - sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( - self.cluster_shape_mnk, self.tiled_mma.thr_id - ) + sfb_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group) 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, @@ -672,9 +736,15 @@ class GemmSm100(GemmSm90): self.cluster_layout_sfb_vmnk.shape, internal_type=cutlass.Int16, ) - if const_expr( - self.cta_tile_shape_mnk[1] == 192 and self.sf_dtype is cutlass.Float8E8M0FNU - ): + # tile_n=192 spans 1.5 SF atoms, so consecutive N-tiles straddle atom + # boundaries and a TMA box can't be placed at a half-atom offset. + # Instead each tile loads a 2-atom (256-wide) window: remap the gmem + # atom sequence to [a0 a1 | a1 a2 | a3 a4 | a4 a5 | ...] (groups of 4 + # presented atoms at offsets (0, x, x, 2x), advancing by 3 atoms) so + # that tile j's window lands on atoms (3j//2, 3j//2 + 1). Odd tiles + # start 64 into their first atom; the mma warp corrects for that via + # the sfb_n_atom_misaligned tmem offset. + if 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) tma_tensor_sfb = cute.make_tensor( @@ -702,28 +772,24 @@ class GemmSm100(GemmSm90): self.num_tma_load_bytes += sfa_copy_size + sfb_copy_size self.num_tma_load_bytes *= atom_thr_size - # Setup TMA store for D - tma_atom_d, tma_tensor_d = None, None - if const_expr(mD is not None): - tma_atom_d, tma_tensor_d = self._make_tma_epi_atoms_and_tensors( - copy_utils.create_ragged_tensor_for_tma(mD, ragged_dim=0, ptr_shift=True) - if varlen_m - else mD, - self.epi_smem_layout_staged, - self.epi_tile, - op_type="store" - if not (hasattr(epilogue_args, "add_to_output") and epilogue_args.add_to_output) - else "add", - ) - tma_atom_c, tma_tensor_c = None, None - if const_expr(mC is not None): - tma_atom_c, tma_tensor_c = self._make_tma_epi_atoms_and_tensors( - mC, self.epi_c_smem_layout_staged, self.epi_tile, op_type="load" - ) + # Setup TMA store for D and TMA load for C. + tma_atom_d, tma_tensor_d, tma_atom_c, tma_tensor_c = ( + self.make_tma_epilogue_atoms_and_tensors(mD, mC, epilogue_args, varlen_m) + ) epilogue_params = self.epi_to_underlying_arguments(epilogue_args) varlen_params = VarlenManager.to_underlying_arguments(varlen_args) + self.epi_load_bytes_per_stage = self.epi_smem_bytes( + epilogue_args, + self.cta_tile_shape_mnk, + self.epi_tile, + self.epi_smem_warp_shape_mnk(), + ).c_stage + if const_expr(mC is not None): + c_smem_layout = cute.slice_(self.epi_c_smem_layout_staged, (None, None, 0)) + self.epi_load_bytes_per_stage += cute.size_in_bytes(self.c_dtype, c_smem_layout) + TileSchedulerCls = self.get_scheduler_class(varlen_m=varlen_m) tile_sched_args = self.get_scheduler_arguments( mA, mB, mD, scheduler_args, varlen_args, epilogue_args @@ -750,19 +816,21 @@ class GemmSm100(GemmSm90): self.cta_tile_shape_mnk[0] if varlen_m else self.cta_tile_shape_mnk[2] ) - # Define shared storage for kernel - @cute.struct + # Define shared storage for kernel. sched_data lives in the RESERVED + # smem partition (with the pipeline mbarriers / TMEM holding buf): a + # small buffer before the 1024-byte aligned epilogue tensors would add + # a 1 KiB pad; CLC responses use i128 copies, so it stays 16-byte + # aligned. + # 4 Int32 per stage, shared by the two (mode-exclusive) users: STATIC/DYNAMIC + # store the STAS-broadcast (pid_m, pid_n, batch_idx, is_valid); CLC stores the + # 16-byte try_cancel response (16B-aligned since each stage slot is 16 bytes). + sched_smem_size = 4 * self.sched_stage if self.is_persistent else 0 + + @partitioned_struct class SharedStorage: - ab_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.ab_stage * 2] - epi_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.epi_c_stage * 2] - acc_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] - sched_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.sched_stage * 2] - a_prefetch_pipeline_array_ptr: cute.struct.MemRange[ - cutlass.Int64, self.a_prefetch_stage * 2 + sched_data: Reserved[ + cute.struct.Align[cute.struct.MemRange[Int32, sched_smem_size], 16] ] - sched_data: cute.struct.MemRange[Int32, self.sched_stage * 12] - tmem_dealloc_mbar_ptr: cutlass.Int64 - tmem_holding_buf: Int32 sAIdx: cute.struct.Align[cute.struct.MemRange[Int32, a_idx_smem_size], 16] # (EPI_TILE_M, EPI_TILE_N, STAGE) sD: cute.struct.Align[ @@ -831,7 +899,6 @@ class GemmSm100(GemmSm90): self.epi_tile, tile_sched_params, TileSchedulerCls, - trace_ptr, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -874,16 +941,11 @@ class GemmSm100(GemmSm90): epi_tile: cute.Tile, tile_sched_params, TileSchedulerCls: cutlass.Constexpr[Callable], - trace_ptr: Optional[cutlass.Int64] = None, ): """ GPU device kernel performing the Persistent batched GEMM computation. """ - from .trace import TraceContext - - tctx = TraceContext.create(trace_ptr) - varlen_m = const_expr(varlen_params.cu_seqlens_m is not None) varlen_k = const_expr(varlen_params.cu_seqlens_k is not None) assert not (varlen_m and varlen_k) @@ -891,6 +953,7 @@ class GemmSm100(GemmSm90): assert varlen_m or varlen_k has_D = const_expr(mD_mnl is not None) has_C = const_expr(mC_mnl is not None) + has_epi_load = const_expr(self.epi_c_stage > 0) warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -920,39 +983,26 @@ class GemmSm100(GemmSm90): # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(self.shared_storage) + storage = self.shared_storage.allocate(smem) # Initialize pipelines and states ab_pipeline = self.make_ab_pipeline( tiled_mma=tiled_mma, cluster_layout_vmnk=cluster_layout_vmnk, - ab_pipeline_mbar_ptr=storage.ab_pipeline_array_ptr.data_ptr(), is_leader_cta=is_leader_cta, ) epi_pipeline = None - if const_expr(has_C): - epi_pipeline = self.make_epi_pipeline( - c_smem_layout=cute.slice_(epi_c_smem_layout, (None, None, 0)), - epi_pipeline_mbar_ptr=storage.epi_pipeline_array_ptr.data_ptr(), - ) - acc_pipeline = self.make_acc_pipeline( - cluster_layout_vmnk=cluster_layout_vmnk, - acc_pipeline_mbar_ptr=storage.acc_pipeline_array_ptr.data_ptr(), - ) + if const_expr(has_epi_load): + epi_pipeline = self.make_epi_pipeline(tx_count=self.epi_load_bytes_per_stage) + acc_pipeline = self.make_acc_pipeline(cluster_layout_vmnk=cluster_layout_vmnk) sched_pipeline = None sched_data = None if const_expr(self.is_persistent): - sched_pipeline = self.make_sched_pipeline( - self.cluster_shape_mnk, - sched_pipeline_mbar_ptr=storage.sched_pipeline_array_ptr.data_ptr(), - has_C=has_C, - ) - sched_data = storage.sched_data.get_tensor((12, self.sched_stage)) + sched_pipeline = self.make_sched_pipeline(self.cluster_shape_mnk, has_C=has_epi_load) + sched_data = storage.sched_data.get_tensor(cute.make_layout((4, self.sched_stage))) a_prefetch_pipeline = None if const_expr(self.gather_A): - a_prefetch_pipeline = self.make_a_prefetch_pipeline( - storage.a_prefetch_pipeline_array_ptr.data_ptr(), - ) + a_prefetch_pipeline = self.make_a_prefetch_pipeline() tmem_alloc_barrier = pipeline.NamedBarrier( barrier_id=int(NamedBarrierGemm.TmemPtr), @@ -960,11 +1010,9 @@ class GemmSm100(GemmSm90): ) # Tensor memory dealloc barrier init tmem = cutlass.utils.TmemAllocator( - storage.tmem_holding_buf, barrier_for_retrieve=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 @@ -1020,13 +1068,18 @@ class GemmSm100(GemmSm90): ) TileSchedulerCls = partial( - TileSchedulerCls.create, tile_sched_params, sched_data, sched_pipeline + TileSchedulerCls.create, + tile_sched_params, + sched_data, + sched_pipeline, + throttle_barrier=self.clc_throttle_barrier, ) epi_load_barrier = None - if const_expr(has_C): + if const_expr(has_epi_load): epi_load_barrier = pipeline.NamedBarrier( - barrier_id=int(NamedBarrierGemm.EpilogueLoad), num_threads=2 * cute.arch.WARP_SIZE + barrier_id=int(NamedBarrierGemm.EpilogueLoad), + num_threads=(self.num_ab_load_warps + 1) * cute.arch.WARP_SIZE, ) # Cluster wait before tensor memory alloc @@ -1042,30 +1095,6 @@ class GemmSm100(GemmSm90): cute.arch.griddepcontrol_wait() if const_expr(self.gather_A): cute.arch.setmaxregister_decrease(self.num_regs_other) - # Compute multicast mask for A/B buffer full - block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) - block_in_cluster_coord_sfb_vmnk = None - if const_expr(self.blockscaled): - block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( - cta_rank_in_cluster - ) - a_mcast_mask, b_mcast_mask = None, None - sfa_mcast_mask, sfb_mcast_mask = None, None - if const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): - a_mcast_mask = cpasync.create_tma_multicast_mask( - cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 - ) - b_mcast_mask = cpasync.create_tma_multicast_mask( - cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 - ) - if const_expr(self.blockscaled): - sfa_mcast_mask = cpasync.create_tma_multicast_mask( - cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 - ) - sfb_mcast_mask = cpasync.create_tma_multicast_mask( - cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1 - ) - # Persistent tile scheduling loop tile_scheduler = TileSchedulerCls() work_tile = tile_scheduler.initial_work_tile_info() @@ -1076,7 +1105,16 @@ class GemmSm100(GemmSm90): pipeline.PipelineUserType.Consumer, self.a_prefetch_stage ) do_epi_load_barrier_arrive = Boolean(True) + # CLC throttle producer: only the first load warp of CTA 0 in the + # cluster signals; commit once per work tile started, or the scheduler + # warp starves of credits. + is_throttle_producer = Boolean(warp_idx == self.ab_load_warp_id) + if const_expr(cute.size(cluster_layout_vmnk) > 1): + is_throttle_producer = is_throttle_producer & Boolean( + cute.arch.block_idx_in_cluster() == 0 + ) while work_tile.is_valid_tile: + tile_scheduler.throttle_producer_commit(is_throttle_producer) tile_coord_mnkl = work_tile.tile_idx batch_idx = tile_coord_mnkl[3] # Local_tile partition global tensors @@ -1102,7 +1140,7 @@ class GemmSm100(GemmSm90): ) if const_expr(self.blockscaled): # (bM, bK) - # SFA uses padded per-expert offset (dQaccum format), not + # SFA uses the tile-aligned per-batch offset (padded SF layout), not # the A-data offset — allows varlen_m seqlens that aren't # multiples of 128. gSFA_mkl = cute.local_tile( @@ -1111,38 +1149,36 @@ class GemmSm100(GemmSm90): (mma_tile_coord_mnl[0], None), ) # (bN, bK) - # SFB uses padded per-expert K offset in varlen_k (dQaccum format). + # SFB uses the tile-aligned per-batch K offset in varlen_k (padded SF layout). + # N-tiles sharing one 128-wide SF atom (tile_n=64) load the same + # atom, so the gmem N-tile coord is divided by sfb_tiles_per_atom. gSFB_nkl = cute.local_tile( varlen_manager.offset_batch_SFB(mSFB_nkl, batch_idx), cute.select(self.mma_tiler_sfb, [1, 2]), - ( - ( - mma_tile_coord_mnl[1] // 2 - if self.cta_tile_shape_mnk[1] == 64 - else mma_tile_coord_mnl[1] - ), - None, - ), + (mma_tile_coord_mnl[1] // self.sfb_tiles_per_atom, None), ) # Partition global tensor for TiledMMA_A/B/D # Then partition global/shared tensor for TMA load A/B len_k = varlen_manager.len_k(batch_idx) - # TMA load A partition_S/D - a_cta_layout = cute.make_layout( - cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape - ) + # block_copy's lowering wants the coordinate held fixed by the + # multicast mask: A/SFA are same-M across N peers, while B/SFB + # are same-N across M peers. Degenerate cluster dimensions are + # left for the compiler lowering to simplify. + a_tma_multicast = { + "cluster_shape": self.cluster_shape_mnk[:2], + "multicast_dim": "M", + } + b_tma_multicast = { + "cluster_shape": self.cluster_shape_mnk[:2], + "multicast_dim": "N", + } copy_A, prefetch_A = None, None if const_expr(not self.gather_A): # (MMA, MMA_M, MMA_K, RestK) tCgA = thr_mma.partition_A(gA_mk) - copy_A, _, _ = copy_utils.tma_get_copy_fn( - tma_atom_a, - cta_coord=block_in_cluster_coord_vmnk[2], - cta_layout=a_cta_layout, - src_tensor=tCgA, - dst_tensor=sA, - mcast_mask=a_mcast_mask, + copy_A = copy_utils.tma_get_block_copy_fn( + tma_atom_a, src_tensor=tCgA, dst_tensor=sA, tma_multicast=a_tma_multicast ) else: # For varlen_m paths (TMA or cp.async): consume indices from @@ -1162,9 +1198,7 @@ class GemmSm100(GemmSm90): warp_idx, ) if const_expr(varlen_m): - cute.arch.sync_warp() - with cute.arch.elect_one(): - a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state) + a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state) a_prefetch_consumer_state.advance() if const_expr(prefetch_A is not None): prefetch_A = partial(prefetch_A, a_prefetch_pipeline) @@ -1176,52 +1210,33 @@ class GemmSm100(GemmSm90): # (MMA, MMA_N, MMA_K) tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) # TMA load B partition_S/D - copy_B, _, _ = copy_utils.tma_get_copy_fn( - tma_atom_b, - cta_coord=block_in_cluster_coord_vmnk[1], - cta_layout=cute.make_layout( - cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape - ), - src_tensor=tCgB, - dst_tensor=sB, - mcast_mask=b_mcast_mask, + copy_B = copy_utils.tma_get_block_copy_fn( + tma_atom_b, src_tensor=tCgB, dst_tensor=sB, tma_multicast=b_tma_multicast ) copy_SFA, copy_SFB = None, None if const_expr(self.blockscaled): # TMA load SFA partition_S/D - copy_SFA, _, _ = copy_utils.tma_get_copy_fn( + copy_SFA = copy_utils.tma_get_block_copy_fn( tma_atom_sfa, - cta_coord=block_in_cluster_coord_vmnk[2], - cta_layout=a_cta_layout, src_tensor=tCgSFA, dst_tensor=sSFA, - filter_zeros=True, - mcast_mask=sfa_mcast_mask, + tma_multicast=a_tma_multicast, ) # TMA load SFB partition_S/D - sfb_cta_layout = cute.make_layout( - cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape - ) - copy_SFB, _, _ = copy_utils.tma_get_copy_fn( + copy_SFB = copy_utils.tma_get_block_copy_fn( tma_atom_sfb, - cta_coord=block_in_cluster_coord_sfb_vmnk[1], - cta_layout=sfb_cta_layout, src_tensor=tCgSFB, dst_tensor=sSFB, - filter_zeros=True, - mcast_mask=sfb_mcast_mask, + tma_multicast=b_tma_multicast, ) k_tile_cnt = cute.ceil_div(len_k, self.cta_tile_shape_mnk[2]) - tctx.b("tma_load") + iket.range_push("tma_load") if const_expr(not self.gather_A): - ab_producer_state = self.load_AB( + ab_producer_state = self.load_tma( ab_pipeline, ab_producer_state, - copy_A, - copy_B, + [copy_A, copy_B, copy_SFA, copy_SFB], k_tile_cnt, - copy_SFA, - copy_SFB, ) elif const_expr(self.use_tma_gather): ab_producer_state, a_prefetch_consumer_state = self.load_AB_tma_gather( @@ -1243,7 +1258,7 @@ class GemmSm100(GemmSm90): copy_B, k_tile_cnt, ) - tctx.e("tma_load") + iket.range_pop() if const_expr(epi_load_barrier is not None): # In the first work tile, the epi load warp will wait for the signal # from the mainloop load warp to start loading C, to avoid interfering @@ -1252,8 +1267,10 @@ class GemmSm100(GemmSm90): epi_load_barrier.arrive() do_epi_load_barrier_arrive = Boolean(False) # Advance to next tile + iket.range_push("sched_fetch") tile_scheduler.advance_to_next_work() work_tile = tile_scheduler.get_current_work() + iket.range_pop() # Wait A/B buffer empty if warp_idx == self.ab_load_warp_id: ab_pipeline.producer_tail(ab_producer_state) @@ -1274,11 +1291,19 @@ class GemmSm100(GemmSm90): work_tile = tile_scheduler.initial_work_tile_info() while work_tile.is_valid_tile: # Advance to next tile + iket.range_push("clc_produce") tile_scheduler.advance_to_next_work(is_scheduler_warp=is_scheduler_warp) + iket.range_pop() + iket.range_push("clc_consume") work_tile = tile_scheduler.get_current_work() + iket.range_pop() # End of persistent scheduler loop if is_scheduler_warp: tile_scheduler.producer_tail() + # Drain the pending-cluster tail (varlen padding) with unobserved + # cancels so it never launches; see cancel_pending_tail for the + # grant-monotonicity assumption this relies on. + tile_scheduler.cancel_pending_tail() # Specialized A-index prefetch warp (gather_A only) if const_expr(self.gather_A): @@ -1359,9 +1384,9 @@ class GemmSm100(GemmSm90): if const_expr(self.gather_A): cute.arch.setmaxregister_decrease(self.num_regs_other) # PDL: wait for prior kernel before any C TMA loads (matches cutlass C++ epi_load) - if const_expr(self.use_pdl and mC_mnl is not None): + if const_expr(self.use_pdl and has_epi_load): cute.arch.griddepcontrol_wait() - if const_expr(mC_mnl is not None): + if const_expr(has_epi_load): epi_producer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Producer, self.epi_c_stage ) @@ -1373,22 +1398,41 @@ class GemmSm100(GemmSm90): # Get tile coord from tile scheduler tile_coord_mnkl = work_tile.tile_idx batch_idx = tile_coord_mnkl[3] - copy_C_fn, _, bGS_gC = self.epilog_gmem_copy_and_partition( - tma_atom_c, - varlen_manager.offset_batch_epi(mC_mnl, batch_idx), - self.cta_tile_shape_mnk[:2], - epi_tile, - sC, + copy_C = None + if const_expr(has_C): + copy_C_fn, _, _ = self.epilog_gmem_copy_and_partition( + tma_atom_c, + varlen_manager.offset_batch_epi(mC_mnl, batch_idx), + self.cta_tile_shape_mnk[:2], + epi_tile, + sC, + tile_coord_mnkl, + ) + copy_C = copy_utils.tma_producer_copy_fn(copy_C_fn, epi_pipeline) + tile_load_copy_fns = self.epi_tile_load_g2s_copy_fns( + epilogue_params, + epi_smem_tensors, tile_coord_mnkl, + varlen_manager, + epi_pipeline, + ) + copy_epi_load = copy_utils.chain_tma_producer_copy_fns( + (copy_C, *tile_load_copy_fns) ) - copy_C = copy_utils.tma_producer_copy_fn(copy_C_fn, epi_pipeline) if do_epi_load_barrier_wait: epi_load_barrier.arrive_and_wait() do_epi_load_barrier_wait = Boolean(False) - epi_tile_num = const_expr(cute.size(bGS_gC, mode=[1])) + epi_tile_num = const_expr( + cute.size( + cute.zipped_divide( + cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile + ), + mode=[1], + ) + ) for epi_idx in cutlass.range(epi_tile_num, unroll=1): epi_pipeline.producer_acquire(epi_producer_state) - copy_C(src_idx=epi_idx, producer_state=epi_producer_state) + copy_epi_load(src_idx=epi_idx, producer_state=epi_producer_state) # Epi pipeline's producer commit is a NOP epi_pipeline.producer_commit(epi_producer_state) epi_producer_state.advance() @@ -1446,21 +1490,8 @@ class GemmSm100(GemmSm90): cute.slice_(sfb_smem_layout, (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) else: tCtSFA, tCtSFB = None, None - tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = None, None, None - tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = None, None, None # Persistent tile scheduling loop tile_scheduler = TileSchedulerCls() @@ -1486,7 +1517,10 @@ class GemmSm100(GemmSm90): ) tCtAcc = tCtAcc_base[None, None, None, acc_stage_idx] tCtSFB_mma = tCtSFB - if const_expr(self.blockscaled and self.mma_inst_shape_mnk[1] in (64, 192)): + if const_expr(self.blockscaled and self.sfb_n_atom_misaligned): + # Odd N-tiles start 64 into a 128-wide SF atom: shift the SFB + # tmem base by 2 columns (in the atom layout (32,4):(16,4), + # N+64 is element offset 8 = 2 tmem columns). tCtSFB_mma = cute.make_tensor( cute.recast_ptr( sfb_tmem_base_ptr + Int32((tile_coord_mnkl[1] % 2) * 2), @@ -1494,7 +1528,25 @@ class GemmSm100(GemmSm90): ), tCtSFB.layout, ) - tctx.b("mma") + copy_s2t_sfa, copy_s2t_sfb = None, None + sf_valid_insts = None + if const_expr(self.blockscaled): + copy_s2t_sfa = copy_utils.s2t_get_copy_fn(sSFA, tCtSFA, self.cta_group) + copy_s2t_sfb = copy_utils.s2t_get_copy_fn(sSFB, tCtSFB, self.cta_group) + # Exploits the fact that for mxfp8 the MMA instruction K size + # equals the SF vec size (== 32), so one instruction consumes + # exactly one SF block and the mma loop can skip the + # instructions for SF pad blocks on a ragged-K last tile (see + # the comment in self.mma). fp4 has inst_k 64 spanning + # multiple SF blocks, but we don't do varlen_k for + # mxfp4/nvfp4. Valid instructions in that tile; % maps + # "aligned or full tile" to 0 = nothing to skip. + if const_expr(self.mma_inst_shape_mnk[2] == self.sf_vec_size): + num_insts = self.mma_tiler[2] // self.mma_inst_shape_mnk[2] + sf_valid_insts = ( + cute.ceil_div(k_len % self.mma_tiler[2], self.sf_vec_size) % num_insts + ) + iket.range_push("mma") ab_consumer_state, acc_producer_state, tiled_mma = self.mma( ab_pipeline, acc_pipeline, @@ -1509,17 +1561,14 @@ class GemmSm100(GemmSm90): cta_rank_in_cluster, tCtSFA, tCtSFB_mma, - tiled_copy_s2t_sfa, - tiled_copy_s2t_sfb, - tCsSFA_compact_s2t, - tCsSFB_compact_s2t, - tCtSFA_compact_s2t, - tCtSFB_compact_s2t, + copy_s2t_sfa, + copy_s2t_sfb, + sf_valid_insts, ) if const_expr(self.overlap_accum_sf): # After iter 0, 2, ..., shift tmem ptr by -256. # After iter 1, 3, ..., shift tmem ptr by 256. - tCtSFA, tCtSFB, tCtSFA_compact_s2t, tCtSFB_compact_s2t = [ + tCtSFA, tCtSFB = [ cute.make_tensor( cute.recast_ptr( # Doing tmem ptr arithmetic requires 32-bit type, wrong otherwise @@ -1532,9 +1581,9 @@ class GemmSm100(GemmSm90): ), mT.layout, ) - for mT in [tCtSFA, tCtSFB, tCtSFA_compact_s2t, tCtSFB_compact_s2t] + for mT in [tCtSFA, tCtSFB] ] - tctx.e("mma") + iket.range_pop() # Advance to next tile tile_scheduler.advance_to_next_work() work_tile = tile_scheduler.get_current_work() @@ -1592,6 +1641,13 @@ class GemmSm100(GemmSm90): pipeline.PipelineUserType.Consumer, self.epi_c_stage ) while work_tile.is_valid_tile: + # Prefetch the next work tile before the epilogue: the response is + # already in smem (3-stage sched pipeline), and consuming it here + # hides the ~300ns decode (swizzle + async fence) behind this tile's + # epilogue — the pacing chain for small-K / double-store epilogues. + # advance_to_next_work stays after the body: num_tiles_executed must + # count completed tiles during the body (sD stage cycling). + next_work_tile = tile_scheduler.get_current_work() # Get tile coord from tile scheduler tile_coord_mnkl = work_tile.tile_idx batch_idx = tile_coord_mnkl[3] @@ -1635,10 +1691,10 @@ class GemmSm100(GemmSm90): acc_release_idx=self.iter_acc_early_release if const_expr(self.overlap_accum_sf) else epi_tile_num - 1, - clear_acc=varlen_k and k_len == 0, + clear_acc=(varlen_k and k_len == 0), ) - tctx.b("epilogue") + iket.range_push("epilogue") epi_read_state, _ = self.epilogue( epilogue_params, epi_smem_tensors, @@ -1667,11 +1723,11 @@ class GemmSm100(GemmSm90): ) # acc_pipeline.consumer_release was already called in self.epi_load_acc_subtile acc_consumer_state.advance() - tctx.e("epilogue") + iket.range_pop() # Advance to next tile tile_scheduler.advance_to_next_work() - work_tile = tile_scheduler.get_current_work() + work_tile = next_work_tile # Wait for D store complete if is_tma_warp: @@ -1682,8 +1738,6 @@ class GemmSm100(GemmSm90): tmem_alloc_barrier.arrive_and_wait() tmem.free(acc_tmem_ptr) - tctx.flush() - @cute.jit def _make_gather_A_copy( self, @@ -1876,19 +1930,15 @@ class GemmSm100(GemmSm90): cta_rank_in_cluster: Int32, tCtSFA: Optional[cute.Tensor] = None, tCtSFB: Optional[cute.Tensor] = None, - tiled_copy_s2t_sfa: Optional[cute.TiledCopy] = None, - tiled_copy_s2t_sfb: Optional[cute.TiledCopy] = None, - tCsSFA_compact_s2t: Optional[cute.Tensor] = None, - tCsSFB_compact_s2t: Optional[cute.Tensor] = None, - tCtSFA_compact_s2t: Optional[cute.Tensor] = None, - tCtSFB_compact_s2t: Optional[cute.Tensor] = None, + copy_s2t_sfa: Optional[Callable] = None, + copy_s2t_sfb: Optional[Callable] = None, + sf_valid_insts_last_tile: Optional[Int32] = None, ) -> Tuple[cutlass.pipeline.PipelineState, cutlass.pipeline.PipelineState, cute.TiledMma]: - blockscaled = const_expr(tiled_copy_s2t_sfa is not None) + blockscaled = const_expr(copy_s2t_sfa is not None) if const_expr(blockscaled): assert all(x is not None for x in (tCtSFA, tCtSFB)) - assert all(x is not None for x in (tiled_copy_s2t_sfa, tiled_copy_s2t_sfb)) - assert all(x is not None for x in (tCsSFA_compact_s2t, tCsSFB_compact_s2t)) - assert all(x is not None for x in (tCtSFA_compact_s2t, tCtSFB_compact_s2t)) + assert copy_s2t_sfb is not None + skip_sf_pad_insts = const_expr(sf_valid_insts_last_tile is not None) # If gather_A and use_2cta_instrs, the cp.async for the non-leader CTA will # arrive at an mbarrier on the non-leader CTA side, then the mma warp of the non-leader # CTA will wait for that then arrive at the mbarrier on the leader CTA. @@ -1911,29 +1961,67 @@ class GemmSm100(GemmSm90): if not is_leader_cta: ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) with cute.arch.elect_one(): - # The odd CTA signals the even CTA - ab_pipeline.sync_object_full.arrive_mbarrier( - ab_consumer_state.index, dst_rank=cta_rank_in_cluster & 0xFE + # The odd CTA signals the even CTA. The arrive must release this + # CTA's cp.async smem writes at cluster scope so that the leader's + # 2-CTA MMA, which reads our smem over DSMEM, is guaranteed to + # observe them; a plain mbarrier.arrive is only release.cta + # (https://github.com/Dao-AILab/quack/issues/63). + mbarrier_arrive_release_cluster( + ab_pipeline.sync_object_full.get_barrier(ab_consumer_state.index), + cta_rank_in_cluster & 0xFE, ) if is_leader_cta: # Conditionally wait for AB buffer full ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + if const_expr(need_nonleader_cta): + # consumer_wait acquires at cta scope only; pair the non-leader's + # cluster-scope release with a cluster-scope acquire of the (already + # completed) phase before the MMA reads the peer CTA's smem. + mbarrier_acquire_cluster( + ab_pipeline.sync_object_full.get_barrier(ab_consumer_state.index), + ab_consumer_state.phase, + ) # Copy SFA/SFB from smem to tmem if const_expr(blockscaled): - 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) + copy_s2t_sfa(ab_consumer_state.index) + copy_s2t_sfb(ab_consumer_state.index) + # Ragged K: the last k-tile's SF atom holds pad bytes beyond the + # valid scale blocks. We exploit the fact that for mxfp8 the MMA + # instruction K size equals the SF vec size (both 32), i.e. each + # instruction consumes exactly one SF block: skipping the + # instructions for pad blocks — whose A/B values are + # TMA-zero-filled and contribute nothing — means the pad scales + # are never consumed and the gmem pad may be arbitrary (e8m0 0xFF + # = NaN would otherwise poison the accumulator via 0-value x + # NaN-scale products). Instruction issue is a leader-only + # decision, so this covers 2-CTA MMA too. fp4 has inst_k 64 (2 + # SF blocks for mxfp4, 4 for nvfp4), but we don't do varlen_k + # for those formats. + # (The set/gemm sequence is duplicated below because the DSL + # rejects closures capturing staged values inside a dynamic if.) + if const_expr(skip_sf_pad_insts): + num_mma_insts = Int32(num_k_blocks) + if sf_valid_insts_last_tile > 0 and k_tile == k_tile_cnt - 1: + num_mma_insts = sf_valid_insts_last_tile for k_blk_idx in cutlass.range(num_k_blocks, unroll_full=True): k_blk_coord = (None, None, k_blk_idx, ab_consumer_state.index) if const_expr(blockscaled): # Set SFA/SFB tensor to tiled_mma sf_kblock_coord = (None, None, k_blk_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, acc, tCrA[k_blk_coord], tCrB[k_blk_coord], acc) - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + if const_expr(skip_sf_pad_insts): + if k_blk_idx < num_mma_insts: + 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, acc, tCrA[k_blk_coord], tCrB[k_blk_coord], acc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + 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, acc, tCrA[k_blk_coord], tCrB[k_blk_coord], acc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + cute.gemm(tiled_mma, acc, tCrA[k_blk_coord], tCrB[k_blk_coord], acc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) # Async arrive AB buffer empty ab_pipeline.consumer_release(ab_consumer_state) ab_consumer_state.advance() @@ -1957,7 +2045,7 @@ class GemmSm100(GemmSm90): tTR_tAcc: cute.Tensor, tTR_rAcc: cute.Tensor, tRS_rD: cute.Tensor, - epi_idx: int, + epi_coord: [int, int], acc_pipeline: pipeline.PipelineAsync, acc_consumer_state: pipeline.PipelineState, acc_release_idx: int, @@ -1965,50 +2053,15 @@ class GemmSm100(GemmSm90): ): if not clear_acc: # Load accumulator from tensor memory buffer to register - cute.copy(tiled_copy_t2r, tTR_tAcc[None, None, None, epi_idx], tTR_rAcc) + cute.copy(tiled_copy_t2r, tTR_tAcc[None, None, None, epi_coord], tTR_rAcc) tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc) tRS_rD.store(tRS_rAcc.load()) else: tRS_rD.fill(0.0) - if epi_idx == acc_release_idx: + assert epi_coord[0] == 0 # For Sm100, we assume epi_M = 1 + if epi_coord[1] == acc_release_idx: cute.arch.fence_view_async_tmem_load() - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - - 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 + acc_pipeline.consumer_release(acc_consumer_state) def epilog_tmem_copy_and_partition( self, @@ -2142,7 +2195,6 @@ class GemmSm100(GemmSm90): self, tiled_mma: cute.TiledMma, cluster_layout_vmnk: cute.Layout, - ab_pipeline_mbar_ptr: cute.Pointer, is_leader_cta: Boolean, ) -> pipeline.PipelineAsync: # If gather_A and use_2cta_instrs, the cp.async for the non-leader CTA will @@ -2155,9 +2207,11 @@ class GemmSm100(GemmSm90): if const_expr(not self.gather_A or self.use_tma_gather): producer_cnt = 1 else: - producer_cnt = self.num_ab_load_warps * 32 + ( - 1 if const_expr(not self.use_2cta_instrs) else 2 - ) + producer_cnt = self.num_ab_load_warps * cute.arch.WARP_SIZE + if const_expr(not self.use_2cta_instrs): + producer_cnt += 1 + else: + producer_cnt += Int32(2) if is_leader_cta else Int32(0) ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt) # Each warp will contribute to the arrive count with the number of mcast size mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 @@ -2165,19 +2219,8 @@ class GemmSm100(GemmSm90): ab_pipeline_consumer_group = pipeline.CooperativeGroup( pipeline.Agent.Thread, consumer_arrive_cnt ) - if const_expr(not self.gather_A): - pipeline_ab = pipeline.PipelineTmaUmma.create( - barrier_storage=ab_pipeline_mbar_ptr, - num_stages=self.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, - ) - elif const_expr(self.use_tma_gather): + if const_expr(not self.gather_A or self.use_tma_gather): pipeline_ab = PipelineTmaUmma.create( - barrier_storage=ab_pipeline_mbar_ptr, num_stages=self.ab_stage, producer_group=ab_pipeline_producer_group, consumer_group=ab_pipeline_consumer_group, @@ -2187,40 +2230,35 @@ class GemmSm100(GemmSm90): ) else: pipeline_ab = PipelineTmaCpAsyncUmma.create( - barrier_storage=ab_pipeline_mbar_ptr, num_stages=self.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, - producer_drop_count=None - if not self.use_2cta_instrs - else (2 if not is_leader_cta else 0), defer_sync=True, ) return pipeline_ab - def make_acc_pipeline( - self, cluster_layout_vmnk: cute.Layout, acc_pipeline_mbar_ptr: cute.Pointer - ) -> pipeline.PipelineAsync: + def make_acc_pipeline(self, cluster_layout_vmnk: cute.Layout) -> pipeline.PipelineAsync: acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) num_acc_consumer_threads = self.num_epi_warps * (2 if self.use_2cta_instrs else 1) acc_pipeline_consumer_group = pipeline.CooperativeGroup( pipeline.Agent.Thread, num_acc_consumer_threads ) - return pipeline.PipelineUmmaAsync.create( - barrier_storage=acc_pipeline_mbar_ptr, + return PipelineUmmaAsync.create( 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, + elect_one_release=True, + # TMEM load consumers are already ordered by fence_view_async_tmem_load() + syncwarp_before_release=False, ) def make_sched_pipeline( self, cluster_layout_mnk: cute.Layout, - sched_pipeline_mbar_ptr: cute.Pointer, has_C: bool = False, ) -> pipeline.PipelineAsync: # Threads/warps participating in this pipeline @@ -2237,32 +2275,40 @@ class GemmSm100(GemmSm90): sched_pipeline_consumer_group = pipeline.CooperativeGroup( pipeline.Agent.Thread, consumer_arrive_cnt ) - return pipeline.PipelineAsync.create( - barrier_storage=sched_pipeline_mbar_ptr, + # Plain PipelineAsync on purpose (vs the DSL example's PipelineClcFetchAsync): + # expect_tx is per-phase mbarrier state, so each mode's producer arms the full + # barrier as a transaction barrier itself — CLC's multicast try_cancel or + # STATIC/DYNAMIC's STAS st.async, both arrive_and_expect_tx(16) per CTA — and + # only the consumer protocol (wait full, elect-one arrive at CTA 0's empty + # barrier) is shared across modes. A CLC-specific pipeline would hardwire the + # producer and still need this one for STATIC/DYNAMIC. + return QuackPipelineAsync.create( num_stages=self.sched_stage, producer_group=sched_pipeline_producer_group, consumer_group=sched_pipeline_consumer_group, # If there's cluster, the consumers must arrive at the mbar of CTA 0 in the cluster. consumer_mask=None if const_expr(cluster_size == 1) else 0, defer_sync=True, + # One arrive per consumer warp (consumer_arrive_cnt counts warps): syncwarp + # so every lane's slot read is complete, then one elected lane signals. + elect_one_release=True, ) @cute.jit - def make_a_prefetch_pipeline( - self, a_prefetch_pipeline_mbar_ptr: cute.Pointer - ) -> pipeline.PipelineAsync: + def make_a_prefetch_pipeline(self) -> pipeline.PipelineAsync: producer_cnt = 32 a_prefetch_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt) consumer_arrive_cnt = self.num_ab_load_warps a_prefetch_consumer_group = pipeline.CooperativeGroup( pipeline.Agent.Thread, consumer_arrive_cnt ) - return pipeline.PipelineCpAsync.create( - barrier_storage=a_prefetch_pipeline_mbar_ptr, + return PipelineCpAsync.create( num_stages=self.a_prefetch_stage, producer_group=a_prefetch_producer_group, consumer_group=a_prefetch_consumer_group, defer_sync=True, + elect_one_release=True, + syncwarp_before_release=True, ) @classmethod @@ -2284,6 +2330,7 @@ class GemmSm100(GemmSm90): prefetch_A_idx: Literal[None, "varlen_m", "varlen_k"], smem_capacity: int, occupancy: int, + warp_shape_mnk: Tuple[int, int, int] | None = None, ) -> Tuple[int, int, int]: """Computes the number of stages for A/B/C operands based on heuristics. @@ -2319,7 +2366,15 @@ class GemmSm100(GemmSm90): # Default D stages epi_stage = 4 if cute.size(epi_tile[1]) <= 16 else 2 - epi_c_stage = 0 if c_dtype is None else (4 if cute.size(epi_tile[1]) <= 16 else 2) + epi_smem_bytes = cls.epi_smem_bytes( + epilogue_args, cta_tile_shape_mnk, epi_tile, warp_shape_mnk + ) + has_tile_load = epi_smem_bytes.c_stage > 0 + epi_c_stage = ( + 0 + if c_dtype is None and not has_tile_load + else (4 if cute.size(epi_tile[1]) <= 16 else 2) + ) # Calculate smem layout and size for one stage of A, B, and C a_smem_layout_staged_one = sm100_utils.make_smem_layout_a( @@ -2373,13 +2428,13 @@ class GemmSm100(GemmSm90): d_bytes_per_stage = ( cute.size_in_bytes(d_dtype, d_smem_layout_staged_one) if d_dtype is not None else 0 ) - epi_bytes_per_stage = d_bytes_per_stage + cls.epi_smem_bytes_per_stage( - epilogue_args, cta_tile_shape_mnk, epi_tile - ) - epi_bytes = epi_bytes_per_stage * epi_stage + epi_bytes_per_stage = d_bytes_per_stage + epi_smem_bytes.d_stage + epi_bytes = epi_smem_bytes.unstaged + epi_bytes_per_stage * epi_stage if const_expr(c_dtype is not None): c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) epi_bytes += c_bytes_per_stage * epi_c_stage + if const_expr(has_tile_load): + epi_bytes += epi_smem_bytes.c_stage * epi_c_stage # Calculate A/B/SFA/SFB stages: # Start with total smem per CTA (capacity / occupancy) @@ -2391,7 +2446,8 @@ class GemmSm100(GemmSm90): # Refine epilogue stages: # Calculate remaining smem after allocating for A/B stages and reserved bytes # Add remaining unused smem to epilogue - epi_stage += (remaining_bytes - ab_bytes_per_stage * ab_stage) // (epi_bytes_per_stage) + if epi_bytes_per_stage > 0: + epi_stage += (remaining_bytes - ab_bytes_per_stage * ab_stage) // epi_bytes_per_stage return num_acc_stage, ab_stage, epi_stage, epi_c_stage @staticmethod @@ -2553,15 +2609,15 @@ class GemmSm100(GemmSm90): @staticmethod def is_valid_mma_tiler_and_cluster_shape( - mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]], + mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]], cluster_shape_mn: Tuple[int, int], blockscaled: bool, ) -> 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 mma_tiler_mnk: The (M, N) or (M, N, K) shape of the MMA instruction tiler + :type mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]] :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster :type cluster_shape_mn: Tuple[int, int] @@ -2571,20 +2627,20 @@ class GemmSm100(GemmSm90): is_valid = True # Skip invalid mma tile shape if not blockscaled: - if mma_tiler_mn[0] not in [64, 128, 256]: + if mma_tiler_mnk[0] not in [64, 128, 256]: is_valid = False else: - if mma_tiler_mn[0] not in [128, 256]: + if mma_tiler_mnk[0] not in [128, 256]: is_valid = False - mma_inst_n = mma_tiler_mn[1] if mma_tiler_mn[1] <= 256 else mma_tiler_mn[1] // 2 + mma_inst_n = mma_tiler_mnk[1] if mma_tiler_mnk[1] <= 256 else mma_tiler_mnk[1] // 2 if not blockscaled: if mma_inst_n not in range(32, 257, 32): is_valid = False else: # Blockscaled currently supports tile_n in {64, 128, 192, 256}. - if mma_tiler_mn[1] not in [64, 128, 192, 256]: + if mma_tiler_mnk[1] not in [64, 128, 192, 256]: is_valid = False - if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0: + if cluster_shape_mn[0] % (2 if mma_tiler_mnk[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 @@ -2662,7 +2718,7 @@ class GemmSm100(GemmSm90): sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, d_dtype: Type[cutlass.Numeric], - mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]], + mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]], cluster_shape_mn: Tuple[int, int], m: int, n: int, @@ -2680,13 +2736,9 @@ class GemmSm100(GemmSm90): if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"): can_implement = False if not GemmSm100.is_valid_mma_tiler_and_cluster_shape( - mma_tiler_mn, cluster_shape_mn, blockscaled=True + mma_tiler_mnk, cluster_shape_mn, blockscaled=True ): can_implement = False - # Multi-tile N iteration with an asymmetric SFB atom size needs the same - # kind of special-case layout rewriting as tile_n==192. - if mma_tiler_mn[1] == 224 and n > 224: - can_implement = False if not GemmSm100.is_valid_tensor_alignment( m, n, k, l, ab_dtype, d_dtype, a_major, b_major, d_major ): @@ -2698,7 +2750,7 @@ class GemmSm100(GemmSm90): ab_dtype: Type[cutlass.Numeric], acc_dtype: Type[cutlass.Numeric], d_dtype: Type[cutlass.Numeric], - mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]], + mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]], cluster_shape_mn: Tuple[int, int], m: int, n: int, @@ -2717,8 +2769,8 @@ class GemmSm100(GemmSm90): :type acc_dtype: Type[cutlass.Numeric] :param d_dtype: The data type of the output tensor :type d_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 mma_tiler_mnk: The (M, N) or (M, N, K) shape of the MMA instruction tiler + :type mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]] :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster :type cluster_shape_mn: Tuple[int, int] :param m: The number of rows in the A tensor @@ -2745,7 +2797,7 @@ class GemmSm100(GemmSm90): can_implement = False # Skip invalid mma tile shape and cluster shape if not GemmSm100.is_valid_mma_tiler_and_cluster_shape( - mma_tiler_mn, cluster_shape_mn, blockscaled=False + mma_tiler_mnk, cluster_shape_mn, blockscaled=False ): can_implement = False # Skip illegal problem shape for load/store alignment diff --git a/build/torch-cuda/quack/gemm_sm120.py b/build/torch-cuda/quack/gemm_sm120.py index 64cf0bc64b40acda60a8a5b4be384f89ed3d432f..bfe2493993a59538e7d49e9fe314a1ebeff54175 100644 --- a/build/torch-cuda/quack/gemm_sm120.py +++ b/build/torch-cuda/quack/gemm_sm120.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025-2026, Tri Dao. +# Copyright (c) 2025-2026, QuACK team. # Based on the cute-dsl example: # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py # SM120-style GEMM using warp-level MMA (MmaF16BF16Op) + ldmatrix. @@ -17,6 +17,7 @@ import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, warp from cutlass import Int32, Boolean, const_expr +from cutlass.utils import SmemPartition from .varlen_utils import VarlenManager from .pipeline import make_pipeline_state @@ -42,11 +43,12 @@ class GemmSm120(GemmSm90): self, acc_dtype: Type[cutlass.Numeric], a_dtype: Type[cutlass.Numeric], - tile_shape_mn: Tuple[int, int], + tile_shape_mnk: Tuple[int, int] | Tuple[int, int, int], cluster_shape_mnk: Tuple[int, int, int], pingpong: bool = False, is_persistent: bool = True, gather_A: bool = False, + concat_layout: tuple | None = None, use_pdl: bool = True, ): # Don't call super().__init__ — we set up our own config @@ -57,22 +59,24 @@ class GemmSm120(GemmSm90): self.use_pdl = use_pdl self.fp8_slow_accum = False self.gather_A = gather_A + self.concat_layout = concat_layout or () if self.pingpong: assert self.is_persistent, "Pingpong gemm requires persistent scheduler" if gather_A: assert cluster_shape_mnk[1] == 1 self.cluster_shape_mnk = cluster_shape_mnk - tile_M, tile_N = tile_shape_mn - self.cta_tile_shape_mnk = (tile_M, tile_N, 1) + assert len(tile_shape_mnk) in [2, 3], "CTA tile shape must be (M, N) or (M, N, K)" + # K dimension: if user provides 3 values, use their K; otherwise default in _setup_tiled_mma. + self.cta_tile_shape_mnk = ( + tuple(tile_shape_mnk) if len(tile_shape_mnk) == 3 else (*tile_shape_mnk, 0) + ) + tile_M, tile_N = self.cta_tile_shape_mnk[:2] # Pingpong: 2 warp groups each with (2,2,1) atom layout # Non-pingpong: 1 group of 8 warps with (4,2,1) atom layout self.mma_inst_mnk = (16, 8, 16) - if not self.pingpong: - self.atom_layout_mnk = (4, 2, 1) - else: - self.atom_layout_mnk = (2, 2, 1) + self.atom_layout_mnk = (4, 2, 1) if not self.pingpong else (2, 2, 1) # num_mma_warps = total warps doing MMA (both warp groups in pingpong) self.num_mma_warps = math.prod(self.atom_layout_mnk) * (1 if not self.pingpong else 2) # For compatibility with SM90 code that uses warp groups @@ -113,6 +117,7 @@ class GemmSm120(GemmSm90): self.ab_stage = None self.epi_stage = None + self.epi_m_major = True self.a_smem_layout_staged = None self.b_smem_layout_staged = None self.epi_smem_layout_staged = None @@ -120,22 +125,33 @@ class GemmSm120(GemmSm90): self.shared_storage = None self.buffer_align_bytes = 1024 + def epi_smem_warp_shape_mnk(self): + return self.atom_layout_mnk + def _setup_tiled_mma(self): """Set up warp-level MMA (MmaF16BF16Op) and tile K dimension.""" op = warp.MmaF16BF16Op(self.a_dtype, self.acc_dtype, self.mma_inst_mnk) tC = cute.make_layout(self.atom_layout_mnk) + atom_m, atom_n, atom_k = self.atom_layout_mnk + # We want each warp to have 16 consecutive elements in the N direction, for STSM + # and for gated epilogue. + permutation_n = cute.make_ordered_layout((self.mma_inst_mnk[1], atom_n, 2), order=(0, 2, 1)) permutation_mnk = ( - self.atom_layout_mnk[0] * self.mma_inst_mnk[0], - self.atom_layout_mnk[1] * self.mma_inst_mnk[1] * 2, - self.atom_layout_mnk[2] * self.mma_inst_mnk[2], + atom_m * self.mma_inst_mnk[0], + permutation_n, + atom_k * self.mma_inst_mnk[2], ) self.tiled_mma = cute.make_tiled_mma(op, tC, permutation_mnk=permutation_mnk) - tile_k = self.mma_inst_mnk[2] * 4 - self.cta_tile_shape_mnk = ( - self.cta_tile_shape_mnk[0], - self.cta_tile_shape_mnk[1], - tile_k, + tile_k = ( + self.cta_tile_shape_mnk[2] + if self.cta_tile_shape_mnk[2] > 0 + else self.mma_inst_mnk[2] * 4 + ) + assert tile_k > 0, "CTA tile K must be positive" + assert tile_k % self.mma_inst_mnk[2] == 0, ( + f"CTA tile K ({tile_k}) must be divisible by MMA instruction K ({self.mma_inst_mnk[2]})" ) + self.cta_tile_shape_mnk = (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1], tile_k) # __call__, _setup_attributes, make_ab_pipeline, make_epi_store_pipeline, # make_sched_pipeline, epilogue are all inherited from GemmSm90. @@ -161,11 +177,8 @@ class GemmSm120(GemmSm90): epi_c_smem_layout: cute.ComposedLayout, tile_sched_params, TileSchedulerCls: cutlass.Constexpr[Callable], - trace_ptr: Optional[cutlass.Int64] = None, ): - from .trace import TraceContext - - tctx = TraceContext.create(trace_ptr) + from cutlass.cute.experimental import iket varlen_m = const_expr(varlen_params.cu_seqlens_m is not None) varlen_k = const_expr(varlen_params.cu_seqlens_k is not None) @@ -189,23 +202,24 @@ class GemmSm120(GemmSm90): ab_pipeline = self.make_ab_pipeline( tiled_mma=tiled_mma, cluster_layout_vmnk=cute.make_layout((1, *cluster_layout_mnk.shape)), - ab_pipeline_mbar_ptr=storage.ab_pipeline_array_ptr.data_ptr(), ) epi_pipeline = None - if const_expr(has_C): - epi_pipeline = self.make_epi_pipeline( - c_smem_layout=cute.slice_(epi_c_smem_layout, (None, None, 0)), - epi_pipeline_mbar_ptr=storage.epi_pipeline_array_ptr.data_ptr(), - ) + has_epi_load = const_expr(self.epi_c_stage > 0) + if const_expr(has_epi_load): + epi_pipeline = self.make_epi_pipeline(tx_count=self.epi_load_bytes_per_stage) sched_pipeline = None sched_data = None if const_expr(self.is_persistent): - sched_pipeline = self.make_sched_pipeline( - cluster_layout_mnk, - sched_pipeline_mbar_ptr=storage.sched_pipeline_array_ptr.data_ptr(), - varlen_k=varlen_k, + sched_pipeline = self.make_sched_pipeline(cluster_layout_mnk, varlen_k=varlen_k) + # Keep scheduler scratch out of SharedStorage. A small buffer before + # the 1024-byte aligned epilogue tensors can add a 1 KiB pad; CLC + # responses also use i128 copies, so this stays 16-byte aligned. + sched_data = smem.allocate_tensor( + Int32, + cute.make_layout((4, self.sched_stage)), + byte_alignment=16, + partition=SmemPartition.RESERVED, ) - sched_data = storage.sched_data.get_tensor((4, self.sched_stage)) # Cluster sync pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk[:-1], is_relaxed=True) @@ -244,17 +258,18 @@ class GemmSm120(GemmSm90): warp_idx >= self.ab_load_warp_id and warp_idx < self.ab_load_warp_id + self.num_ab_load_warps ): - # Get mcast mask - cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) - block_in_cluster_coord_mnk = cluster_layout_mnk.get_flat_coord(cta_rank_in_cluster) - a_mcast_mask = cute.make_layout_image_mask( - cluster_layout_mnk, block_in_cluster_coord_mnk, mode=1 - ) - b_mcast_mask = cute.make_layout_image_mask( - cluster_layout_mnk, block_in_cluster_coord_mnk, mode=0 - ) - a_mcast_mask = a_mcast_mask if self.is_a_mcast else 0 - b_mcast_mask = b_mcast_mask if self.is_b_mcast else 0 + # block_copy's lowering wants the coordinate held fixed by the + # multicast mask: A is same-M across N peers, while B is + # same-N across M peers. Degenerate cluster dimensions are + # left for the compiler lowering to simplify. + a_tma_multicast = { + "cluster_shape": self.cluster_shape_mnk[:2], + "multicast_dim": "M", + } + b_tma_multicast = { + "cluster_shape": self.cluster_shape_mnk[:2], + "multicast_dim": "N", + } # Persistent tile scheduling loop is_scheduler_warp = self.num_ab_load_warps == 1 or warp_idx == self.ab_load_warp_id @@ -266,7 +281,7 @@ class GemmSm120(GemmSm90): pipeline.PipelineUserType.Producer, self.ab_stage ) while work_tile.is_valid_tile: - tctx.b("tma_load") + iket.range_push("tma_load") tile_coord_mnkl = work_tile.tile_idx batch_idx = tile_coord_mnkl[3] # Local_tile partition global tensors @@ -280,15 +295,11 @@ class GemmSm120(GemmSm90): (tile_coord_mnkl[0], None), ) # TMA load A partition_S/D - copy_A, _, _ = copy_utils.tma_get_copy_fn( + copy_A = copy_utils.tma_get_block_copy_fn( tma_atom_a, - cta_coord=block_in_cluster_coord_mnk[1], - cta_layout=cute.make_layout( - cute.slice_(cluster_layout_mnk, (0, None, 0)).shape - ), src_tensor=gA_mk, dst_tensor=sA, - mcast_mask=a_mcast_mask, + tma_multicast=a_tma_multicast, ) else: copy_A, prefetch_A = self._make_gather_A_copy( @@ -301,21 +312,17 @@ class GemmSm120(GemmSm90): (tile_coord_mnkl[1], None), ) # TMA load B partition_S/D - copy_B, _, _ = copy_utils.tma_get_copy_fn( + copy_B = copy_utils.tma_get_block_copy_fn( tma_atom_b, - cta_coord=block_in_cluster_coord_mnk[0], - cta_layout=cute.make_layout( - cute.slice_(cluster_layout_mnk, (None, 0, 0)).shape - ), src_tensor=gB_nk, dst_tensor=sB, - mcast_mask=b_mcast_mask, + tma_multicast=b_tma_multicast, ) len_k = varlen_manager.len_k(batch_idx) k_tile_cnt = cute.ceil_div(len_k, self.cta_tile_shape_mnk[2]) if const_expr(not self.gather_A): - ab_producer_state = self.load_AB( - ab_pipeline, ab_producer_state, copy_A, copy_B, k_tile_cnt + ab_producer_state = self.load_tma( + ab_pipeline, ab_producer_state, [copy_A, copy_B], k_tile_cnt ) else: ab_producer_state = self.load_AB_gather_A( @@ -327,7 +334,7 @@ class GemmSm120(GemmSm90): k_tile_cnt, varlen_m=varlen_m, ) - tctx.e("tma_load") + iket.range_pop() tile_scheduler.advance_to_next_work(is_scheduler_warp=is_scheduler_warp) work_tile = tile_scheduler.get_current_work() # End of persistent scheduler loop @@ -420,7 +427,7 @@ class GemmSm120(GemmSm90): acc.fill(0.0) if const_expr(self.pingpong): self.pingpong_barrier_sync(warp_group_idx, stage="mma") - tctx.b("mma") + iket.range_push("mma") ab_read_state = self.mma( ab_pipeline, ab_read_state, @@ -437,14 +444,14 @@ class GemmSm120(GemmSm90): if const_expr(self.pingpong): # Cue for next WG's MMA to start self.pingpong_barrier_arrive(1 - warp_group_idx, stage="mma") - tctx.e("mma") + iket.range_pop() # ============================================================ # EPILOGUE — reuse SM90's epilogue flow # ============================================================ if const_expr(self.pingpong): self.pingpong_barrier_sync(warp_group_idx, "epi") - tctx.b("epilogue") + iket.range_push("epilogue") copy_D = None if const_expr(has_D): @@ -467,12 +474,22 @@ class GemmSm120(GemmSm90): tile_coord_mnkl, ) copy_C = copy_utils.tma_producer_copy_fn(copy_C_fn, epi_pipeline) + if const_expr(has_epi_load): + tile_load_copy_fns = self.epi_tile_load_g2s_copy_fns( + epilogue_params, + epi_smem_tensors, + tile_coord_mnkl, + varlen_manager, + epi_pipeline, + ) + copy_C = copy_utils.chain_tma_producer_copy_fns((copy_C, *tile_load_copy_fns)) d_dtype_for_layout = self.d_dtype if self.d_dtype is not None else cutlass.BFloat16 tiled_copy_r2s, tRS_rD, tRS_sD = self.epilog_smem_store_and_partition( tiled_mma, self.d_layout, d_dtype_for_layout, sD, tidx ) - tRS_rAcc = self.epi_retile_acc(acc, tRS_rD, tiled_copy_r2s, tidx) + # (R2S, R2S_M, R2S_N, (epi_M, epi_N)) + tRS_rAcc = self.epi_retile_acc(acc, tRS_rD, tiled_copy_r2s) load_acc_subtile = partial(self.epi_load_acc_subtile, tRS_rAcc) if const_expr(has_C): tiled_copy_s2r, tRS_rC, tSR_rC, tSR_sC = self.epilog_smem_load_and_partition( @@ -517,7 +534,7 @@ class GemmSm120(GemmSm90): if is_tma_warp: epi_store_pipeline.producer_tail() self.pingpong_barrier_arrive(1 - warp_group_idx, stage="epi") - tctx.e("epilogue") + iket.range_pop() if const_expr(not self.pingpong): tile_scheduler.advance_to_next_work() @@ -546,8 +563,6 @@ class GemmSm120(GemmSm90): if is_tma_warp: epi_store_pipeline.producer_tail() - tctx.flush() - @cute.jit def mma( self, @@ -585,7 +600,12 @@ class GemmSm120(GemmSm90): for k in cutlass.range_constexpr(num_k_blocks): k_next = 0 if k + 1 == num_k_blocks else k + 1 if const_expr(k == num_k_blocks - 1): - # Don't need to sync_warp: the previous instruction was mma.sync from cute.gemm + # TMA writes this smem stage through the async proxy, while ldmatrix + # reads it through the generic proxy. Fence before release so the + # producer's next async-proxy write cannot race those reads; sync the + # warp because only one lane signals the empty mbarrier. + cute.arch.fence_view_async_shared() + cute.arch.sync_warp() ab_pipeline.consumer_release(ab_read_state) ab_read_state.advance() peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_read_state) @@ -601,6 +621,12 @@ class GemmSm120(GemmSm90): for k in cutlass.range_constexpr(num_k_blocks): k_next = 0 if k + 1 == num_k_blocks else k + 1 if const_expr(k == num_k_blocks - 1): + # TMA writes this smem stage through the async proxy, while ldmatrix + # reads it through the generic proxy. Fence before release so the + # producer's next async-proxy write cannot race those reads; sync the + # warp because only one lane signals the empty mbarrier. + cute.arch.fence_view_async_shared() + cute.arch.sync_warp() ab_pipeline.consumer_release(ab_read_state) ab_read_state.advance() if const_expr(k_next > 0): @@ -610,17 +636,28 @@ class GemmSm120(GemmSm90): return ab_read_state - def epi_retile_acc(self, acc, tRS_rD, tiled_copy_r2s, tidx=None): - """Retile accumulator for epilogue. Warp-level MMA uses tiled_copy_r2s.retile.""" - if tidx is None: - tidx = cute.arch.thread_idx()[0] - thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) - self._epi_size_tRS_rD = cute.size(tRS_rD) - return thr_copy_r2s.retile(acc) - - @cute.jit - def epi_load_acc_subtile(self, tRS_rAcc, tRS_rD, epi_idx): - """Load acc subtile using retile-based flat indexing (warp-level MMA layout).""" - size_rD = self._epi_size_tRS_rD - for i in cutlass.range_constexpr(size_rD): - tRS_rD[i] = tRS_rAcc[epi_idx * size_rD + i] + @staticmethod + def _compute_tile_shape_or_override( + cta_tile_shape_mnk: Tuple[int, int, int], + atom_layout_mnk: Tuple[int, int, int], + element_type: Optional[Type[cutlass.Numeric]] = None, + epi_tile_override: Tuple[int, int] | None = None, + ) -> Tuple[int, int]: + """Compute the epilogue tile shape or use override if provided. + + :param cta_tile_shape_mnk: CTA tile shape (M,N,K) + :type cta_tile_shape_mnk: Tuple[int, int, int] + :param element_type: Data type of elements + :type element_type: type[cutlass.Numeric] + :param epi_tile_override: Optional override for epilogue tile shape + :type epi_tile_override: Tuple[int, int] or None + + :return: Computed epilogue tile shape + :rtype: Tuple[int, int] + """ + if epi_tile_override is not None: + return epi_tile_override + n_perf = 64 if element_type is not None and element_type.width == 8 else 32 + tile_m = math.gcd(64, cute.size(cta_tile_shape_mnk, mode=[0])) + tile_n = math.gcd(n_perf, cute.size(cta_tile_shape_mnk, mode=[1])) + return (tile_m, tile_n) diff --git a/build/torch-cuda/quack/gemm_sm80.py b/build/torch-cuda/quack/gemm_sm80.py new file mode 100644 index 0000000000000000000000000000000000000000..d26121642f6c62039e7a9ff5fe034053546bff59 --- /dev/null +++ b/build/torch-cuda/quack/gemm_sm80.py @@ -0,0 +1,149 @@ +# Copyright (c) 2026, Tri Dao. +# +# Ampere GEMM using warp-level MMA and cp.async global-to-shared loads. +# All CTA threads participate in cp.async, MMA, and epilogue. + +from typing import Optional, Tuple, Type, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.cute.nvgpu import warp + +from .gemm_base import GemmBase, NamedBarrierGemm +from .tile_scheduler import TileSchedulerOptions +from .varlen_utils import VarlenArguments + + +class GemmSm80(GemmBase): + """SM80 GEMM with cp.async loads and warp-level tensor-core MMA. + + SM80 has no TMA, so both the A/B mainloop and the epilogue global memory + movement use per-thread copies. The epilogue still reuses the standard + composable epilogue hooks used by the SM90/SM120 classes. + """ + + arch = 80 + _supported_archs = (80, 86, 87, 89) + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + tile_shape_mnk: Union[Tuple[int, int], Tuple[int, int, int]], + cluster_shape_mnk: Tuple[int, int, int], + pingpong: bool = False, + is_persistent: bool = False, + gather_A: bool = False, + concat_layout: tuple | None = None, + use_pdl: bool = False, + num_warps: Optional[int] = None, + occupancy: Optional[int] = None, + arch: int = 80, + ): + if arch not in self._supported_archs: + raise ValueError( + f"SM80-family GEMM supports arch in {self._supported_archs}, got {arch}" + ) + self.arch = arch + self.acc_dtype = acc_dtype + self.pingpong = False + self.is_persistent = is_persistent + self.use_clc_persistence = False + self.use_pdl = use_pdl + self.fp8_slow_accum = False + self.gather_A = gather_A + self.concat_layout = concat_layout or () + + assert not pingpong, "SM8x GEMM does not use pingpong scheduling" + assert cluster_shape_mnk == (1, 1, 1), "SM8x GEMM does not support CTA clusters" + # The shared tile scheduler API still takes a cluster shape. SM8x only launches + # independent CTAs, so keep this fixed at a single-CTA cluster. + self.cluster_shape_mnk = (1, 1, 1) + self.mma_inst_mnk = (16, 8, 16) + if len(tile_shape_mnk) == 3: + tile_m, tile_n, tile_k = tile_shape_mnk + elif len(tile_shape_mnk) == 2: + tile_m, tile_n = tile_shape_mnk + tile_k = 4 * self.mma_inst_mnk[2] + else: + raise ValueError("SM80 tile shape must be (M, N) or (M, N, K)") + if tile_m % 16 != 0 or tile_n % 16 != 0: + raise ValueError("SM80 tile shape M/N must be divisible by 16") + if tile_k <= 0 or tile_k % self.mma_inst_mnk[2] != 0: + raise ValueError("SM80 tile shape K must be a positive multiple of MMA instruction K") + self.cta_tile_shape_mnk = (tile_m, tile_n, tile_k) + + self.num_warps = ( + num_warps + if num_warps is not None + else (8 if (tile_m, tile_n) in ((128, 256), (128, 192)) else 4) + ) + if self.num_warps not in (4, 8): + raise ValueError("SM80 GEMM supports num_warps=4 or 8") + self.atom_layout_mnk = (2, self.num_warps // 2, 1) + self.mma_inst_tile_k = tile_k // self.mma_inst_mnk[2] + self.threads_per_cta = self.num_warps * cute.arch.WARP_SIZE + + self.smem_capacity = self._smem_capacity_for_arch(self.arch) + self.buffer_align_bytes = 1024 + default_occupancy = 1 if self.num_warps == 8 else 2 + ab_bytes_per_stage = (tile_m + tile_n) * tile_k * a_dtype.width // 8 + if ( + 3 * ab_bytes_per_stage + > self.smem_capacity // default_occupancy - self.buffer_align_bytes + ): + default_occupancy = 1 + self.occupancy = occupancy if occupancy is not None else default_occupancy + self.num_epi_warps = self.num_warps + self.epilogue_barrier = pipeline.NamedBarrier( + barrier_id=int(NamedBarrierGemm.Epilogue), + num_threads=self.num_epi_warps * cute.arch.WARP_SIZE, + ) + + self.ab_stage = None + self.epi_stage = None + self.epi_c_stage = None + self.epi_m_major = True + self.a_smem_layout_staged = None + self.b_smem_layout_staged = None + self.epi_smem_layout_staged = None + self.epi_c_smem_layout_staged = None + self.epi_tile = None + self.shared_storage = None + + @staticmethod + def _smem_capacity_for_arch(arch: int) -> int: + # CUDA documents CC 8.7 with the same shared-memory capacity as CC 8.0, + # but this CUTLASS helper build does not accept "sm_87". + if arch == 87: + arch = 80 + return cutlass.utils.get_smem_capacity_in_bytes(f"sm_{arch}") + + def _setup_tiled_mma(self): + op = warp.MmaF16BF16Op(self.a_dtype, self.acc_dtype, self.mma_inst_mnk) + atom_m, atom_n, atom_k = self.atom_layout_mnk + tC = cute.make_layout(self.atom_layout_mnk) + permutation_mnk = ( + atom_m * self.mma_inst_mnk[0], + atom_n * self.mma_inst_mnk[1] * 2, + atom_k * self.mma_inst_mnk[2], + ) + self.tiled_mma = cute.make_tiled_mma(op, tC, permutation_mnk=permutation_mnk) + self.tiled_mma_gated_postact = self.tiled_mma + + @cute.jit + def __call__( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mD: Optional[cute.Tensor], + mC: Optional[cute.Tensor], + epilogue_args: tuple, + scheduler_args: TileSchedulerOptions, + varlen_args: Optional[VarlenArguments], + stream: cuda.CUstream, + ): + raise NotImplementedError("Gemm Sm80 is not implemented yet") diff --git a/build/torch-cuda/quack/gemm_sm90.py b/build/torch-cuda/quack/gemm_sm90.py index ea3078b721c92d514239d1ed4fbbba6e6dd43354..947603579d560c227724f80e584a188e7524747a 100644 --- a/build/torch-cuda/quack/gemm_sm90.py +++ b/build/torch-cuda/quack/gemm_sm90.py @@ -1,9 +1,9 @@ -# Copyright (c) 2025-2026, Tri Dao. +# Copyright (c) 2025-2026, QuACK team. + # Based on the cute-dsl example: # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/hopper/dense_gemm.py -import enum -from typing import Tuple, Type, Callable, Optional, Union, Literal +from typing import Tuple, Type, Callable, Optional from functools import partial import math @@ -17,28 +17,18 @@ from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, warp, warpgroup import cutlass.utils.hopper_helpers as sm90_utils from cutlass import Int32, Float32, Float16, Boolean, const_expr -from cutlass.utils import LayoutEnum - +from cutlass.utils import LayoutEnum, SmemPartition -from dataclasses import dataclass -from .cute_dsl_utils import ParamsBase from . import layout_utils -from .tile_scheduler import ( - TileSchedulerOptions, - TileSchedulerArguments, - TileScheduler, - VarlenMTileSchedulerArguments, - VarlenMTileScheduler, - PersistenceMode, -) +from .gemm_base import GemmTmaBase, NamedBarrierGemm +from .tile_scheduler import TileSchedulerOptions from .varlen_utils import VarlenArguments, VarlenManager # return PipelineStateWAdvance instead of PipelineState -from .pipeline import make_pipeline_state, PipelineTmaCpAsync +from .pipeline import PipelineAsync as QuackPipelineAsync, make_pipeline_state from . import copy_utils as copy_utils from . import sm90_utils as quack_sm90_utils -from .rounding import RoundingMode """ A high-performance batched dense GEMM (C = A * B) example for the NVIDIA Hopper architecture @@ -78,27 +68,16 @@ Constraints: """ -class NamedBarrierGemm(enum.IntEnum): - Epilogue = enum.auto() # starts from 1 as barrier 0 is reserved for sync_threads() - # For mainloop load warps to signal that the epilogue load warp can start. - # This is to avoid loading C too early, interfering with loading A and B. - EpilogueLoad = enum.auto() - MmaWG0 = enum.auto() - MmaWG1 = enum.auto() - EpiWG0 = enum.auto() - EpiWG1 = enum.auto() - TmemPtr = enum.auto() - - -class GemmSm90: +class GemmSm90(GemmTmaBase): """ This class implements batched matrix multiplication (C = A x B) with support for various data types and architectural features specific to Hopper GPUs with persistent tile scheduling and warp specialization. :param acc_dtype: Data type for accumulation during computation :type acc_dtype: type[cutlass.Numeric] - :param tile_shape_mn: Shape of the CTA tile (M,N) - :type tile_shape_mn: Tuple[int, int, int] + :param tile_shape_mnk: Shape of the CTA tile. Pass (M, N) to default K to + 4 MMA instructions, or (M, N, K) to set K explicitly. + :type tile_shape_mnk: Tuple[int, int] | Tuple[int, int, int] :param cluster_shape_mnk: Cluster dimensions (M,N,K) for parallel processing :type cluster_shape_mnk: Tuple[int, int, int] @@ -121,25 +100,21 @@ class GemmSm90: Example: >>> gemm = GemmSm90( ... acc_dtype=Float32, - ... tile_shape_mn=(128, 256), + ... tile_shape_mnk=(128, 256), ... cluster_shape_mnk=(1, 1, 1) ... ) >>> gemm(a_tensor, b_tensor, c_tensor, stream) """ arch = 90 - - @dataclass - class EpilogueArguments: - pass - - EpilogueParams = ParamsBase + EpilogueArguments = GemmTmaBase.EpilogueArguments + EpilogueParams = GemmTmaBase.EpilogueParams def __init__( self, acc_dtype: Type[cutlass.Numeric], a_dtype: Type[cutlass.Numeric], - tile_shape_mn: Tuple[int, int], + tile_shape_mnk: Tuple[int, int] | Tuple[int, int, int], cluster_shape_mnk: Tuple[int, int, int], pingpong: bool = False, is_persistent: bool = True, @@ -157,8 +132,8 @@ class GemmSm90: :param acc_dtype: Data type for accumulation during computation :type acc_dtype: type[cutlass.Numeric] - :param tile_shape_mn: Shape of the CTA tile (M,N) - :type tile_shape_mn: Tuple[int, int] + :param tile_shape_mnk: Shape of the CTA tile (M,N) or (M,N,K) + :type tile_shape_mnk: Tuple[int, int] | Tuple[int, int, int] :param cluster_shape_mnk: Cluster dimensions (M,N,K) for parallel processing :type cluster_shape_mnk: Tuple[int, int, int] """ @@ -179,8 +154,11 @@ class GemmSm90: assert cluster_shape_mnk[1] == 1, "Cluster shape N must be 1 for gather A " self.cluster_shape_mnk = cluster_shape_mnk - # K dimension is deferred in _setup_attributes - self.cta_tile_shape_mnk = (*tile_shape_mn, 1) + assert len(tile_shape_mnk) in [2, 3], "CTA tile shape must be (M, N) or (M, N, K)" + # K dimension: if user provides 3 values, use their K; otherwise default in _setup_tiled_mma. + self.cta_tile_shape_mnk = ( + tuple(tile_shape_mnk) if len(tile_shape_mnk) == 3 else (*tile_shape_mnk, 0) + ) tile_M, tile_N = self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1] # check the cta tile shape if not self.pingpong: @@ -268,6 +246,7 @@ class GemmSm90: self.ab_stage = None self.epi_stage = None + self.epi_m_major = True self.a_smem_layout_staged = None self.b_smem_layout_staged = None @@ -277,6 +256,10 @@ class GemmSm90: self.shared_storage = None self.buffer_align_bytes = 1024 + def epi_smem_warp_shape_mnk(self): + atom_m, atom_n, atom_k = self.atom_layout_mnk + return (atom_m * 4, atom_n, atom_k) + def _setup_tiled_mma(self): """Set up tiled MMA and tile K dimension. Override for different MMA types.""" self.tiled_mma = sm90_utils.make_trivial_tiled_mma( @@ -304,11 +287,17 @@ class GemmSm90: permutation_mnk=(None, permutation_n, None), ) mma_inst_shape_k = cute.size(self.tiled_mma.shape_mnk, mode=[2]) - mma_inst_tile_k = 4 + tile_k = ( + self.cta_tile_shape_mnk[2] if self.cta_tile_shape_mnk[2] > 0 else mma_inst_shape_k * 4 + ) + assert tile_k > 0, "CTA tile K must be positive" + assert tile_k % mma_inst_shape_k == 0, ( + f"CTA tile K ({tile_k}) must be divisible by MMA instruction K ({mma_inst_shape_k})" + ) self.cta_tile_shape_mnk = ( self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1], - mma_inst_shape_k * mma_inst_tile_k, + tile_k, ) def _setup_attributes(self, epilogue_args: EpilogueArguments): @@ -325,14 +314,16 @@ class GemmSm90: - Computing A/B/C shared memory layout """ self._setup_tiled_mma() + self.epi_m_major = self.resolve_epi_m_major(epilogue_args) self.cluster_layout_mnk = cute.make_layout(self.cluster_shape_mnk) - self.epi_tile = self._sm90_compute_tile_shape_or_override( + self.epi_tile = self._compute_tile_shape_or_override( self.cta_tile_shape_mnk, self.atom_layout_mnk, self.d_dtype, ) + self.epi_tile_shape = cute.ceil_div(self.cta_tile_shape_mnk[:2], self.epi_tile) # Compute stage before compute smem layout self.ab_stage, self.epi_stage, self.epi_c_stage = self._compute_stages( @@ -345,6 +336,7 @@ class GemmSm90: epilogue_args, cutlass.utils.get_smem_capacity_in_bytes(f"sm_{self.arch}"), # smem_capacity self.occupancy, + self.epi_smem_warp_shape_mnk(), ) self.sched_stage = 2 if self.pingpong else 1 @@ -380,7 +372,6 @@ class GemmSm90: scheduler_args: TileSchedulerOptions, varlen_args: Optional[VarlenArguments], stream: cuda.CUstream, - trace_ptr: Optional[cutlass.Int64] = None, ): """Execute the GEMM operation in steps: - Setup static attributes @@ -434,52 +425,31 @@ class GemmSm90: a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, 0)) b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, 0)) - tma_atom_a, tma_tensor_a = None, None - if const_expr(not self.gather_A): - tma_atom_a, tma_tensor_a = self._make_tma_atoms_and_tensors( - copy_utils.create_ragged_tensor_for_tma(mA, ragged_dim=1) - if varlen_k and not self.gather_A - else mA, - a_smem_layout, - (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[2]), - self.cluster_shape_mnk[1], - ) - tma_atom_b, tma_tensor_b = self._make_tma_atoms_and_tensors( - copy_utils.create_ragged_tensor_for_tma(mB, ragged_dim=1) if varlen_k else mB, - b_smem_layout, - (self.cta_tile_shape_mnk[1], self.cta_tile_shape_mnk[2]), - self.cluster_shape_mnk[0], + tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b = self.make_tma_load_atoms_and_tensors( + mA, mB, a_smem_layout, b_smem_layout, varlen_k ) self.num_tma_load_bytes = cute.size_in_bytes(self.b_dtype, b_smem_layout) if const_expr(not self.gather_A): self.num_tma_load_bytes += cute.size_in_bytes(self.a_dtype, a_smem_layout) - tma_atom_d, tma_tensor_d = None, None - if const_expr(mD is not None): - tma_atom_d, tma_tensor_d = self._make_tma_epi_atoms_and_tensors( - copy_utils.create_ragged_tensor_for_tma( - mD, - ragged_dim=0, - ptr_shift=True, - ) - if varlen_m - else mD, - self.epi_smem_layout_staged, - self.epi_tile, - op_type="store" - if not (hasattr(epilogue_args, "add_to_output") and epilogue_args.add_to_output) - else "add", - ) - tma_atom_c, tma_tensor_c = None, None - if const_expr(mC is not None): - tma_atom_c, tma_tensor_c = self._make_tma_epi_atoms_and_tensors( - mC, self.epi_c_smem_layout_staged, self.epi_tile, op_type="load" - ) + tma_atom_d, tma_tensor_d, tma_atom_c, tma_tensor_c = ( + self.make_tma_epilogue_atoms_and_tensors(mD, mC, epilogue_args, varlen_m) + ) epilogue_params = self.epi_to_underlying_arguments(epilogue_args) varlen_params = VarlenManager.to_underlying_arguments(varlen_args) + self.epi_load_bytes_per_stage = self.epi_smem_bytes( + epilogue_args, + self.cta_tile_shape_mnk, + self.epi_tile, + self.epi_smem_warp_shape_mnk(), + ).c_stage + if const_expr(mC is not None): + c_smem_layout = cute.slice_(self.epi_c_smem_layout_staged, (None, None, 0)) + self.epi_load_bytes_per_stage += cute.size_in_bytes(self.c_dtype, c_smem_layout) + TileSchedulerCls = self.get_scheduler_class(varlen_m=varlen_m) tile_sched_args = self.get_scheduler_arguments( mA, mB, mD, scheduler_args, varlen_args, epilogue_args @@ -494,10 +464,6 @@ class GemmSm90: @cute.struct class SharedStorage: - ab_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.ab_stage * 2] - epi_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.epi_c_stage * 2] - sched_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.sched_stage * 2] - sched_data: cute.struct.MemRange[Int32, self.sched_stage * 4] sD: cute.struct.Align[ cute.struct.MemRange[ self.d_dtype if self.d_dtype is not None else Int32, epi_smem_size @@ -542,7 +508,6 @@ class GemmSm90: self.epi_c_smem_layout_staged, tile_sched_params, TileSchedulerCls, - trace_ptr, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -575,7 +540,6 @@ class GemmSm90: epi_c_smem_layout: cute.ComposedLayout, tile_sched_params, TileSchedulerCls: cutlass.Constexpr[Callable], - trace_ptr: Optional[cutlass.Int64] = None, ): """ GPU device kernel performing the batched GEMM computation. @@ -604,9 +568,7 @@ class GemmSm90: :type epi_smem_layout: cute.ComposedLayout """ - from .trace import TraceContext - - tctx = TraceContext.create(trace_ptr) + from cutlass.cute.experimental import iket varlen_m = const_expr(varlen_params.cu_seqlens_m is not None) varlen_k = const_expr(varlen_params.cu_seqlens_k is not None) @@ -615,6 +577,7 @@ class GemmSm90: assert varlen_m or varlen_k has_D = const_expr(mD_mnl is not None) has_C = const_expr(mC_mnl is not None) + has_epi_load = const_expr(self.epi_c_stage > 0) warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -631,23 +594,23 @@ class GemmSm90: ab_pipeline = self.make_ab_pipeline( tiled_mma=tiled_mma, cluster_layout_vmnk=cute.make_layout((1, *cluster_layout_mnk.shape)), - ab_pipeline_mbar_ptr=storage.ab_pipeline_array_ptr.data_ptr(), ) epi_pipeline = None - if const_expr(has_C): - epi_pipeline = self.make_epi_pipeline( - c_smem_layout=cute.slice_(epi_c_smem_layout, (None, None, 0)), - epi_pipeline_mbar_ptr=storage.epi_pipeline_array_ptr.data_ptr(), - ) + if const_expr(has_epi_load): + epi_pipeline = self.make_epi_pipeline(tx_count=self.epi_load_bytes_per_stage) sched_pipeline = None sched_data = None if const_expr(self.is_persistent): - sched_pipeline = self.make_sched_pipeline( - cluster_layout_mnk, - sched_pipeline_mbar_ptr=storage.sched_pipeline_array_ptr.data_ptr(), - varlen_k=varlen_k, + sched_pipeline = self.make_sched_pipeline(cluster_layout_mnk, varlen_k=varlen_k) + # Keep scheduler scratch out of SharedStorage. A small buffer before + # the 1024-byte aligned epilogue tensors can add a 1 KiB pad; CLC + # responses also use i128 copies, so this stays 16-byte aligned. + sched_data = smem.allocate_tensor( + Int32, + cute.make_layout((4, self.sched_stage)), + byte_alignment=16, + partition=SmemPartition.RESERVED, ) - sched_data = storage.sched_data.get_tensor((4, self.sched_stage)) # Cluster arrive after barrier init pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk[:-1], is_relaxed=True) @@ -690,17 +653,18 @@ class GemmSm90: # PDL: wait for prior kernel before any TMA loads (matches cutlass C++ sm90 mainloop producer) if const_expr(self.use_pdl): cute.arch.griddepcontrol_wait() - # Get mcast mask - cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) - block_in_cluster_coord_mnk = cluster_layout_mnk.get_flat_coord(cta_rank_in_cluster) - a_mcast_mask = cute.make_layout_image_mask( - cluster_layout_mnk, block_in_cluster_coord_mnk, mode=1 - ) - b_mcast_mask = cute.make_layout_image_mask( - cluster_layout_mnk, block_in_cluster_coord_mnk, mode=0 - ) - a_mcast_mask = a_mcast_mask if self.is_a_mcast else 0 - b_mcast_mask = b_mcast_mask if self.is_b_mcast else 0 + # block_copy's lowering wants the coordinate held fixed by the + # multicast mask: A is same-M across N peers, while B is + # same-N across M peers. Degenerate cluster dimensions are + # left for the compiler lowering to simplify. + a_tma_multicast = { + "cluster_shape": self.cluster_shape_mnk[:2], + "multicast_dim": "M", + } + b_tma_multicast = { + "cluster_shape": self.cluster_shape_mnk[:2], + "multicast_dim": "N", + } # Persistent tile scheduling loop is_scheduler_warp = self.num_ab_load_warps == 1 or warp_idx == self.ab_load_warp_id @@ -712,7 +676,7 @@ class GemmSm90: pipeline.PipelineUserType.Producer, self.ab_stage ) while work_tile.is_valid_tile: - tctx.b("tma_load") + iket.range_push("tma_load") tile_coord_mnkl = work_tile.tile_idx batch_idx = tile_coord_mnkl[3] # Local_tile partition global tensors @@ -726,15 +690,11 @@ class GemmSm90: (tile_coord_mnkl[0], None), ) # TMA load A partition_S/D - copy_A, _, _ = copy_utils.tma_get_copy_fn( + copy_A = copy_utils.tma_get_block_copy_fn( tma_atom_a, - cta_coord=block_in_cluster_coord_mnk[1], - cta_layout=cute.make_layout( - cute.slice_(cluster_layout_mnk, (0, None, 0)).shape - ), src_tensor=gA_mk, dst_tensor=sA, - mcast_mask=a_mcast_mask, + tma_multicast=a_tma_multicast, ) else: copy_A, prefetch_A = self._make_gather_A_copy( @@ -747,21 +707,17 @@ class GemmSm90: (tile_coord_mnkl[1], None), ) # TMA load B partition_S/D - copy_B, _, _ = copy_utils.tma_get_copy_fn( + copy_B = copy_utils.tma_get_block_copy_fn( tma_atom_b, - cta_coord=block_in_cluster_coord_mnk[0], - cta_layout=cute.make_layout( - cute.slice_(cluster_layout_mnk, (None, 0, 0)).shape - ), src_tensor=gB_nk, dst_tensor=sB, - mcast_mask=b_mcast_mask, + tma_multicast=b_tma_multicast, ) len_k = varlen_manager.len_k(batch_idx) k_tile_cnt = cute.ceil_div(len_k, self.cta_tile_shape_mnk[2]) if const_expr(not self.gather_A): - ab_producer_state = self.load_AB( - ab_pipeline, ab_producer_state, copy_A, copy_B, k_tile_cnt + ab_producer_state = self.load_tma( + ab_pipeline, ab_producer_state, [copy_A, copy_B], k_tile_cnt ) else: ab_producer_state = self.load_AB_gather_A( @@ -773,7 +729,7 @@ class GemmSm90: k_tile_cnt, varlen_m=varlen_m, ) - tctx.e("tma_load") + iket.range_pop() tile_scheduler.advance_to_next_work(is_scheduler_warp=is_scheduler_warp) work_tile = tile_scheduler.get_current_work() # End of persistent scheduler loop @@ -824,7 +780,7 @@ class GemmSm90: k_tile_cnt_static = cute.ceil_div( cute.size(mA_mkl, mode=[1]), self.cta_tile_shape_mnk[2] ) - c_tile_cnt = cute.size(cute.ceil_div(self.cta_tile_shape_mnk[:2], self.epi_tile)) + c_tile_cnt = cute.size(self.epi_tile_shape) ab_read_state = make_pipeline_state(pipeline.PipelineUserType.Consumer, self.ab_stage) epi_store_pipeline = self.make_epi_store_pipeline() @@ -857,19 +813,19 @@ class GemmSm90: k_tile_cnt = cute.ceil_div(len_k, self.cta_tile_shape_mnk[2]) if const_expr(self.pingpong): self.pingpong_barrier_sync(warp_group_idx, stage="mma") - tctx.b("mma") + iket.range_push("mma") ab_read_state = self.mma( ab_pipeline, ab_read_state, mma_fn, acc, acc_slow, k_tile_cnt, warp_group_idx ) if const_expr(varlen_k): if k_tile_cnt == 0: acc.fill(0.0) - tctx.e("mma") + iket.range_pop() # EPILOGUE if const_expr(self.pingpong): self.pingpong_barrier_sync(warp_group_idx, "epi") - tctx.b("epilogue") + iket.range_push("epilogue") copy_D = None if const_expr(has_D): @@ -892,12 +848,21 @@ class GemmSm90: tile_coord_mnkl, ) copy_C = copy_utils.tma_producer_copy_fn(copy_C_fn, epi_pipeline) + if const_expr(has_epi_load): + tile_load_copy_fns = self.epi_tile_load_g2s_copy_fns( + epilogue_params, + epi_smem_tensors, + tile_coord_mnkl, + varlen_manager, + epi_pipeline, + ) + copy_C = copy_utils.chain_tma_producer_copy_fns((copy_C, *tile_load_copy_fns)) d_dtype_for_layout = self.d_dtype if self.d_dtype is not None else cutlass.BFloat16 tiled_copy_r2s, tRS_rD, tRS_sD = self.epilog_smem_store_and_partition( tiled_mma, self.d_layout, d_dtype_for_layout, sD, tidx ) - # (R2S, R2S_M, R2S_N, num_epi) + # (R2S, R2S_M, R2S_N, (epi_M, epi_N)) tRS_rAcc = self.epi_retile_acc(acc, tRS_rD, tiled_copy_r2s) load_acc_subtile = partial(self.epi_load_acc_subtile, tRS_rAcc) if const_expr(has_C): @@ -943,7 +908,7 @@ class GemmSm90: if is_tma_warp: epi_store_pipeline.producer_tail() self.pingpong_barrier_arrive(1 - warp_group_idx, stage="epi") - tctx.e("epilogue") + iket.range_pop() if const_expr(not self.pingpong): tile_scheduler.advance_to_next_work() @@ -977,48 +942,6 @@ class GemmSm90: if is_tma_warp: epi_store_pipeline.producer_tail() - tctx.flush() - - @cute.jit - def load_AB( - self, - ab_pipeline: cutlass.pipeline.PipelineAsync, - ab_producer_state: cutlass.pipeline.PipelineState, - copy_A: Optional[Callable], - copy_B: Callable, - k_tile_cnt: Int32, - # These are for Sm100 blockscaled gemm - copy_SFA: Optional[Callable] = None, - copy_SFB: Optional[Callable] = None, - ) -> cutlass.pipeline.PipelineState: - blockscaled = const_expr(copy_SFA is not None) - if const_expr(blockscaled): - assert copy_SFB is not None - # Peek (try_wait) AB buffer empty for k_block = prefetch_k_tile_cnt - peek_ab_empty_status = Boolean(True) - if 0 < k_tile_cnt: - peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) - # TMA load - for k_tile in cutlass.range(k_tile_cnt, unroll=1): - # Wait for A/B buffers to be empty before loading into them - # Also sets the transaction barrier for the A/B buffers - ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) - tma_bar_ptr = ab_pipeline.producer_get_barrier(ab_producer_state) - smem_idx = ab_producer_state.index - if const_expr(copy_A is not None): - copy_A(k_tile, smem_idx, tma_bar_ptr=tma_bar_ptr) - copy_B(k_tile, smem_idx, tma_bar_ptr=tma_bar_ptr) - if const_expr(blockscaled): - copy_SFA(k_tile, smem_idx, tma_bar_ptr=tma_bar_ptr) - copy_SFB(k_tile, smem_idx, tma_bar_ptr=tma_bar_ptr) - # Mainloop pipeline's producer commit is a NOP - ab_pipeline.producer_commit(ab_producer_state) - ab_producer_state.advance() - peek_ab_empty_status = Boolean(True) - if k_tile + 1 < k_tile_cnt: - peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) - return ab_producer_state - @cute.jit def load_AB_gather_A( self, @@ -1189,338 +1112,21 @@ class GemmSm90: acc.store(acc_slow.load()) return ab_read_state - @cute.jit - def epilogue( - self, - params: EpilogueParams, - epi_smem_tensors: Tuple[cute.Tensor, ...], - epi_pipeline: cutlass.pipeline.PipelineAsync, - epi_store_pipeline: cutlass.pipeline.PipelineAsync, - epi_read_state: cutlass.pipeline.PipelineState, - epi_producer_state: Optional[cutlass.pipeline.PipelineState], - epi_tile: cute.Tile, - load_acc_subtile: Callable, - tRS_rD: cute.Tensor, - tRS_rC: Optional[cute.Tensor], - tiled_copy_t2r: Optional[cute.TiledCopy], # Only for Sm100 - tiled_copy_r2s: cute.TiledCopy, - tRS_sD: cute.Tensor, - tiled_copy_s2r: Optional[cute.ThrCopy], - tSR_rC: Optional[cute.Tensor], - tSR_sC: Optional[cute.Tensor], - copy_D: Optional[Callable], - copy_C: Optional[Callable], - tile_coord_mnkl: cute.Coord, - varlen_manager: VarlenManager, - epilogue_barrier: cutlass.pipeline.NamedBarrier, - tile_scheduler, - tidx: Int32, - is_tma_warp: Boolean, - ) -> Tuple[cutlass.pipeline.PipelineState, cutlass.pipeline.PipelineState]: - has_C = const_expr(tRS_rC is not None) - has_D = const_expr(copy_D is not None) - - # Setup postact output (returns None for default epilogue, context tuple for Act) - postact_ctx = self.epi_setup_postact( - params, - epi_smem_tensors, - tiled_copy_r2s, - tiled_copy_t2r, - tile_coord_mnkl, - varlen_manager, - tidx, - ) - - epi_tile_shape = cute.zipped_divide( - cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile - ).shape[1] - # We iterate over epi tiles in the N dimension first before the M dimension - epi_tile_layout = cute.make_ordered_layout(epi_tile_shape, order=(1, 0)) - epi_tile_num = cute.size(epi_tile_shape) - num_prev_subtiles = tile_scheduler.num_tiles_executed * epi_tile_num - - epi_tensors = self.epi_begin( - params, - epi_smem_tensors, - epi_tile, - tiled_copy_t2r, - tiled_copy_r2s, - tile_coord_mnkl, - varlen_manager, - epilogue_barrier, - tidx, - ) - - if const_expr(copy_C is not None): - for epi_idx in cutlass.range(min(epi_tile_num, self.epi_c_stage), unroll=1): - gmem_coord_C = epi_tile_layout.get_hier_coord(epi_idx) - if is_tma_warp: - epi_pipeline.producer_acquire(epi_producer_state) - copy_C(src_idx=gmem_coord_C, producer_state=epi_producer_state) - epi_pipeline.producer_commit(epi_producer_state) - epi_producer_state.advance() - - for epi_idx in cutlass.range_constexpr(epi_tile_num): - # The global memory coordinate for the current epi tile - gmem_coord = epi_tile_layout.get_hier_coord(epi_idx) - # Copy from acc to D registers - load_acc_subtile(tRS_rD, epi_idx) - epi_loop_tensors = self.epi_begin_loop(params, epi_tensors, gmem_coord) - if const_expr(has_C): - epi_pipeline.consumer_wait(epi_read_state) - cute.copy(tiled_copy_s2r, tSR_sC[None, None, None, epi_read_state.index], tSR_rC) - # Fence to make sure shared memory read is visible to TMA load - cute.arch.fence_view_async_shared() - cute.arch.sync_warp() - with cute.arch.elect_one(): - epi_pipeline.consumer_release(epi_read_state) - epi_read_state.advance() - if const_expr(copy_C is not None and epi_idx + self.epi_c_stage < epi_tile_num): - gmem_coord_C = epi_tile_layout.get_hier_coord(epi_idx + self.epi_c_stage) - if is_tma_warp: - epi_pipeline.producer_acquire(epi_producer_state) - copy_C(src_idx=gmem_coord_C, producer_state=epi_producer_state) - epi_pipeline.producer_commit(epi_producer_state) - epi_producer_state.advance() - tRS_rPostAct = self.epi_visit_subtile(params, epi_loop_tensors, tRS_rD, tRS_rC) - # Convert and store postact if this epilogue produces one - if const_expr(postact_ctx is not None): - tRS_rPostAct_out = self.epi_convert_postact( - tRS_rPostAct, - epi_loop_tensors["sr_seed"], - tidx, - tile_coord_mnkl, - num_prev_subtiles, - epi_idx, - ) - if is_tma_warp: - epi_store_pipeline.producer_acquire() - epilogue_barrier.arrive_and_wait() - # Copy from D registers to shared memory - epi_buffer = (num_prev_subtiles + epi_idx) % self.epi_stage - if const_expr(has_D): - if const_expr( - self.rounding_mode == RoundingMode.RS - and self.acc_dtype == cutlass.Float32 - and self.d_dtype == cutlass.BFloat16 - ): - seed = epi_loop_tensors["sr_seed"] + ( - tile_coord_mnkl[0] * 65537 - + tile_coord_mnkl[1] * 257 - + tile_coord_mnkl[3] * 17 - + (num_prev_subtiles + epi_idx) * 7 - ) - copy_utils.sr_cvt_copy( - tiled_copy_r2s, - tRS_rD, - tRS_sD[None, None, None, epi_buffer], - seed, - tidx, - ) - else: - copy_utils.cvt_copy( - tiled_copy_r2s, tRS_rD, tRS_sD[None, None, None, epi_buffer] - ) - # Copy postact from registers to shared memory - if const_expr(postact_ctx is not None): - tiled_copy_postact_r2s, tRS_sPostAct, copy_postact = postact_ctx - cute.copy( - tiled_copy_postact_r2s, - tiled_copy_postact_r2s.retile(tRS_rPostAct_out), - tRS_sPostAct[None, None, None, epi_buffer], - ) - # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_view_async_shared() - epilogue_barrier.arrive_and_wait() - # Copy from shared memory to global memory - if is_tma_warp: - if const_expr(has_D): - copy_D(src_idx=epi_buffer, dst_idx=gmem_coord) - if const_expr(postact_ctx is not None): - copy_postact(src_idx=epi_buffer, dst_idx=gmem_coord) - epi_store_pipeline.producer_commit() - - self.epi_end( - params, - epi_tensors, - epi_tile, - tiled_copy_t2r, - tiled_copy_r2s, - tile_coord_mnkl, - varlen_manager, - tidx, - ) - - return epi_read_state, epi_producer_state - - def get_scheduler_class(self, varlen_m: bool = False): - """Return the scheduler class to use. Override in subclasses for custom schedulers.""" - return TileScheduler if not varlen_m else VarlenMTileScheduler - - def get_scheduler_arguments( - self, - mA: cute.Tensor, - mB: cute.Tensor, - mD: Optional[cute.Tensor], - scheduler_args, - varlen_args, - epilogue_args, - ): - """Create scheduler arguments. Override in subclasses for custom schedulers.""" - if const_expr(not self.is_persistent): - persistence_mode = PersistenceMode.NONE - else: - if const_expr(self.arch >= 100 and self.use_clc_persistence): - persistence_mode = PersistenceMode.CLC - elif const_expr(scheduler_args.tile_count_semaphore is not None): - persistence_mode = PersistenceMode.DYNAMIC - else: - persistence_mode = PersistenceMode.STATIC - if const_expr(varlen_args.mCuSeqlensM is None): - num_problems = ( - mD.shape[2] - if mD is not None - else ( - mB.shape[2] - if varlen_args.mCuSeqlensK is None - else varlen_args.mCuSeqlensK.shape[0] - 1 - ) - ) - problem_shape_ntile_mnl = ( - cute.ceil_div(cute.size(mA, mode=[0]), self.cta_tile_shape_mnk[0]), - cute.ceil_div(cute.size(mB, mode=[0]), self.cta_tile_shape_mnk[1]), - num_problems, - ) - tile_sched_args = TileSchedulerArguments( - problem_shape_ntile_mnl=problem_shape_ntile_mnl, - raster_order=scheduler_args.raster_order, - group_size=scheduler_args.max_swizzle_size, - cluster_shape_mnk=self.cluster_shape_mnk, - tile_count_semaphore=scheduler_args.tile_count_semaphore, - batch_idx_permute=scheduler_args.batch_idx_permute, - persistence_mode=persistence_mode, - ) - else: - assert (mD is not None) or (epilogue_args.mPostAct is not None) or (not self.gather_A) - problem_shape_ntile_mnl = ( - None, - cute.ceil_div(cute.size(mB, mode=[0]), self.cta_tile_shape_mnk[1]), - varlen_args.mCuSeqlensM.shape[0] - 1, - ) - tile_sched_args = VarlenMTileSchedulerArguments( - problem_shape_ntile_mnl=problem_shape_ntile_mnl, - total_m=mD.shape[0] if mD is not None else varlen_args.mAIdx.shape[0], - cu_seqlens_m=varlen_args.mCuSeqlensM, - raster_order=scheduler_args.raster_order, - group_size=scheduler_args.max_swizzle_size, - tile_shape_mn=self.cta_tile_shape_mnk[:2], - cluster_shape_mnk=self.cluster_shape_mnk, - tile_count_semaphore=scheduler_args.tile_count_semaphore, - persistence_mode=persistence_mode, - ) - return tile_sched_args - def epi_retile_acc(self, acc, tRS_rD, tiled_copy_r2s): - """Retile accumulator for epilogue subtile access. SM90 uses flat_divide.""" - return cute.flat_divide(acc, tRS_rD.layout) - - @cute.jit - def epi_load_acc_subtile(self, tRS_rAcc: cute.Tensor, tRS_rD: cute.Tensor, epi_idx: int): - cute.autovec_copy(tRS_rAcc[None, None, None, epi_idx], tRS_rD) - - @cute.jit - def epi_begin( - self, - params: EpilogueParams, - epi_smem_tensors: Tuple[cute.Tensor, ...], - epi_tile: cute.Tile, - tiled_copy_t2r: Optional[cute.TiledCopy], - tiled_copy_r2s: cute.TiledCopy, - tile_coord_mnkl: cute.Coord, - varlen_manager: VarlenManager, - epilogue_barrier: cutlass.pipeline.NamedBarrier, - tidx: Int32, - ) -> Tuple[cute.Tensor, ...]: - return () - - def epi_begin_loop( - self, params: EpilogueParams, epi_tensors: Tuple[cute.Tensor, ...], epi_coord: cute.Coord - ) -> Tuple[cute.Tensor, ...]: - return () - - def epi_visit_subtile( - self, - params: EpilogueParams, - epi_loop_tensors: Tuple[cute.Tensor, ...], - tRS_rD: cute.Tensor, - tRS_rC: Optional[cute.Tensor] = None, - ) -> Optional[cute.Tensor]: - return None - - def epi_visit_acc( - self, - params: EpilogueParams, - acc: cute.Tensor, - tiled_mma: cute.TiledMma, - tile_coord_mnkl: cute.Coord, - tidx: Int32, - ) -> None: - pass - - @cute.jit - def epi_end( - self, - params: EpilogueParams, - epi_tensors: Tuple[cute.Tensor, ...], - epi_tile: cute.Tile, - tiled_copy_t2r: Optional[cute.TiledCopy], - tiled_copy_r2s: cute.TiledCopy, - tile_coord_mnkl: cute.Coord, - varlen_manager, - tidx, - ) -> None: - pass - - def epi_to_underlying_arguments( - self, args: EpilogueArguments, *, loc=None, ip=None - ) -> EpilogueParams: - return self.EpilogueParams() - - def epi_get_tma_atoms( - self, params: EpilogueParams, *, loc=None, ip=None - ) -> list[cute.CopyAtom]: - """Subclasses can override this""" - return [] - - @staticmethod - def epi_smem_bytes_per_stage( - args: Optional[EpilogueArguments], - cta_tile_shape_mnk: Tuple[int, int, int], - epi_tile: cute.Tile, - ) -> int: - return 0 - - def epi_get_smem_struct(self, params: EpilogueParams): - return cute.struct.MemRange[Int32, 0] # Dummy struct - - def epi_get_smem_tensors(self, params: EpilogueParams, storage) -> Tuple[cute.Tensor, ...]: - return tuple() - - def pingpong_barrier_sync(self, warp_group_idx: Int32, stage: Literal["mma", "epi"]): - assert stage in ["mma", "epi"] - barrier = NamedBarrierGemm.MmaWG0 if stage == "mma" else NamedBarrierGemm.EpiWG0 - cute.arch.barrier( - barrier_id=int(barrier) + warp_group_idx, - number_of_threads=2 * self.num_threads_per_warp_group, - ) - - def pingpong_barrier_arrive(self, warp_group_idx: Int32, stage: Literal["mma", "epi"]): - assert stage in ["mma", "epi"] - barrier = NamedBarrierGemm.MmaWG0 if stage == "mma" else NamedBarrierGemm.EpiWG0 - cute.arch.barrier_arrive( - barrier_id=int(barrier) + warp_group_idx, - number_of_threads=2 * self.num_threads_per_warp_group, + """Retile accumulator for epilogue subtile access.""" + acc_reshaped = layout_utils.reshape_acc_to_frgA(acc) # ((2, 2, 2), MMA_M, MMA_N) + # ((2, 2, 2), MMA_M / epi_M, MMA_N / epi_N) + epi_acc_shape = ( + acc_reshaped.shape[0], + *cute.ceil_div(acc_reshaped.shape[1:], self.epi_tile_shape), ) + # ((2, 2, 2), MMA_M / epi_M, MMA_N / epi_N, (1, 1, 1), epi_M, epi_N) + acc_divide = cute.flat_divide(acc_reshaped, epi_acc_shape) + assert cute.size(acc_divide, mode=[3]) == 1 + # ((2, 2, 2), MMA_M / epi_M, MMA_N / epi_N, (epi_M, epi_N)) + tRS_rAcc = cute.group_modes(acc_divide[None, None, None, 0, None, None], 3, 5) + # (((2,2,2),1), MMA_M / epi_M, MMA_N / epi_N, (epi_M, epi_N)) + return tiled_copy_r2s.retile(tRS_rAcc) def epilog_smem_copy_atom(self, tiled_mma: cute.TiledMma) -> cute.TiledCopy: copy_atom_C = cute.make_copy_atom( @@ -1576,90 +1182,23 @@ class GemmSm90: tSR_rC = thr_copy_s2r.retile(tRS_rC) return tiled_copy_s2r, tRS_rC, tSR_rC, tSR_sC - def epilog_gmem_copy_and_partition( - self, - atom: Union[cute.CopyAtom, cute.TiledCopy], - mD_mn: cute.Tensor, - tile_shape_mn: cute.Tile, - epi_tile: cute.Tile, - sD: cute.Tensor, - tile_coord_mnkl: cute.Coord, - ) -> Tuple[cute.Tensor, cute.Tensor]: - # (bM, bN) - gD = cute.local_tile(mD_mn, tile_shape_mn, tile_coord_mnkl[:2]) - tDgD_for_tma_partition = cute.zipped_divide(gD, epi_tile) - is_s2g = isinstance( - atom.op, (cpasync.CopyBulkTensorTileS2GOp, cpasync.CopyReduceBulkTensorTileS2GOp) - ) - src_tensor, dst_tensor = ( - (sD, tDgD_for_tma_partition) if is_s2g else (tDgD_for_tma_partition, sD) - ) - return copy_utils.tma_get_copy_fn( - atom, - cta_coord=0, - cta_layout=cute.make_layout(1), - src_tensor=src_tensor, - dst_tensor=dst_tensor, - ) - - def make_ab_pipeline( - self, - tiled_mma: cute.TiledMma, - cluster_layout_vmnk: cute.Layout, - ab_pipeline_mbar_ptr: cute.Pointer, - ): - # Threads/warps participating in this pipeline - producer_cnt = 1 if const_expr(not self.gather_A) else 1 + self.num_ab_load_warps * 32 - ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt) - # Each warp will contribute to the arrive count with the number of mcast size - mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 - consumer_arrive_cnt = mcast_size * tiled_mma.size // cute.arch.WARP_SIZE - ab_pipeline_consumer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, consumer_arrive_cnt - ) - pipeline_cls = pipeline.PipelineTmaAsync if not self.gather_A else PipelineTmaCpAsync - return pipeline_cls.create( - barrier_storage=ab_pipeline_mbar_ptr, - num_stages=self.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, - ) - - def make_epi_pipeline( - self, c_smem_layout: cute.Layout | cute.ComposedLayout, epi_pipeline_mbar_ptr: cute.Pointer - ): - # Threads/warps participating in this pipeline - epi_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - # Each warp will contribute 1 to the arrive count - consumer_arrive_cnt = self.num_epi_warps - epi_pipeline_consumer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, consumer_arrive_cnt - ) - tma_copy_c_bytes = cute.size_in_bytes(self.c_dtype, c_smem_layout) - return pipeline.PipelineTmaAsync.create( - barrier_storage=epi_pipeline_mbar_ptr, - num_stages=self.epi_c_stage, - producer_group=epi_pipeline_producer_group, - consumer_group=epi_pipeline_consumer_group, - tx_count=tma_copy_c_bytes, - defer_sync=True, + def pingpong_barrier_sync(self, warp_group_idx: Int32, stage: str): + assert stage in ["mma", "epi"] + barrier = NamedBarrierGemm.MmaWG0 if stage == "mma" else NamedBarrierGemm.EpiWG0 + cute.arch.barrier( + barrier_id=int(barrier) + warp_group_idx, + number_of_threads=2 * self.num_threads_per_warp_group, ) - def make_epi_store_pipeline(self): - # Threads/warps participating in tma store pipeline - num_epi_threads = self.num_epi_warps * cute.arch.WARP_SIZE - epi_store_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_epi_threads) - return pipeline.PipelineTmaStore.create( - num_stages=self.epi_stage, producer_group=epi_store_producer_group + def pingpong_barrier_arrive(self, warp_group_idx: Int32, stage: str): + assert stage in ["mma", "epi"] + barrier = NamedBarrierGemm.MmaWG0 if stage == "mma" else NamedBarrierGemm.EpiWG0 + cute.arch.barrier_arrive( + barrier_id=int(barrier) + warp_group_idx, + number_of_threads=2 * self.num_threads_per_warp_group, ) - def make_sched_pipeline( - self, cluster_layout_mnk: cute.Layout, sched_pipeline_mbar_ptr: cute.Pointer, varlen_k: bool - ): - # Threads/warps participating in this pipeline + def make_sched_pipeline(self, cluster_layout_mnk: cute.Layout, varlen_k: bool): sched_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) cluster_size = cute.size(cluster_layout_mnk) # Each warp will contribute 1 to the arrive count @@ -1672,14 +1211,16 @@ class GemmSm90: sched_pipeline_consumer_group = pipeline.CooperativeGroup( pipeline.Agent.Thread, consumer_arrive_cnt ) - return pipeline.PipelineAsync.create( - barrier_storage=sched_pipeline_mbar_ptr, + return QuackPipelineAsync.create( num_stages=self.sched_stage, producer_group=sched_pipeline_producer_group, consumer_group=sched_pipeline_consumer_group, # If there's cluster, the consumers must arrive at the mbar of CTA 0 in the cluster. consumer_mask=None if const_expr(cluster_size == 1) else 0, defer_sync=True, + # One arrive per consumer warp (consumer_arrive_cnt counts warps): syncwarp + # so every lane's slot read is complete, then one elected lane signals. + elect_one_release=True, ) @classmethod @@ -1694,6 +1235,7 @@ class GemmSm90: epilogue_args: EpilogueArguments, smem_capacity: int, occupancy: int, + warp_shape_mnk: Tuple[int, int, int] | None = None, ) -> Tuple[int, int]: """Computes the number of stages for A/B/C operands based on heuristics. @@ -1714,14 +1256,21 @@ class GemmSm90: """ epi_stage = 4 if epi_tile[1] <= 16 else 2 - d_bytes_per_stage = cute.size(epi_tile) * d_dtype.width // 8 if d_dtype is not None else 0 - epi_bytes_per_stage = d_bytes_per_stage + cls.epi_smem_bytes_per_stage( - epilogue_args, cta_tile_shape_mnk, epi_tile + epi_smem_bytes = cls.epi_smem_bytes( + epilogue_args, cta_tile_shape_mnk, epi_tile, warp_shape_mnk + ) + has_tile_load = epi_smem_bytes.c_stage > 0 + epi_tile_elems = cute.size(cute.shape(epi_tile)) + d_bytes_per_stage = epi_tile_elems * d_dtype.width // 8 if d_dtype is not None else 0 + epi_bytes_per_stage = d_bytes_per_stage + epi_smem_bytes.d_stage + epi_bytes = epi_smem_bytes.unstaged + epi_bytes_per_stage * epi_stage + epi_c_stage = ( + 0 if c_dtype is None and not has_tile_load else (4 if epi_tile[1] <= 16 else 2) ) - epi_bytes = epi_bytes_per_stage * epi_stage - epi_c_stage = 0 if c_dtype is None else (4 if epi_tile[1] <= 16 else 2) if c_dtype is not None: - epi_bytes += cute.size(epi_tile) * c_dtype.width // 8 * epi_c_stage + epi_bytes += epi_tile_elems * c_dtype.width // 8 * epi_c_stage + if has_tile_load: + epi_bytes += epi_smem_bytes.c_stage * epi_c_stage a_shape = cute.slice_(cta_tile_shape_mnk, (None, 0, None)) b_shape = cute.slice_(cta_tile_shape_mnk, (0, None, None)) @@ -1741,7 +1290,7 @@ class GemmSm90: return ab_stage, epi_stage, epi_c_stage @staticmethod - def _sm90_compute_tile_shape_or_override( + def _compute_tile_shape_or_override( cta_tile_shape_mnk: Tuple[int, int, int], atom_layout_mnk: Tuple[int, int, int], element_type: Optional[Type[cutlass.Numeric]] = None, @@ -1827,8 +1376,8 @@ class GemmSm90: """ a_smem_shape = cute.slice_(cta_tile_shape_mnk, (None, 0, None)) - a_is_k_major = a_layout.sm90_mma_major_mode() == warpgroup.OperandMajorMode.K - b_is_k_major = b_layout.sm90_mma_major_mode() == warpgroup.OperandMajorMode.K + a_is_k_major = a_layout.sm90_mma_major_mode() == cute.nvgpu.OperandMajorMode.K + b_is_k_major = b_layout.sm90_mma_major_mode() == cute.nvgpu.OperandMajorMode.K a_major_mode_size = cta_tile_shape_mnk[2 if a_is_k_major else 0] a_smem_layout_atom = warpgroup.make_smem_layout_atom( sm90_utils.get_smem_layout_atom(a_layout, a_dtype, a_major_mode_size), @@ -1873,105 +1422,6 @@ class GemmSm90: epi_c_smem_layout_staged, ) - @staticmethod - def _make_tma_epi_atoms_and_tensors( - tensor_d: cute.Tensor, - epi_smem_layout_staged: cute.ComposedLayout, - epi_tile: Tuple[int, int], - op_type: Literal["store", "load", "add"], - ) -> Tuple[cute.CopyAtom, cute.Tensor]: - """Create TMA atoms and tensors for storing D or loading C. - - :param tensor_d: Output tensor D - :type tensor_d: cute.Tensor - :param epi_smem_layout_staged: Shared memory layout for epilogue - :type epi_smem_layout_staged: cute.ComposedLayout - :param epi_tile: Epilogue tile shape - :type epi_tile: Tuple[int, int] - - :return: TMA atom and tensor for C - :rtype: Tuple[cute.CopyAtom, cute.Tensor] - """ - assert op_type in ["load", "store", "add"] - epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0)) - d_cta_v_layout = cute.composition(cute.make_identity_layout(tensor_d.shape), epi_tile) - op = ( - cpasync.CopyBulkTensorTileG2SOp() - if op_type == "load" - else cpasync.CopyBulkTensorTileS2GOp() - if op_type == "store" - else cpasync.CopyReduceBulkTensorTileS2GOp(cute.ReductionOp.ADD) - ) - tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom( - op, tensor_d, epi_smem_layout, d_cta_v_layout - ) - return tma_atom_d, tma_tensor_d - - @staticmethod - def _make_tma_atoms_and_tensors( - tensor: cute.Tensor, - smem_layout: cute.ComposedLayout, - smem_tile: Tuple[int, int], - mcast_dim: int, - ) -> Tuple[cute.CopyAtom, cute.Tensor]: - """Create TMA atoms and tensors for input tensors. - - :param tensor: Input tensor (A or B) - :type tensor: cute.Tensor - :param smem_layout: Shared memory layout for the tensor - :type smem_layout: cute.ComposedLayout - :param smem_tile: Shared memory tile shape - :type smem_tile: Tuple[int, int] - :param mcast_dim: Multicast dimension - :type mcast_dim: int - - :return: TMA atom and tensor - :rtype: Tuple[cute.CopyAtom, cute.Tensor] - """ - op = ( - cpasync.CopyBulkTensorTileG2SOp() - if mcast_dim == 1 - else cpasync.CopyBulkTensorTileG2SMulticastOp() - ) - tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( - op, - tensor, - smem_layout, - smem_tile, - num_multicast=mcast_dim, - ) - return tma_atom, tma_tensor - - def _make_gmem_tiled_copy_A(self, dtype, major_mode, num_threads, copy_bits=128): - atom_async_copy = cute.make_copy_atom( - cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), - dtype, - num_bits_per_copy=copy_bits, - ) - copy_elems = copy_bits // dtype.width - loads_per_cache_line = 128 * 8 // copy_bits # 128 bytes per cache line - shape_dim_1 = cute.size(self.cta_tile_shape_mnk[2]) // copy_elems - if shape_dim_1 > loads_per_cache_line: - shape_dim_1 = math.gcd(shape_dim_1, loads_per_cache_line) - # thread layout for copy - thread_layout = cute.make_layout( - (num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1) - ) - if major_mode != LayoutEnum.ROW_MAJOR: - shape_dim_0 = cute.size(self.cta_tile_shape_mnk[0]) // copy_elems - if shape_dim_0 > loads_per_cache_line: - shape_dim_0 = math.gcd(shape_dim_0, loads_per_cache_line) - thread_layout = cute.make_layout( - (shape_dim_0, num_threads // shape_dim_0), stride=(1, shape_dim_0) - ) - # Value layout for copy - value_layout = ( - cute.make_layout((1, copy_elems)) - if major_mode == LayoutEnum.ROW_MAJOR - else cute.make_layout((copy_elems, 1)) - ) - return cute.make_tiled_copy_tv(atom_async_copy, thread_layout, value_layout) - @staticmethod def is_valid_dtypes( a_dtype: Type[cutlass.Numeric], diff --git a/build/torch-cuda/quack/gemm_sq_reduce.py b/build/torch-cuda/quack/gemm_sq_reduce.py index 1dac5d25c974f08175e45ba703e92bf0fcf94cab..9a088d65261fd281b1effba78d4a3eb2dfee0a9f 100644 --- a/build/torch-cuda/quack/gemm_sq_reduce.py +++ b/build/torch-cuda/quack/gemm_sq_reduce.py @@ -8,7 +8,7 @@ from torch import Tensor import cutlass import cutlass.cute as cute -from cutlass import Float32, const_expr +from cutlass import Int32, Float32, const_expr from .cute_dsl_utils import ( mlir_namedtuple, @@ -16,18 +16,29 @@ from .cute_dsl_utils import ( get_device_capacity, get_max_active_clusters, ) -from .epi_ops import ColVecReduce, colvec_reduce_accumulate, vec_multiply +from .epi_ops import ( + ColVecReduce, + RowVecLoad, + Scalar, + TileStore, + colvec_reduce_accumulate, + vec_multiply, +) +from .gemm_act import GemmActMixin +from .gemm_sm80 import GemmSm80 from .gemm_sm90 import GemmSm90 from .gemm_sm100 import GemmSm100 from .gemm_sm120 import GemmSm120 -from .gemm_default_epi import GemmDefaultEpiMixin from .rounding import RoundingMode from .compile_utils import make_fake_tensor as fake_tensor -from .cache_utils import jit_cache +from .cache import jit_cache from .gemm_tvm_ffi_utils import ( + div_for_dtype, + get_major, get_majors, get_dtypes, perm3d, + perm3d_single, make_scheduler_args, make_varlen_args, make_fake_scheduler_args, @@ -38,37 +49,50 @@ from .gemm_tvm_ffi_utils import ( from . import utils as utils -class GemmSqReduceMixin(GemmDefaultEpiMixin): +class GemmSqReduceMixin(GemmActMixin): """GEMM + sq_reduce + optional rowvec scaling. D_raw = A @ B (+ C), reduce[m] = sum_n(D_raw[m,n]^2), D_out = D_raw * rowvec. - The sq_sum is computed BEFORE the rowvec scaling. + The sq_sum is computed BEFORE the rowvec scaling. If mAuxOut is provided, the + pre-rowvec value (D_raw, after alpha/beta/C) is written to it. """ - _epi_ops = (*GemmDefaultEpiMixin._epi_ops, ColVecReduce("mColVecReduce")) + _epi_ops = ( + Scalar("alpha"), + Scalar("beta"), + Scalar("sr_seed", dtype=Int32), + RowVecLoad("mRowVecBroadcast"), + ColVecReduce("mColVecReduce"), + TileStore("mAuxOut"), + ) + _extra_param_fields = () # no act_fn @mlir_namedtuple class EpilogueArguments(NamedTuple): alpha: Optional[Float32 | cute.Tensor] = None beta: Optional[Float32 | cute.Tensor] = None mRowVecBroadcast: Optional[cute.Tensor] = None - mColVecBroadcast: Optional[cute.Tensor] = None mColVecReduce: Optional[cute.Tensor] = None + mAuxOut: Optional[cute.Tensor] = None add_to_output: cutlass.Constexpr[bool] = False rounding_mode: cutlass.Constexpr[int] = RoundingMode.RN - sr_seed: None = None + sr_seed: Optional[Int32 | cute.Tensor] = None # EpilogueParams auto-generated from _epi_ops def epi_to_underlying_arguments(self, args, *, loc=None, ip=None): self.rounding_mode = args.rounding_mode + if args.mAuxOut is not None: + self.aux_out_dtype = args.mAuxOut.element_type + self.aux_out_layout = cutlass.utils.LayoutEnum.from_tensor(args.mAuxOut) + self.cta_tile_shape_aux_out_mn = self.cta_tile_shape_mnk[:2] d = self._epi_ops_to_params_dict(args) return self.EpilogueParams(**d) @cute.jit def epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC=None): - tDrColVecReduce = epi_loop_tensors["mColVecReduce"] - tDrRowVec = epi_loop_tensors["mRowVecBroadcast"] + tDrColVecReduce = epi_loop_tensors.get("mColVecReduce") + tDrRowVec = epi_loop_tensors.get("mRowVecBroadcast") # Load accumulator, apply alpha/beta/C (skip rowvec/colvec — we handle rowvec below) rD = tRS_rD.load() if const_expr(hasattr(params, "alpha") and params.alpha is not None): @@ -83,15 +107,26 @@ class GemmSqReduceMixin(GemmDefaultEpiMixin): tRS_rD.store(rD) # Accumulate sq_sum BEFORE rowvec scaling: reduce[m] += sum_n(D[m,n]^2) colvec_reduce_accumulate(self, tDrColVecReduce, tRS_rD, rScale=tRS_rD) + # Snapshot pre-rowvec value if the caller wants the aux output written. + if const_expr(getattr(params, "mAuxOut", None) is not None): + tRS_rAuxOut = cute.make_rmem_tensor_like(tRS_rD) + tRS_rAuxOut.store(tRS_rD.load()) + tRS_rAuxOuts = (tRS_rAuxOut,) + else: + tRS_rAuxOuts = () # Multiply by rowvec (norm_weight) AFTER sq_sum vec_multiply(self, tRS_rD, None, tDrRowVec) - return None + return tRS_rAuxOuts class GemmSqReduceSm90(GemmSqReduceMixin, GemmSm90): pass +class GemmSqReduceSm80(GemmSqReduceMixin, GemmSm80): + pass + + class GemmSqReduceSm100(GemmSqReduceMixin, GemmSm100): pass @@ -118,9 +153,12 @@ def _compile_gemm_sq_reduce( colvec_reduce_dtype, colvec_reduce_ndim, rowvec_dtype, + aux_out_dtype, + aux_out_major, device_capacity, ): sm_to_cls = { + 8: GemmSqReduceSm80, 9: GemmSqReduceSm90, 10: GemmSqReduceSm100, 11: GemmSqReduceSm100, @@ -153,9 +191,20 @@ def _compile_gemm_sq_reduce( divisibility=1, ) mRowVec = fake_tensor(rowvec_dtype, (l, n), leading_dim=1, divisibility=4) + if aux_out_dtype is not None: + aux_leading = 1 if aux_out_major == "n" else 0 + mAuxOut = fake_tensor( + aux_out_dtype, + (m, n, l), + leading_dim=aux_leading, + divisibility=div_for_dtype(aux_out_dtype), + ) + else: + mAuxOut = None epi_args = GemmCls.EpilogueArguments( mRowVecBroadcast=mRowVec, mColVecReduce=mColVecReduce, + mAuxOut=mAuxOut, ) scheduler_args = make_fake_scheduler_args( (is_dynamic_persistent and device_capacity[0] == 9), False, l @@ -192,24 +241,36 @@ def gemm_sq_reduce( tile_N: int, cluster_M: int, cluster_N: int, + tile_K: int | None = None, pingpong: bool = False, persistent: bool = True, is_dynamic_persistent: bool = False, max_swizzle_size: int = 8, rowvec: Optional[Tensor] = None, # (l, n) — norm_weight + aux_out: Optional[Tensor] = None, # (l, m, n) — pre-rowvec output snapshot ) -> None: """GEMM + sq_reduce + optional rowvec scaling. D_raw = A @ B (+ C), colvec_reduce[m] = sum_n(D_raw[m,n]^2), D_out = D_raw * rowvec. + If aux_out is provided, the pre-rowvec value (D_raw, after alpha/beta/C) is also + written to it. """ device_capacity = get_device_capacity(A.device) - assert device_capacity[0] in [9, 10, 11, 12], "Only SM90, SM100, SM110, and SM120 are supported" - if device_capacity[0] == 12: - raise NotImplementedError("SM120 GEMM sq reduce epilogue is not yet supported") + assert device_capacity[0] in [8, 9, 10, 11, 12], ( + "Only SM8x, SM90, SM100, SM110, and SM120 are supported" + ) A_p, B_p, D_p, C_p = perm3d(A, B, D, C) a_major, b_major, d_major, c_major = get_majors(A_p, B_p, D_p, C_p) a_dtype, b_dtype, d_dtype, c_dtype = get_dtypes(A, B, D, C) + if aux_out is not None: + AuxOut_p = perm3d_single(aux_out) + aux_out_dtype = torch2cute_dtype_map[aux_out.dtype] + aux_out_major = get_major(AuxOut_p, "m", "n") + else: + AuxOut_p = None + aux_out_dtype = None + aux_out_major = None if is_dynamic_persistent and device_capacity[0] == 9: assert tile_count_semaphore is not None, ( @@ -225,7 +286,7 @@ def gemm_sq_reduce( b_major, d_major, c_major, - (tile_M, tile_N), + (tile_M, tile_N, tile_K) if tile_K is not None else (tile_M, tile_N), (cluster_M, cluster_N, 1), pingpong, persistent, @@ -233,18 +294,16 @@ def gemm_sq_reduce( torch2cute_dtype_map[colvec_reduce.dtype], colvec_reduce.ndim, torch2cute_dtype_map[rowvec.dtype] if rowvec is not None else None, + aux_out_dtype, + aux_out_major, device_capacity, ) - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY: - return - max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0 epi_args = GemmSqReduceMixin.EpilogueArguments( mRowVecBroadcast=rowvec, mColVecReduce=colvec_reduce, + mAuxOut=AuxOut_p, add_to_output=None, # Constexpr, pass None at runtime rounding_mode=None, # Constexpr, pass None at runtime ) @@ -254,6 +313,6 @@ def gemm_sq_reduce( varlen_args = make_varlen_args(None, None, None) if device_capacity[0] in [10, 11]: - compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None, None) + compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None) else: - compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None) + compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args) diff --git a/build/torch-cuda/quack/gemm_symmetric.py b/build/torch-cuda/quack/gemm_symmetric.py index 9467efe662a368fae33698d3935b8787a0a9b29c..6821bbd3a4174c6c293fac0585ce24c45d1baf26 100644 --- a/build/torch-cuda/quack/gemm_symmetric.py +++ b/build/torch-cuda/quack/gemm_symmetric.py @@ -1,4 +1,4 @@ -from typing import Tuple, Optional, Callable +from typing import Dict, Tuple, Optional, Callable from torch import Tensor @@ -11,6 +11,7 @@ from .compile_utils import make_fake_tensor as fake_tensor from .cute_dsl_utils import get_device_capacity, get_max_active_clusters, torch2cute_dtype_map from .activation import act_fn_map from .gemm_act import GemmActMixin +from .gemm_sm80 import GemmSm80 from .gemm_sm90 import GemmSm90 from .gemm_sm100 import GemmSm100 from .gemm_sm120 import GemmSm120 @@ -23,11 +24,11 @@ from .gemm_tvm_ffi_utils import ( make_fake_scheduler_args, compile_gemm_kernel, ) -from .cache_utils import jit_cache +from .cache import jit_cache from .tile_scheduler import TriangularTileScheduler from .varlen_utils import VarlenManager from . import copy_utils as copy_utils -from .rounding import RoundingMode +from .rounding import RoundingMode, epilogue_sr_seed class GemmSymmetricMixin(GemmActMixin): @@ -38,11 +39,11 @@ class GemmSymmetricMixin(GemmActMixin): def epilogue( self, params: GemmActMixin.EpilogueParams, - epi_smem_tensors: Tuple[cute.Tensor, ...], - epi_pipeline: cutlass.pipeline.PipelineAsync, - epi_store_pipeline: cutlass.pipeline.PipelineAsync, - epi_read_state: cutlass.pipeline.PipelineState, - epi_producer_state: cutlass.pipeline.PipelineState, + epi_smem_tensors: Dict[str, cute.Tensor], + epi_pipeline: Optional[cutlass.pipeline.PipelineAsync], + epi_store_pipeline: Optional[cutlass.pipeline.PipelineAsync], + epi_read_state: Optional[cutlass.pipeline.PipelineState], + epi_producer_state: Optional[cutlass.pipeline.PipelineState], epi_tile: cute.Tile, load_acc_subtile: Callable, tRS_rD: cute.Tensor, @@ -50,7 +51,7 @@ class GemmSymmetricMixin(GemmActMixin): tiled_copy_t2r: Optional[cute.TiledCopy], # Only for Sm100 tiled_copy_r2s: cute.TiledCopy, tRS_sD: cute.Tensor, - tiled_copy_s2r: Optional[cute.TiledCopy], + tiled_copy_s2r: Optional[cute.ThrCopy], tSR_rC: Optional[cute.Tensor], tSR_sC: Optional[cute.Tensor], copy_D: Optional[Callable], @@ -63,9 +64,21 @@ class GemmSymmetricMixin(GemmActMixin): is_tma_warp: Boolean, ) -> Tuple[cutlass.pipeline.PipelineState, cutlass.pipeline.PipelineState]: has_C = const_expr(tRS_rC is not None) + has_epi_load = const_expr(self.epi_c_stage > 0) has_D = const_expr(copy_D is not None) + use_tma_epi = const_expr(epi_store_pipeline is not None) + use_tma_c = const_expr(epi_pipeline is not None) + inline_epi_load = const_expr(copy_C is not None) + use_stochastic_rounding = const_expr( + self.rounding_mode == RoundingMode.RS + and self.acc_dtype == cutlass.Float32 + and self.d_dtype == cutlass.BFloat16 + ) - tiled_copy_postact_r2s, tRS_sPostAct, copy_postact = self.epi_setup_postact( + # Setup aux outputs. Returns a tuple of ``(tiled_copy_r2s, + # tRS_sAuxOut, copy_aux_out)`` triples — empty when no aux output + # was requested, one entry per aux output otherwise. + aux_out_ctxs = self.epi_setup_aux_out( params, epi_smem_tensors, tiled_copy_r2s, @@ -75,14 +88,20 @@ class GemmSymmetricMixin(GemmActMixin): tidx, ) - # We iterate over epi tiles in the N dimension first before the M dimension epi_tile_shape = cute.zipped_divide( cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile ).shape[1] - epi_tile_layout = cute.make_layout(epi_tile_shape, stride=(epi_tile_shape[1], 1)) + epi_tile_layout = cute.make_ordered_layout( + epi_tile_shape, order=(0, 1) if const_expr(self.epi_m_major) else (1, 0) + ) epi_tile_num = cute.size(epi_tile_shape) num_prev_subtiles = tile_scheduler.num_tiles_executed * epi_tile_num + # Symmetric guard: skip the mirrored aux write on the diagonal, + # otherwise we'd write the same gmem location twice. + square_tile_m = tile_coord_mnkl[0] // self.cluster_shape_mnk[0] + square_tile_n = tile_coord_mnkl[1] // self.cluster_shape_mnk[1] + epi_tensors = self.epi_begin( params, epi_smem_tensors, @@ -93,95 +112,133 @@ class GemmSymmetricMixin(GemmActMixin): varlen_manager, epilogue_barrier, tidx, + tRS_rD.layout, ) - if const_expr(copy_C is not None): + if const_expr(inline_epi_load): for epi_idx in cutlass.range(min(epi_tile_num, self.epi_c_stage), unroll=1): - gmem_coord_C = epi_tile_layout.get_hier_coord(epi_idx) - if is_tma_warp: - epi_pipeline.producer_acquire(epi_producer_state) - copy_C(src_idx=gmem_coord_C, producer_state=epi_producer_state) - epi_pipeline.producer_commit(epi_producer_state) - epi_producer_state.advance() + epi_coord_C = epi_tile_layout.get_hier_coord(epi_idx) + if const_expr(use_tma_c): + if is_tma_warp: + epi_pipeline.producer_acquire(epi_producer_state) + copy_C(src_idx=epi_coord_C, producer_state=epi_producer_state) + epi_pipeline.producer_commit(epi_producer_state) + epi_producer_state.advance() + else: + # TODO: turn this to cp.async instead of direct G2R copy + copy_C(src_idx=epi_coord_C, dst_idx=epi_idx % self.epi_c_stage) + if const_expr(use_tma_c): + epilogue_barrier.arrive_and_wait() for epi_idx in cutlass.range_constexpr(epi_tile_num): - # The global memory coordinate for the current epi tile - gmem_coord = epi_tile_layout.get_hier_coord(epi_idx) + epi_coord = epi_tile_layout.get_hier_coord(epi_idx) # (epi_m, epi_n) # Copy from acc to D registers - load_acc_subtile(tRS_rD, epi_idx) - epi_loop_tensors = self.epi_begin_loop(params, epi_tensors, gmem_coord) - if const_expr(has_C): - epi_pipeline.consumer_wait(epi_read_state) - cute.copy(tiled_copy_s2r, tSR_sC[None, None, None, epi_read_state.index], tSR_rC) - # Fence to make sure shared memory read is visible to TMA load - cute.arch.fence_view_async_shared() - cute.arch.sync_warp() - with cute.arch.elect_one(): + load_acc_subtile(tRS_rD, epi_coord) + if const_expr(has_epi_load): + if const_expr(use_tma_c): + epi_pipeline.consumer_wait(epi_read_state) + if const_expr(has_C): + cute.copy( + tiled_copy_s2r, tSR_sC[None, None, None, epi_read_state.index], tSR_rC + ) + self.epi_tile_load_s2r(params, epi_tensors, epi_read_state.index) + cute.arch.fence_view_async_shared() epi_pipeline.consumer_release(epi_read_state) - epi_read_state.advance() - if const_expr(copy_C is not None and epi_idx + self.epi_c_stage < epi_tile_num): - gmem_coord_C = epi_tile_layout.get_hier_coord(epi_idx + self.epi_c_stage) - if is_tma_warp: - epi_pipeline.producer_acquire(epi_producer_state) - copy_C(src_idx=gmem_coord_C, producer_state=epi_producer_state) - epi_pipeline.producer_commit(epi_producer_state) - epi_producer_state.advance() - tRS_rPostAct = self.epi_visit_subtile(params, epi_loop_tensors, tRS_rD, tRS_rC) - tRS_rPostAct_out = self.epi_convert_postact( - tRS_rPostAct, - epi_loop_tensors["sr_seed"], - tidx, + epi_read_state.advance() + else: + c_buffer = epi_idx % self.epi_c_stage + cute.copy(tiled_copy_s2r, tSR_sC[None, None, None, c_buffer], tSR_rC) + # TODO: cp.async wait once we switch to cp.async + epilogue_barrier.arrive_and_wait() + epi_loop_tensors = self.epi_begin_loop(params, epi_tensors, epi_coord) + if const_expr(inline_epi_load and epi_idx + self.epi_c_stage < epi_tile_num): + epi_coord_C = epi_tile_layout.get_hier_coord(epi_idx + self.epi_c_stage) + if const_expr(use_tma_c): + if is_tma_warp: + epi_pipeline.producer_acquire(epi_producer_state) + copy_C(src_idx=epi_coord_C, producer_state=epi_producer_state) + epi_pipeline.producer_commit(epi_producer_state) + epi_producer_state.advance() + else: + epilogue_barrier.arrive_and_wait() + copy_C( + src_idx=epi_coord_C, + dst_idx=(epi_idx + self.epi_c_stage) % self.epi_c_stage, + ) + # Returns a tuple of register tensors — one per aux output. + # Length matches ``aux_out_ctxs``. + tRS_rAuxOuts = self.epi_visit_subtile(params, epi_loop_tensors, tRS_rD, tRS_rC) + self.epi_end_loop( + params, + epi_tensors, + epi_coord, + epi_tile, + tiled_copy_t2r, + tiled_copy_r2s, tile_coord_mnkl, - num_prev_subtiles, - epi_idx, + varlen_manager, + tidx, + ) + # Convert each output to its storage dtype. + tRS_rAuxOuts_out = tuple( + self.epi_convert_aux_out( + i, + tRS_rAuxOuts[i], + epi_loop_tensors.get("sr_seed"), + tidx, + tile_coord_mnkl, + num_prev_subtiles, + epi_idx, + ) + for i in range(len(aux_out_ctxs)) ) - if is_tma_warp: - epi_store_pipeline.producer_acquire() - epilogue_barrier.arrive_and_wait() - # Copy from D registers to shared memory + if const_expr(use_tma_epi): + if is_tma_warp: + epi_store_pipeline.producer_acquire() + else: + epilogue_barrier.arrive_and_wait() + if const_expr(use_tma_epi): + epilogue_barrier.arrive_and_wait() epi_buffer = (num_prev_subtiles + epi_idx) % self.epi_stage if const_expr(has_D): - if const_expr( - self.rounding_mode == RoundingMode.RS - and self.acc_dtype == cutlass.Float32 - and self.d_dtype == cutlass.BFloat16 - ): - seed = epi_loop_tensors["sr_seed"] + ( - tile_coord_mnkl[0] * 65537 - + tile_coord_mnkl[1] * 257 - + tile_coord_mnkl[3] * 17 - + (num_prev_subtiles + epi_idx) * 7 - ) - copy_utils.sr_cvt_copy( - tiled_copy_r2s, - tRS_rD, - tRS_sD[None, None, None, epi_buffer], - seed, - tidx, + tRS_sD_cur = tRS_sD[None, None, None, epi_buffer] + if const_expr(use_stochastic_rounding): + seed = epilogue_sr_seed( + epi_loop_tensors.get("sr_seed"), + tile_coord_mnkl, + num_prev_subtiles + epi_idx, ) + copy_utils.sr_cvt_copy(tiled_copy_r2s, tRS_rD, tRS_sD_cur, seed, tidx) else: - copy_utils.cvt_copy( - tiled_copy_r2s, tRS_rD, tRS_sD[None, None, None, epi_buffer] - ) - cute.copy( - tiled_copy_postact_r2s, - tiled_copy_postact_r2s.retile(tRS_rPostAct_out), - tRS_sPostAct[None, None, None, epi_buffer], - ) - pid_m = tile_coord_mnkl[0] - pid_n = tile_coord_mnkl[1] - # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_view_async_shared() - epilogue_barrier.arrive_and_wait() - # Copy from shared memory to global memory - if is_tma_warp: - square_tile_m = pid_m // self.cluster_shape_mnk[0] - square_tile_n = pid_n // self.cluster_shape_mnk[1] + copy_utils.cvt_copy(tiled_copy_r2s, tRS_rD, tRS_sD_cur) + for i in cutlass.range_constexpr(len(aux_out_ctxs)): + tiled_copy_aux_out_r2s, tRS_sAuxOut, _ = aux_out_ctxs[i] + cute.copy( + tiled_copy_aux_out_r2s, + # Need contiguous for Sm80 and Sm120 where acc layout is ((2, 2), MMA_M, MMA_N) + tiled_copy_aux_out_r2s.retile(tRS_rAuxOuts_out[i]).contiguous(), + tRS_sAuxOut[None, None, None, epi_buffer], + ) + if const_expr(use_tma_epi): + cute.arch.fence_view_async_shared() + epilogue_barrier.arrive_and_wait() + if is_tma_warp: + if const_expr(has_D): + copy_D(src_idx=epi_buffer, dst_idx=epi_coord) + for i in cutlass.range_constexpr(len(aux_out_ctxs)): + _, _, copy_aux_out = aux_out_ctxs[i] + if square_tile_m != square_tile_n: # don't write twice on the diagonal + copy_aux_out(src_idx=epi_buffer, dst_idx=epi_coord) + epi_store_pipeline.producer_commit() + else: + epilogue_barrier.arrive_and_wait() if const_expr(has_D): - copy_D(src_idx=epi_buffer, dst_idx=gmem_coord) - if square_tile_m != square_tile_n: # don't write twice to the same tile - copy_postact(src_idx=epi_buffer, dst_idx=gmem_coord) - epi_store_pipeline.producer_commit() + copy_D(src_idx=epi_buffer, dst_idx=epi_coord) + for i in cutlass.range_constexpr(len(aux_out_ctxs)): + _, _, copy_aux_out = aux_out_ctxs[i] + if square_tile_m != square_tile_n: # don't write twice on the diagonal + copy_aux_out(src_idx=epi_buffer, dst_idx=epi_coord) + epilogue_barrier.arrive_and_wait() self.epi_end( params, @@ -197,6 +254,10 @@ class GemmSymmetricMixin(GemmActMixin): return epi_read_state, epi_producer_state +class GemmSymmetricSm80(GemmSymmetricMixin, GemmSm80): + pass + + class GemmSymmetricSm90(GemmSymmetricMixin, GemmSm90): pass @@ -231,6 +292,7 @@ def _compile_gemm_symmetric( device_capacity, ): sm_to_cls = { + 8: GemmSymmetricSm80, 9: GemmSymmetricSm90, 10: GemmSymmetricSm100, 11: GemmSymmetricSm100, @@ -252,7 +314,7 @@ def _compile_gemm_symmetric( # PostAct = D.mT, so it has the opposite major from D (m↔n swapped) div_pa = div_for_dtype(postact_dtype) postact_leading = 1 if postact_major == "n" else 0 - mPostAct = fake_tensor( + mAuxOut = fake_tensor( postact_dtype, (m, m, l), leading_dim=postact_leading, divisibility=div_pa ) @@ -267,7 +329,7 @@ def _compile_gemm_symmetric( activation = None # identity act_fn = act_fn_map[activation] epi_args = GemmCls.EpilogueArguments( - mPostAct, + mAuxOut, act_fn, alpha=fake_scalar(alpha_mode), beta=fake_scalar(beta_mode), @@ -306,6 +368,7 @@ def gemm_symmetric( tile_N: int, cluster_M: int, cluster_N: int, + tile_K: int | None = None, pingpong: bool = False, persistent: bool = True, is_dynamic_persistent: bool = False, @@ -325,14 +388,16 @@ def gemm_symmetric( postact_major = "n" if PostAct_p.stride(1) == 1 else "m" device_capacity = get_device_capacity(A.device) - assert device_capacity[0] in [9, 10, 11, 12], "Only SM90, SM100, SM110, and SM120 are supported" + assert device_capacity[0] in [8, 9, 10, 11, 12], ( + "Only SM8x, SM90, SM100, SM110, and SM120 are supported" + ) - if is_dynamic_persistent and device_capacity[0] == 9: + if is_dynamic_persistent and device_capacity[0] <= 9: assert tile_count_semaphore is not None, ( "Dynamic persistent tile scheduler in SM90 requires a semaphore in GMEM" ) - tile_shape_mn = (tile_M, tile_N) + tile_shape_mn = (tile_M, tile_N, tile_K) if tile_K is not None else (tile_M, tile_N) cluster_shape_mnk = (cluster_M, cluster_N, 1) alpha_mode = 2 if isinstance(alpha, Tensor) else (1 if alpha != 1.0 else 0) beta_mode = 2 if isinstance(beta, Tensor) else (1 if beta != 1.0 else 0) @@ -358,12 +423,10 @@ def gemm_symmetric( device_capacity, ) - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY: - return - - max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0 + cluster_size = cluster_M * cluster_N + max_active_clusters = ( + get_max_active_clusters(cluster_size, device_capacity=device_capacity) if persistent else 0 + ) def scalar_arg(scalar, mode): if mode == 0: @@ -389,6 +452,6 @@ def gemm_symmetric( varlen_args = None if device_capacity[0] in [10, 11]: - compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None, None) + compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None) else: - compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None) + compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args) diff --git a/build/torch-cuda/quack/gemm_tvm_ffi_utils.py b/build/torch-cuda/quack/gemm_tvm_ffi_utils.py index 1d59c3c3cb335f79416e77da675128c55cdb4b1e..a99388f5eb7badc800b3f63b622d06cb5d76e935 100644 --- a/build/torch-cuda/quack/gemm_tvm_ffi_utils.py +++ b/build/torch-cuda/quack/gemm_tvm_ffi_utils.py @@ -3,9 +3,10 @@ from functools import partial +import torch import cutlass.cute as cute -from cutlass import Int32, Int64, Float32 +from cutlass import Int32, Float32 from cutlass.cute.runtime import make_ptr from .compile_utils import make_fake_tensor as fake_tensor @@ -13,6 +14,99 @@ from .cute_dsl_utils import torch2cute_dtype_map from .tile_scheduler import TileSchedulerOptions from .varlen_utils import VarlenArguments +# Blockscaled scale-factor dtype determines the quantization block size along K: +# e8m0 -> MX formats (32-element blocks), e4m3 -> NVFP4 (16-element blocks). +SF_DTYPE_TO_VEC_SIZE = { + torch.float8_e8m0fnu: 32, + torch.float8_e4m3fn: 16, +} + + +def validate_blockscaled_sf(A, B, SFA, SFB, device_capacity, num_batches=None, varlen_k=False): + """Validate blockscaled scale factors against kernel-layout operands. + + A is (l, m, k[/2 if fp4]) and B is (l, n, k[/2]); SFA/SFB are + (l, rm/rn, rk, 32, 4, 4) with the inner (32, 4, 4) block contiguous + (strides (16, 4, 1) — one 512 B atom per 128 rows x 4 K-blocks). + + When num_batches is not None and varlen_k is False (varlen_m), A is + (total_m, k) and SFA must be a single M-padded buffer (tile-aligned + per-batch padding) (1, total_padded_rm, rk, 32, 4, 4) with + total_padded_rm >= ceil(total_m/128) + (num_batches - 1) — the bound from + AI/varlen_blockscaled_sf_layout.md that suffices for any per-batch split + of total_m. SFB stays per-batch: (num_batches, rn, rk, 32, 4, 4). + + When varlen_k, A is (m, total_k) m-major and B is (n, total_k) n-major + (MXFP8 only — fp4 operands must be K-major), and BOTH SF buffers are + K-padded with tile-aligned per-batch padding: + (1, rm/rn, total_padded_rk, 32, 4, 4) with + total_padded_rk >= ceil(total_k/128) + (num_batches - 1). + SF pad bytes inside each batch's last atom are loaded by the kernel but + never consumed: the mma loop skips the MMA instructions for pad k-blocks + (one instruction per SF block for mxfp8; see GemmSm100.mma), so the pad + may be arbitrary bytes — torch.empty buffers are fine. + Returns (sf_dtype, sf_vec_size) as (cutlass dtype, int). + """ + varlen_m = num_batches is not None and not varlen_k + assert not varlen_k or num_batches is not None, "varlen_k requires num_batches" + assert SFB is not None, "SFA and SFB must be provided together" + assert device_capacity[0] in [10, 11], "Blockscaled GEMM requires SM100/SM110" + assert SFA.dtype == SFB.dtype, f"SF dtype mismatch: {SFA.dtype} vs {SFB.dtype}" + assert SFA.dtype in SF_DTYPE_TO_VEC_SIZE, f"unsupported SF dtype: {SFA.dtype}" + sf_vec_size = SF_DTYPE_TO_VEC_SIZE[SFA.dtype] + sf_dtype = torch2cute_dtype_map[SFA.dtype] + # A.shape[-1] is packed K for fp4 (two elements per byte) while dlpack presents + # the logical extent to the kernel, so validate rk against logical K. + k_logical = A.shape[-1] * (2 if A.dtype == torch.float4_e2m1fn_x2 else 1) + rk = (k_logical + 4 * sf_vec_size - 1) // (4 * sf_vec_size) + if varlen_k: + assert A.dtype != torch.float4_e2m1fn_x2, ( + "varlen_k blockscaled supports MXFP8 only: fp4 operands must be K-major, " + "but varlen_k requires m-major A / n-major B" + ) + assert A.ndim == 2 and B.ndim == 2, ( + f"varlen_k expects A (m, total_k) and B (n, total_k), " + f"got shapes {tuple(A.shape)} / {tuple(B.shape)}" + ) + # rk here = ceil(total_k/128); K-padded buffers need one extra atom + # column per additional batch. + min_rk = rk + (num_batches - 1) + for name, SF, mn in (("SFA", SFA, A.shape[0]), ("SFB", SFB, B.shape[0])): + r_mn = (mn + 127) // 128 + assert SF.shape[0] == 1 and SF.shape[1] == r_mn and tuple(SF.shape[3:]) == (32, 4, 4), ( + f"{name} shape {tuple(SF.shape)} != (1, {r_mn}, total_padded_rk, 32, 4, 4)" + ) + assert SF.shape[2] >= min_rk, ( + f"{name} padded rk {SF.shape[2]} < ceil(total_k/128) + (L-1) = {min_rk}" + ) + shapes = [] + elif varlen_m: + assert A.ndim == 2, f"varlen_m expects A as (total_m, k), got shape {tuple(A.shape)}" + assert B.shape[0] == num_batches, ( + f"B batch dim {B.shape[0]} != len(cu_seqlens_m) - 1 = {num_batches}" + ) + min_rm = (A.shape[0] + 127) // 128 + (num_batches - 1) + assert SFA.shape[0] == 1 and tuple(SFA.shape[2:]) == (rk, 32, 4, 4), ( + f"SFA shape {tuple(SFA.shape)} != (1, total_padded_rm, {rk}, 32, 4, 4)" + ) + assert SFA.shape[1] >= min_rm, ( + f"SFA padded rm {SFA.shape[1]} < ceil(total_m/128) + (L-1) = {min_rm}" + ) + shapes = [("SFB", SFB, (num_batches, (B.shape[-2] + 127) // 128, rk, 32, 4, 4))] + else: + shapes = [ + (name, SF, (A.shape[0], (mn + 127) // 128, rk, 32, 4, 4)) + for name, SF, mn in (("SFA", SFA, A.shape[-2]), ("SFB", SFB, B.shape[-2])) + ] + for name, SF, expected in shapes: + assert tuple(SF.shape) == expected, f"{name} shape {tuple(SF.shape)} != {expected}" + for name, SF in (("SFA", SFA), ("SFB", SFB)): + assert SF.stride()[-3:] == (16, 4, 1), ( + f"{name}: inner (32, 4, 4) block must be contiguous with strides (16, 4, 1), " + f"got {SF.stride()[-3:]}" + ) + return sf_dtype, sf_vec_size + def div_for_dtype(dtype): """16-byte alignment: divisibility in elements = 128 // dtype_width_bits.""" @@ -136,7 +230,11 @@ def make_fake_gemm_tensors( b_leading = 1 if b_major == "k" else 0 d_leading = 1 if d_major == "n" else 0 c_leading = 1 if c_major == "n" else 0 - m, n, k, l = cute.sym_int(), cute.sym_int(), cute.sym_int(), cute.sym_int() + m, n, l = cute.sym_int(), cute.sym_int(), cute.sym_int() + # Sub-byte (fp4) tensors need their contiguous extent statically divisible by the + # packing factor; fp4 operands are k-major, so mark k. Harmless for 8-bit+ dtypes. + k_div = div_for_dtype(a_dtype) if a_dtype.width < 8 else 1 + k = cute.sym_int(divisibility=k_div) div_a = div_for_dtype(a_dtype) div_b = div_for_dtype(b_dtype) div_d = div_for_dtype(d_dtype) if d_dtype is not None else 1 @@ -165,6 +263,32 @@ def make_fake_gemm_tensors( return mA, mB, mD, mC, m, n, k, l +def make_fake_sf_tensor(sf_dtype, l): + """Fake (l, rm, rk, 32, 4, 4) blockscaled scale-factor tensor. + + The inner (32, 4, 4) block has static strides (16, 4, 1) — one contiguous + 512 B atom per 128 rows x 4 K-blocks, so TMA loads it as one box. The + kernel only consumes the base pointer and the outer (l, rm, rk) strides + (the atom layout is hardware-fixed); outer strides are dynamic but + atom-granular, so slices of larger scale buffers are accepted without a + copy. + """ + rm, rk = cute.sym_int(), cute.sym_int() + return cute.runtime.make_fake_tensor( + sf_dtype, + (l, rm, rk, 32, 4, 4), + stride=( + cute.sym_int64(divisibility=512), + cute.sym_int64(divisibility=512), + cute.sym_int64(divisibility=512), + 16, + 4, + 1, + ), + assumed_align=16, + ) + + def compile_gemm_kernel( GemmCls, a_dtype, @@ -185,18 +309,24 @@ def compile_gemm_kernel( post_init=None, mSFA=None, mSFB=None, - has_trace_ptr=False, use_tma_gather=False, concat_layout=None, + num_warps=None, + sf_vec_size=None, ): """Build GemmCls instance, apply SM90 partial, and cute.compile with TVM-FFI.""" - if device_capacity[0] in [9, 12]: + if device_capacity[0] == 8: + sm8x_kwargs = {"is_persistent": persistent, "num_warps": num_warps} + sm8x_kwargs["arch"] = device_capacity[0] * 10 + device_capacity[1] + GemmCls = partial(GemmCls, **sm8x_kwargs) + elif device_capacity[0] in [9, 12]: GemmCls = partial(GemmCls, pingpong=pingpong, is_persistent=persistent) elif device_capacity[0] in [10, 11]: GemmCls = partial( GemmCls, use_clc_persistence=is_dynamic_persistent, use_tma_gather=use_tma_gather, + sf_vec_size=sf_vec_size, ) gemm_obj = GemmCls( Float32, @@ -209,10 +339,7 @@ def compile_gemm_kernel( if post_init: post_init(gemm_obj) stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - sf_args = () if device_capacity[0] in (9, 12) else (mSFA, mSFB) - # Trace pointer: Optional[Int64]. Compile with Int64(0) when tracing is - # requested, None otherwise. TVM-FFI caches each variant separately. - trace_ptr = Int64(0) if has_trace_ptr else None + sf_args = () if device_capacity[0] in (8, 9, 12) else (mSFA, mSFB) return cute.compile( gemm_obj, mA, @@ -224,6 +351,5 @@ def compile_gemm_kernel( varlen_args, stream, *sf_args, - trace_ptr, options="--enable-tvm-ffi", ) diff --git a/build/torch-cuda/quack/jax_utils.py b/build/torch-cuda/quack/jax_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5d7e7dcd7331b9057e7f5b109f7803473d392e90 --- /dev/null +++ b/build/torch-cuda/quack/jax_utils.py @@ -0,0 +1,142 @@ +"""Shared helpers for optional JAX bindings.""" + +from __future__ import annotations + +from typing import Any, Callable + +import cutlass +import jax +import jax.numpy as jnp +import numpy as np + + +JAX_TO_CUTLASS_DTYPE: dict[jnp.dtype, type[cutlass.Numeric]] = { + jnp.dtype(jnp.float16): cutlass.Float16, + jnp.dtype(jnp.bfloat16): cutlass.BFloat16, + jnp.dtype(jnp.float32): cutlass.Float32, +} + + +def require_jax_tvm_ffi(): + try: + import jax_tvm_ffi + except ImportError as e: + raise ImportError( + "This QuACK JAX TVM-FFI path requires jax-tvm-ffi. Install it with " + "`pip install jax-tvm-ffi`." + ) from e + return jax_tvm_ffi + + +def dtype_name(dtype: jnp.dtype) -> str: + return str(jnp.dtype(dtype)).replace("float", "f").replace("bfloat", "bf") + + +def cutlass_dtype(dtype: jnp.dtype) -> type[cutlass.Numeric]: + dtype = jnp.dtype(dtype) + if dtype not in JAX_TO_CUTLASS_DTYPE: + raise TypeError(f"Unsupported dtype: {dtype}") + return JAX_TO_CUTLASS_DTYPE[dtype] + + +def check_rank(name: str, x: Any, rank: int) -> None: + if len(x.shape) != rank: + raise ValueError(f"{name} must be {rank}D, got shape {x.shape}") + + +def shape_dtype_like(x: Any) -> jax.ShapeDtypeStruct: + return jax.ShapeDtypeStruct(x.shape, x.dtype) + + +def tvm_ffi_call( + target: str, + *args, + output_shape_dtype, + vmap_method: str = "broadcast_all", + **kwargs, +): + call = jax.ffi.ffi_call( + target, + output_shape_dtype, + vmap_method=vmap_method, + **kwargs, + ) + return call(*args) + + +def _tvm_ffi_key_part(value: Any) -> str: + if isinstance(value, np.dtype): + return dtype_name(value) + if isinstance(value, type) and issubclass(value, cutlass.Numeric): + return value.__name__.lower() + return str(value).replace(" ", "").replace("/", "_").replace(".", "p") + + +class TvmFfiKernel: + """Lazy JAX FFI wrapper for a family of TVM-FFI compiled kernels.""" + + def __init__( + self, + name: str, + compile_fn: Callable[..., Any], + *, + target_name: Callable[..., str] | None = None, + platform: str = "gpu", + arg_spec: list[str] | None = None, + allow_cuda_graph: bool = False, + pass_owned_tensor: bool = False, + use_last_output_for_alloc_workspace: bool = False, + ) -> None: + self.name = name + self.compile_fn = compile_fn + self.target_name = target_name + self.platform = platform + self.arg_spec = arg_spec + self.allow_cuda_graph = allow_cuda_graph + self.pass_owned_tensor = pass_owned_tensor + self.use_last_output_for_alloc_workspace = use_last_output_for_alloc_workspace + self._targets: dict[tuple[Any, ...], str] = {} + self._compiled: dict[tuple[Any, ...], Any] = {} + + def _target_for_key(self, key: tuple[Any, ...]) -> str: + if self.target_name is not None: + return self.target_name(*key) + if not key: + return self.name + suffix = "_".join(_tvm_ffi_key_part(part) for part in key) + return f"{self.name}_{suffix}" + + def target(self, key: tuple[Any, ...]) -> str: + key = tuple(key) + if key in self._targets: + return self._targets[key] + target = self._target_for_key(key) + compiled = self.compile_fn(*key) + require_jax_tvm_ffi().register_ffi_target( + target, + compiled, + arg_spec=self.arg_spec, + platform=self.platform, + allow_cuda_graph=self.allow_cuda_graph, + pass_owned_tensor=self.pass_owned_tensor, + use_last_output_for_alloc_workspace=self.use_last_output_for_alloc_workspace, + ) + self._targets[key] = target + self._compiled[key] = compiled + return target + + def __call__( + self, + *args, + key: tuple[Any, ...], + output_shape_dtype, + vmap_method: str = "broadcast_all", + **kwargs, + ): + return tvm_ffi_call( + self.target(key), + *args, + output_shape_dtype=output_shape_dtype, + vmap_method=vmap_method, + **kwargs, + ) diff --git a/build/torch-cuda/quack/layout_utils.py b/build/torch-cuda/quack/layout_utils.py index 5ad26397979f0a7f48e2f69b96411c0b6deac906..ac799d7d04d25287962ff981b8aa5193a3dace92 100644 --- a/build/torch-cuda/quack/layout_utils.py +++ b/build/torch-cuda/quack/layout_utils.py @@ -230,15 +230,16 @@ def convert_layout_acc_frgA(acc_layout: cute.Layout) -> cute.Layout: ) else: # Sm80 # (4, MMA_M, MMA_N) -> (4, MMA_M, (2, MMA_N / 2)) + assert acc_layout.shape[2] % 2 == 0 l = cute.logical_divide(acc_layout, (None, None, 2)) rA_mma_view = cute.make_layout( ( - (l.shape[0], l.shape[2][0]), + (l.shape[0][0], l.shape[0][1], l.shape[2][0]), l.shape[1], l.shape[2][1], ), stride=( - (l.stride[0], l.stride[2][0]), + (l.stride[0][0], l.stride[0][1], l.stride[2][0]), l.stride[1], l.stride[2][1], ), @@ -277,69 +278,41 @@ def convert_layout_zero_stride( def mma_partition_C_vec( - sVec: cute.Tensor, thr_mma: cute.core.ThrMma, expand_shape: int, is_colvec: bool + sVec: cute.Tensor, thr_mma: cute.ThrMma, expand_shape: int, is_colvec: bool ) -> cute.Tensor: assert cute.rank(sVec) == 2 assert sVec.stride[0] == 1 - stage = sVec.shape[1] - shape = ( - (sVec.shape[0], expand_shape, stage) - if const_expr(is_colvec) - else (expand_shape, sVec.shape[0], stage) - ) - stride = (1, 0, sVec.stride[1]) if const_expr(is_colvec) else (0, 1, sVec.stride[1]) - sVec_mma = cute.make_tensor(sVec.iterator, cute.make_layout(shape, stride=stride)) - tC_sVec = make_acc_tensor_mn_view(thr_mma.partition_C(sVec_mma)) + sVec_mma = expand(sVec, 1 if const_expr(is_colvec) else 0, expand_shape) + tC_sVec = reshape_acc_to_mn(thr_mma.partition_C(sVec_mma)) return tC_sVec[None, 0, None] if const_expr(is_colvec) else tC_sVec[0, None, None] def mma_partition_A_vec( - sVec: cute.Tensor, thr_mma: cute.core.ThrMma, expand_shape: int, is_colvec: bool + sVec: cute.Tensor, thr_mma: cute.ThrMma, expand_shape: int, is_colvec: bool ) -> cute.Tensor: assert cute.rank(sVec) == 2 assert sVec.stride[0] == 1 - stage = sVec.shape[1] - shape = ( - (sVec.shape[0], expand_shape, stage) - if const_expr(is_colvec) - else (expand_shape, sVec.shape[0], stage) - ) - stride = (1, 0, sVec.stride[1]) if const_expr(is_colvec) else (0, 1, sVec.stride[1]) - sVec_mma = cute.make_tensor(sVec.iterator, cute.make_layout(shape, stride=stride)) - tC_sVec = make_acc_tensor_mn_view(thr_mma.partition_A(sVec_mma)) + sVec_mma = expand(sVec, 1 if const_expr(is_colvec) else 0, expand_shape) + tC_sVec = reshape_acc_to_mn(thr_mma.partition_A(sVec_mma)) return tC_sVec[None, 0, None] if const_expr(is_colvec) else tC_sVec[0, None, None] def copy_partition_S_vec( - sVec: cute.Tensor, thr_copy: cute.core.ThrCopy, expand_shape: int, is_colvec: bool + sVec: cute.Tensor, thr_copy: cute.ThrCopy, expand_shape: int, is_colvec: bool ) -> cute.Tensor: assert cute.rank(sVec) == 2 assert sVec.stride[0] == 1 - stage = sVec.shape[1] - shape = ( - (sVec.shape[0], expand_shape, stage) - if const_expr(is_colvec) - else (expand_shape, sVec.shape[0], stage) - ) - stride = (1, 0, sVec.stride[1]) if const_expr(is_colvec) else (0, 1, sVec.stride[1]) - sVec_thr = cute.make_tensor(sVec.iterator, cute.make_layout(shape, stride=stride)) + sVec_thr = expand(sVec, 1 if const_expr(is_colvec) else 0, expand_shape) tC_sVec = reshape_acc_to_mn(thr_copy.partition_S(sVec_thr)) return tC_sVec[None, 0, None] if const_expr(is_colvec) else tC_sVec[0, None, None] def copy_partition_D_vec( - sVec: cute.Tensor, thr_copy: cute.core.ThrCopy, expand_shape: int, is_colvec: bool + sVec: cute.Tensor, thr_copy: cute.ThrCopy, expand_shape: int, is_colvec: bool ) -> cute.Tensor: assert cute.rank(sVec) == 2 assert sVec.stride[0] == 1 - stage = sVec.shape[1] - shape = ( - (sVec.shape[0], expand_shape, stage) - if const_expr(is_colvec) - else (expand_shape, sVec.shape[0], stage) - ) - stride = (1, 0, sVec.stride[1]) if const_expr(is_colvec) else (0, 1, sVec.stride[1]) - sVec_thr = cute.make_tensor(sVec.iterator, cute.make_layout(shape, stride=stride)) + sVec_thr = expand(sVec, 1 if const_expr(is_colvec) else 0, expand_shape) tC_sVec = reshape_acc_to_mn(thr_copy.partition_D(sVec_thr)) return tC_sVec[None, 0, None] if const_expr(is_colvec) else tC_sVec[0, None, None] @@ -366,11 +339,11 @@ def tile_atom_to_shape_SF_strided( shape: A/B operand shape. Rank-3 `(m/n, k, l)` or rank-2 `(total_mn, k)` (varlen_m). sf_vec_size: Scale factor vector size (16 or 32). - sf_strides: Strides of the scale tensor, which has logical shape - `(L, rmn, rk, 512)` (rank 4). Only `sf_strides[0..2]` are used: - `sf_strides[1]` as the rmn stride, `sf_strides[2]` as the rk - stride, and `sf_strides[0]` as the L stride (only for rank-3 - `shape`). + sf_strides: Strides of the `(L, rmn, rk, 32, 4, 4)` scale tensor. + Only `sf_strides[0..2]` are used: `sf_strides[1]` as the rmn + stride, `sf_strides[2]` as the rk stride, and `sf_strides[0]` as + the L stride (only for rank-3 `shape`); the inner atom layout is + hardware-fixed. """ from cutlass.utils.blockscaled_layout import BlockScaledBasicChunk diff --git a/build/torch-cuda/quack/linear.py b/build/torch-cuda/quack/linear.py index e5477433f8eeaeb0cf537e6534a011f97c484f34..3351f63d7d22805b5b4da88a1bb26810aa146204 100644 --- a/build/torch-cuda/quack/linear.py +++ b/build/torch-cuda/quack/linear.py @@ -48,7 +48,7 @@ def linear_bwd_compute_input_grad(ctx, dout, weight, matmul_fn): def linear_bwd_compute_weight_grad(ctx, dout, x, weight_og, matmul_fn, matmul_inplace_fn): if ctx.needs_input_grad[1]: assert x is not None - x = x.reshape(-1, x.shape[-1]) + x = x.flatten(0, -2) # fuse_grad_accum is not compatible with torch.compile if not ctx.fuse_grad_accum or weight_og.grad is None or torch.compiler.is_compiling(): dweight = matmul_fn(dout.T, x, out_dtype=ctx.weight_dtype) @@ -88,6 +88,7 @@ class _LinearUntunedOps(_LinearOps): matmul_fwd_fn = partial(gemm, tuned=False) matmul_bwd_dx = partial(gemm, dynamic_scheduler=True, tuned=False) matmul_bwd_dw = partial(gemm, dynamic_scheduler=True, tuned=False) + matmul_bwd_dw_inplace = partial(gemm_add_inplace, dynamic_scheduler=True, tuned=False) class _LinearActOps(_LinearOps): @@ -167,7 +168,7 @@ class LinearFunc(torch.autograd.Function): ctx.ops = ops weight_og = weight batch_shape = x.shape[:-1] - x = x.reshape(-1, x.shape[-1]) + x = x.flatten(0, -2) out = ops.matmul_fwd_fn(x, weight.T, bias=bias) linear_fwd_postprocess( ctx, x, weight, weight_og, needs_x_w_grad=ctx.needs_input_grad[:2] @@ -185,7 +186,7 @@ class LinearFunc(torch.autograd.Function): ops = ctx.ops x, weight, weight_og = ctx.saved_tensors # weight_og is None if not ctx.fuse_grad_accum batch_shape = dout.shape[:-1] - dout = _ensure_contiguous(dout.reshape(-1, dout.shape[-1])) + dout = _ensure_contiguous(dout.flatten(0, -2)) dbias = dout.sum(0, dtype=ctx.bias_dtype) if ctx.compute_dbias else None dx = linear_bwd_compute_input_grad(ctx, dout, weight, ops.matmul_bwd_dx) dx = dx.reshape(*batch_shape, dx.shape[-1]) if dx is not None else None @@ -217,7 +218,7 @@ class LinearActFunc(torch.autograd.Function): ctx.ops = ops weight_og = weight batch_shape = x.shape[:-1] - x = x.reshape(-1, x.shape[-1]) + x = x.flatten(0, -2) out, postact = ops.matmul_fwd_fn( x, weight.T, bias=bias, activation=activation, store_preact=store_preact ) @@ -238,7 +239,7 @@ class LinearActFunc(torch.autograd.Function): ops = ctx.ops x, weight, weight_og = ctx.saved_tensors batch_shape = dout.shape[:-1] - dout = _ensure_contiguous(dout.reshape(-1, dout.shape[-1])) + dout = _ensure_contiguous(dout.flatten(0, -2)) dbias = dout.sum(0, dtype=ctx.bias_dtype) if ctx.compute_dbias else None dx = linear_bwd_compute_input_grad(ctx, dout, weight, ops.matmul_bwd_dx) dx = dx.reshape(*batch_shape, dx.shape[-1]) if dx is not None else None @@ -289,7 +290,7 @@ class DActLinearFunc(torch.autograd.Function): ctx.ops = ops weight_og = weight batch_shape = x.shape[:-1] - x = x.reshape(-1, x.shape[-1]) + x = x.flatten(0, -2) out = ops.matmul_fwd_fn(x, weight.T, bias=bias) # Store preact instead of x, we will recompute x (postact) in backward. # dpreact needs gemm_dact(dout, weight, preact) → needs both weight and preact. @@ -315,16 +316,16 @@ class DActLinearFunc(torch.autograd.Function): # weight_og is None if not ctx.fuse_grad_accum preact, weight, weight_og = ctx.saved_tensors batch_shape = dout.shape[:-1] - dout = _ensure_contiguous(dout.reshape(-1, dout.shape[-1])) + dout = _ensure_contiguous(dout.flatten(0, -2)) dbias = dout.sum(0, dtype=ctx.bias_dtype) if ctx.compute_dbias else None if ctx.needs_input_grad[0]: # Need dpreact: gemm_dact(dout, weight, preact) → (dpreact, postact) - preact = preact.reshape(-1, preact.shape[-1]) + preact = preact.flatten(0, -2) assert weight is not None dpreact, x = ops.matmul_bwd_dx(dout, weight, preact, activation=ctx.activation) elif ctx.needs_input_grad[1]: # Only need dweight: recompute postact from preact cheaply (no GEMM needed) - preact = preact.reshape(-1, preact.shape[-1]) + preact = preact.flatten(0, -2) x = ops.recompute_postact(preact, ctx.activation) dpreact = None else: diff --git a/build/torch-cuda/quack/linear_cross_entropy.py b/build/torch-cuda/quack/linear_cross_entropy.py index 98ad3c470232c8f1633d2637af9a51046246c295..b50140f6ace219904c862405e4ec991ada007519 100644 --- a/build/torch-cuda/quack/linear_cross_entropy.py +++ b/build/torch-cuda/quack/linear_cross_entropy.py @@ -92,6 +92,7 @@ def chunked_linear_cross_entropy_fwd( loss=loss_chunk, lse=None, # we don't need lse here dx=dlogits_chunk, + weight=None, ignore_index=ignore_index, ) # Compute dx for this chunk: dlogits @ weight diff --git a/build/torch-cuda/quack/mlp.py b/build/torch-cuda/quack/mlp.py index 4316285db87241321b3bc88c2d65ab889befafba..06614c69551d523b1194b91b9bb1ef74cd2a5f1c 100644 --- a/build/torch-cuda/quack/mlp.py +++ b/build/torch-cuda/quack/mlp.py @@ -26,10 +26,14 @@ from .gemm_interface import ( Activation = Literal[ "gelu_tanh_approx", + "silu", + "silu-tanh", "relu", "relu_sq", "swiglu", + "swiglu-tanh", "swiglu_oai", + "swiglu_oai-tanh", "reglu", "geglu", "glu", @@ -86,7 +90,7 @@ class _MLPGatedConcatUntunedOps(_MLPGatedUntunedOps): matmul_bwd_dx = partial(gemm, dynamic_scheduler=True, tuned=False, concat_layout=("B",)) matmul_bwd_dw1 = partial(gemm, dynamic_scheduler=True, tuned=False, concat_layout=("out",)) matmul_bwd_dw1_inplace = partial( - gemm_add_inplace, dynamic_scheduler=True, tuned=False, concat_layout=("out",) + gemm_add_inplace, dynamic_scheduler=True, tuned=False, concat_layout=("C", "out") ) recompute_fwd = partial(gemm, tuned=False, concat_layout=("B",)) diff --git a/build/torch-cuda/quack/pipeline.py b/build/torch-cuda/quack/pipeline.py index 7589ff555c24357626e0473d0fb9d34c83546d0f..ec470d99e57d3358712176cf9e1ee3e72fc42d25 100644 --- a/build/torch-cuda/quack/pipeline.py +++ b/build/torch-cuda/quack/pipeline.py @@ -6,13 +6,14 @@ from dataclasses import dataclass import cutlass.cute as cute from cutlass import Boolean, Int32, const_expr from cutlass.cutlass_dsl import if_generate, and_, dsl_user_op -from cutlass.pipeline import MbarrierArray, CooperativeGroup, PipelineOp +from cutlass._mlir.dialects import nvvm as _nvvm, llvm +from cutlass.cute.typing import AddressSpace, Int, Pointer from cutlass.pipeline import PipelineState, PipelineUserType -from cutlass.pipeline import Agent, agent_sync from cutlass.pipeline import NamedBarrier as NamedBarrierOg from cutlass.pipeline import PipelineAsync as PipelineAsyncOg from cutlass.pipeline import PipelineCpAsync as PipelineCpAsyncOg from cutlass.pipeline import PipelineTmaAsync as PipelineTmaAsyncOg +from cutlass.pipeline import PipelineTmaStore as PipelineTmaStoreOg from cutlass.pipeline import PipelineTmaUmma as PipelineTmaUmmaOg from cutlass.pipeline import PipelineUmmaAsync as PipelineUmmaAsyncOg from cutlass.pipeline import PipelineAsyncUmma as PipelineAsyncUmmaOg @@ -34,6 +35,72 @@ def _override_create(parent_cls, child_cls): return create +@dsl_user_op +def mbarrier_arrive_release_cluster( + mbar_ptr: Pointer, peer_cta_rank_in_cluster: Int, *, loc=None, ip=None +) -> None: + """Arrive on a peer CTA's mbarrier with cluster-scope release semantics. + + cute.arch.mbarrier_arrive with a peer rank emits mbarrier.arrive with the default + .release.cta semantics, which does not order this CTA's smem writes with the peer + CTA's reads of them (e.g. a 2-CTA tcgen05.mma reading our smem over DSMEM). Emit + fence.release.sync_restrict::shared::cta.cluster followed by + mbarrier.arrive.relaxed.cluster.shared::cluster instead: the fence releases all smem + writes this thread has observed at cluster scope, and a cluster-scope acquire of the + barrier phase on the consumer side (mbarrier_test_wait_acquire_cluster) completes the + release-acquire relation. This is the cheaper form of mbarrier.arrive.release.cluster, + which lowers to MEMBAR.GPU. + """ + _nvvm.fence_sync_restrict(_nvvm.MemOrderKind.RELEASE, loc=loc, ip=ip) + remote_ptr = _nvvm.mapa( + llvm.PointerType.get(AddressSpace.dsmem), + mbar_ptr.to_llvm_ptr(loc=loc, ip=ip), + Int32(peer_cta_rank_in_cluster).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + _nvvm.mbarrier_arrive( + None, + remote_ptr, + count=Int32(1).ir_value(loc=loc, ip=ip), + scope=_nvvm.MemScopeKind.CLUSTER, + relaxed=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def mbarrier_acquire_cluster(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> None: + """Cluster-scope acquire observation of an already-completed mbarrier phase. + + Emits mbarrier.test_wait.parity.acquire.cluster. The regular waits + (mbarrier.try_wait.parity) only acquire at cta scope, which cannot pair with a + cluster-scope release from another CTA (mbarrier_arrive_release_cluster). Call this + after a regular wait on the same (barrier, phase) has already succeeded: the test_wait + then observes the completed phase and upgrades the observation to cluster scope. + The result must feed control flow, otherwise the test_wait is dead-code-eliminated; + branch to a (never-taken) blocking wait so the observation cannot be dropped. + """ + status = Boolean( + _nvvm.mbarrier_wait_parity( + mbar_ptr.to_llvm_ptr(loc=loc, ip=ip), + Int32(phase).ir_value(loc=loc, ip=ip), + _nvvm.MBarrierWaitKind.TEST, + scope=_nvvm.MBarrierScopeKind.CLUSTER, + order=_nvvm.MemOrderKind.ACQUIRE, + loc=loc, + ip=ip, + ) + ) + if_generate( + status == 0, + lambda: cute.arch.mbarrier_wait(mbar_ptr, phase, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + + def _make_state(index: Int32, phase: Int32) -> PipelineState: """Construct a PipelineState from index and phase (count/stages unused by callers).""" return PipelineState(stages=0, count=Int32(0), index=index, phase=phase) @@ -238,11 +305,12 @@ class PipelineCpAsync(_PipelineIndexPhaseMixin, PipelineCpAsyncOg): @staticmethod def create( *args, + barrier_storage: Optional[cute.Pointer] = None, elect_one_release: bool = False, syncwarp_before_release: bool = True, **kwargs, ): - obj = PipelineCpAsyncOg.create(*args, **kwargs) + obj = PipelineCpAsyncOg.create(*args, barrier_storage=barrier_storage, **kwargs) object.__setattr__(obj, "__class__", PipelineCpAsync) object.__setattr__(obj, "_elect_one_release", elect_one_release) object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release) @@ -268,7 +336,23 @@ class PipelineCpAsync(_PipelineIndexPhaseMixin, PipelineCpAsyncOg): @dataclass(frozen=True) class PipelineTmaAsync(_PipelineIndexPhaseMixin, PipelineTmaAsyncOg): - """Override producer_acquire to take in extra_tx_count parameter.""" + """PipelineTmaAsync with extra_tx_count plus optional elected consumer release.""" + + _elect_one_release: bool = False + _syncwarp_before_release: bool = True + + @staticmethod + def create( + *args, + elect_one_release: bool = False, + syncwarp_before_release: bool = True, + **kwargs, + ): + obj = PipelineTmaAsyncOg.create(*args, **kwargs) + object.__setattr__(obj, "__class__", PipelineTmaAsync) + object.__setattr__(obj, "_elect_one_release", elect_one_release) + object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release) + return obj @dsl_user_op def producer_acquire( @@ -289,14 +373,48 @@ class PipelineTmaAsync(_PipelineIndexPhaseMixin, PipelineTmaAsyncOg): loc=loc, ip=ip, ) - if const_expr(extra_tx_count == 0): + if const_expr(isinstance(extra_tx_count, int) and extra_tx_count == 0): self.sync_object_full.arrive(state.index, self.producer_mask, loc=loc, ip=ip) else: tx_count = self.sync_object_full.tx_count + extra_tx_count self.sync_object_full.arrive_and_expect_tx(state.index, tx_count, loc=loc, ip=ip) + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineTmaAsyncOg.consumer_release, + self, + state, + self._elect_one_release, + self._syncwarp_before_release, + loc, + ip, + ) + + +# ── PipelineTmaStore ──────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineTmaStore(PipelineTmaStoreOg): + """PipelineTmaStore with configurable cp.async.bulk wait read flag.""" + + _read: bool = True + + @staticmethod + def create(*args, read: bool = True, **kwargs): + obj = PipelineTmaStoreOg.create(*args, **kwargs) + object.__setattr__(obj, "__class__", PipelineTmaStore) + object.__setattr__(obj, "_read", read) + return obj + + @dsl_user_op + def producer_acquire(self, *, loc=None, ip=None) -> None: + cute.arch.cp_async_bulk_wait_group(self.num_stages - 1, read=self._read, loc=loc, ip=ip) -PipelineTmaAsync.create = _override_create(PipelineTmaAsyncOg, PipelineTmaAsync) + @dsl_user_op + def producer_tail(self, *, loc=None, ip=None) -> None: + cute.arch.cp_async_bulk_wait_group(0, read=self._read, loc=loc, ip=ip) # ── PipelineTmaUmma ───────────────────────────────────────────────────────── @@ -357,10 +475,56 @@ PipelineTmaUmma.create = _override_create(PipelineTmaUmmaOg, PipelineTmaUmma) @dataclass(frozen=True) class PipelineUmmaAsync(_PipelineIndexPhaseMixin, PipelineUmmaAsyncOg): - pass + """ + PipelineUmmaAsync with optional elect_one for producer_commit and + consumer_release, mirroring PipelineAsync. + """ + + _elect_one_commit: bool = False + _syncwarp_before_commit: bool = True + _elect_one_release: bool = False + _syncwarp_before_release: bool = True + + @staticmethod + def create( + *args, + elect_one_commit: bool = False, + syncwarp_before_commit: bool = True, + elect_one_release: bool = False, + syncwarp_before_release: bool = True, + **kwargs, + ): + obj = PipelineUmmaAsyncOg.create(*args, **kwargs) + object.__setattr__(obj, "__class__", PipelineUmmaAsync) + object.__setattr__(obj, "_elect_one_commit", elect_one_commit) + object.__setattr__(obj, "_syncwarp_before_commit", syncwarp_before_commit) + object.__setattr__(obj, "_elect_one_release", elect_one_release) + object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release) + return obj + @dsl_user_op + def producer_commit(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineUmmaAsyncOg.producer_commit, + self, + state, + self._elect_one_commit, + self._syncwarp_before_commit, + loc, + ip, + ) -PipelineUmmaAsync.create = _override_create(PipelineUmmaAsyncOg, PipelineUmmaAsync) + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineUmmaAsyncOg.consumer_release, + self, + state, + self._elect_one_release, + self._syncwarp_before_release, + loc, + ip, + ) # ── PipelineAsyncUmma ─────────────────────────────────────────────────────── @@ -368,10 +532,56 @@ PipelineUmmaAsync.create = _override_create(PipelineUmmaAsyncOg, PipelineUmmaAsy @dataclass(frozen=True) class PipelineAsyncUmma(_PipelineIndexPhaseMixin, PipelineAsyncUmmaOg): - pass + """ + PipelineAsyncUmma with optional elect_one for producer_commit and + consumer_release, mirroring PipelineAsync. + """ + + _elect_one_commit: bool = False + _syncwarp_before_commit: bool = True + _elect_one_release: bool = False + _syncwarp_before_release: bool = True + + @staticmethod + def create( + *args, + elect_one_commit: bool = False, + syncwarp_before_commit: bool = True, + elect_one_release: bool = False, + syncwarp_before_release: bool = True, + **kwargs, + ): + obj = PipelineAsyncUmmaOg.create(*args, **kwargs) + object.__setattr__(obj, "__class__", PipelineAsyncUmma) + object.__setattr__(obj, "_elect_one_commit", elect_one_commit) + object.__setattr__(obj, "_syncwarp_before_commit", syncwarp_before_commit) + object.__setattr__(obj, "_elect_one_release", elect_one_release) + object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release) + return obj + @dsl_user_op + def producer_commit(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineAsyncUmmaOg.producer_commit, + self, + state, + self._elect_one_commit, + self._syncwarp_before_commit, + loc, + ip, + ) -PipelineAsyncUmma.create = _override_create(PipelineAsyncUmmaOg, PipelineAsyncUmma) + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineAsyncUmmaOg.consumer_release, + self, + state, + self._elect_one_release, + self._syncwarp_before_release, + loc, + ip, + ) # ── PipelineTmaCpAsync ────────────────────────────────────────────────────── @@ -420,54 +630,6 @@ class PipelineTmaCpAsync(_PipelineIndexPhaseMixin, PipelineTmaAsyncOg): PipelineTmaCpAsync.create = _override_create(PipelineTmaAsyncOg, PipelineTmaCpAsync) -# ── MbarrierArrayWDropCount ───────────────────────────────────────────────── - - -class MbarrierArrayWDropCount(MbarrierArray): - @dsl_user_op - def __init__( - self, - barrier_storage: cute.Pointer, - num_stages: int, - agent: tuple[PipelineOp, CooperativeGroup], - tx_count: int = 0, - drop_count: Optional[Int32] = None, - *, - loc=None, - ip=None, - ) -> None: - self.barrier_storage = barrier_storage - self.tx_count = tx_count - self.num_stages = num_stages - self.op_type, self.cg = agent - self.arrive_count = self.cg.size - self.drop_count = drop_count - - if self.num_stages <= 0: - raise ValueError("Error: Mbarrier stage count must be greater than 0.") - if self.arrive_count <= 0: - raise ValueError("Error: Mbarrier arrive count must be greater than 0.") - if self.op_type is PipelineOp.TmaLoad and self.tx_count < 0: - raise ValueError("Error: Mbarrier tx count must not be less than 0 for TMA ops.") - - if const_expr(drop_count is not None): - self.arrive_count = self.arrive_count - drop_count - - # Store mbarrier base pointer - self.mbarrier_base = self.barrier_storage - - # Mbarrier initialization in constructor - self.mbarrier_init(loc=loc, ip=ip) - - def __extract_mlir_values__(self): - return [self.barrier_storage, self.drop_count] - - def __new_from_mlir_values__(self, values): - return MbarrierArrayWDropCount( - values[0], self.num_stages, (self.op_type, self.cg), self.tx_count, values[1] - ) - - # ── PipelineTmaCpAsyncUmma ────────────────────────────────────────────────── @@ -478,108 +640,6 @@ class PipelineTmaCpAsyncUmma(PipelineTmaUmmaOg): (e.g. Blackwell mainloops) """ - @dsl_user_op - @staticmethod - def create( - *, - num_stages: int, - producer_group: CooperativeGroup, - consumer_group: CooperativeGroup, - tx_count: int, - barrier_storage: cute.Pointer = None, - cta_layout_vmnk: Optional[cute.Layout] = None, - mcast_mode_mn: tuple[int, int] = (1, 1), - defer_sync: bool = False, - producer_drop_count: Optional[Int32] = None, - loc=None, - ip=None, - ): - """Creates and initializes a new PipelineTmaUmma instance. - - :param num_stages: Number of buffer stages for this pipeline - :type num_stages: int - :param producer_group: CooperativeGroup for the producer agent - :type producer_group: CooperativeGroup - :param consumer_group: CooperativeGroup for the consumer agent - :type consumer_group: CooperativeGroup - :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage - :type tx_count: int - :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers - :type barrier_storage: cute.Pointer, optional - :param cta_layout_vmnk: Layout of the cluster shape - :type cta_layout_vmnk: cute.Layout, optional - :param mcast_mode_mn: Tuple specifying multicast modes for m and n dimensions (each 0 or 1) - :type mcast_mode_mn: tuple[int, int], optional - :raises ValueError: If barrier_storage is not a cute.Pointer instance - :return: A new PipelineTmaUmma instance configured with the provided parameters - :rtype: PipelineTmaUmma - """ - if not isinstance(barrier_storage, cute.Pointer): - raise TypeError( - f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" - ) - - producer_type = PipelineOp.TmaLoad - consumer_type = PipelineOp.TCGen05Mma - - producer = (producer_type, producer_group) - consumer = (consumer_type, consumer_group) - - sync_object_full = MbarrierArrayWDropCount( - barrier_storage.align(min_align=8), - num_stages, - producer, - tx_count, - drop_count=producer_drop_count, - loc=loc, - ip=ip, - ) - sync_object_empty = PipelineTmaUmmaOg._make_sync_object( - barrier_storage.align(min_align=8) + num_stages, - num_stages, - consumer, - loc=loc, - ip=ip, - ) - - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1: - # No mcast mask if not using clusters - producer_mask = None - # All threadblocks are leaders if not using clusters - is_leader_cta = True - else: - producer_mask = PipelineTmaUmmaOg._compute_mcast_arrival_mask( - cta_layout_vmnk, mcast_mode_mn, loc=loc, ip=ip - ) - is_leader_cta = PipelineTmaUmmaOg._compute_is_leader_cta( - cta_layout_vmnk, loc=loc, ip=ip - ) - - cta_group = ( - cute.nvgpu.tcgen05.CtaGroup.ONE - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip) == 1 - else cute.nvgpu.tcgen05.CtaGroup.TWO - ) - - consumer_mask = producer_mask - - if not defer_sync: - cute.arch.mbarrier_init_fence() - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1: - agent_sync(Agent.ThreadBlock) - else: - agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) - - return PipelineTmaCpAsyncUmma( - sync_object_full, - sync_object_empty, - num_stages, - producer_mask, - consumer_mask, - is_leader_cta, - cta_group, - ) - @dsl_user_op def producer_acquire( self, @@ -617,3 +677,6 @@ class PipelineTmaCpAsyncUmma(PipelineTmaUmmaOg): cute.arch.cp_async_mbarrier_arrive_noinc( self.producer_get_barrier(state, loc=loc, ip=ip), loc=loc, ip=ip ) + + +PipelineTmaCpAsyncUmma.create = _override_create(PipelineTmaUmmaOg, PipelineTmaCpAsyncUmma) diff --git a/build/torch-cuda/quack/reduce.py b/build/torch-cuda/quack/reduce.py index ea7c1fa3b9af9be7bc2034f00ec8f0ba6d8e254a..eb9eac5a2bfd9e1a59aae5109397c7bae7ab67f9 100644 --- a/build/torch-cuda/quack/reduce.py +++ b/build/torch-cuda/quack/reduce.py @@ -7,13 +7,52 @@ from typing import Callable, Optional import cutlass import cutlass.cute as cute from cutlass import Int32, Int64, Float32, Boolean, const_expr +from cutlass.base_dsl.arch import Arch from . import utils as utils +_operator_max = getattr(operator, "max", None) +_operator_min = getattr(operator, "min", None) +_cutlass_min = getattr(cutlass, "min", None) + + +@cute.jit +def warp_reduce( + val: cute.Numeric, + op: Callable, + threads_in_group: cutlass.Constexpr[int] = cute.arch.WARP_SIZE, + dtype: cutlass.Constexpr = None, +) -> cute.Numeric: + arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum() + if const_expr(threads_in_group == cute.arch.WARP_SIZE): + val_dtype = dtype if const_expr(dtype is not None) else getattr(val, "dtype", None) + if const_expr(val_dtype == Int32): + if const_expr(op is operator.add): + return cute.arch.warp_redux_sync(val, "add") + if const_expr(op is max or op is cutlass.max or op is _operator_max): + return cute.arch.warp_redux_sync(val, "max") + if const_expr(op is min or op is _cutlass_min or op is _operator_min): + return cute.arch.warp_redux_sync(val, "min") + if const_expr(val_dtype == Float32 and arch.is_family_of(Arch.sm_100f)): + if const_expr( + op is max or op is cutlass.max or op is cute.arch.fmax or op is _operator_max + ): + return cute.arch.warp_redux_sync(val, "fmax") + if const_expr( + op is min or op is _cutlass_min or op is cute.arch.fmin or op is _operator_min + ): + return cute.arch.warp_redux_sync(val, "fmin") + return cute.arch.warp_reduction(val, op, threads_in_group=threads_in_group) + + @cute.jit def block_reduce( - val: cute.Numeric, op: Callable, reduction_buffer: cute.Tensor, init_val: cute.Numeric = 0.0 + val: cute.Numeric, + op: Callable, + reduction_buffer: cute.Tensor, + init_val: cute.Numeric = 0.0, + dtype: cutlass.Constexpr = None, ) -> cute.Numeric: """reduction_buffer has shape (num_warps / warp_per_row, warps_per_row)""" lane_idx, warp_idx = cute.arch.lane_idx(), cute.arch.warp_idx() @@ -25,7 +64,7 @@ def block_reduce( block_reduce_val = init_val if lane_idx < warps_per_row: block_reduce_val = reduction_buffer[row_idx, lane_idx] - return cute.arch.warp_reduction(block_reduce_val, op) + return warp_reduce(block_reduce_val, op, dtype=dtype) @cute.jit @@ -36,6 +75,7 @@ def cluster_reduce( mbar_ptr: cute.Pointer, init_val: cute.Numeric = 0.0, phase: Optional[Int32] = None, + dtype: cutlass.Constexpr = None, ) -> cute.Numeric: """reduction_buffer has shape (num_warps / warps_per_row, (warps_per_row, cluster_n))""" cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -63,7 +103,7 @@ def cluster_reduce( idx = lane_idx + i * cute.arch.WARP_SIZE if idx < cute.size(reduction_buffer, mode=[1]): block_reduce_val = op(block_reduce_val, reduction_buffer[row_idx, idx]) - return cute.arch.warp_reduction(block_reduce_val, op) + return warp_reduce(block_reduce_val, op, dtype=dtype) @cute.jit @@ -74,12 +114,27 @@ def block_or_cluster_reduce( mbar_ptr: Optional[cute.Pointer], phase: Optional[Int32] = None, init_val: cute.Numeric = 0.0, + dtype: cutlass.Constexpr = None, ) -> cute.Numeric: """Perform either block or cluster reduction based on whether mbar_ptr is provided.""" if const_expr(mbar_ptr is None): - return block_reduce(val, op, reduction_buffer, init_val=init_val) + return block_reduce( + val, + op, + reduction_buffer, + init_val=init_val, + dtype=dtype, + ) else: - return cluster_reduce(val, op, reduction_buffer, mbar_ptr, phase=phase, init_val=init_val) + return cluster_reduce( + val, + op, + reduction_buffer, + mbar_ptr, + phase=phase, + init_val=init_val, + dtype=dtype, + ) @cute.jit @@ -101,13 +156,14 @@ def row_reduce( warp_op = { cute.ReductionOp.ADD: operator.add, cute.ReductionOp.MAX: cute.arch.fmax if const_expr(x.dtype == Float32) else max, - cute.ReductionOp.MIN: min, + cute.ReductionOp.MIN: cute.arch.fmin if const_expr(x.dtype == Float32) else min, cute.ReductionOp.MUL: operator.mul, }[op] - val = cute.arch.warp_reduction( + val = warp_reduce( val, warp_op, threads_in_group=min(threads_per_row, cute.arch.WARP_SIZE), + dtype=x.dtype, ) if const_expr(hook_fn is not None): hook_fn() @@ -118,7 +174,13 @@ def row_reduce( ) if const_expr(warps_per_row > 1 or cluster_n > 1): val = block_or_cluster_reduce( - val, warp_op, reduction_buffer, mbar_ptr, phase=phase, init_val=init_val + val, + warp_op, + reduction_buffer, + mbar_ptr, + phase=phase, + init_val=init_val, + dtype=x.dtype, ) return val @@ -135,10 +197,11 @@ def online_softmax_reduce( ) -> [Float32, Float32, Optional[cute.TensorSSA]]: assert x.dtype == Float32, "x must be of type Float32" """reduction_buffer must have shape (num_warps / warps_per_row, (warps_per_row, cluster_n), 2)""" - max_x = cute.arch.warp_reduction( + max_x = warp_reduce( x.reduce(cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0), cute.arch.fmax, threads_in_group=min(threads_per_row, cute.arch.WARP_SIZE), + dtype=Float32, ) log2_e = math.log2(math.e) exp_x = cute.math.exp2(x * log2_e - (max_x * log2_e), fastmath=True) @@ -170,7 +233,7 @@ def online_softmax_reduce( max_x_single_warp, sum_exp_x = utils.i64_to_f32x2( reduction_buffer[row_idx, lane_idx] ) - max_x_final = cute.arch.warp_reduction(max_x_single_warp, cute.arch.fmax) + max_x_final = warp_reduce(max_x_single_warp, cute.arch.fmax, dtype=Float32) sum_exp_x *= cute.math.exp(max_x_single_warp - max_x_final, fastmath=True) sum_exp_x = cute.arch.warp_reduction(sum_exp_x, operator.add) if const_expr(return_exp_x): @@ -209,7 +272,7 @@ def online_softmax_reduce( max_x_final = max_x_single_warp.load().reduce( cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0 ) - max_x_final = cute.arch.warp_reduction(max_x_final, cute.arch.fmax) + max_x_final = warp_reduce(max_x_final, cute.arch.fmax, dtype=Float32) sum_exp_x = 0.0 for i in cutlass.range_constexpr(num_iter): sum_exp_x += sum_exp_x_single_warp[i] * cute.math.exp( diff --git a/build/torch-cuda/quack/reduction_base.py b/build/torch-cuda/quack/reduction_base.py index 9139f512375179121eab12a4ab76922d10d281db..c0894d2687031eb956e18b19f074fd8f25636d7f 100644 --- a/build/torch-cuda/quack/reduction_base.py +++ b/build/torch-cuda/quack/reduction_base.py @@ -25,6 +25,20 @@ class ReductionBase: def _set_cluster_n(self): self.cluster_n = 1 + def _cap_cluster_n(self, vecsize: int) -> None: + """Cap ``cluster_n`` so every peer CTA owns a distinct, non-empty N-tile. + + A clustered launch splits the row across ``cluster_n`` CTAs. If + ``threads_per_row * cluster_n`` exceeds the number of vector blocks in the + row (``N // vecsize``), one CTA tile already spans the whole row + (``tiler_mn[1] >= N``); local_tile then collapses every peer onto tile 0, + so the peers re-reduce the same columns and double-count in the cluster + reduction. Capping to ``(N // vecsize) // threads_per_row`` guarantees + ``tiler_mn[1] < N`` whenever the resulting ``cluster_n > 1``. + """ + max_cluster_n = max(1, (self.N // vecsize) // self._threads_per_row()) + self.cluster_n = min(self.cluster_n, max_cluster_n) + def _get_tiled_copy(self, vecsize: int = 1): assert self.N % vecsize == 0, f"Input N {self.N} is not divisible by vector size {vecsize}" threads_per_row = self._threads_per_row() diff --git a/build/torch-cuda/quack/rms_final_reduce.py b/build/torch-cuda/quack/rms_final_reduce.py index 1b65d95ed6e43ef21d38e291d6f40a977fdee901..24505d3ef9c06906c22a78f77b1670f53df26f5e 100644 --- a/build/torch-cuda/quack/rms_final_reduce.py +++ b/build/torch-cuda/quack/rms_final_reduce.py @@ -13,15 +13,16 @@ import cutlass.cute as cute from cutlass import Float32, const_expr import torch -from ._ops_compat import add_quack_op_namespace_prefix +from ._ops_compat import add_op_namespace_prefix from torch import Tensor from . import copy_utils as copy_utils from .compile_utils import make_fake_tensor as fake_tensor from .reduce import row_reduce from .reduction_base import ReductionBase -from .cache_utils import jit_cache +from .cache import jit_cache from .cute_dsl_utils import torch2cute_dtype_map +from .dsl import cute_op class RmsFinalReduce(ReductionBase): @@ -135,8 +136,8 @@ def _compile_rms_final_reduce(dtype, N): ) -@torch.library.custom_op( - add_quack_op_namespace_prefix("rms_final_reduce_out"), +@cute_op( + add_op_namespace_prefix("rms_final_reduce_out"), mutates_args=("rstd",), device_types="cuda", ) @@ -153,15 +154,6 @@ def _rms_final_reduce_out( compiled_fn(x, rstd, scale, eps) -@_rms_final_reduce_out.register_fake -def _rms_final_reduce_out_fake(x, rstd, scale, eps): - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY and not isinstance(x.shape[0], torch.SymInt): - x_dtype = torch2cute_dtype_map[x.dtype] - _compile_rms_final_reduce(x_dtype, x.shape[1]) - - def rms_final_reduce( x: Tensor, # (M, N) partial squared sums scale: float, # typically 1.0 / total_columns @@ -171,11 +163,5 @@ def rms_final_reduce( assert x.ndim == 2 M = x.shape[0] rstd = torch.empty(M, dtype=torch.float32, device=x.device) - - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY: - return rstd - _rms_final_reduce_out(x, rstd, scale, eps) return rstd diff --git a/build/torch-cuda/quack/rmsnorm.py b/build/torch-cuda/quack/rmsnorm.py index b3bcfc362fa2697392d756b7baa974ea93fa2d1c..9e840f9fc93ffbc9818ce54d0db85a4a8b543798 100644 --- a/build/torch-cuda/quack/rmsnorm.py +++ b/build/torch-cuda/quack/rmsnorm.py @@ -8,21 +8,49 @@ import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute -from cutlass import Float32, Int32, const_expr +import cutlass.pipeline as pipeline +from cutlass import Float32, Int32, Int64, const_expr +from cutlass.cute.nvgpu import cpasync import torch -from ._ops_compat import add_quack_op_namespace_prefix +from ._ops_compat import add_op_namespace_prefix from torch import Tensor from . import utils as utils from . import copy_utils as copy_utils from . import layout_utils as layout_utils from .compile_utils import make_fake_tensor as fake_tensor +from .dsl import cute_op from .reduce import row_reduce from .reduction_base import ReductionBase -from .cache_utils import jit_cache +from .cache import jit_cache from .cute_dsl_utils import torch2cute_dtype_map -from cutlass.base_dsl import Arch +from .autotuner import autotune, AutotuneConfig +from .rmsnorm_config import ( + RmsNormBwdConfig, + RmsNormFwdConfig, + get_all_bwd_configs, + get_all_fwd_configs, + get_sm_count, + prune_invalid_rmsnorm_bwd_configs, + prune_invalid_rmsnorm_fwd_configs, +) +from cutlass.base_dsl.arch import Arch + + +def _bucket_T_hint(T_hint: int) -> int: + """Round ``T_hint`` up to the next power of 2 to bucket the JIT cache key. + + Each base-2 order of magnitude becomes a single bucket so adjacent T + values share a compiled binary instead of triggering a recompile per row + count. Buckets align with powers of 2 (..., 512, 1024, 2048, ...), which + keeps the analytical heuristic thresholds (e.g. ``T_hint <= 1024``) exact + at their power-of-2 boundaries. ``T_hint <= 0`` (the "unknown shape" + SymInt sentinel) is preserved. + """ + if T_hint <= 0: + return 0 + return 1 << (T_hint - 1).bit_length() def _ensure_contiguous(t): @@ -35,18 +63,31 @@ def _ensure_contiguous(t): class RMSNorm(ReductionBase): - def __init__(self, dtype: Type[cutlass.Numeric], N: int, is_layernorm: bool = False): + def __init__( + self, + dtype: Type[cutlass.Numeric], + N: int, + is_layernorm: bool = False, + config: Optional["RmsNormFwdConfig"] = None, + ): super().__init__(dtype, N, stage=2 if is_layernorm else 1) self.is_layernorm = is_layernorm - self.reload_from = None if N <= (16384 if is_layernorm else 8192) else "smem" - self.delay_w_load = False + if config is None: + config = RmsNormFwdConfig.from_analytical_heuristic( + N, dtype.width, is_layernorm=is_layernorm + ) + self.config = config + self.reload_from = config.reload_from + self.delay_w_load = config.delay_w_load + self._num_threads_val = config.num_threads + self._threads_per_row_val = config.threads_per_row + self._cluster_n_val = config.cluster_n + + def _num_threads(self): + return self._num_threads_val def _threads_per_row(self): - N = self.N - for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]: - if N <= limit: - return threads - return 256 + return self._threads_per_row_val def _set_cluster_n(self): arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum() @@ -56,21 +97,7 @@ class RMSNorm(ReductionBase): return # SM12x supports cluster up to 8 max_cluster = 8 if arch.major == 12 else 16 - N = self.N - # cluster_n = 4 is faster and cluster_n = 2 for N=64k for some reason - # Similarly cluster_n = 8 is faster for N=128k - if arch.major == 12 and const_expr(self.dtype.width >= 32): - # SM12x 99 KB SMEM: fp32 needs tighter clustering (conservative for residual case) - thresholds = [(8 * 1024, 1), (16 * 1024, 2), (32 * 1024, 4), (64 * 1024, 8)] - elif const_expr(self.dtype.width == 16): - thresholds = [(16 * 1024, 1), (32 * 1024, 2), (64 * 1024, 4), (128 * 1024, 8)] - else: - thresholds = [(32 * 1024, 1), (64 * 1024, 2), (128 * 1024, 4), (256 * 1024, 8)] - for limit, cluster in thresholds: - if N <= limit: - self.cluster_n = cluster - return - self.cluster_n = max_cluster + self.cluster_n = min(self._cluster_n_val, max_cluster) @cute.jit def __call__( @@ -92,6 +119,7 @@ class RMSNorm(ReductionBase): max(*(t.element_type.width for t in [mX, mRes, mW, mB, mO, mResO] if t is not None)) ) vecsize = math.gcd(self.N, 128 // largest_dtype_width) + self._cap_cluster_n(vecsize) tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) num_threads = tiled_copy.size mW, mB = [ @@ -259,8 +287,16 @@ class RMSNorm(ReductionBase): if const_expr(mRes is not None): copy(tXgRes, tXrRes) x += tXrRes.load().to(cute.Float32) + x_centered = x - mean + if const_expr(not is_even_N): + # OOB lanes are zero-filled for the mean pass, but they must contribute zero + # to the variance pass (not mean^2 from (0 - mean)^2). + tXrX_centered = cute.make_rmem_tensor_like(tXrX, Float32) + tXrX_centered.store(x_centered) + utils.fill_oob(tXrX_centered, tXpX, fill_value=Float32.zero) + x_centered = tXrX_centered.load() sum_sq_x_sub_mean = row_reduce( - (x - mean) * (x - mean), + x_centered * x_centered, cute.ReductionOp.ADD, threads_per_row, reduction_buffer[None, None, 1], @@ -317,8 +353,8 @@ class RMSNorm(ReductionBase): copy(tXrO, tXgO) -@torch.library.custom_op( - add_quack_op_namespace_prefix("_rmsnorm_fwd"), +@cute_op( + add_op_namespace_prefix("_rmsnorm_fwd"), mutates_args=("out", "rstd", "mean", "residual_out"), device_types="cuda", # We need to specify the schema manually since we're mutating an optional tensor @@ -345,13 +381,15 @@ def _rmsnorm_fwd( Returns: Normalized output tensor of same shape as x """ - # Don't need to check is_cuda since torch.library ensures that + # TVM FFI validates tensor devices at runtime. supported_types = {torch.float16, torch.bfloat16, torch.float32} assert x.dtype in supported_types, "Unsupported dtype" if weight is not None: assert weight.dtype in supported_types, "Weight must be float32, float16 or bfloat16" if residual is not None: assert residual.dtype in supported_types, "Residual must be float16, bfloat16, or float32" + if x.numel() == 0: + return N = x.size(-1) per_head = (weight is not None and weight.dim() == 2) or (bias is not None and bias.dim() == 2) @@ -374,58 +412,6 @@ def _rmsnorm_fwd( )(x, weight, bias, residual, out, residual_out, rstd, mean, eps) -@_rmsnorm_fwd.register_fake -def _rmsnorm_fwd_fake( - x: Tensor, - weight: Optional[Tensor], - out: Tensor, - bias: Optional[Tensor] = None, - rstd: Optional[Tensor] = None, - mean: Optional[Tensor] = None, - residual: Optional[Tensor] = None, - residual_out: Optional[Tensor] = None, - eps: float = 1e-6, - is_layernorm: bool = False, -) -> None: - # See softmax.py _softmax_fwd_fake for why register_fake is needed. - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY and not isinstance(x.size(-1), torch.SymInt): - N = x.size(-1) - per_head = (weight is not None and weight.dim() == 2) or ( - bias is not None and bias.dim() == 2 - ) - dtype, out_dtype, weight_dtype, bias_dtype, res_dtype, res_out_dtype = [ - torch2cute_dtype_map[t.dtype] if t is not None else None - for t in [x, out, weight, bias, residual, residual_out] - ] - _compile_rmsnorm_fwd( - dtype, - out_dtype, - res_dtype, - weight_dtype, - bias_dtype, - res_out_dtype, - N, - rstd is not None, - mean is not None, - is_layernorm, - per_head, - ) - _compile_rmsnorm_bwd( - N, - dtype, - dtype, - dtype, - weight_dtype, - bias is not None, - res_dtype, - res_out_dtype, - weight is not None, - per_head, - ) - - @jit_cache def _compile_rmsnorm_fwd( dtype, @@ -439,6 +425,7 @@ def _compile_rmsnorm_fwd( has_mean, is_layernorm, per_head, + config: Optional[RmsNormFwdConfig] = None, ): batch_sym = cute.sym_int() head_sym = cute.sym_int() if per_head else None @@ -456,7 +443,7 @@ def _compile_rmsnorm_fwd( rstd_cute = fake_tensor(Float32, batch_shape) if has_rstd else None mean_cute = fake_tensor(Float32, batch_shape) if has_mean else None return cute.compile( - RMSNorm(dtype, N, is_layernorm=is_layernorm), + RMSNorm(dtype, N, is_layernorm=is_layernorm, config=config), x_cute, weight_cute, bias_cute, @@ -487,7 +474,7 @@ def rmsnorm_fwd( out_dtype = x.dtype if out_dtype is None else out_dtype out = torch.empty_like(x, dtype=out_dtype) rstd = torch.empty(*x.shape[:-1], device=x.device, dtype=torch.float32) if store_rstd else None - if residual is not None: + if residual is not None and residual_dtype is None: residual_dtype = residual.dtype if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype): residual_out = torch.empty_like( @@ -502,6 +489,58 @@ def rmsnorm_fwd( return out, residual_out, rstd +@autotune( + configs=[AutotuneConfig(config=c) for c in get_all_fwd_configs()], + key=["is_layernorm", "per_head"], + prune_configs_by={"early_config_prune": prune_invalid_rmsnorm_fwd_configs}, +) +def rmsnorm_fwd_tuned( + x: Tensor, + weight: Optional[Tensor], + out: Tensor, + bias: Optional[Tensor] = None, + rstd: Optional[Tensor] = None, + mean: Optional[Tensor] = None, + residual: Optional[Tensor] = None, + residual_out: Optional[Tensor] = None, + eps: float = 1e-6, + is_layernorm: bool = False, + per_head: bool = False, + config: Optional[RmsNormFwdConfig] = None, +) -> None: + """Autotuned RMSNorm/LayerNorm forward dispatch. + + The ``@autotune`` decorator injects ``config`` from the exhaustive search + space at first call for a given (shape, dtype, ``is_layernorm``, ``per_head``) + and caches the winner for subsequent calls. The un-tuned counterpart is + :func:`rmsnorm_fwd`, which uses the analytical heuristic. + """ + if config is None: + raise RuntimeError( + "rmsnorm_fwd_tuned requires a config (provided automatically by " + "the @autotune decorator). Use rmsnorm_fwd for the un-tuned path." + ) + N = x.size(-1) + dtype, out_dtype, weight_dtype, bias_dtype, res_dtype, res_out_dtype = [ + torch2cute_dtype_map[t.dtype] if t is not None else None + for t in [x, out, weight, bias, residual, residual_out] + ] + _compile_rmsnorm_fwd( + dtype, + out_dtype, + res_dtype, + weight_dtype, + bias_dtype, + res_out_dtype, + N, + rstd is not None, + mean is not None, + is_layernorm, + per_head, + config=config, + )(x, weight, bias, residual, out, residual_out, rstd, mean, eps) + + def rmsnorm_ref(x, w=None, bias=None, residual=None, eps=1e-6): x_f32 = x.float() if residual is not None: @@ -537,23 +576,45 @@ def rmsnorm_bwd_ref(x, w, dout, rstd, eps=1e-6): class RMSNormBackward(ReductionBase): - def __init__(self, dtype: cutlass.Numeric, N: int): + def __init__( + self, + dtype: cutlass.Numeric, + N: int, + dout_dtype: Optional[Type[cutlass.Numeric]] = None, + T_hint: int = 0, + per_head: bool = False, + config: Optional["RmsNormBwdConfig"] = None, + ): # 2 stages for double buffering when computing mean of x_hat * wdy super().__init__(dtype, N, stage=2, reduction_dtype=Float32) - self.reload_wdy = None if N <= 16 * 1024 else "smem" + dout_width = dout_dtype.width if dout_dtype is not None else dtype.width + if config is None: + config = RmsNormBwdConfig.from_analytical_heuristic( + N, dtype.width, dout_width, T_hint=T_hint + ) + self.config = config + self.reload_wdy = config.reload_wdy + self.reload_x = config.reload_x + self.per_head = per_head + self._num_threads_val = config.num_threads + self._threads_per_row_val = config.threads_per_row + self._cluster_n_val = config.cluster_n + tile_n = N // max(1, config.cluster_n) + row_bytes_x = tile_n * dtype.width // 8 + row_bytes_do = tile_n * dout_width // 8 + self.USE_TMA = ( + config.use_tma and not per_head and row_bytes_x % 16 == 0 and row_bytes_do % 16 == 0 + ) + self._can_use_tma = self.USE_TMA if self.N > 128 * 1024 and self.dtype.width >= 32: # Not enough smem raise ValueError("RMSNormBackward does not support N > 128k with dtype >= 32 bits") def _num_threads(self): - return 128 if self.N <= 4096 else 256 + return self._num_threads_val def _threads_per_row(self): - N = self.N - for limit, threads in [(64, 8), (128, 16), (256, 32), (512, 64), (4096, 128)]: - if N <= limit: - return threads - return 256 + return self._threads_per_row_val def _set_cluster_n(self): arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum() @@ -561,19 +622,8 @@ class RMSNormBackward(ReductionBase): if arch < Arch.sm_90: self.cluster_n = 1 return - # SM12x supports cluster up to 8 max_cluster = 8 if arch.major == 12 else 16 - N = self.N - if arch.major == 12 and const_expr(self.dtype.width >= 32): - # SM12x 99 KB SMEM: fp32 bwd double-buffers 2 tensors, needs much tighter clustering - thresholds = [(1024, 1), (8 * 1024, 2), (16 * 1024, 4), (32 * 1024, 8)] - else: - thresholds = [(8 * 1024, 1), (16 * 1024, 2), (32 * 1024, 4), (64 * 1024, 8)] - for limit, cluster in thresholds: - if N <= limit: - self.cluster_n = cluster - return - self.cluster_n = max_cluster + self.cluster_n = min(self._cluster_n_val, max_cluster) @cute.jit def __call__( @@ -596,15 +646,41 @@ class RMSNormBackward(ReductionBase): max(*(t.element_type.width for t in [mX, mW, mdO, mdResO, mdX, mdRes] if t is not None)) ) vecsize = math.gcd(self.N, 128 // largest_dtype_width) + self._cap_cluster_n(vecsize) tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) num_threads = tiled_copy.size mW = ( layout_utils.expand(mW, dim=0, size=tiler_mn[0]) if const_expr(mW is not None) else None ) + use_tma = const_expr(self.USE_TMA) + if const_expr(use_tma): + tma_smem_layout = cute.make_ordered_layout(tiler_mn, order=(1, 0)) + tma_op = cpasync.CopyBulkTensorTileG2SOp() + tma_atom_X, mX_tma = cpasync.make_tiled_tma_atom(tma_op, mX, tma_smem_layout, tiler_mn) + tma_atom_dO, mdO_tma = cpasync.make_tiled_tma_atom( + tma_op, mdO, tma_smem_layout, tiler_mn + ) + else: + tma_atom_X, mX_tma, tma_atom_dO, mdO_tma = None, None, None, None num_blocks = sm_count - num_heads = mX.shape[1] if const_expr(cute.rank(mX) == 3) else 1 + num_heads = mX.shape[1] if const_expr(self.per_head) else 1 self.kernel( - mX, mW, mdO, mdResO, mRstd, mdX, mdW, mdB, mdRes, tiler_mn, tiled_copy, threads_per_row + mX, + mW, + mdO, + mdResO, + mRstd, + mdX, + mdW, + mdB, + mdRes, + tma_atom_X, + mX_tma, + tma_atom_dO, + mdO_tma, + tiler_mn, + tiled_copy, + threads_per_row, ).launch( grid=[num_blocks, self.cluster_n, num_heads], block=[num_threads, 1, 1], @@ -624,18 +700,26 @@ class RMSNormBackward(ReductionBase): mdW: Optional[cute.Tensor], mdB: Optional[cute.Tensor], mdRes: Optional[cute.Tensor], + tma_atom_X: Optional[cute.CopyAtom], + mX_tma: Optional[cute.Tensor], + tma_atom_dO: Optional[cute.CopyAtom], + mdO_tma: Optional[cute.Tensor], tiler_mn: cute.Shape, tiled_copy: cute.TiledCopy, threads_per_row: cutlass.Constexpr[int], ): tidx, _, _ = cute.arch.thread_idx() - bidx_start, _, bidz = cute.arch.block_idx() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if const_expr(self.per_head): + bidx_start, _, bidz = cute.arch.block_idx() + else: + bidx_start, _, _ = cute.arch.block_idx() gdim, _, _ = cute.arch.grid_dim() cluster_y = const_expr(0) if const_expr(self.cluster_n == 1) else cute.arch.block_idx()[1] tv_layout = tiled_copy.layout_tv_tiled # Slice per head - if const_expr(cute.rank(mX) == 3): + if const_expr(self.per_head): mX, mW, mdO, mdResO, mdX, mdW, mdB, mdRes = [ mT[None, bidz, None] if const_expr(mT is not None) else None for mT in (mX, mW, mdO, mdResO, mdX, mdW, mdB, mdRes) @@ -643,15 +727,20 @@ class RMSNormBackward(ReductionBase): mRstd = mRstd[None, bidz] shape = mX.shape - M, N = shape[0], shape[1] + M = shape[0] is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) idX = cute.make_identity_tensor(shape) smem = cutlass.utils.SmemAllocator() - smem_layout = cute.make_ordered_layout((tiler_mn[0], tiler_mn[1], 2), order=(1, 0, 2)) - sX = smem.allocate_tensor(mX.element_type, smem_layout, byte_alignment=16) - sdO = smem.allocate_tensor(mdO.element_type, smem_layout, byte_alignment=16) + USE_TMA = const_expr(self.USE_TMA) + n_smem_stages = const_expr(self.config.smem_stages) + smem_layout = cute.make_ordered_layout( + (tiler_mn[0], tiler_mn[1], n_smem_stages), order=(1, 0, 2) + ) + smem_align = const_expr(128 if USE_TMA else 16) + sX = smem.allocate_tensor(mX.element_type, smem_layout, byte_alignment=smem_align) + sdO = smem.allocate_tensor(mdO.element_type, smem_layout, byte_alignment=smem_align) reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar( smem, tv_layout, is_persistent=True ) @@ -716,6 +805,10 @@ class RMSNormBackward(ReductionBase): tXrdB = cute.make_rmem_tensor_like(tXgdB, Float32) num_warps = cute.size(tiled_copy) // cute.arch.WARP_SIZE + NUM_PIPE_STAGES = const_expr(self.config.smem_stages) + + if const_expr(USE_TMA): + tma_mbar_ptr = smem.allocate_array(Int64, num_elems=NUM_PIPE_STAGES * 2) self._initialize_cluster(tidx, mbar_ptr, num_warps, is_persistent=True) @@ -727,23 +820,49 @@ class RMSNormBackward(ReductionBase): if const_expr(not is_even_N): tXrW.fill(0.0) copy(tXgW, tXrW) - - # Prefetch the first batch - row = tXcX[None, None, None, bidx_start][0][0] - if row < M: - copy(tXgX[None, None, None, bidx_start], tXsX[None, None, None, 0], is_async=True) - copy(tXgdO[None, None, None, bidx_start], tXsdO[None, None, None, 0], is_async=True) - else: - if const_expr(tiler_mn[0] > 1): - # Fill with zero, otherwise smem will be uninitialized, and we could read this back - # later into registers, causing wrong dW. - utils.fill_oob(tXsX[None, None, None, 0], None, fill_value=mX.element_type.zero) - utils.fill_oob(tXsdO[None, None, None, 0], None, fill_value=mdO.element_type.zero) - cute.arch.cp_async_commit_group() + # No-op fp32 round-trip; pins tXrW into stable registers across the loop. + tXrW.store((tXrW.load().to(Float32) + Float32(0.0)).to(tXrW.element_type)) if const_expr(self.cluster_n > 1): cute.arch.cluster_wait() + producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, NUM_PIPE_STAGES + ) + consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, NUM_PIPE_STAGES + ) + if const_expr(USE_TMA): + num_threads_total = cute.size(tiled_copy) + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_threads_total) + tma_bytes_x = const_expr(cute.size(tiler_mn) * mX.element_type.width // 8) + tma_bytes_do = const_expr(cute.size(tiler_mn) * mdO.element_type.width // 8) + tma_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=tma_mbar_ptr, + num_stages=NUM_PIPE_STAGES, + producer_group=producer_group, + consumer_group=consumer_group, + tx_count=tma_bytes_x + tma_bytes_do, + cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), + ) + gX_tma = cute.local_tile(mX_tma, tiler_mn, (None, cluster_y)) + gdO_tma = cute.local_tile(mdO_tma, tiler_mn, (None, cluster_y)) + tXsX_tma, tXgX_tma = cpasync.tma_partition( + tma_atom_X, + 0, + cute.make_layout(1), + cute.group_modes(sX, 0, 2), + cute.group_modes(gX_tma, 0, 2), + ) + tXsdO_tma, tXgdO_tma = cpasync.tma_partition( + tma_atom_dO, + 0, + cute.make_layout(1), + cute.group_modes(sdO, 0, 2), + cute.group_modes(gdO_tma, 0, 2), + ) + if const_expr(mdW is not None): tXrdW.fill(0.0) if const_expr(mdB is not None): @@ -751,28 +870,115 @@ class RMSNormBackward(ReductionBase): stage = Int32(0) producer_phase = Int32(1) consumer_phase = Int32(0) + next_wave_work_id = (NUM_PIPE_STAGES - 1) * gdim + next_wave_row_id = next_wave_work_id * tiler_mn[0] + + M_ceil = cute.ceil_div(M, tiler_mn[0]) + if const_expr(USE_TMA): + for prefetch_iter in cutlass.range_constexpr(const_expr(NUM_PIPE_STAGES - 1)): + init_bidx = bidx_start + prefetch_iter * gdim + if warp_id == 0: + if init_bidx < M_ceil: + tma_pipeline.producer_acquire(producer_state) + pipe_bar = tma_pipeline.producer_get_barrier(producer_state) + cute.copy( + tma_atom_X, + tXgX_tma[None, init_bidx], + tXsX_tma[None, producer_state.index], + tma_bar_ptr=pipe_bar, + ) + cute.copy( + tma_atom_dO, + tXgdO_tma[None, init_bidx], + tXsdO_tma[None, producer_state.index], + tma_bar_ptr=pipe_bar, + ) + tma_pipeline.producer_commit(producer_state) + producer_state.advance() + else: + # Pre-issue NUM_PIPE_STAGES-1 prefetches into smem stages 0..N-2. + # The bidx loop then maintains exactly NUM_PIPE_STAGES groups in flight + # via cp_async_wait_group(NUM_PIPE_STAGES - 1). + for prefetch_iter in cutlass.range_constexpr(const_expr(NUM_PIPE_STAGES - 1)): + init_bidx = bidx_start + prefetch_iter * gdim + init_row = tXcX[None, None, None, init_bidx][0][0] + if init_row < M: + copy( + tXgX[None, None, None, init_bidx], + tXsX[None, None, None, producer_state.index], + is_async=True, + ) + copy( + tXgdO[None, None, None, init_bidx], + tXsdO[None, None, None, producer_state.index], + is_async=True, + ) + else: + if const_expr(tiler_mn[0] > 1): + utils.fill_oob( + tXsX[None, None, None, producer_state.index], + None, + fill_value=mX.element_type.zero, + ) + utils.fill_oob( + tXsdO[None, None, None, producer_state.index], + None, + fill_value=mdO.element_type.zero, + ) + cute.arch.cp_async_commit_group() + producer_state.advance() + for bidx in cutlass.range(bidx_start, cute.ceil_div(M, tiler_mn[0]), gdim): row = tXcX[None, None, None, bidx][0][0] - if row + gdim * tiler_mn[0] < M: # Prefetch the next batch - copy( - tXgX[None, None, None, bidx + gdim], - tXsX[None, None, None, stage ^ 1], - is_async=True, - ) - copy( - tXgdO[None, None, None, bidx + gdim], - tXsdO[None, None, None, stage ^ 1], - is_async=True, - ) + if const_expr(USE_TMA): + ahead_bidx = bidx + next_wave_work_id + if warp_id == 0: + if ahead_bidx < M_ceil: + tma_pipeline.producer_acquire(producer_state) + pipe_bar = tma_pipeline.producer_get_barrier(producer_state) + cute.copy( + tma_atom_X, + tXgX_tma[None, ahead_bidx], + tXsX_tma[None, producer_state.index], + tma_bar_ptr=pipe_bar, + ) + cute.copy( + tma_atom_dO, + tXgdO_tma[None, ahead_bidx], + tXsdO_tma[None, producer_state.index], + tma_bar_ptr=pipe_bar, + ) + tma_pipeline.producer_commit(producer_state) + producer_state.advance() else: - if const_expr(tiler_mn[0] > 1): - utils.fill_oob( - tXsX[None, None, None, stage ^ 1], None, fill_value=mX.element_type.zero + # cp.async: prefetch the (NUM_PIPE_STAGES-1)-ahead batch into the + # smem slot we're about to free up. + ahead_bidx = bidx + next_wave_work_id + if row + next_wave_row_id < M: + copy( + tXgX[None, None, None, ahead_bidx], + tXsX[None, None, None, producer_state.index], + is_async=True, ) - utils.fill_oob( - tXsdO[None, None, None, stage ^ 1], None, fill_value=mdO.element_type.zero + copy( + tXgdO[None, None, None, ahead_bidx], + tXsdO[None, None, None, producer_state.index], + is_async=True, ) - cute.arch.cp_async_commit_group() + else: + if const_expr(tiler_mn[0] > 1): + utils.fill_oob( + tXsX[None, None, None, producer_state.index], + None, + fill_value=mX.element_type.zero, + ) + utils.fill_oob( + tXsdO[None, None, None, producer_state.index], + None, + fill_value=mdO.element_type.zero, + ) + cute.arch.cp_async_commit_group() + producer_state.advance() rstd = cutlass.Float.zero if row < M or tiler_mn[0] == 1: rstd = mRstd[row] @@ -781,10 +987,14 @@ class RMSNormBackward(ReductionBase): copy(tXgdResO[None, None, None, bidx], tXrdResO) elif tiler_mn[0] > 1: tXrdResO.fill(0.0) - cute.arch.cp_async_wait_group(1) - cute.autovec_copy(tXsX[None, None, None, stage], tXrX) + if const_expr(USE_TMA): + tma_pipeline.consumer_wait(consumer_state) + else: + cute.arch.cp_async_wait_group(const_expr(NUM_PIPE_STAGES - 1)) + smem_stage = consumer_state.index + cute.autovec_copy(tXsX[None, None, None, smem_stage], tXrX) x = tXrX.load().to(cute.Float32) - cute.autovec_copy(tXsdO[None, None, None, stage], tXrdO) + cute.autovec_copy(tXsdO[None, None, None, smem_stage], tXrdO) dout = tXrdO.load().to(cute.Float32) x_hat = x * rstd wdy = dout @@ -818,12 +1028,17 @@ class RMSNormBackward(ReductionBase): ) if const_expr(self.reload_wdy == "smem"): - cute.autovec_copy(tXsdO[None, None, None, stage], tXrdO) + cute.autovec_copy(tXsdO[None, None, None, smem_stage], tXrdO) dout = tXrdO.load().to(cute.Float32) wdy = dout if const_expr(mW is not None): wdy *= tXrW.load().to(Float32) + if const_expr(self.reload_x == "smem"): + cute.autovec_copy(tXsX[None, None, None, smem_stage], tXrX) + x = tXrX.load().to(cute.Float32) + x_hat = x * rstd + dx = (wdy - x_hat * mean_xhat_wdy) * rstd if const_expr(mdResO is not None): dx += tXrdResO.load().to(cute.Float32) @@ -839,6 +1054,12 @@ class RMSNormBackward(ReductionBase): if const_expr(mdB is not None): tXrdB.store(tXrdB.load() + dout) + if const_expr(USE_TMA): + tma_pipeline.sync_object_empty.arrive( + consumer_state.index, tma_pipeline.consumer_mask + ) + consumer_state.advance() + stage ^= 1 if stage == 0: consumer_phase ^= 1 @@ -903,25 +1124,8 @@ class RMSNormBackward(ReductionBase): cute.arch.mbarrier_wait(mbar_empty_ptr + stage, producer_phase) -def _get_sm_count(N: int, device: torch.device) -> int: - # This should be tuned on how many CTAs can be launched on each SM - sm_count_multiple = ( - 16 if N <= 256 else (8 if N <= 1024 else (4 if N <= 2048 else (2 if N <= 4096 else 1))) - ) - sm_count = torch.cuda.get_device_properties(device).multi_processor_count - # By right, if we're using cluster, this should be cluster_count not sm_count. - # But for cluster >= 4, due to quantization we would need to query active max cluster. - # Instead we just do sm_count * 2, which is reasonably larger than active_cluster_count to - # avoid wave quantization. - sm_count = ( - sm_count * sm_count_multiple if N <= 8192 else sm_count // 2 if N <= 16384 else sm_count * 2 - ) - - return sm_count - - -@torch.library.custom_op( - add_quack_op_namespace_prefix("_rmsnorm_bwd"), +@cute_op( + add_op_namespace_prefix("_rmsnorm_bwd"), mutates_args={"dx", "dw_partial", "db_partial", "dresidual"}, device_types="cuda", # We need to specify the schema manually since we're mutating an optional tensor @@ -951,23 +1155,21 @@ def _rmsnorm_bwd( - dw: Weight gradients tensor of same shape as weight (or None if weight is None) """ assert x.dim() in (2, 3), "Input must be 2D or 3D" - assert x.is_cuda, "Input tensor must be on CUDA device" supported_types = {torch.float16, torch.bfloat16, torch.float32} assert x.dtype in supported_types, "Unsupported dtype" per_head = x.dim() == 3 if weight is not None: - assert weight.is_cuda, "Weight tensor must be on CUDA device" assert weight.dtype in supported_types, "Weight must be float32, float16 or bfloat16" if dresidual_out is not None: assert dresidual_out.shape == x.shape - assert dresidual_out.is_cuda assert dresidual_out.dtype in supported_types, ( "Residual must be float16, bfloat16, or float32" ) if dresidual is not None: assert dresidual.shape == x.shape - assert dresidual.is_cuda assert dresidual.dtype in supported_types, "Residual must be float16, bfloat16, or float32" + if x.numel() == 0: + return N = x.size(-1) if dw_partial is None and db_partial is None: @@ -978,6 +1180,7 @@ def _rmsnorm_bwd( torch2cute_dtype_map[t.dtype] if t is not None else None for t in [x, dout, dx, weight, dresidual, dresidual_out] ] + T_hint = _bucket_T_hint(int(x.size(0)) if not isinstance(x.size(0), torch.SymInt) else 0) _compile_rmsnorm_bwd( N, dtype, @@ -989,48 +1192,10 @@ def _rmsnorm_bwd( dres_out_dtype, dw_partial is not None, per_head, + T_hint=T_hint, )(x, weight, dout, dresidual_out, rstd, dx, dw_partial, dresidual, db_partial, sm_count) -@_rmsnorm_bwd.register_fake -def _rmsnorm_bwd_fake( - x: Tensor, - weight: Optional[Tensor], - dout: Tensor, - rstd: Tensor, - dx: Tensor, - dw_partial: Optional[Tensor], - db_partial: Optional[Tensor] = None, - dresidual_out: Optional[Tensor] = None, - dresidual: Optional[Tensor] = None, - sm_count: Optional[int] = None, -) -> None: - # See softmax.py _softmax_fwd_fake for why register_fake is needed. - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY and not isinstance(x.size(-1), torch.SymInt): - N = x.size(-1) - per_head = x.dim() == 3 - if dw_partial is None and db_partial is None and sm_count is None: - return - dtype, dout_dtype, dx_dtype, weight_dtype, dres_dtype, dres_out_dtype = [ - torch2cute_dtype_map[t.dtype] if t is not None else None - for t in [x, dout, dx, weight, dresidual, dresidual_out] - ] - _compile_rmsnorm_bwd( - N, - dtype, - dout_dtype, - dx_dtype, - weight_dtype, - db_partial is not None, - dres_dtype, - dres_out_dtype, - dw_partial is not None, - per_head, - ) - - @jit_cache def _compile_rmsnorm_bwd( N, @@ -1043,6 +1208,8 @@ def _compile_rmsnorm_bwd( dres_out_dtype, has_dw_partial, per_head=False, + T_hint=0, + config: Optional[RmsNormBwdConfig] = None, ): batch_sym, batch_partial_sym = cute.sym_int(), cute.sym_int() head_sym = cute.sym_int() if per_head else None @@ -1060,7 +1227,14 @@ def _compile_rmsnorm_bwd( dw_partial_cute = fake_tensor(Float32, dw_shape, div) if has_dw_partial else None db_partial_cute = fake_tensor(Float32, dw_shape, div) if has_db_partial else None return cute.compile( - RMSNormBackward(dtype, N), + RMSNormBackward( + dtype, + N, + dout_dtype=dout_dtype, + T_hint=T_hint, + per_head=per_head, + config=config, + ), x_cute, weight_cute, dout_cute, @@ -1093,7 +1267,7 @@ def rmsnorm_bwd( dresidual = torch.empty_like(x, dtype=dresidual_out.dtype) else: dresidual = None - sm_count = _get_sm_count(N, device) + sm_count = get_sm_count(N, device) if per_head: H = x.size(1) sm_count = max(round(sm_count / H), 1) @@ -1108,19 +1282,90 @@ def rmsnorm_bwd( db_shape = (sm_count, H, N) if per_head else (sm_count, N) db_partial = torch.empty(db_shape, device=device, dtype=torch.float32) if has_bias else None - _rmsnorm_bwd( - x, weight, dout, rstd, dx, dw_partial, db_partial, dresidual_out, dresidual, sm_count - ) - - # we have summed the partial gradients in fp32, now we convert back to the weight dtype - dw = dw_partial.sum(dim=0).to(weight.dtype) if weight is not None else None - db = db_partial.sum(dim=0).to(weight.dtype) if has_bias else None + if x.numel() > 0: + _rmsnorm_bwd( + x, weight, dout, rstd, dx, dw_partial, db_partial, dresidual_out, dresidual, sm_count + ) + # we have summed the partial gradients in fp32, now we convert back to the weight dtype + dw = dw_partial.sum(dim=0).to(weight.dtype) if weight is not None else None + db = db_partial.sum(dim=0).to(weight.dtype) if has_bias else None + else: + dw = torch.zeros_like(weight) if weight is not None else None + db = torch.zeros_like(weight) if has_bias else None # dresidual is the same as dx in this case if has_residual and dresidual is None: dresidual = dx return dx, dw, db, dresidual +@autotune( + configs=[AutotuneConfig(config=c) for c in get_all_bwd_configs()], + key=["per_head", "has_dw_partial", "has_db_partial"], + prune_configs_by={"early_config_prune": prune_invalid_rmsnorm_bwd_configs}, +) +def rmsnorm_bwd_tuned( + x: Tensor, + weight: Optional[Tensor], + dout: Tensor, + rstd: Tensor, + dx: Tensor, + dw_partial: Optional[Tensor] = None, + db_partial: Optional[Tensor] = None, + dresidual_out: Optional[Tensor] = None, + dresidual: Optional[Tensor] = None, + sm_count: Optional[int] = None, + per_head: bool = False, + has_dw_partial: bool = False, + has_db_partial: bool = False, + config: Optional[RmsNormBwdConfig] = None, +) -> None: + """Autotuned RMSNorm backward dispatch. + + The ``@autotune`` decorator injects ``config`` from the exhaustive search + space at first call for a given (shape, dtype, ``per_head``, has_*) and + caches the winner for subsequent calls. The un-tuned counterpart is + :func:`rmsnorm_bwd`, which uses the analytical heuristic. + """ + if config is None: + raise RuntimeError( + "rmsnorm_bwd_tuned requires a config (provided automatically by " + "the @autotune decorator). Use rmsnorm_bwd for the un-tuned path." + ) + # The persistent grid size is encoded in the partial-accumulator shape + # (dw_partial / db_partial have shape (sm_count, ..., N)). Derive + # sm_count from there when available; require it explicitly only when + # neither buffer is provided. Mirrors the _rmsnorm_bwd torch-op + # contract. + if dw_partial is not None: + sm_count = dw_partial.shape[0] + elif db_partial is not None: + sm_count = db_partial.shape[0] + elif sm_count is None: + raise ValueError( + "rmsnorm_bwd_tuned: sm_count is required when neither dw_partial " + "nor db_partial is provided." + ) + N = x.size(-1) + dtype, dout_dtype, dx_dtype, weight_dtype, dres_dtype, dres_out_dtype = [ + torch2cute_dtype_map[t.dtype] if t is not None else None + for t in [x, dout, dx, weight, dresidual, dresidual_out] + ] + _compile_rmsnorm_bwd( + N, + dtype, + dout_dtype, + dx_dtype, + weight_dtype, + has_db_partial, + dres_dtype, + dres_out_dtype, + has_dw_partial, + per_head, + T_hint=0, + config=config, + )(x, weight, dout, dresidual_out, rstd, dx, dw_partial, dresidual, db_partial, sm_count) + + class RMSNormFunction(torch.autograd.Function): """Autograd wrapper for rmsnorm. diff --git a/build/torch-cuda/quack/rmsnorm_config.py b/build/torch-cuda/quack/rmsnorm_config.py new file mode 100644 index 0000000000000000000000000000000000000000..5f801f7fc10238a8e2caa1ae980ba6c4689f28d6 --- /dev/null +++ b/build/torch-cuda/quack/rmsnorm_config.py @@ -0,0 +1,508 @@ +# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. + +"""Launch configuration for the RMSNorm forward and backward kernels. + +Mirrors :mod:`quack.gemm_config`: frozen dataclasses that capture the launch +knobs, plus arch-specific factories that own the heuristics. +""" + +import itertools +from dataclasses import dataclass +from typing import List, Optional + +import torch + + +@dataclass(frozen=True) +class RmsNormFwdConfig: + num_threads: int + threads_per_row: int + cluster_n: int + # None = compute once into registers; "smem" / "gmem" = reload x (and + # residual) before the post-reduction epilogue. + reload_from: Optional[str] + # Defer the weight/bias load until after the row reduction. + delay_w_load: bool = False + + @classmethod + def from_analytical_heuristic( + cls, + N: int, + dtype_width: int, + arch_major: Optional[int] = None, + is_layernorm: bool = False, + ) -> "RmsNormFwdConfig": + """Pick a launch config from the hand-tuned analytical heuristic. + + ``arch_major`` defaults to the current device's capability. The same + ladder is used for Hopper, Blackwell, and SM12x today; a future + ``_for_blackwell_fwd`` factory can be added and dispatched on + ``arch_major >= 10``. For autotuning, use :func:`get_all_fwd_configs`. + """ + if arch_major is None: + arch_major = _detect_arch_major() + return _for_hopper_fwd(N, dtype_width, arch_major, is_layernorm) + + +def _for_hopper_fwd( + N: int, dtype_width: int, arch_major: int, is_layernorm: bool +) -> RmsNormFwdConfig: + num_threads = 128 if N <= 16 * 1024 else 256 + + threads_per_row = 256 + for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]: + if N <= limit: + threads_per_row = threads + break + + if arch_major < 9: + cluster_n = 1 + else: + max_cluster = 8 if arch_major == 12 else 16 + # cluster_n=4 is faster than cluster_n=2 for N=64k; cluster_n=8 is + # faster for N=128k. + if arch_major == 12 and dtype_width >= 32: + # SM12x 99 KB SMEM: fp32 needs tighter clustering (conservative for residual case) + thresholds = [(8 * 1024, 1), (16 * 1024, 2), (32 * 1024, 4), (64 * 1024, 8)] + elif dtype_width == 16: + thresholds = [(16 * 1024, 1), (32 * 1024, 2), (64 * 1024, 4), (128 * 1024, 8)] + elif is_layernorm: + # fp32 layernorm: bump cluster earlier than fp16/bf16. The 2-pass path's + # single-CTA tile is bandwidth-limited at N=16k/32k; cluster_n=2 splits + # the row across two CTAs and recovers ~3-14% at those sizes. + thresholds = [(8 * 1024, 1), (64 * 1024, 2), (128 * 1024, 4), (256 * 1024, 8)] + else: + # fp32 rmsnorm (1-pass) is already saturated at cluster_n=1 for N<=32k; + # bumping to cluster_n=2 there regresses ~3%. + thresholds = [(32 * 1024, 1), (64 * 1024, 2), (128 * 1024, 4), (256 * 1024, 8)] + cluster_n = max_cluster + for limit, cluster in thresholds: + if N <= limit: + cluster_n = cluster + break + + reload_threshold = 16 * 1024 if is_layernorm else 8 * 1024 + return RmsNormFwdConfig( + num_threads=num_threads, + threads_per_row=threads_per_row, + cluster_n=cluster_n, + reload_from=None if N <= reload_threshold else "smem", + delay_w_load=False, + ) + + +@dataclass(frozen=True) +class RmsNormBwdConfig: + num_threads: int + threads_per_row: int + cluster_n: int + # None = recompute from registers; "smem" = reload from shared memory. + reload_wdy: Optional[str] + reload_x: Optional[str] + use_tma: bool + # Number of smem stages used by the prefetch pipeline. Drives both the + # cp.async (use_tma=False) and TMA (use_tma=True) paths, which lead by + # ``smem_stages - 1`` batches. Larger depths hide more latency at the cost + # of smem footprint. + smem_stages: int = 2 + + @classmethod + def from_analytical_heuristic( + cls, + N: int, + dtype_width: int, + dout_width: int, + arch_major: Optional[int] = None, + T_hint: int = 0, + ) -> "RmsNormBwdConfig": + """Pick a launch config from the hand-tuned analytical heuristic. + + ``arch_major`` defaults to the current device's capability. + ``arch_major >= 10`` selects the Blackwell heuristic; anything else + uses the legacy/default path tuned on Hopper. For autotuning, use + :func:`get_all_bwd_configs`. + """ + if arch_major is None: + arch_major = _detect_arch_major() + if arch_major >= 10: + return _for_blackwell_bwd(N, dtype_width, dout_width, T_hint) + return _for_hopper_bwd(N, dtype_width, arch_major) + + +def _for_hopper_bwd(N: int, dtype_width: int, arch_major: int) -> RmsNormBwdConfig: + num_threads = 128 if N <= 4096 else 256 + for limit, threads in [(64, 8), (128, 16), (256, 32), (512, 64), (4096, 128)]: + if N <= limit: + threads_per_row = threads + break + else: + threads_per_row = 256 + + if arch_major < 9: + cluster_n = 1 + else: + max_cluster = 8 if arch_major == 12 else 16 + if arch_major == 12 and dtype_width >= 32: + thresholds = [(1024, 1), (8 * 1024, 2), (16 * 1024, 4), (32 * 1024, 8)] + else: + thresholds = [(8 * 1024, 1), (16 * 1024, 2), (32 * 1024, 4), (64 * 1024, 8)] + cluster_n = max_cluster + for limit, cluster in thresholds: + if N <= limit: + cluster_n = cluster + break + + return RmsNormBwdConfig( + num_threads=num_threads, + threads_per_row=threads_per_row, + cluster_n=cluster_n, + reload_wdy=None if N <= 16 * 1024 else "smem", + reload_x=None, + use_tma=False, + ) + + +def _for_blackwell_bwd( + N: int, dtype_width: int, dout_width: int, T_hint: int = 0 +) -> RmsNormBwdConfig: + """Pick a launch config for RMSNorm bwd on Blackwell. + + All thresholds are expressed in ``row_bytes = N * max(x, dout)`` so a + single ladder handles bf16, fp32, and mixed-dtype combinations. ``x_bytes`` + governs the X tile (loads, smem footprint); ``max_bytes`` is the wider + side which sets register pressure for the per-thread fragments. + """ + # Safety floor for very narrow rows: keep tpr below 128 so we don't over- + # parallelise tiny problems. + if N <= 64: + threads_per_row = 8 + elif N <= 128: + threads_per_row = 16 + elif N <= 256: + threads_per_row = 32 + elif N <= 512: + threads_per_row = 64 + else: + threads_per_row = None + + if threads_per_row is not None: + return RmsNormBwdConfig( + num_threads=128, + threads_per_row=threads_per_row, + cluster_n=1, + reload_wdy=None, + reload_x=None, + use_tma=False, + ) + + max_bytes = max(dtype_width, dout_width) // 8 + row_bytes = N * max_bytes + + if row_bytes >= 48 * 1024: + # Spread the row across a CTA cluster. Step back to cluster_n=4 only + # when T is tiny AND the row isn't extreme — otherwise cn=8 keeps each + # CTA's tile small enough to fit comfortably in registers. + cluster_n = 4 if 0 < T_hint <= 1024 and row_bytes <= 64 * 1024 else 8 + num_threads, threads_per_row = 128, 128 + # Override if this cluster_n would overflow the device's smem budget. + cluster_n = _bump_cluster_n_for_smem( + cluster_n, + N, + smem_stages=2, + sum_bytes=(dtype_width + dout_width) // 8, + max_cluster=_max_cluster_for(10), # Blackwell + ) + elif row_bytes > 16 * 1024: + # Wider than 128 threads can comfortably handle at cluster_n=1; bump + # threads/row to keep per-thread fragments small. + cluster_n = 1 + num_threads, threads_per_row = 256, 256 + else: + cluster_n = 1 + num_threads, threads_per_row = 128, 128 + + bytes_per_thread_frag = (N // cluster_n) // threads_per_row * max_bytes + + # TMA pays off when the cluster needs prefetch (cn>=4), and also for + # fp32-class single-CTA wide rows where TMA's wider descriptors amortise + # setup. Pure bf16 single-CTA cases mostly don't benefit and can lose ~5%. + use_tma = cluster_n >= 4 or (max_bytes >= 4 and row_bytes >= 16 * 1024) + # reload_x: wide end of the cluster ladder, plus fp32-class single-CTA + # cases where the wider X fragments crowd registers across the row + # reduction barrier. + reload_x = ( + "smem" + if (cluster_n >= 8 and N >= 32 * 1024) + or (cluster_n == 1 and max_bytes >= 4 and bytes_per_thread_frag >= 64) + else None + ) + # reload_wdy: cluster cases get it for free, plus single-CTA cases where + # each thread holds ≥64 bytes of fragment (the wdy register count is then + # large enough to spill). + reload_wdy = "smem" if cluster_n >= 4 or bytes_per_thread_frag >= 64 else None + + return RmsNormBwdConfig( + num_threads=num_threads, + threads_per_row=threads_per_row, + cluster_n=cluster_n, + reload_wdy=reload_wdy, + reload_x=reload_x, + use_tma=use_tma, + ) + + +def _get_sm_count_hopper(N: int, sm_count: int) -> int: + # This should be tuned on how many CTAs can be launched on each SM. + sm_count_multiple = ( + 16 if N <= 256 else (8 if N <= 1024 else (4 if N <= 2048 else (2 if N <= 4096 else 1))) + ) + # By right, if we're using cluster, this should be cluster_count not sm_count. + # But for cluster >= 4, due to quantization we would need to query active max cluster. + # Instead we just do sm_count * 2, which is reasonably larger than active_cluster_count to + # avoid wave quantization. + return ( + sm_count * sm_count_multiple if N <= 8192 else sm_count // 2 if N <= 16384 else sm_count * 2 + ) + + +def _get_sm_count_blackwell(N: int, sm_count: int) -> int: + if N <= 256: + return sm_count * 16 + if N <= 1024: + return sm_count * 8 + if N <= 2048: + return sm_count * 4 + return sm_count * 2 + + +def get_sm_count(N: int, device: torch.device) -> int: + props = torch.cuda.get_device_properties(device) + if props.major >= 10: + return _get_sm_count_blackwell(N, props.multi_processor_count) + return _get_sm_count_hopper(N, props.multi_processor_count) + + +_CTA_THREAD_SIZE = (128, 256) +# Full launch-knob menu lives here; the per-call pruner in +# ``prune_invalid_rmsnorm_{fwd,bwd}_configs`` drops layouts that don't fit the +# current row. The tiny widths (8, 16, 32) are kept for the analytical +# heuristic's narrow-row safety floor (N <= 512) but are dropped from the +# autotune search space below since they're only optimal for that floor. +_THREADS_PER_REDUCTION_DIM = (8, 16, 32, 64, 128, 256) +_AUTOTUNE_THREADS_PER_REDUCTION_DIM = (64, 128, 256) +# smem_stages=4 doubles the data-buffer footprint vs stages=2 and crashes at +# launch for fp32 N>=64K on Blackwell (227 KB opt-in smem). Stages=3 still +# offers a meaningful prefetch depth without that risk. +_AUTOTUNE_SMEM_STAGES = (2, 3) + + +def _max_dynamic_smem_bytes() -> int: + """Per-CTA opt-in dynamic smem capacity for the current device. + + Returns 0 when CUDA is unavailable (callers should treat this as "no + smem-budget guard"). Falls back to ``shared_memory_per_block`` on older + PyTorch builds that lack the ``_optin`` field. + """ + if not torch.cuda.is_available(): + return 0 + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + return getattr(props, "shared_memory_per_block_optin", props.shared_memory_per_block) + + +def _bump_cluster_n_for_smem( + cluster_n: int, + N: int, + smem_stages: int, + sum_bytes: int, + max_cluster: int, +) -> int: + """Raise ``cluster_n`` to the smallest power of 2 such that the bwd's + sX+sdO data buffers fit under the device's opt-in dynamic smem. + + Footprint per CTA is ``(N / cluster_n) * smem_stages * sum_bytes`` where + ``sum_bytes = x_bytes + dout_bytes``. Snaps up to the next power of 2, + floors at the input ``cluster_n`` (never lowers the tuning's choice), and + caps at ``max_cluster``. If the required cluster_n exceeds ``max_cluster`` + the runtime guard in ``RMSNormBackward.__init__`` raises a precise overflow + error. Returns the input unchanged if CUDA is unavailable. + """ + # Reserved for row-reduction buffer, mbars, and smem alignment overhead. + _BWD_SMEM_RESERVED_BYTES = 4 * 1024 + smem_max = _max_dynamic_smem_bytes() + if smem_max <= 0: + return cluster_n + budget = max(smem_max - _BWD_SMEM_RESERVED_BYTES, 1) + needed = (N * smem_stages * sum_bytes + budget - 1) // budget # ceil-div + pow2 = 1 + while pow2 < needed: + pow2 *= 2 + return min(max(pow2, cluster_n), max_cluster) + + +def _detect_arch_major() -> int: + """Return the major device capability of the current CUDA device. + + Honors the ``QUACK_ARCH`` override (via ``get_device_capacity``) so + GPU-blind processes — compile-pool workers, CPU-only boxes — never + initialize CUDA here. This function runs at ``import quack`` time (the + module-level ``get_all_fwd_configs()`` call in ``rmsnorm.py``), so a raw + ``torch.cuda.current_device()`` would create a CUDA context on import, + which both slows imports and poisons forked children (torch's + "Cannot re-initialize CUDA in forked subprocess" guard). + + Falls back to 0 (no-cluster, no-TMA) when CUDA is unavailable so the + autotune search space stays well-defined for CPU-only imports. + """ + import os + + if os.environ.get("QUACK_ARCH") is None and not torch.cuda.is_available(): + return 0 + from .cute_dsl_utils import get_device_capacity + + return get_device_capacity()[0] + + +def _max_cluster_for(arch_major: int) -> int: + """Maximum cluster_n supported on this arch.""" + if arch_major < 9: + return 1 + # SM12x (RTX 50) supports up to 8; Hopper/Blackwell up to 16. + return 8 if arch_major == 12 else 16 + + +def get_all_fwd_configs() -> List[RmsNormFwdConfig]: + """Exhaustive search space of RMSNorm fwd configs for the current device. + + The search space is over launch knobs only — ``device_capacity`` is not a + tunable parameter, so the current device's capability is queried once and + used to bound ``cluster_n``. + """ + arch_major = _detect_arch_major() + max_cluster = _max_cluster_for(arch_major) + cluster_vals = tuple(c for c in (1, 2, 4, 8, 16) if c <= max_cluster) + reload_from_vals = (None, "smem", "gmem") + delay_w_load_vals = (False,) + + configs: List[RmsNormFwdConfig] = [] + for num_threads, threads_per_row, cluster_n, reload_from, delay_w_load in itertools.product( + _CTA_THREAD_SIZE, + _AUTOTUNE_THREADS_PER_REDUCTION_DIM, + cluster_vals, + reload_from_vals, + delay_w_load_vals, + ): + if threads_per_row > num_threads: + continue + if num_threads % threads_per_row != 0: + continue + configs.append( + RmsNormFwdConfig( + num_threads=num_threads, + threads_per_row=threads_per_row, + cluster_n=cluster_n, + reload_from=reload_from, + delay_w_load=delay_w_load, + ) + ) + return configs + + +def get_all_bwd_configs() -> List[RmsNormBwdConfig]: + """Exhaustive search space of RMSNorm bwd configs for the current device. + + Like :func:`get_all_fwd_configs`, the current device's capability bounds + ``cluster_n`` and gates ``use_tma`` (TMA requires SM90+). ``smem_stages`` + sweeps the safe depths in :data:`_AUTOTUNE_SMEM_STAGES`. + """ + arch_major = _detect_arch_major() + max_cluster = _max_cluster_for(arch_major) + cluster_vals = tuple(c for c in (1, 2, 4, 8, 16) if c <= max_cluster) + use_tma_vals = (False, True) if arch_major >= 9 else (False,) + reload_vals = (None, "smem") + + configs: List[RmsNormBwdConfig] = [] + for ( + num_threads, + threads_per_row, + cluster_n, + reload_wdy, + reload_x, + use_tma, + smem_stages, + ) in itertools.product( + _CTA_THREAD_SIZE, + _AUTOTUNE_THREADS_PER_REDUCTION_DIM, + cluster_vals, + reload_vals, + reload_vals, + use_tma_vals, + _AUTOTUNE_SMEM_STAGES, + ): + if threads_per_row > num_threads: + continue + if num_threads % threads_per_row != 0: + continue + configs.append( + RmsNormBwdConfig( + num_threads=num_threads, + threads_per_row=threads_per_row, + cluster_n=cluster_n, + reload_wdy=reload_wdy, + reload_x=reload_x, + use_tma=use_tma, + smem_stages=smem_stages, + ) + ) + return configs + + +def prune_invalid_rmsnorm_fwd_configs(configs, named_args: dict, **kwargs): + """Drop configs whose CTA layout doesn't fit the row width. + + The search space (see :func:`get_all_fwd_configs`) is already restricted to + the current device's capability, so all that's left is a per-call shape + check: ``threads_per_row * cluster_n > N`` would leave cluster CTAs with no + work to do, so skip those. + """ + kwargs = named_args | kwargs + x = kwargs["x"] + N = int(x.size(-1)) + pruned = [] + for ac in configs: + c = ac.kwargs["config"] + if c.threads_per_row * c.cluster_n > N: + continue + pruned.append(ac) + return pruned + + +def prune_invalid_rmsnorm_bwd_configs(configs, named_args: dict, **kwargs): + """Same per-call shape filter as the fwd, plus three ``use_tma`` drops + that mirror the runtime ``USE_TMA`` guard in :class:`RMSNormBackward` + (``USE_TMA = use_tma and not per_head and row_bytes_x % 16 == 0 and + row_bytes_do % 16 == 0``). Configs that would silently fall back to the + cp.async path are dropped here so the autotune bench doesn't time the + same kernel twice and pick whichever happens to bench faster by noise. + """ + kwargs = named_args | kwargs + x = kwargs["x"] + dout = kwargs.get("dout") + N = int(x.size(-1)) + per_head = bool(kwargs.get("per_head", x.dim() == 3)) + x_bytes = x.element_size() + dout_bytes = dout.element_size() if dout is not None else x_bytes + pruned = [] + for ac in configs: + c = ac.kwargs["config"] + if c.threads_per_row * c.cluster_n > N: + continue + if c.use_tma: + if per_head: + continue + tile_n = N // max(1, c.cluster_n) + row_bytes_x = tile_n * x_bytes + row_bytes_do = tile_n * dout_bytes + if row_bytes_x % 16 != 0 or row_bytes_do % 16 != 0: + continue + pruned.append(ac) + return pruned diff --git a/build/torch-cuda/quack/rotary.py b/build/torch-cuda/quack/rotary.py new file mode 100644 index 0000000000000000000000000000000000000000..f7765b737a7d6c516118006ae51120a811362dbf --- /dev/null +++ b/build/torch-cuda/quack/rotary.py @@ -0,0 +1,834 @@ +# Copyright (c) 2026, Tri Dao. + +import math +from functools import partial +from typing import Optional + +import torch +from ._ops_compat import add_op_namespace_prefix +from torch import Tensor + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr + +from . import copy_utils +from .cache import jit_cache +from .compile_utils import make_fake_tensor as fake_tensor +from .cute_dsl_utils import torch2cute_dtype_map +from .dsl import cute_op + + +def _ensure_last_dim_contiguous(t: Tensor) -> Tensor: + """Ensure last-dim stride is 1 while avoiding copies for strided-but-row-contiguous inputs.""" + if torch.compiler.is_compiling(): + return t.contiguous() + return t if t.stride(-1) == 1 else t.contiguous() + + +class RotaryKernel: + def __init__( + self, + dtype: type[cutlass.Numeric], + dim: int, + interleaved: bool = False, + conjugate: bool = False, + ): + self.dtype = dtype + self.dim = dim + self.interleaved = interleaved + self.conjugate = conjugate + self.num_threads = 128 + self.tile_h = 2 if self.dim <= 96 else 1 + multiple = 32 if dim <= 128 else 64 + self.tile_d = (dim + multiple - 1) // multiple * multiple + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mCos: cute.Tensor, + mSin: cute.Tensor, + mSeqlenOffsets: Optional[cute.Tensor], + mCuSeqlens: Optional[cute.Tensor], + mO: cute.Tensor, + max_seqlen: Int32, + stream: cuda.CUstream, + ): + assert mX.element_type == self.dtype + assert mO.element_type == self.dtype + assert mCos.element_type == mSin.element_type + assert mCos.shape[1] == mSin.shape[1] + assert mCos.shape[1] * 2 == self.dim + + self.is_varlen = const_expr(mCuSeqlens is not None) + + # If not self.interleaved, we use cp.async for copying X from gmem -> smem, sync, then + # smem -> rmem with a different thread layout. + # If self.interleaved, then we directly copy X from gmem -> rmem, so the layout for X + # has to be compatible w the layout for cos/sin. + vecsize = math.gcd(128 // mX.element_type.width, self.dim) + if const_expr(not self.interleaved): + vecsize_cs = math.gcd(128 // mCos.element_type.width, self.dim // 2) + else: + vecsize_cs = vecsize // 2 + assert (128 // mCos.element_type.width) % vecsize_cs == 0 + vecs_per_row = self.tile_d // vecsize + vecs_per_row_cs = self.tile_d // 2 // vecsize_cs + threads_per_row = math.gcd(32, vecs_per_row) + threads_per_row_cs = math.gcd(32, vecs_per_row_cs) + if const_expr(not self.interleaved): + # Multiply so that all threads can fetch 1 cos and 1 sin. + multiple = max(mX.element_type.width * 2 // mCos.element_type.width, 1) + else: + multiple = 1 + tiler_mn = (self.num_threads // threads_per_row * multiple, self.tile_d) + tiled_copy = copy_utils.tiled_copy_2d( + mX.element_type, threads_per_row, self.num_threads, vecsize + ) + tiled_copy_cs = copy_utils.tiled_copy_2d( + mCos.element_type, threads_per_row_cs, self.num_threads, vecsize_cs + ) + assert tiler_mn[0] % (self.num_threads // threads_per_row_cs) == 0 + + # (b, s, h, d) -> (s, d, h, b); (s, h, d) -> (s, d, h) + x_layout_transpose = [0, 2, 1] if const_expr(self.is_varlen) else [1, 3, 2, 0] + mX, mO = [ + cute.make_tensor(t.iterator, cute.select(t.layout, mode=x_layout_transpose)) + for t in (mX, mO) + ] + assert cute.rank(mX) == (3 if const_expr(self.is_varlen) else 4) + if const_expr(self.is_varlen): + batch = mCuSeqlens.shape[0] - 1 + seqlen = max_seqlen + nheads = mX.shape[2] + else: + batch = mX.shape[3] + seqlen = mX.shape[0] + nheads = mX.shape[2] + self.kernel( + mX, + mCos, + mSin, + mSeqlenOffsets, + mCuSeqlens, + mO, + max_seqlen, + tiler_mn, + tiled_copy, + tiled_copy_cs, + ).launch( + grid=[ + cute.ceil_div(nheads, self.tile_h), + cute.ceil_div(seqlen, tiler_mn[0]), + batch, + ], + block=[self.num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mCos: cute.Tensor, + mSin: cute.Tensor, + mSeqlenOffsets: Optional[cute.Tensor], + mCuSeqlens: Optional[cute.Tensor], + mO: cute.Tensor, + max_seqlen: Int32, + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + tiled_copy_cs: cute.TiledCopy, + ): + tidx, _, _ = cute.arch.thread_idx() + head_idx, m_idx, batch_idx = cute.arch.block_idx() + + tiler_mnh = (tiler_mn[0], tiler_mn[1], self.tile_h) + tiler_cossin = (tiler_mn[0], tiler_mn[1] // 2) + + smem = cutlass.utils.SmemAllocator() + sX = None + if const_expr(not self.interleaved): + sX = smem.allocate_tensor( + mX.element_type, + cute.make_ordered_layout(tiler_mnh, order=(1, 0, 2)), + byte_alignment=16, + ) + + offset = Int32(0) + if const_expr(mSeqlenOffsets is not None): + offset = Int32(mSeqlenOffsets[batch_idx]) + + cX_cols = const_expr(max(self.dim, tiler_mn[1])) + if const_expr(self.is_varlen): + seq_start = Int32(mCuSeqlens[batch_idx]) + seq_len = Int32(mCuSeqlens[batch_idx + 1]) - seq_start + nheads = mX.shape[2] + cX_shape = (max_seqlen, cX_cols) + max_seqlen_or_x = max_seqlen + else: + seq_len = mX.shape[0] + nheads = mX.shape[2] + cX_shape = (mX.shape[0], cX_cols) + max_seqlen_or_x = mX.shape[0] + + if const_expr(self.is_varlen): + mX_batch = cute.domain_offset((seq_start, None, None), mX) + mO_batch = cute.domain_offset((seq_start, None, None), mO) + else: + mX_batch = mX[None, None, None, batch_idx] + mO_batch = mO[None, None, None, batch_idx] + gX = cute.local_tile(mX_batch, tiler_mnh, (m_idx, 0, head_idx)) + gO = cute.local_tile(mO_batch, tiler_mnh, (m_idx, 0, head_idx)) + gCos = cute.local_tile(cute.domain_offset((offset, None), mCos), tiler_cossin, (m_idx, 0)) + gSin = cute.local_tile(cute.domain_offset((offset, None), mSin), tiler_cossin, (m_idx, 0)) + cX = cute.local_tile(cute.make_identity_tensor(cX_shape), tiler_mn, (m_idx, 0)) + cCosSin = cute.local_tile( + cute.make_identity_tensor((max_seqlen_or_x, mCos.shape[1])), + tiler_cossin, + (m_idx, 0), + ) + + thr_copy = tiled_copy.get_slice(tidx) + tXcX_full = thr_copy.partition_S(cX) + thr_copy_cs = tiled_copy_cs.get_slice(tidx) + tScCosSin_full = thr_copy_cs.partition_S(cCosSin) + tXcX = tXcX_full[(0, None), None, None] + tXgX = thr_copy.partition_S(gX) + tXsX = thr_copy.partition_D(sX) if const_expr(sX is not None) else None + tXgO = thr_copy.partition_D(gO) + tCSgCos = thr_copy_cs.partition_S(gCos) + tCSgSin = thr_copy_cs.partition_S(gSin) + tCSrCos = cute.make_rmem_tensor_like(tCSgCos) + tCSrSin = cute.make_rmem_tensor_like(tCSgSin) + tXrX_g2r = cute.make_rmem_tensor_like(tXgX) + + is_even_dim = const_expr(tiler_mn[1] == self.dim) + pred, pred_cs = None, None + if const_expr(not is_even_dim): + pred = copy_utils.predicate_k(tXcX_full, limit=self.dim) + pred_cs = copy_utils.predicate_k(tScCosSin_full, limit=self.dim // 2) + copy = partial(copy_utils.copy, pred=pred[None, 0, None] if not is_even_dim else None) + copy_cs = partial(copy_utils.copy, pred=pred_cs[None, 0, None] if not is_even_dim else None) + + tScCosSin = tScCosSin_full[(0, None), None, None] + for m in cutlass.range(cute.size(tCSgCos, mode=[1]), unroll_full=True): + row_cs = tScCosSin[0, m, 0][0] + sincos_is_valid = row_cs < seq_len + if const_expr(mSeqlenOffsets is not None): + sincos_is_valid = sincos_is_valid and row_cs + offset < mCos.shape[0] + if sincos_is_valid: + copy_cs(tCSgCos[None, m, None], tCSrCos[None, m, None]) + copy_cs(tCSgSin[None, m, None], tCSrSin[None, m, None]) + + for h in cutlass.range_constexpr(self.tile_h): + if self.tile_h == 1 or h < nheads - head_idx * self.tile_h: + for m in cutlass.range(cute.size(tXgX, mode=[1]), unroll_full=True): + if tXcX[0, m, 0][0] < seq_len: + if const_expr(not self.interleaved): + copy(tXgX[None, m, None, h], tXsX[None, m, None, h], is_async=True) + else: + copy(tXgX[None, m, None, h], tXrX_g2r[None, m, None, h]) + if const_expr(not self.interleaved): + cute.arch.cp_async_commit_group() + + cos_vals = tCSrCos.load().to(Float32) + sin_vals = tCSrSin.load().to(Float32) + if const_expr(self.conjugate): + sin_vals = -sin_vals + rCos = cute.make_rmem_tensor(tCSrCos.shape, Float32) + rCos.store(cos_vals) + rSin = cute.make_rmem_tensor(tCSrSin.shape, Float32) + rSin.store(sin_vals) + if const_expr(not self.interleaved): + sX0 = cute.composition(sX, (tiler_mn[0], tiler_mn[1] // 2, self.tile_h)) + sX1 = cute.domain_offset((None, self.dim // 2, None), sX0) + tCsX0 = thr_copy_cs.partition_D(sX0) + tCsX1 = thr_copy_cs.partition_D(sX1) + for h in cutlass.range_constexpr(self.tile_h): + cute.arch.cp_async_wait_group(self.tile_h - h - 1) + cute.arch.sync_threads() + tCrX0 = copy_utils.load_s2r(tCsX0[None, None, None, h]) + tCrX1 = copy_utils.load_s2r(tCsX1[None, None, None, h]) + x0_vals = tCrX0.load().to(Float32) + x1_vals = tCrX1.load().to(Float32) + tCrX0.store((x0_vals * cos_vals - x1_vals * sin_vals).to(tCrX0.element_type)) + tCrX1.store((x0_vals * sin_vals + x1_vals * cos_vals).to(tCrX1.element_type)) + if const_expr(is_even_dim): + cute.autovec_copy(tCrX0, tCsX0[None, None, None, h]) + else: + for k in cutlass.range(cute.size(tCrX0, mode=[2]), unroll_full=True): + if pred_cs[0, 0, k]: # Need predication to avoid overwriting tCsX1 + cute.autovec_copy(tCrX0[None, None, k], tCsX0[None, None, k, h]) + cute.autovec_copy(tCrX1, tCsX1[None, None, None, h]) + cute.arch.sync_threads() + else: + for h in cutlass.range_constexpr(self.tile_h): + tCrX_f32 = tXrX_g2r[None, None, None, h].to(Float32) + assert cute.size(tCrX_f32.shape) == cute.size(rCos) * 2 + for i in cutlass.range(cute.size(tCrX_f32.shape) // 2, unroll_full=True): + x0, x1 = tCrX_f32[2 * i], tCrX_f32[2 * i + 1] + tCrX_f32[2 * i] = x0 * rCos[i] - x1 * rSin[i] + tCrX_f32[2 * i + 1] = x0 * rSin[i] + x1 * rCos[i] + tXrX_g2r[None, None, None, h].store(tCrX_f32.load().to(tXrX_g2r.element_type)) + + for h in cutlass.range_constexpr(self.tile_h): + if const_expr(not self.interleaved): + tXrX = copy_utils.load_s2r(tXsX[None, None, None, h]) + else: + tXrX = tXrX_g2r[None, None, None, h] + if self.tile_h == 1 or h < nheads - head_idx * self.tile_h: + for m in cutlass.range(cute.size(tXgO, mode=[1]), unroll_full=True): + if tXcX[0, m, 0][0] < seq_len: + copy(tXrX[None, m, None], tXgO[None, m, None, h]) + + @staticmethod + @jit_cache + def compile( + dtype, + cossin_dtype, + seqlen_offsets_dtype, + cu_seqlens_dtype, + dim, + interleaved, + conjugate, + ): + is_varlen = cu_seqlens_dtype is not None + has_seqlen_offsets = seqlen_offsets_dtype is not None + batch_sym = cute.sym_int() + batch_p1_sym = cute.sym_int() + seqlen_sym = cute.sym_int() + total_seqlen_sym = cute.sym_int() + nheads_sym = cute.sym_int() + x_dim_sym = cute.sym_int() + seqlen_ro_sym = cute.sym_int() + x_shape = ( + (total_seqlen_sym, nheads_sym, x_dim_sym) + if is_varlen + else (batch_sym, seqlen_sym, nheads_sym, x_dim_sym) + ) + x_divby = math.gcd(128 // dtype.width, dim) + cossin_divby = math.gcd(128 // cossin_dtype.width, dim // 2) + x_cute = fake_tensor(dtype, x_shape, x_divby) + out_cute = fake_tensor(dtype, x_shape, x_divby) + cos_cute = fake_tensor(cossin_dtype, (seqlen_ro_sym, dim // 2), cossin_divby) + sin_cute = fake_tensor(cossin_dtype, (seqlen_ro_sym, dim // 2), cossin_divby) + seqlen_offsets_cute = ( + cute.runtime.make_fake_tensor( + seqlen_offsets_dtype, (batch_sym,), stride=(cute.sym_int64(divisibility=1),) + ) + if has_seqlen_offsets + else None + ) + cu_seqlens_cute = fake_tensor(cu_seqlens_dtype, (batch_p1_sym,)) if is_varlen else None + return cute.compile( + RotaryKernel(dtype, dim, interleaved=interleaved, conjugate=conjugate), + x_cute, + cos_cute, + sin_cute, + seqlen_offsets_cute, + cu_seqlens_cute, + out_cute, + Int32(0), # max_seqlen, just for compilation + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _launch_rotary( + x: Tensor, + cos: Tensor, + sin: Tensor, + seqlen_offsets: Optional[Tensor], + cu_seqlens: Optional[Tensor], + out: Tensor, + max_seqlen: int, + interleaved: bool, + conjugate: bool, +) -> None: + assert x.stride(-1) == 1 and out.stride(-1) == 1, ( + "Rotary vectorized path requires last-dim stride 1" + ) + assert cos.dtype == sin.dtype and cos.shape == sin.shape + if x.numel() == 0: + return + dtype = torch2cute_dtype_map[x.dtype] + cossin_dtype = torch2cute_dtype_map[cos.dtype] + dim_half = cos.size(1) + dim = dim_half * 2 + seqlen_offsets_dtype = ( + torch2cute_dtype_map[seqlen_offsets.dtype] if seqlen_offsets is not None else None + ) + cu_seqlens_dtype = Int32 if cu_seqlens is not None else None + RotaryKernel.compile( + dtype, + cossin_dtype, + seqlen_offsets_dtype, + cu_seqlens_dtype, + dim, + interleaved, + conjugate, + )(x, cos, sin, seqlen_offsets, cu_seqlens, out, max_seqlen) + + +@cute_op( + add_op_namespace_prefix("_rotary_fwd_out"), + mutates_args=("out",), + device_types="cuda", + schema="(Tensor x, Tensor cos, Tensor sin, Tensor? seqlen_offsets, Tensor? cu_seqlens, Tensor(a!) out, int max_seqlen, bool interleaved, bool conjugate) -> ()", +) +def _rotary_fwd_out( + x: Tensor, + cos: Tensor, + sin: Tensor, + seqlen_offsets: Optional[Tensor], + cu_seqlens: Optional[Tensor], + out: Tensor, + max_seqlen: int, + interleaved: bool, + conjugate: bool, +) -> None: + _launch_rotary(x, cos, sin, seqlen_offsets, cu_seqlens, out, max_seqlen, interleaved, conjugate) + + +@cute_op( + add_op_namespace_prefix("_rotary_fwd_inplace"), + mutates_args=("x",), + device_types="cuda", + schema="(Tensor(a!) x, Tensor cos, Tensor sin, Tensor? seqlen_offsets, Tensor? cu_seqlens, int max_seqlen, bool interleaved, bool conjugate) -> ()", +) +def _rotary_fwd_inplace( + x: Tensor, + cos: Tensor, + sin: Tensor, + seqlen_offsets: Optional[Tensor], + cu_seqlens: Optional[Tensor], + max_seqlen: int, + interleaved: bool, + conjugate: bool, +) -> None: + _launch_rotary(x, cos, sin, seqlen_offsets, cu_seqlens, x, max_seqlen, interleaved, conjugate) + + +# CustomOpDef.register_effect() is only public in PyTorch 2.10+ (pytorch#163284). +# On 2.8 / 2.9 fall back to the private torch._higher_order_ops.effects API, +# which takes the underlying OpOverload. Collapse this once 2.10 is the floor. +def _register_ordered_effect(op) -> None: + try: + from torch._library.custom_ops import EffectType + + op.register_effect(EffectType.ORDERED) + except ImportError: + from torch._higher_order_ops.effects import _EffectType, _register_effectful_op + + _register_effectful_op(op._opoverload, _EffectType.ORDERED) + + +@cute_op( + add_op_namespace_prefix("_rotary_inplace_bwd"), + mutates_args=(), + device_types="cuda", + schema="(Tensor dout, Tensor cos, Tensor sin, Tensor? seqlen_offsets, Tensor? cu_seqlens, int? max_seqlen, bool interleaved) -> ()", +) +def _rotary_inplace_bwd( + dout: Tensor, + cos: Tensor, + sin: Tensor, + seqlen_offsets: Optional[Tensor], + cu_seqlens: Optional[Tensor], + max_seqlen: Optional[int], + interleaved: bool, +) -> None: + # In-place forward can invert the rotation by mutating grad_output directly. + # Declaring this as a normal mutating op makes AOTAutograd clone grad_output + # first, so we register it as an ordered effect to keep Dynamo to one kernel. + max_seqlen = dout.shape[1] if cu_seqlens is None else max_seqlen + assert max_seqlen is not None + _launch_rotary( + dout, + cos, + sin, + seqlen_offsets, + cu_seqlens, + dout, + max_seqlen, + interleaved, + conjugate=True, + ) + + +_register_ordered_effect(_rotary_inplace_bwd) + + +def apply_rotary( + x: Tensor, + cos: Tensor, + sin: Tensor, + seqlen_offsets: Optional[Tensor] = None, + cu_seqlens: Optional[Tensor] = None, + max_seqlen: Optional[int] = None, + interleaved: bool = False, + inplace: bool = False, + conjugate: bool = False, +) -> Tensor: + """ + Apply rotary embedding to the first rotary_dim dimensions of x. + + x is (batch, seqlen, nheads, headdim) when cu_seqlens is None, otherwise + (total_seqlen, nheads, headdim). cos/sin are (seqlen_ro, rotary_dim / 2). + """ + supported_types = {torch.float16, torch.bfloat16, torch.float32} + assert x.dtype in supported_types, "Unsupported x dtype" + assert cos.dtype == sin.dtype, "cos and sin must have the same dtype" + assert cos.dtype in supported_types and sin.dtype in supported_types, ( + "Unsupported cos/sin dtype" + ) + is_varlen = cu_seqlens is not None + if not is_varlen: + _batch, seqlen, _nheads, headdim = x.shape + launch_max_seqlen = seqlen + else: + assert max_seqlen is not None, "If cu_seqlens is passed, max_seqlen must be passed" + assert cu_seqlens.dtype == torch.int32, "cu_seqlens must have dtype torch.int32" + total_seqlen, _nheads, headdim = x.shape + seqlen = max_seqlen + launch_max_seqlen = int(max_seqlen) + seqlen_ro, rotary_dim_half = cos.shape + rotary_dim = rotary_dim_half * 2 + assert rotary_dim <= headdim, "rotary_dim must be <= headdim" + assert headdim <= 512, "Only support headdim <= 512" + assert headdim % 8 == 0, "headdim must be divisible by 8" + assert rotary_dim % 8 == 0, "rotary_dim must be divisible by 8" + assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" + + cos, sin = _ensure_last_dim_contiguous(cos), _ensure_last_dim_contiguous(sin) + out = x if inplace else torch.empty_like(x) + if rotary_dim < headdim and not inplace: + out[..., rotary_dim:].copy_(x[..., rotary_dim:]) + if inplace: + _rotary_fwd_inplace( + x, + cos, + sin, + seqlen_offsets, + cu_seqlens, + launch_max_seqlen, + interleaved, + conjugate, + ) + else: + _rotary_fwd_out( + x, + cos, + sin, + seqlen_offsets, + cu_seqlens, + out, + launch_max_seqlen, + interleaved, + conjugate, + ) + return out + + +class ApplyRotaryEmb(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x: Tensor, + cos: Tensor, + sin: Tensor, + interleaved: bool = False, + inplace: bool = False, + seqlen_offsets: Optional[Tensor] = None, + cu_seqlens: Optional[Tensor] = None, + max_seqlen: Optional[int] = None, + ): + out = apply_rotary( + x, + cos, + sin, + seqlen_offsets=seqlen_offsets, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + interleaved=interleaved, + inplace=inplace, + ) + ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) + ctx.interleaved = interleaved + ctx.inplace = inplace + ctx.max_seqlen = max_seqlen + return out if not inplace else x + + @staticmethod + def backward(ctx, do): + cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors + if ctx.inplace: + cos, sin = _ensure_last_dim_contiguous(cos), _ensure_last_dim_contiguous(sin) + _rotary_inplace_bwd( + do, cos, sin, seqlen_offsets, cu_seqlens, ctx.max_seqlen, ctx.interleaved + ) + dx = do + else: + dx = apply_rotary( + do, + cos, + sin, + seqlen_offsets=seqlen_offsets, + cu_seqlens=cu_seqlens, + max_seqlen=ctx.max_seqlen, + interleaved=ctx.interleaved, + inplace=False, + conjugate=True, + ) + return dx, None, None, None, None, None, None, None + + +def apply_rotary_emb( + x: Tensor, + cos: Tensor, + sin: Tensor, + interleaved: bool = False, + inplace: bool = False, + seqlen_offsets: Optional[Tensor] = None, + cu_seqlens: Optional[Tensor] = None, + max_seqlen: Optional[int] = None, +) -> Tensor: + return ApplyRotaryEmb.apply( + x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen + ) + + +apply_rotary_emb_func = apply_rotary_emb + + +def _mark_dirty(ctx, tensor: Tensor) -> None: + # This custom autograd Function returns the same storage it mutates, so + # mark_dirty is the correct contract. Eager PyTorch rejects mark_dirty on a + # leaf requiring grad ("a leaf Variable ... used in an in-place operation"), + # but the public in-place rotary API has historically accepted that case. + if torch.compiler.is_compiling() or not (tensor.requires_grad and tensor.is_leaf): + ctx.mark_dirty(tensor) + + +def _apply_rotary_qkv_inplace( + qkv: Tensor, + cos: Tensor, + sin: Tensor, + seqlen_offsets: Optional[Tensor], + num_heads_q: int, + interleaved: bool, + conjugate: bool, +) -> None: + if qkv.dim() == 5: + batch, seqlen, three, nheads, headdim = qkv.shape + assert three == 3 + if qkv.is_contiguous(): + qk = qkv[:, :, :2].reshape(batch, seqlen, 2 * nheads, headdim) + apply_rotary( + qk, + cos, + sin, + seqlen_offsets=seqlen_offsets, + interleaved=interleaved, + inplace=True, + conjugate=conjugate, + ) + else: + # The packed QK reshape is only a view for contiguous QKV. Keep + # non-contiguous inputs correct by rotating Q and K separately. + apply_rotary( + qkv[:, :, 0], + cos, + sin, + seqlen_offsets=seqlen_offsets, + interleaved=interleaved, + inplace=True, + conjugate=conjugate, + ) + apply_rotary( + qkv[:, :, 1], + cos, + sin, + seqlen_offsets=seqlen_offsets, + interleaved=interleaved, + inplace=True, + conjugate=conjugate, + ) + else: + assert qkv.dim() == 4 + num_heads_k = (qkv.shape[2] - num_heads_q) // 2 + assert qkv.shape[2] == num_heads_q + 2 * num_heads_k + qk = qkv[:, :, : num_heads_q + num_heads_k] + apply_rotary( + qk, + cos, + sin, + seqlen_offsets=seqlen_offsets, + interleaved=interleaved, + inplace=True, + conjugate=conjugate, + ) + + +# Keep the QKV view/reshape work behind this mutating custom op. Dynamo's +# custom-autograd tracing still fails when the packed-QK GQA path is inlined. +@cute_op( + add_op_namespace_prefix("_rotary_qkv_inplace"), + mutates_args=("qkv",), + device_types="cuda", + schema="(Tensor(a!) qkv, Tensor cos, Tensor sin, Tensor? seqlen_offsets, int num_heads_q, bool interleaved, bool conjugate) -> ()", +) +def _rotary_qkv_inplace( + qkv: Tensor, + cos: Tensor, + sin: Tensor, + seqlen_offsets: Optional[Tensor], + num_heads_q: int, + interleaved: bool, + conjugate: bool, +) -> None: + _apply_rotary_qkv_inplace(qkv, cos, sin, seqlen_offsets, num_heads_q, interleaved, conjugate) + + +@cute_op( + add_op_namespace_prefix("_rotary_qkv_inplace_bwd"), + mutates_args=(), + device_types="cuda", + schema="(Tensor dqkv, Tensor cos, Tensor sin, Tensor? seqlen_offsets, int num_heads_q, bool interleaved) -> ()", +) +def _rotary_qkv_inplace_bwd( + dqkv: Tensor, + cos: Tensor, + sin: Tensor, + seqlen_offsets: Optional[Tensor], + num_heads_q: int, + interleaved: bool, +) -> None: + _apply_rotary_qkv_inplace( + dqkv, + cos, + sin, + seqlen_offsets, + num_heads_q, + interleaved, + conjugate=True, + ) + + +# The backward consumes and mutates grad_output in place. If this op is declared +# as a normal mutating custom op, AOTAutograd functionalizes it by cloning dqkv +# first. Mark it as an ordered effect instead so Dynamo keeps the call without +# inserting that clone; the returned grad input is the same mutated dqkv tensor. +_register_ordered_effect(_rotary_qkv_inplace_bwd) + + +class ApplyRotaryEmbQKV_(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv: Tensor, + cos: Tensor, + sin: Tensor, + interleaved=False, + seqlen_offsets=None, + num_heads_q=0, + ): + num_heads_q = int(num_heads_q) + _rotary_qkv_inplace( + qkv, + cos, + sin, + seqlen_offsets, + num_heads_q, + interleaved, + conjugate=False, + ) + ctx.save_for_backward(cos, sin, seqlen_offsets) + ctx.interleaved = interleaved + ctx.num_heads_q = num_heads_q + _mark_dirty(ctx, qkv) + return qkv + + @staticmethod + def backward(ctx, dqkv): + cos, sin, seqlen_offsets = ctx.saved_tensors + _rotary_qkv_inplace_bwd(dqkv, cos, sin, seqlen_offsets, ctx.num_heads_q, ctx.interleaved) + return dqkv, None, None, None, None, None + + +def apply_rotary_emb_qkv_( + qkv: Tensor, + cos: Tensor, + sin: Tensor, + interleaved: bool = False, + seqlen_offsets: Optional[Tensor] = None, + num_heads_q: Optional[int] = None, +) -> Tensor: + if qkv.dim() == 5: + return ApplyRotaryEmbQKV_.apply(qkv, cos, sin, interleaved, seqlen_offsets, 0) + + assert qkv.dim() == 4 + assert num_heads_q is not None + num_heads_q = int(num_heads_q) + return ApplyRotaryEmbQKV_.apply( + qkv, + cos, + sin, + interleaved, + seqlen_offsets, + num_heads_q, + ) + + +class ApplyRotaryEmbKV_(torch.autograd.Function): + @staticmethod + def forward( + ctx, + kv: Tensor, + cos: Tensor, + sin: Tensor, + interleaved: bool = False, + seqlen_offsets: Optional[Tensor] = None, + ): + batch, seqlen, two, nheads, headdim = kv.shape + assert two == 2 + apply_rotary( + kv[:, :, 0], + cos, + sin, + seqlen_offsets=seqlen_offsets, + interleaved=interleaved, + inplace=True, + ) + ctx.save_for_backward(cos, sin, seqlen_offsets) + ctx.interleaved = interleaved + return kv + + @staticmethod + def backward(ctx, dkv): + cos, sin, seqlen_offsets = ctx.saved_tensors + dk = apply_rotary( + dkv[:, :, 0], + cos, + sin, + seqlen_offsets=seqlen_offsets, + interleaved=ctx.interleaved, + inplace=False, + conjugate=True, + ) + dkv = torch.stack([dk, dkv[:, :, 1]], dim=2) + return dkv, None, None, None, None + + +def apply_rotary_emb_kv_( + kv: Tensor, + cos: Tensor, + sin: Tensor, + interleaved: bool = False, + seqlen_offsets: Optional[Tensor] = None, +) -> Tensor: + return ApplyRotaryEmbKV_.apply(kv, cos, sin, interleaved, seqlen_offsets) diff --git a/build/torch-cuda/quack/rounding.py b/build/torch-cuda/quack/rounding.py index 0886252b6ce8ae32fa5f356642c8f97df56147db..15f3541376a1d26920f4602031abae6a5cd27387 100644 --- a/build/torch-cuda/quack/rounding.py +++ b/build/torch-cuda/quack/rounding.py @@ -10,9 +10,10 @@ instruction and is only supported on Blackwell (SM100+) GPUs. from enum import IntEnum import cutlass +import cutlass.cute as cute from cutlass import Float32, Uint32 from cutlass._mlir import ir -from cutlass._mlir.dialects import arith, llvm, vector +from cutlass._mlir.dialects import llvm, vector from cutlass.cutlass_dsl import dsl_user_op, Int32, T @@ -27,6 +28,13 @@ class RoundingMode(IntEnum): RS = 1 +# Odd strides used to derive distinct Philox counters across output tiles and subtiles. +EPILOGUE_SR_SEED_M_STRIDE = 65537 +EPILOGUE_SR_SEED_N_STRIDE = 257 +EPILOGUE_SR_SEED_BATCH_STRIDE = 17 +EPILOGUE_SR_SEED_SUBTILE_STRIDE = 7 +EPILOGUE_SR_SEED_AUX_OUT_SALT = 0x9E3779B1 + PHILOX_N_ROUNDS_DEFAULT = 7 PHILOX_ROUND_A = 0xD2511F53 @@ -35,27 +43,52 @@ PHILOX_KEY_A = 0x9E3779B9 PHILOX_KEY_B = 0xBB67AE85 +@dsl_user_op +def epilogue_sr_seed( + base_seed: Int32, + tile_coord_mnkl: cute.Coord, + subtile_idx, + *, + loc=None, + ip=None, +) -> Int32: + return base_seed + ( + tile_coord_mnkl[0] * EPILOGUE_SR_SEED_M_STRIDE + + tile_coord_mnkl[1] * EPILOGUE_SR_SEED_N_STRIDE + + tile_coord_mnkl[3] * EPILOGUE_SR_SEED_BATCH_STRIDE + + subtile_idx * EPILOGUE_SR_SEED_SUBTILE_STRIDE + ) + + +@dsl_user_op +def epilogue_aux_out_sr_seed( + base_seed: Int32, + tile_coord_mnkl: cute.Coord, + subtile_idx, + *, + loc=None, + ip=None, +) -> Int32: + return epilogue_sr_seed( + base_seed + EPILOGUE_SR_SEED_AUX_OUT_SALT, + tile_coord_mnkl, + subtile_idx, + loc=loc, + ip=ip, + ) + + @dsl_user_op def mul_wide_u32(a: Uint32, b: Uint32, *, loc=None, ip=None) -> tuple: """Unsigned 32b x 32b -> 64 wide multiply via PTX `mul.wide.u32`. Returns (hi, lo) as a pair of Uint32 values. """ - struct_ty = ir.Type.parse("!llvm.struct<(i32, i32)>") - result = llvm.inline_asm( - struct_ty, - [ - Uint32(a).ir_value(loc=loc, ip=ip), - Uint32(b).ir_value(loc=loc, ip=ip), - ], - "{\n .reg .u64 prod;\n mul.wide.u32 prod, $2, $3;\n mov.b64 {$1, $0}, prod;\n}", - "=r,=r,r,r", - has_side_effects=False, - is_align_stack=False, - ) - i32_ty = T.i32() - hi = cutlass.Uint32(llvm.extractvalue(i32_ty, result, [0], loc=loc, ip=ip)) - lo = cutlass.Uint32(llvm.extractvalue(i32_ty, result, [1], loc=loc, ip=ip)) + prod = cute.arch.mul_wide(Uint32(a), Uint32(b), loc=loc, ip=ip) + # ptxas folds the shift and to(Uint32) into the same IMAD.WIDE.U32 register pair + # that the previous inline PTX mov.b64 split produced. + hi = (prod >> 32).to(Uint32, loc=loc, ip=ip) + lo = prod.to(Uint32, loc=loc, ip=ip) return hi, lo @@ -160,15 +193,17 @@ def convert_f32_to_bf16_sr( lo_idx = pair_idx * 2 hi_idx = pair_idx * 2 + 1 - src_lo = vector.extractelement( + src_lo = vector.extract( src_vec, - position=arith.constant(Int32.mlir_type, lo_idx, loc=loc, ip=ip), + dynamic_position=[], + static_position=[lo_idx], loc=loc, ip=ip, ) - src_hi = vector.extractelement( + src_hi = vector.extract( src_vec, - position=arith.constant(Int32.mlir_type, hi_idx, loc=loc, ip=ip), + dynamic_position=[], + static_position=[hi_idx], loc=loc, ip=ip, ) @@ -183,10 +218,11 @@ def convert_f32_to_bf16_sr( packed_i32 = cvt_f32x2_bf16x2_rs(Float32(src_lo), Float32(src_hi), entropy, loc=loc, ip=ip) packed_i32_val = cutlass.Int32(packed_i32).ir_value(loc=loc, ip=ip) - i32_vec = vector.insertelement( + i32_vec = vector.insert( packed_i32_val, i32_vec, - position=arith.constant(Int32.mlir_type, pair_idx, loc=loc, ip=ip), + dynamic_position=[], + static_position=[pair_idx], loc=loc, ip=ip, ) diff --git a/build/torch-cuda/quack/sm100_utils.py b/build/torch-cuda/quack/sm100_utils.py index 4911a88e38d0f45c4552eb8a25ceb89a67478c25..9f01fee6cc739af9ddbcc2eeddc810c9467579eb 100644 --- a/build/torch-cuda/quack/sm100_utils.py +++ b/build/torch-cuda/quack/sm100_utils.py @@ -4,8 +4,8 @@ from typing import Type, Union import cutlass.cute as cute import cutlass.utils.blackwell_helpers as sm100_utils_og -from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode from cutlass.cutlass_dsl import Numeric, dsl_user_op +from cutlass.cute.nvgpu import OperandMajorMode @dsl_user_op diff --git a/build/torch-cuda/quack/sm90_utils.py b/build/torch-cuda/quack/sm90_utils.py index ebabfe6509ffa2152aa90baedeb5982b76c95750..a20549f91e0e71b6478431eb48f308976d63f9ce 100644 --- a/build/torch-cuda/quack/sm90_utils.py +++ b/build/torch-cuda/quack/sm90_utils.py @@ -1,6 +1,6 @@ # Copyright (c) 2025, Tri Dao. -from typing import Type, Union, Optional +from typing import Literal, Type, Union, Optional import cutlass import cutlass.cute as cute @@ -42,6 +42,90 @@ def make_smem_layout( make_smem_layout_epi = make_smem_layout +def choose_sm90_wgmma_layout_mn( + tile_m: int, + tile_n: int, + num_wg: int, + *, + allow_swap_ab: bool = True, +) -> tuple[bool, int]: + """Return ``(swap_AB, AtomLayoutM)`` minimizing Hopper SS WGMMA SMEM traffic. + + The logical MMA is ``(tile_m, tile_n)``. Hopper's physical WGMMA M mode is + 64, and for ``num_wg`` in {1, 2, 3} the only useful warp-group layouts are + all WGs along physical M or all WGs along physical N. The returned + ``AtomLayoutM`` is in the caller's logical coordinate system. Callers that + pass a physical ``atom_layout_mnk`` directly to a lower-level MMA builder + should swap the first two atom-layout modes when ``swap_AB`` is true. + """ + if num_wg not in (1, 2, 3): + raise ValueError(f"SM90 WGMMA layout chooser expects num_wg in {{1, 2, 3}}, got {num_wg}") + if tile_m <= 0 or tile_n <= 0: + raise ValueError(f"tile_m and tile_n must be positive, got {(tile_m, tile_n)}") + + def best_physical_layout(x: int, y: int) -> tuple[int, int] | None: + # Prefer split-M when valid: it has wg_n=1, strictly lower traffic than + # split-N for the same orientation when num_wg > 1. + if x % (64 * num_wg) == 0 and y % 8 == 0: + return (num_wg, 1) + if x % 64 == 0 and y % (8 * num_wg) == 0: + return (1, num_wg) + return None + + best: tuple[int, bool, int] | None = None + for swap_ab in (False, True) if allow_swap_ab else (False,): + physical_m, physical_n = (tile_n, tile_m) if swap_ab else (tile_m, tile_n) + layout = best_physical_layout(physical_m, physical_n) + if layout is None: + continue + physical_atom_m, physical_atom_n = layout + score = physical_m * physical_atom_n + atom_layout_m = physical_atom_n if swap_ab else physical_atom_m + candidate = (score, swap_ab, atom_layout_m) + if best is None or candidate < best: + best = candidate + if best is None: + raise ValueError( + "no valid SM90 WGMMA layout for " + f"tile_m={tile_m}, tile_n={tile_n}, num_wg={num_wg}, " + f"allow_swap_ab={allow_swap_ab}" + ) + _, swap_ab, atom_layout_m = best + return swap_ab, atom_layout_m + + +def make_tiled_mma( + a_dtype: Type[Numeric], + a_major: Literal["K", "MN"], + b_major: Literal["K", "MN"], + tiler_n: int, + source: Literal["SS", "RS"] = "SS", + atom_layout_mnk: tuple = (1, 1, 1), + swap_AB: bool = False, + b_dtype: Optional[Type[Numeric]] = None, + acc_dtype: Type[Numeric] = Float32, +) -> cute.TiledMma: + """`b_dtype` defaults to `a_dtype`; pass it for mixed-precision MMAs (e.g. fp8). + `acc_dtype` defaults to Float32.""" + if b_dtype is None: + b_dtype = a_dtype + mode = {"K": cute.nvgpu.OperandMajorMode.K, "MN": cute.nvgpu.OperandMajorMode.MN} + a_mode, b_mode = mode[a_major], mode[b_major] + if swap_AB: + a_mode, b_mode = b_mode, a_mode + a_source = warpgroup.OperandSource.RMEM if source == "RS" else warpgroup.OperandSource.SMEM + return sm90_utils_og.make_trivial_tiled_mma( + a_dtype, + b_dtype, + a_mode, + b_mode, + acc_dtype, + atom_layout_mnk=atom_layout_mnk, + tiler_mn=(64, tiler_n), + a_source=a_source, + ) + + @dsl_user_op def partition_for_epilogue( cT: cute.Tensor, diff --git a/build/torch-cuda/quack/softmax.py b/build/torch-cuda/quack/softmax.py index c32e6e4c69e827048dfab9fb54668efb7ae9080c..64ef62d93b2ee187789038df6ff2917958e953f2 100644 --- a/build/torch-cuda/quack/softmax.py +++ b/build/torch-cuda/quack/softmax.py @@ -6,7 +6,7 @@ from functools import partial import torch -from ._ops_compat import add_quack_op_namespace_prefix +from ._ops_compat import add_op_namespace_prefix import cuda.bindings.driver as cuda import cutlass @@ -16,11 +16,12 @@ from cutlass import Int64, Float32, const_expr from . import utils as utils from . import copy_utils as copy_utils from .compile_utils import make_fake_tensor as fake_tensor +from .dsl import cute_op from .reduce import row_reduce, online_softmax_reduce from .reduction_base import ReductionBase -from .cache_utils import jit_cache +from .cache import jit_cache from .cute_dsl_utils import torch2cute_dtype_map -from cutlass.base_dsl import Arch +from cutlass.base_dsl.arch import Arch class Softmax(ReductionBase): @@ -99,7 +100,7 @@ class Softmax(ReductionBase): bidx, _, _ = cute.arch.block_idx() cluster_y = const_expr(0) if const_expr(self.cluster_n == 1) else cute.arch.block_idx()[1] - shape = mX.shape + shape = (cute.size(mX, mode=[0]), self.N) idX = cute.make_identity_tensor(shape) # slice for CTAs gX, gO, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, mO, idX)] @@ -175,23 +176,22 @@ class Softmax(ReductionBase): if tXcX[0][0] < shape[0]: copy(tXrO, tXgO) - -@jit_cache -def _compile_softmax_fwd(dtype, out_dtype, N): - batch_sym = cute.sym_int() - div = math.gcd(128 // dtype.width, N) - x_cute, out_cute = [fake_tensor(dt, (batch_sym, N), div) for dt in [dtype, out_dtype]] - softmax_op = Softmax(dtype, N) - return cute.compile( - softmax_op, - x_cute, - out_cute, - cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), - options="--enable-tvm-ffi", - ) + @staticmethod + @jit_cache + def compile(dtype, out_dtype, N): + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + x_cute, out_cute = [fake_tensor(dt, (batch_sym, N), div) for dt in [dtype, out_dtype]] + return cute.compile( + Softmax(dtype, N), + x_cute, + out_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) -@torch.library.custom_op(add_quack_op_namespace_prefix("_softmax_fwd"), mutates_args={"out"}) +@cute_op(add_op_namespace_prefix("_softmax_fwd"), mutates_args={"out"}) def _softmax_fwd(x: torch.Tensor, out: torch.Tensor) -> None: """Softmax forward pass. Args: @@ -200,28 +200,12 @@ def _softmax_fwd(x: torch.Tensor, out: torch.Tensor) -> None: Softmax output tensor of same shape as x """ assert x.dim() == 2, "Input must be 2D" - assert x.is_cuda, "Tensor must be on CUDA device" assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype" + if x.numel() == 0: + return N = x.size(1) dtype, out_dtype = [torch2cute_dtype_map[t.dtype] for t in [x, out]] - _compile_softmax_fwd(dtype, out_dtype, N)(x, out) - - -@_softmax_fwd.register_fake -def _softmax_fwd_fake(x: torch.Tensor, out: torch.Tensor) -> None: - # This register_fake serves two purposes: - # 1. torch.compile: When dynamo traces with symbolic shapes (SymInt), we must be a no-op. - # Without register_fake, dynamo would trace the real impl which calls _compile_softmax_fwd - # with a SymInt N — crashing @lru_cache since SymInt isn't hashable. - # 2. --compile-only mode: We enter FakeTensorMode with *concrete* shapes to pre-compile - # kernels without GPU memory. Here we trigger both fwd and bwd compilation. - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY and not isinstance(x.size(1), torch.SymInt): - N = x.size(1) - dtype, out_dtype = [torch2cute_dtype_map[t.dtype] for t in [x, out]] - _compile_softmax_fwd(dtype, out_dtype, N) - _compile_softmax_backward(dtype, out_dtype, out_dtype, N) + Softmax.compile(dtype, out_dtype, N)(x, out) def softmax_fwd(x: torch.Tensor) -> torch.Tensor: @@ -304,7 +288,7 @@ class SoftmaxBackward(ReductionBase): cluster_y = const_expr(0) if const_expr(self.cluster_n == 1) else cute.arch.block_idx()[1] tv_layout = tiled_copy.layout_tv_tiled - shape = mdY.shape + shape = (cute.size(mdY, mode=[0]), self.N) idX = cute.make_identity_tensor(shape) # slice for CTAs gdY, gY, gdX, cX = [ @@ -369,26 +353,25 @@ class SoftmaxBackward(ReductionBase): if tXcX[0][0] < shape[0]: copy(tdXrdX, tdXgdX) + @staticmethod + @jit_cache + def compile(dtype, y_dtype, dx_dtype, N): + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + dy_cute, y_cute, dx_cute = [ + fake_tensor(dt, (batch_sym, N), div) for dt in [dtype, y_dtype, dx_dtype] + ] + return cute.compile( + SoftmaxBackward(dtype, N), + dy_cute, + y_cute, + dx_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + -@jit_cache -def _compile_softmax_backward(dtype, y_dtype, dx_dtype, N): - batch_sym = cute.sym_int() - div = math.gcd(128 // dtype.width, N) - dy_cute, y_cute, dx_cute = [ - fake_tensor(dt, (batch_sym, N), div) for dt in [dtype, y_dtype, dx_dtype] - ] - softmax_backward_op = SoftmaxBackward(dtype, N) - return cute.compile( - softmax_backward_op, - dy_cute, - y_cute, - dx_cute, - cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), - options="--enable-tvm-ffi", - ) - - -@torch.library.custom_op(add_quack_op_namespace_prefix("_softmax_backward"), mutates_args={"dx"}) +@cute_op(add_op_namespace_prefix("_softmax_backward"), mutates_args={"dx"}) def _softmax_backward(dy: torch.Tensor, y: torch.Tensor, dx: torch.Tensor) -> None: """Softmax backward pass. Args: @@ -400,23 +383,13 @@ def _softmax_backward(dy: torch.Tensor, y: torch.Tensor, dx: torch.Tensor) -> No assert dy.dim() == 2, "dy must be 2D" assert y.dim() == 2, "y must be 2D" assert dy.shape == y.shape, "dy and y must have same shape" - assert dy.is_cuda and y.is_cuda, "Tensors must be on CUDA device" assert dy.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype" assert y.dtype == dy.dtype, "dy and y must have same dtype" + if dy.numel() == 0: + return N = dy.size(1) dtype, y_dtype, dx_dtype = [torch2cute_dtype_map[t.dtype] for t in [dy, y, dx]] - _compile_softmax_backward(dtype, y_dtype, dx_dtype, N)(dy, y, dx) - - -@_softmax_backward.register_fake -def _softmax_backward_fake(dy: torch.Tensor, y: torch.Tensor, dx: torch.Tensor) -> None: - # See _softmax_fwd_fake for why register_fake is needed. - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY and not isinstance(dy.size(1), torch.SymInt): - N = dy.size(1) - dtype, y_dtype, dx_dtype = [torch2cute_dtype_map[t.dtype] for t in [dy, y, dx]] - _compile_softmax_backward(dtype, y_dtype, dx_dtype, N) + SoftmaxBackward.compile(dtype, y_dtype, dx_dtype, N)(dy, y, dx) def softmax_bwd(dy: torch.Tensor, y: torch.Tensor) -> torch.Tensor: diff --git a/build/torch-cuda/quack/softmax_jax.py b/build/torch-cuda/quack/softmax_jax.py new file mode 100644 index 0000000000000000000000000000000000000000..89f7cceabcf21cabb6c4c69169b95fb69d480b50 --- /dev/null +++ b/build/torch-cuda/quack/softmax_jax.py @@ -0,0 +1,89 @@ +"""JAX bindings for QuACK softmax kernels.""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp + +from .jax_utils import ( + TvmFfiKernel, + check_rank, + cutlass_dtype, + shape_dtype_like, +) +from .softmax import Softmax, SoftmaxBackward + + +def _check_2d(name: str, x) -> None: + check_rank(name, x, 2) + + +_SOFTMAX_FWD = TvmFfiKernel( + "quack_softmax_fwd", + lambda dtype, n_cols: Softmax.compile( + cutlass_dtype(dtype), + cutlass_dtype(dtype), + n_cols, + ), +) + +_SOFTMAX_BWD = TvmFfiKernel( + "quack_softmax_bwd", + lambda dtype, n_cols: SoftmaxBackward.compile( + cutlass_dtype(dtype), + cutlass_dtype(dtype), + cutlass_dtype(dtype), + n_cols, + ), +) + + +def _softmax_fwd(x): + _check_2d("x", x) + cutlass_dtype(x.dtype) + if 0 in x.shape: + return jnp.empty_like(x) + return _SOFTMAX_FWD( + x, + key=(jnp.dtype(x.dtype), x.shape[1]), + output_shape_dtype=shape_dtype_like(x), + ) + + +def _softmax_bwd(dy, y): + _check_2d("dy", dy) + _check_2d("y", y) + if dy.shape != y.shape: + raise ValueError(f"dy and y must have the same shape, got {dy.shape} and {y.shape}") + if dy.dtype != y.dtype: + raise TypeError(f"dy and y must have the same dtype, got {dy.dtype} and {y.dtype}") + cutlass_dtype(dy.dtype) + if 0 in dy.shape: + return jnp.empty_like(dy) + return _SOFTMAX_BWD( + dy, + y, + key=(jnp.dtype(dy.dtype), dy.shape[1]), + output_shape_dtype=shape_dtype_like(dy), + ) + + +@jax.custom_vjp +def softmax(x): + """Apply QuACK softmax with a custom JAX VJP.""" + return _softmax_fwd(x) + + +def _softmax_rule_fwd(x): + y = _softmax_fwd(x) + return y, y + + +def _softmax_rule_bwd(y, dy): + return (_softmax_bwd(dy, y),) + + +softmax.defvjp(_softmax_rule_fwd, _softmax_rule_bwd) + + +__all__ = ["softmax"] diff --git a/build/torch-cuda/quack/sort/bitonic_sort.py b/build/torch-cuda/quack/sort/bitonic_sort.py index edce1b2753b5f906079b7bbc3faec916ce96340b..974adee0034f23cb744f69370b36c4664f139477 100644 --- a/build/torch-cuda/quack/sort/bitonic_sort.py +++ b/build/torch-cuda/quack/sort/bitonic_sort.py @@ -7,7 +7,6 @@ import cutlass import cutlass.cute as cute from cutlass import Int32, Float32, const_expr -from .. import utils as utils from .utils import compare_and_swap from .sorting_networks import optimal_sort @@ -79,7 +78,7 @@ def bitonic_topk_merge( if const_expr(k is None): k = cute.size(arr0.shape) if const_expr(arr0.element_type == Float32): - minmax_fn = utils.fmin if ascending else cute.arch.fmax + minmax_fn = cute.arch.fmin if ascending else cute.arch.fmax else: minmax_fn = min if ascending else max # Write the top k elements to the first half of the array diff --git a/build/torch-cuda/quack/sort/utils.py b/build/torch-cuda/quack/sort/utils.py index 8a73d2ffce12f5a30430c97a79f8d19725fb0c0e..85273e4691f78c5a74717779f1e99cc8c136f734 100644 --- a/build/torch-cuda/quack/sort/utils.py +++ b/build/torch-cuda/quack/sort/utils.py @@ -1,8 +1,6 @@ import cutlass.cute as cute from cutlass import Float32, const_expr -from .. import utils as utils - @cute.jit def compare_and_swap( @@ -23,7 +21,7 @@ def compare_and_swap( # arr[i] = b # arr[j] = a else: - min_fn = min if const_expr(arr.element_type != Float32) else utils.fmin + min_fn = min if const_expr(arr.element_type != Float32) else cute.arch.fmin max_fn = max if const_expr(arr.element_type != Float32) else cute.arch.fmax if const_expr(ascending): arr[i], arr[j] = min_fn(arr[i], arr[j]), max_fn(arr[i], arr[j]) diff --git a/build/torch-cuda/quack/spec/__init__.py b/build/torch-cuda/quack/spec/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..824ee3b52bf8d2ebcb6648414aaf2e152f96141a --- /dev/null +++ b/build/torch-cuda/quack/spec/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2025-2026, Tri Dao. + +"""Spec-layer helpers for TensorSpec, TMA descriptors, and TMEM layouts.""" diff --git a/build/torch-cuda/quack/spec/mma.py b/build/torch-cuda/quack/spec/mma.py new file mode 100644 index 0000000000000000000000000000000000000000..6ffd74d7d0fa77c073717180b518cd16bf6e1ea2 --- /dev/null +++ b/build/torch-cuda/quack/spec/mma.py @@ -0,0 +1,168 @@ +# Copyright (c) 2025-2026, Tri Dao. + +from typing import Callable, Literal, Optional, Tuple, Type + +import cutlass +import cutlass.cute as cute +from cutlass import Float32 +from cutlass.cute.nvgpu import warp, warpgroup, tcgen05 +from cutlass.cute.nvgpu import OperandMajorMode +from cutlass.cutlass_dsl import Numeric +import cutlass.utils.hopper_helpers as sm90_utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.utils import LayoutEnum + + +@cute.jit +def gemm_sm100( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + frag_A: cute.Tensor, + frag_B: cute.Tensor, + stage, + *, + stage_B=None, + zero_init=False, + pre_kblock_fn: Optional[Callable] = None, +) -> None: + """Issue one tcgen05 GEMM over all static K-blocks for a staged A/B view.""" + mma_atom = cute.make_mma_atom(tiled_mma.op) + stage_B = stage if cutlass.const_expr(stage_B is None) else stage_B + for k_blk in cutlass.range_constexpr(cute.size(frag_A, mode=[2])): + if cutlass.const_expr(pre_kblock_fn is not None): + pre_kblock_fn(mma_atom, k_blk) + mma_atom.set(tcgen05.Field.ACCUMULATE, not zero_init or k_blk != 0) + cute.gemm( + mma_atom, + acc, + frag_A[None, None, k_blk, stage], + frag_B[None, None, k_blk, stage_B], + acc, + ) + + +def operand_leading_atom(rows: int, cta_group: int) -> Tuple[int, int]: + """Return `(atom_rows, rest_rows)` for a per-CTA MMA operand tile.""" + assert cta_group in (1, 2), f"MMA operand layouts support cta_group 1 or 2, got {cta_group}" + full_rows = rows * cta_group + inst_rows = full_rows if full_rows <= 256 else full_rows // 2 + assert full_rows % inst_rows == 0, ( + f"full leading dim {full_rows} must be divisible by instruction leading dim {inst_rows}" + ) + assert inst_rows % cta_group == 0, ( + f"instruction leading dim {inst_rows} must be divisible by cta_group {cta_group}" + ) + return inst_rows // cta_group, full_rows // inst_rows + + +def resolve_mma_inst_k(dtype: Type[Numeric], mma_inst_k: Optional[int] = None) -> int: + if mma_inst_k is not None: + assert mma_inst_k > 0, f"mma_inst_k must be positive, got {mma_inst_k}" + return mma_inst_k + return 256 // dtype.width + + +def make_tiled_mma_for_arch( + spec, + source: Literal["SS", "RS", "TS"] = "SS", + atom_layout_mnk: Tuple[int, int, int] = (1, 1, 1), + acc_dtype: Type[cutlass.Numeric] = Float32, + permutation_mnk: Optional[Tuple[int, int, int]] = None, + arch=None, +) -> cute.TiledMma: + """Arch-dispatched TiledMma builder for a MatmulSpec. + + Source modes are arch-specific: SM90 accepts SS/RS, SM100 accepts SS/TS. + cta_group comes off the spec (a storage-distribution property of the + operands); spec.M/N/K are full-tile dims. + """ + if arch is None: + arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum() + cta_group = spec.cta_group + if arch.major == 9: # Hopper — WGMMA + assert cta_group == 1, f"SM90 tiled_mma requires cta_group=1, got {cta_group}" + assert source in ("SS", "RS"), f"SM90 tiled_mma source must be SS or RS, got {source}" + assert permutation_mnk is None, "SM90 tiled_mma does not accept permutation_mnk" + # WGMMA RS source means the physical A operand is a register fragment, + # whose layout convention is K-major. This can differ from the logical + # TensorSpec storage (e.g. a P.T view may have MN-major SMEM backing but + # be fed directly from registers). Treat RS A as K-major regardless of + # the spec's SMEM layout. + a_major = "K" if source == "RS" else spec._operand_major(spec.A, is_A=True) + b_major = spec._operand_major(spec.B, is_A=False) + mode = {"K": cute.nvgpu.OperandMajorMode.K, "MN": cute.nvgpu.OperandMajorMode.MN} + a_source = warpgroup.OperandSource.RMEM if source == "RS" else warpgroup.OperandSource.SMEM + return sm90_utils.make_trivial_tiled_mma( + spec.A.dtype, + spec.B.dtype, + mode[a_major], + mode[b_major], + acc_dtype, + atom_layout_mnk=atom_layout_mnk, + # `atom_layout_mnk[1]` partitions the logical/physical N tile across + # warp-groups. The WGMMA atom's per-warpgroup N extent is therefore + # the full matmul N divided by that atom-layout N factor. + tiler_mn=(64, spec.N // atom_layout_mnk[1]), + a_source=a_source, + ) + elif arch.major in [8, 12]: # SM8x and SM12x — warp-level MMA + assert cta_group == 1, f"warp-level tiled_mma requires cta_group=1, got {cta_group}" + assert source in ("SS", "RS"), f"warp-level tiled_mma source must be SS or RS, got {source}" + mma_inst_mnk = (16, 8, 16) + if spec.A.dtype.width == 16: # fp16 / bf16 + op = warp.MmaF16BF16Op(spec.A.dtype, acc_dtype, mma_inst_mnk) + else: + raise NotImplementedError( + "warp-level MMA backend doesn't yet support " + f"a_dtype={spec.A.dtype} (width={spec.A.dtype.width})" + ) + tC = cute.make_layout(atom_layout_mnk) + if permutation_mnk is None: + atom_m, atom_n, atom_k = atom_layout_mnk + # The N dim is multiplied by 2 to leverage ldmatrix.x4 (matches the reference + # blackwell_geforce/dense_gemm.py). A nested-layout permutation_n adds extra + # modes to partition_A/B output, which breaks the standard mainloop slicing. + permutation_mnk = ( + atom_m * mma_inst_mnk[0], + atom_n * mma_inst_mnk[1] * 2, + atom_k * mma_inst_mnk[2], + ) + return cute.make_tiled_mma(op, tC, permutation_mnk=permutation_mnk) + elif arch.major in [10, 11]: # Blackwell tcgen05 + assert source in ("SS", "TS"), f"SM100 tiled_mma source must be SS or TS, got {source}" + assert permutation_mnk is None, "SM100 tiled_mma does not accept permutation_mnk" + cta_group_enum = tcgen05.CtaGroup.TWO if cta_group == 2 else tcgen05.CtaGroup.ONE + m_full, n_full = spec.M, spec.N + n_inst = n_full if n_full <= 256 else n_full // 2 + if source == "TS": + # TMEM A is a freshly materialized physical operand, not a logical + # view of existing SMEM/TMA storage. Ignore TensorSpec.transposed so + # `S.T` can be stored as a row-major `(D, N)` TS-A tile. + a_major = ( + OperandMajorMode.K if spec.A.layout == LayoutEnum.ROW_MAJOR else OperandMajorMode.MN + ) + else: + a_major = ( + OperandMajorMode.K + if spec._storage_major(spec.A, is_A=True) == "K" + else OperandMajorMode.MN + ) + b_major = ( + OperandMajorMode.K + if spec._storage_major(spec.B, is_A=False) == "K" + else OperandMajorMode.MN + ) + a_source = tcgen05.OperandSource.TMEM if source == "TS" else tcgen05.OperandSource.SMEM + return sm100_utils.make_trivial_tiled_mma( + spec.A.dtype, + spec.B.dtype, + a_major, + b_major, + acc_dtype, + cta_group_enum, + (m_full, n_inst), + a_source, + ) + raise NotImplementedError( + f"make_tiled_mma_for_arch has no backend for {arch.name} (major={arch.major})." + ) diff --git a/build/torch-cuda/quack/spec/smem.py b/build/torch-cuda/quack/spec/smem.py new file mode 100644 index 0000000000000000000000000000000000000000..524e1e1bbf29c494c6bea370d0c096372c26aa1c --- /dev/null +++ b/build/torch-cuda/quack/spec/smem.py @@ -0,0 +1,160 @@ +# Copyright (c) 2025-2026, Tri Dao. + +from typing import Type, Tuple, Union, Optional + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import warpgroup, tcgen05, OperandMajorMode +from cutlass.cutlass_dsl import Numeric, dsl_user_op +from cutlass import const_expr +import cutlass.utils.hopper_helpers as sm90_utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.utils import LayoutEnum + +from . import mma as spec_mma + + +def make_smem_layout_kmajor( + dtype: Type[Numeric], + shape: Tuple[int, int], + stages: int, + cta_group: int = 1, + mma_inst_k: Optional[int] = None, +): + """SM100 operand SMEM layout where K is the fast storage axis. + + `shape = (mn, k)`, per CTA. This matches operands whose storage-contiguous + axis is K, independent of whether the operand is used as A or B. + """ + rows, cols = shape + k_inst = spec_mma.resolve_mma_inst_k(dtype, mma_inst_k) + assert cols % k_inst == 0, f"K-major cols must be divisible by {k_inst}, got {cols}" + atom_rows, rest_rows = spec_mma.operand_leading_atom(rows, cta_group) + atom = sm100_utils.make_smem_layout_atom( + sm100_utils.get_smem_layout_atom_ab(OperandMajorMode.K, dtype, shape), + dtype, + ) + return cute.make_composed_layout( + atom.inner, + 0, + cute.make_layout( + ((atom_rows, k_inst), rest_rows, cols // k_inst, stages), + stride=( + (cols, 1), + atom_rows * cols if rest_rows > 1 else 0, + k_inst, + rows * cols, + ), + ), + ) + + +def make_smem_layout_mnmajor( + dtype: Type[Numeric], + shape: Tuple[int, int], + stages: int, + cta_group: int = 1, + mma_inst_k: Optional[int] = None, +): + """SM100 operand SMEM layout where MN is the fast storage axis. + + `shape = (mn, k)`, per CTA. The logical MN axis is storage-contiguous, while + K is the slow axis in the logical operand view. The constructed nested + layout still factors the K extent with `mma_inst_k`, matching CUTLASS' + role-aware A/B helpers. + """ + rows, cols = shape + k_inst = spec_mma.resolve_mma_inst_k(dtype, mma_inst_k) + assert cols % k_inst == 0, f"MN-major cols must be divisible by {k_inst}, got {cols}" + atom_rows, rest_rows = spec_mma.operand_leading_atom(rows, cta_group) + atom = sm100_utils.make_smem_layout_atom( + sm100_utils.get_smem_layout_atom_ab(OperandMajorMode.MN, dtype, shape), + dtype, + ) + mn_contiguous = atom.outer.shape[0] + assert atom_rows % mn_contiguous == 0, ( + f"MN-major atom rows {atom_rows} must be divisible by contiguous row atom {mn_contiguous}" + ) + if atom_rows == mn_contiguous: + return cute.make_composed_layout( + atom.inner, + 0, + cute.make_layout( + ((atom_rows, k_inst), rest_rows, cols // k_inst, stages), + stride=( + (1, atom_rows), + atom_rows * cols if rest_rows > 1 else 0, + k_inst * atom_rows, + rows * cols, + ), + ), + ) + return cute.make_composed_layout( + atom.inner, + 0, + cute.make_layout( + ( + ((mn_contiguous, atom_rows // mn_contiguous), k_inst), + rest_rows, + cols // k_inst, + stages, + ), + stride=( + ((1, mn_contiguous * cols), mn_contiguous), + atom_rows * cols if rest_rows > 1 else 0, + k_inst * mn_contiguous, + rows * cols, + ), + ), + ) + + +def _sm100_smem_tile_shape(smem_layout): + layout = ( + cute.select(smem_layout, mode=[0, 1, 2]) if cute.rank(smem_layout) == 4 else smem_layout + ) + return layout.outer.shape if hasattr(layout, "outer") else layout.shape + + +@dsl_user_op +def make_smem_layout( + dtype: Type[Numeric], + layout: LayoutEnum, + tile: cute.Tile, + num_stages: Optional[int] = None, + major_mode_size: Optional[int] = None, + *, + loc=None, + ip=None, +) -> Union[cute.Layout, cute.ComposedLayout]: + shape = cute.product_each(cute.shape(tile, loc=loc, ip=ip), loc=loc, ip=ip) + if const_expr(major_mode_size is None): + major_mode_size = shape[1] if layout.is_n_major_c() else shape[0] + arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum() + if arch.major not in [10, 11]: + smem_layout_atom = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom(layout, dtype, major_mode_size), + dtype, + ) + else: # Blackwell + major_mode = OperandMajorMode.MN if layout.is_m_major_c() else OperandMajorMode.K + smem_layout_atom = tcgen05.make_smem_layout_atom( + sm100_utils.get_smem_layout_atom_ab(major_mode, dtype, tile), + dtype, + ) + order = (1, 0, 2) if const_expr(layout.is_m_major_c()) else (0, 1, 2) + smem_layout_staged = cute.tile_to_shape( + smem_layout_atom, + cute.append(shape, num_stages) if const_expr(num_stages is not None) else shape, + order=order if const_expr(num_stages is not None) else order[:2], + ) + # TensorSpec exposes a role-free storage/allocation view. Coalesce removes + # swizzle-atom factoring from the outer modes while preserving addressing, + # so callers see a canonical (M, N[, stage]) layout; MMA-specific nested + # operand views are constructed separately in MatmulSpec. + return cute.coalesce( + smem_layout_staged, + target_profile=(1, 1, 1) if const_expr(num_stages is not None) else (1, 1), + loc=loc, + ip=ip, + ) diff --git a/build/torch-cuda/quack/spec/tensor_spec.py b/build/torch-cuda/quack/spec/tensor_spec.py new file mode 100644 index 0000000000000000000000000000000000000000..39fda683f25906dd9d065ec115fc4af72752ac3d --- /dev/null +++ b/build/torch-cuda/quack/spec/tensor_spec.py @@ -0,0 +1,1389 @@ +# Copyright (c) 2025-2026, Tri Dao. + +"""Spec abstractions for declarative kernel operands. +Note: this is a prototype and the API could change rapidly. + +`TensorSpec` is a declarative description of a staged tile (dtype, shape, SMEM +stage, layout) that drives SMEM layout creation, TMA atom construction, and TMA +pipelines. The spec is **storage-only and MMA/epilogue-role agnostic**: the +physical SMEM layout (swizzle + addressing) is keyed by storage facts (dtype, +tile shape, major-axis pattern, stages). For SM100, TensorSpec's storage/TMA +layout is flat; tcgen05's nested operand layout is derived later by MatmulSpec. + +`MatmulSpec` (returned by `A @ B` on two TensorSpecs) owns everything +role-dependent: operand major modes deduced from `(layout, transposed, is_A)`, +the tiled_mma, and the role-nested SMEM views. TMA atoms are storage-layout +driven: single-CTA paths use the same flat CTA-value map as CuTe's generic tile +TMA helper, while SM100 2-CTA loads use a tcgen05 panel map because each peer +CTA owns instruction panels rather than a contiguous half tile. For SM90/SM120, +`bind_mma(thr)` returns a `BoundMMA` with per-warpgroup partitioned A/B +fragments. For SM100 (tcgen05), use `tiled_mma()` + `smem_view_A/B()` or +`bind_mma(tiled_mma=...)` for fragment views, and `with_tma_load(gmem, ...)` / +`with_tma(op, gmem, ...)` for TMA bindings. + +Shapes are FULL logical tiles (what the MMA computes on). Peer-CTA +distribution is a TensorSpec storage property: `cta_group=2` splits the +storage-leading (MN) mode across the peer pair, so each CTA allocates/loads +`storage_shape = (mn/2, k)`. The split rule is role-free — a 2-CTA MMA splits +A along M and B along N, but both are mode 0 of the `(MN, K)` storage tile — +so SMEM layouts and TMA atoms never need to know operand roles; MMA +construction reads `cta_group` back off the operand specs and validates it. + +Designed to be marshaled across the `@cute.kernel` boundary: `tma_atom` and +`gmem` cross via `__extract_mlir_values__` / `__new_from_mlir_values__`; +`smem` is populated inside the kernel via `with_smem(storage_field)` and +preserved from the host-side template (it lives in JIT-local scope, so cute +doesn't need to marshal it). +""" + +from dataclasses import dataclass, replace +from typing import Literal, Optional, Tuple, Type +from functools import partial + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.utils import LayoutEnum +import cutlass.utils.blackwell_helpers as sm100_utils + + +from . import copy_utils +from . import layout_utils +from . import pipeline +from . import sm90_utils +from . import mma as spec_mma +from . import smem as spec_smem +from . import tmem as spec_tmem +from . import tma as spec_tma + + +@dataclass +class TensorSpec: + """Declarative spec for an operand tile. Owns dtype/shape/layout/stage so + SMEM layouts, TMA atoms, and MMA configs can be derived from it. + + Shape rank: + - 2D `(rows, cols)` — matmul operand tiles. Drives `MatmulSpec`/`bind_mma`, + SM90/SM100 swizzled SMEM layouts, and 2-mode TMA atoms. + - 1D `(vec,)` — vector-with-stage aux operands (Scale, Bias, Gamma, ...). + `smem_layout()` returns a trivial `cute.make_layout((vec, stage))` (no + swizzle); `tma_copy_bytes()` and `_make_tma_atom()` use `mode=[0]`. Cannot + be used as a matmul operand (`__matmul__` / `bind_mma` are 2D-only). + + `shape` is the FULL logical tile. With `cta_group=2` each peer CTA + stores/loads only its shard: `storage_shape` splits the storage-leading + (MN) mode in half, and SMEM layouts / TMA atoms / TMEM layouts are derived + from that per-CTA shard. + + `stage=None` means the tile lives in registers (no SMEM layout, no TMA). + `transposed=True` is a logical .T view of the same storage (2D only). + + After `with_tma(...)`, the returned spec also carries the call-dynamic + `tma_atom` and `gmem` (TMA tensor). The bound spec crosses the `@cute.kernel` + boundary as a single arg — only the cute-object fields are MLIR-marshaled; + static fields are preserved from the host-side template.""" + + dtype: Type[cutlass.Numeric] + shape: Tuple[ + int, ... + ] # 2D for matmul operands; 1D supported for vector-with-stage aux operands + stage: Optional[int] = None + # `layout` is the physical storage order of the backing tile. `transposed` + # is only a logical view flag from `.T`: the same storage can be used as a + # transposed matmul operand without changing its backing layout/bytes. + layout: LayoutEnum = LayoutEnum.ROW_MAJOR + # Optional override for the swizzled major-mode extent used to select the + # SMEM atom. Most kernels can derive this from `shape`, but multi-warpgroup + # SM90 kernels sometimes intentionally split the major mode per warpgroup + # while keeping the same logical tile shape. + major_mode_size: Optional[int] = None + transposed: bool = False + # Peer-CTA distribution for tcgen05 2-CTA MMAs. `shape` stays the full + # logical tile; with cta_group=2 the storage-leading mode (the MN dim of + # the (MN, K) storage convention) is split across the peer pair, so each + # CTA allocates/loads `(mn / 2, k)`. This is a storage fact, not an MMA + # role: A splits along M and B along N, but both are mode 0 of the storage + # tile, so SMEM layouts and TMA atoms stay role-free. MMAs consuming the + # spec must be built with a matching cta_group (validated in MatmulSpec). + cta_group: int = 1 + # TMA binding: one bundled `cpasync.TmaInfo` (atom + TMA coordinate tensor + # + construction smem layout) — what _make_tma_atom returns. `tma_atom`/`gmem` + # are exposed as properties. + tma: Optional[cpasync.TmaInfo] = None + # Plain (non-TMA) gmem tensor, for operands copied with non-TMA helpers. + gmem_raw: Optional[cute.Tensor] = None + smem: Optional[cute.Tensor] = None # populated inside the kernel via with_smem() + tmem: Optional[cute.Tensor] = None # populated inside the kernel via with_tmem() + + def __extract_mlir_values__(self): + # Marshal the TMA binding + plain gmem across the kernel boundary + # (`TmaInfo` implements the marshaling protocol itself). `smem` is + # created inside the kernel via with_smem() and lives in JIT-local + # scope; cute can pass it by reference without marshaling. + values = [] + self._n_tma = 0 + self._n_gmem = 0 + if self.tma is not None: + v = cutlass.extract_mlir_values(self.tma) + values += v + self._n_tma = len(v) + if self.gmem_raw is not None: + v = cutlass.extract_mlir_values(self.gmem_raw) + values += v + self._n_gmem = len(v) + return values + + def __new_from_mlir_values__(self, values): + offset = 0 + new_tma = None + if self.tma is not None: + new_tma = cutlass.new_from_mlir_values(self.tma, values[offset : offset + self._n_tma]) + offset += self._n_tma + new_gmem = None + if self.gmem_raw is not None: + new_gmem = cutlass.new_from_mlir_values( + self.gmem_raw, values[offset : offset + self._n_gmem] + ) + offset += self._n_gmem + return replace(self, tma=new_tma, gmem_raw=new_gmem) + + @property + def tma_atom(self) -> Optional[cute.CopyAtom]: + return self.tma.atom if self.tma is not None else None + + @property + def gmem(self) -> Optional[cute.Tensor]: + """The gmem tensor: the TMA coordinate tensor when TMA-bound, else the + plain tensor attached via with_gmem().""" + if self.tma is not None: + return self.tma.tma_tensor + return self.gmem_raw + + def with_tma( + self, + op, + gmem_tensor: cute.Tensor, + *, + num_multicast: int = 1, + internal_type: Optional[Type[cutlass.Numeric]] = None, + cta_v_map=None, + gmem_raw: Optional[cute.Tensor] = None, + ) -> "TensorSpec": + """Return a new spec carrying a generic tile TMA binding. + + `gmem_tensor` is the TMA coordinate tensor. Use `gmem_raw` when the + kernel also needs ordinary GMEM indexing, e.g. when `gmem_tensor` is a + ragged/role-specific view used only for TMA. + """ + tma = self._make_tma_atom( + op, + gmem_tensor, + num_multicast=num_multicast, + internal_type=internal_type, + cta_v_map=cta_v_map, + ) + return replace( + self, + tma=tma, + gmem_raw=gmem_raw if gmem_raw is not None else self.gmem_raw, + ) + + def with_tma_load( + self, + gmem_tensor: cute.Tensor, + *, + num_multicast: int = 1, + internal_type: Optional[Type[cutlass.Numeric]] = None, + gmem_raw: Optional[cute.Tensor] = None, + ) -> "TensorSpec": + """`with_tma` for G2S loads with the op derived from the spec. + + Role-free: the op's cta_group comes from the spec's storage + distribution and multicast from `num_multicast` — no A/B distinction + (replaces role-named op selectors like `cluster_shape_to_tma_atom_B` + at the kernel level). + """ + cg = tcgen05.CtaGroup.TWO if self.cta_group == 2 else tcgen05.CtaGroup.ONE + op = ( + cpasync.CopyBulkTensorTileG2SMulticastOp(cg) + if num_multicast > 1 + else cpasync.CopyBulkTensorTileG2SOp(cg) + ) + return self.with_tma( + op, + gmem_tensor, + num_multicast=num_multicast, + internal_type=internal_type, + gmem_raw=gmem_raw, + ) + + def with_tma_info( + self, + tma_atom: cute.CopyAtom, + tma_tensor, + smem_layout=None, + *, + gmem_raw: Optional[cute.Tensor] = None, + ) -> "TensorSpec": + """Return a new spec carrying an externally-built TMA binding. + + Use this for kernels that need CUTLASS' role-aware TMA construction but + still want to pass one TensorSpec across the kernel boundary. + """ + return replace( + self, + tma=cpasync.TmaInfo(tma_atom, tma_tensor, smem_layout), + gmem_raw=gmem_raw if gmem_raw is not None else self.gmem_raw, + ) + + def with_gmem(self, gmem) -> "TensorSpec": + """Return a new spec with only a plain GMEM tensor attached. + + Use for operands that should cross the kernel boundary as TensorSpecs but + are copied with non-TMA helpers. + """ + return replace(self, gmem_raw=gmem) + + def with_smem( + self, storage_or_tensor, *, layout=None, single_stage: bool = False + ) -> "TensorSpec": + """Return a new spec with the SMEM tensor attached (call inside the kernel). + Semantics: + - storage field, no layout: derive this spec's `smem_layout()` and recast + the backing pointer to this spec's dtype if the storage field differs. + - storage field, layout: use the supplied layout and recast to this spec's dtype. + - cute.Pointer, no layout: derive this spec's `smem_layout()` and reinterpret the + pointer as this spec's dtype/layout. + - cute.Pointer, layout: reinterpret the pointer as this spec's dtype with the layout. + - cute.Tensor, no layout: bind the tensor exactly as-is; dtype must already match. + - cute.Tensor, layout: unsupported; pass `tensor.iterator` to request reinterpretation.""" + is_storage_field = const_expr(hasattr(storage_or_tensor, "get_tensor")) + is_tensor = const_expr(hasattr(storage_or_tensor, "iterator")) + is_pointer = const_expr(isinstance(storage_or_tensor, cute.Pointer)) + + if layout is not None: + assert not is_tensor, ( + "with_smem(tensor, layout=...) is unsupported; pass tensor.iterator" + ) + if const_expr(is_storage_field): + smem = self.get_smem_tensor(storage_or_tensor, layout) + else: + assert is_pointer, "with_smem(..., layout=...) expects a storage field or pointer" + smem = self._make_smem_tensor_from_ptr(storage_or_tensor, layout) + elif is_storage_field: + smem = self.get_smem_tensor(storage_or_tensor) + elif is_pointer: + smem = self._make_smem_tensor_from_ptr(storage_or_tensor, self.smem_layout()) + else: + assert is_tensor, "with_smem expects a storage field or cute.Tensor" + smem = storage_or_tensor + assert const_expr(smem.element_type == self.dtype), "SMEM tensor dtype mismatch" + if const_expr(single_stage): + assert self.stage == 1, "single_stage=True requires TensorSpec stage=1" + assert const_expr(cute.rank(smem) == self.rank + 1), ( + "single_stage=True expects a staged SMEM tensor" + ) + smem = smem[..., 0] + return replace(self, smem=smem) + + def _make_smem_tensor_from_ptr(self, ptr: cute.Pointer, layout) -> cute.Tensor: + """Reinterpret an SMEM pointer with `layout` and this spec's dtype.""" + if hasattr(layout, "outer"): + return cute.make_tensor( + cute.recast_ptr(ptr, layout.inner, dtype=self.dtype), layout.outer + ) + return cute.make_tensor(cute.recast_ptr(ptr, dtype=self.dtype), layout) + + def tmem_layout(self): + """Role-free flat TMEM storage layout for this spec. + + TMEM MMAs may need nested role layouts (e.g. tcgen05 A operand), but the + TensorSpec-owned storage identity stays `(rows, cols[, stage])`, matching + the SMEM/TMA side of the spec for the all-DP physical M=128 case. The + strides are TMEM-specific, not compact: rows stride by DP lane, columns + are contiguous, and stages advance by the logical column footprint. A + physical M=64 is represented with a nested row mode because tcgen05 uses + half-subpartition DP lanes for that MMA shape. For 2CTA TS-A, the full + MMA M/N live at the tiled_mma level; this local TMEM A view is still per + CTA. Per-CTA M=64 is duplicated into all 128 local DP lanes, and per-CTA + M=128 already occupies all 128 local DP lanes. + """ + assert not self.in_rmem, "register tensor has no TMEM layout" + assert self.rank == 2, "TMEM TensorSpec storage is currently only defined for 2D tiles" + rows, cols = self.storage_shape + if self.cta_group == 2: + assert rows in (64, 128), ( + f"2CTA TS-A TMEM layout expects per-CTA M=64 or 128, got {rows}" + ) + local_tmem_rows = 128 + else: + local_tmem_rows = rows + return spec_tmem.make_tmem_layout(self.dtype, (local_tmem_rows, cols), self.stage) + + def with_tmem( + self, + storage_or_tensor, + *, + m64_partition: Literal["lower", "upper"] = "lower", + ) -> "TensorSpec": + """Return a new spec with a TMEM tensor attached. + + Existing TMEM tensors are attached as-is. This keeps the public API + role-free: `TmemOperandA` fields already carry the spec storage layout, + while `TmemAcc` fields carry the MMA C-fragment layout. + + Raw pointers are interpreted as this spec's dense storage layout. For + 1CTA M=64 storage, `m64_partition` can select the alternate 16-DP + half-subpartition by shifting the base pointer. + """ + if hasattr(storage_or_tensor, "iterator"): + assert m64_partition == "lower", "m64_partition only applies to TS-A TMEM" + return replace(self, tmem=storage_or_tensor) + + ptr = storage_or_tensor + rows, _ = self.storage_shape + assert m64_partition == "lower" or (self.cta_group == 1 and rows == 64), ( + "m64_partition='upper' is only meaningful for 1CTA M=64 TS-A TMEM" + ) + ptr = cute.recast_ptr(ptr, dtype=self.dtype) + if m64_partition == "upper": + ptr = ptr + spec_tmem.m64_half_partition_offset(self.dtype, m64_partition) + tmem = cute.make_tensor(ptr, self.tmem_layout()) + return replace(self, tmem=tmem) + + @property + def T(self) -> "TensorSpec": + # Logical .T view of the same storage — carries over tma_atom/gmem/smem + # so the matmul-spec can default operand smems from `spec.smem` regardless of T. + return replace(self, shape=(self.shape[1], self.shape[0]), transposed=not self.transposed) + + @property + def in_rmem(self) -> bool: + return self.stage is None + + @property + def rank(self) -> int: + """Shape rank: 2 for matmul operand tiles, 1 for vector-with-stage aux operands. + Drives the rank-1 branches in `smem_layout`/`tma_copy_bytes`/`make_tma_atom`/ + `storage_shape` (skip swizzle/matmul-side logic).""" + return len(self.shape) + + def __post_init__(self): + assert self.cta_group in (1, 2), f"cta_group must be 1 or 2, got {self.cta_group}" + if self.cta_group == 2: + assert len(self.shape) == 2, "cta_group=2 requires a 2D matmul-operand spec" + + @property + def full_storage_shape(self) -> Tuple[int, ...]: + """Full-tile shape in storage order — the GMEM tile extent for + `cute.local_tile` at TMA load sites, independent of cta_group.""" + # 1D has nothing to transpose. + if self.rank == 1: + return self.shape + return (self.shape[1], self.shape[0]) if self.transposed else self.shape + + @property + def storage_shape(self) -> Tuple[int, ...]: + """Per-CTA storage shape: the full tile with the leading (MN) storage + mode split across the peer pair when cta_group=2.""" + full = self.full_storage_shape + if self.rank == 1 or self.cta_group == 1: + return full + assert full[0] % self.cta_group == 0, ( + f"leading storage dim {full[0]} not divisible by cta_group {self.cta_group}" + ) + return (full[0] // self.cta_group, full[1]) + + def smem_layout(self): + """Derive the SMEM layout for this operand. + + 1D specs (vector-with-stage aux operands) use a trivial + `cute.make_layout`, with no swizzling. + + 2D specs use the arch-specific SMEM atom selected by + `spec_smem.make_smem_layout`, then expose a coalesced role-free + storage/allocation view. The swizzled axis is the storage-contiguous + axis named by `layout`: ROW_MAJOR means K-contiguous, COL_MAJOR means + MN-contiguous. MMA-specific nested operand views are constructed + separately in MatmulSpec. + + The layout is built FRESH on every call (not cached). This matters because + SMEM layout values are MLIR-region-local: a layout built on the host doesn't + survive crossing into a `@cute.kernel`. The 2D path reconstructs the + layout from primitive (marshaled) inputs on each call — compile-time + only, no runtime cost.""" + assert not self.in_rmem, "register tensor has no SMEM layout" + if self.rank == 1: + # 1D vector-with-stage: no swizzling needed (small + naturally aligned). + return cute.make_layout((self.shape[0], self.stage)) + return spec_smem.make_smem_layout( + self.dtype, + self.layout, + self.storage_shape, + self.stage, + self.major_mode_size, + ) + + def tma_copy_bytes(self, *, full_tile: bool = False) -> int: + """Bytes of one CTA's TMA transfer for a single stage. + + `full_tile=True` scales to the whole logical tile across the peer pair + (× cta_group) — the mbarrier tx count for 2-CTA loads, where both + peers' transactions arrive at the same barrier.""" + # 1D specs have a single non-stage mode; 2D specs have two. + modes = [0] if self.rank == 1 else [0, 1] + per_cta = cute.size_in_bytes(self.dtype, cute.select(self.smem_layout(), mode=modes)) + return per_cta * self.cta_group if full_tile else per_cta + + def _make_tma_atom( + self, + op, + gmem_tensor, + num_multicast: int = 1, + internal_type: Optional[Type[cutlass.Numeric]] = None, + cta_v_map=None, + ): + if cta_v_map is not None: + modes = [0] if self.rank == 1 else [0, 1] + tma_smem_layout = cute.select(self.smem_layout(), mode=modes) + elif self.rank == 2 and self.cta_group == 2: + assert ( + isinstance( + op, + (cpasync.CopyBulkTensorTileG2SOp, cpasync.CopyBulkTensorTileG2SMulticastOp), + ) + and op.cta_group == tcgen05.CtaGroup.TWO + ), f"cta_group=2 spec requires a 2-CTA G2S load op, got {op}" + # The SMEM descriptor can use the normal flat per-CTA storage view. + # Only the GMEM coordinate map is non-contiguous: for a full 512-wide + # leading tile, one CTA owns panels 0..127 and 256..383. + tma_smem_layout = cute.select(self.smem_layout(), mode=[0, 1]) + cta_v_map = spec_tma._sm100_dense_tma_flat_cta_v_map(self.storage_shape, cta_group=2) + else: + assert getattr(op, "cta_group", None) != tcgen05.CtaGroup.TWO, ( + "2-CTA load op on a cta_group=1 spec — declare the spec with cta_group=2" + ) + modes = [0] if self.rank == 1 else [0, 1] + tma_smem_layout = cute.select(self.smem_layout(), mode=modes) + cta_v_map = cute.composition( + cute.make_identity_layout(gmem_tensor.shape), + self.storage_shape, + ) + return spec_tma._make_tiled_tma_atom_from_cta_v_map( + op, + gmem_tensor, + tma_smem_layout, + cta_v_map, + num_multicast, + internal_type=internal_type, + ) + + def smem_struct(self, align: int): + """The aligned SMEM byte allocation for this spec, for inclusion in a SharedStorage class.""" + return cute.struct.Align[ + cute.struct.MemRange[self.dtype, cute.cosize(self.smem_layout())], align + ] + + def get_smem_tensor(self, storage_field, layout=None): + """Materialize the SMEM tensor backed by `storage_field` with this spec's + layout (or the supplied `layout`), recast to this spec's dtype.""" + if layout is None: + layout = self.smem_layout() + if hasattr(layout, "outer"): + smem = storage_field.get_tensor(layout.outer, swizzle=layout.inner, dtype=self.dtype) + else: + smem = storage_field.get_tensor(layout, dtype=self.dtype) + if const_expr(smem.element_type != self.dtype): + smem = cute.make_tensor(cute.recast_ptr(smem.iterator, dtype=self.dtype), smem.layout) + return smem + + @property + def smem_T(self) -> cute.Tensor: + """`transpose_view` of `self.smem` — the layout-transposed view used as + partition_B input when the operand's matmul B-side is MN-major. Hot path + in mamba/linear-attn kernels (sBt = transpose_view(B.smem)).""" + assert self.smem is not None, "smem not bound — call with_smem(...) first" + return layout_utils.transpose_view(self.smem) + + def tma_load_fn(self, g_tile, cta_coord=0, cta_layout=None, *, peer_coord=0, **kwargs): + """Build a TMA load copy fn (gmem → smem) bound to this spec's tma_atom + smem. + + `g_tile` is the FULL-tile gmem slice — typically + `cute.local_tile(m, spec.full_storage_shape, coord)`. For cta_group=2 + this CTA's leading-mode shard is selected internally from `peer_coord` + (the MMA peer rank, `mma_tile_coord_v`), keeping the slice convention + paired with the TMA atom's flat CTA-value map. + Defaults `cta_coord=0`, `cta_layout=cute.make_layout(1)` (no multicast). + Returns the same `(copy_fn, ...)` tuple as `copy_utils.tma_get_copy_fn`.""" + if self.cta_group != 1: + g_tile = spec_tma.slice_tma_tile_by_mma_cta( + g_tile, self.storage_shape[0], peer_coord, self.cta_group + ) + if cta_layout is None: + cta_layout = cute.make_layout(1) + return copy_utils.tma_get_copy_fn( + self.tma_atom, cta_coord, cta_layout, g_tile, self.smem, **kwargs + ) + + def tma_store_fn(self, g_tile, cta_coord=0, cta_layout=None, **kwargs): + """Build a TMA store copy fn (smem → gmem) bound to this spec's tma_atom + smem. + Defaults `cta_coord=0`, `cta_layout=cute.make_layout(1)` (no multicast).""" + if cta_layout is None: + cta_layout = cute.make_layout(1) + return copy_utils.tma_get_copy_fn( + self.tma_atom, cta_coord, cta_layout, self.smem, g_tile, **kwargs + ) + + def tma_pipeline_umma( + self, + producer_group, + consumer_group, + *, + barrier_storage=None, + full_tile: bool = False, + extra_bytes: int = 0, + **kwargs, + ): + """A PipelineTmaUmma sized by this spec — num_stages and tx_count come + from the spec (stage count / per-stage TMA bytes), so the pipeline + cannot drift from the storage ring it guards and kernels don't repeat + the stage/byte bookkeeping per operand. + + Omit `barrier_storage` to let the pipeline allocate reserved smem for + its mbarriers. + + `extra_bytes` is added to tx_count, for one pipeline guarding several + operands loaded per stage (the gemm A+B pattern): + `A.tma_pipeline_umma(..., extra_bytes=B.tma_copy_bytes())`.""" + return pipeline.PipelineTmaUmma.create( + num_stages=self.stage, + producer_group=producer_group, + consumer_group=consumer_group, + tx_count=self.tma_copy_bytes(full_tile=full_tile) + extra_bytes, + barrier_storage=barrier_storage, + defer_sync=True, + **kwargs, + ) + + def tma_pipeline_async( + self, + producer_group, + consumer_group, + *, + barrier_storage=None, + extra_bytes: int = 0, + **kwargs, + ): + """SM90 counterpart of `tma_pipeline_umma`: a PipelineTmaAsync (TMA + producer -> async thread consumers, e.g. WGMMA warpgroups) sized by + this spec's stage count and per-stage TMA bytes. Multicast cluster + shape goes through `cta_layout_vmnk=`. No `full_tile` here — peer-CTA + storage splitting (cta_group=2) is tcgen05-only. + + Omit `barrier_storage` to let the pipeline allocate reserved smem for + its mbarriers. + + `extra_bytes` is added to tx_count, for one pipeline guarding several + operands loaded per stage (the gemm A+B pattern): + `A.tma_pipeline_async(..., extra_bytes=B.tma_copy_bytes())`.""" + return pipeline.PipelineTmaAsync.create( + num_stages=self.stage, + producer_group=producer_group, + consumer_group=consumer_group, + tx_count=self.tma_copy_bytes() + extra_bytes, + barrier_storage=barrier_storage, + defer_sync=True, + **kwargs, + ) + + def __matmul__(self, other: "TensorSpec") -> "MatmulSpec": + return MatmulSpec(self, other) + + +@dataclass +class BoundMMA: + """A tiled_mma plus its partitioned operand fragments and matmul shape (M, N, K). + Bundles the per-MMA boilerplate that follows `(A @ B).bind_mma(...)`. + + `frag_A`/`frag_B` semantics differ per arch: + - WGMMA (Hopper): multi-stage descriptors used with `A_idx`/`B_idx` in gemm. + - Warp-level (SM120): single-stage rmem fragments — the kernel must do its + own ldmatrix SMEM->RMEM step into them before each MMA. + + `tiled_copy_s2r_A`/`B` are the SMEM->RMEM (ldmatrix) `cute.TiledCopy`s — + used by SM120 for the explicit SMEM->RMEM step before each MMA, and useful + on SM90 for non-WGMMA paths that load into register frags. + `tiled_copy_r2s_A`/`B` are the RMEM->SMEM (stmatrix) counterparts — used by + kernels that stage an A-operand transform back through SMEM (e.g. for a + follow-on MMA in a different tiling). + (The per-stage SMEM *partition view* is a kernel mainloop concern and is + not on this object — derive via + `tiled_copy_s2r_A.get_slice(thr).partition_S(sA)` at the use site.)""" + + tiled_mma: cute.TiledMma + frag_A: Optional[cute.Tensor] + frag_B: Optional[cute.Tensor] + M: int # logical M (the user's matmul A side); when swap_AB, physical wgmma sees N here + N: int # logical N (the user's matmul B side) + K: int + tiled_copy_s2r_A: Optional[cute.TiledCopy] = None + tiled_copy_s2r_B: Optional[cute.TiledCopy] = None + tiled_copy_r2s_A: Optional[cute.TiledCopy] = None + tiled_copy_r2s_B: Optional[cute.TiledCopy] = None + # When True, the underlying wgmma was constructed with operand roles swapped + # (logical A → physical B and vice versa) — typically as a wgmma-instruction + # reduction trick when the logical A's M is too large but B's N would fit a + # single instance. The user keeps thinking in logical (A, B) terms; .acc(), + # .fn(), and .r2s_C() handle the physical swap internally. + swap_AB: bool = False + + # MLIR marshaling — without these the cute jit boundary auto-flattens this + # dataclass to its first cute-typed field (`tiled_mma`), losing the rest. + # Static fields (M/N/K) are preserved from the host-side template. + def __extract_mlir_values__(self): + values = [] + self._lengths = {} + for name in ( + "tiled_mma", + "frag_A", + "frag_B", + "tiled_copy_s2r_A", + "tiled_copy_s2r_B", + "tiled_copy_r2s_A", + "tiled_copy_r2s_B", + ): + obj = getattr(self, name) + if obj is not None: + v = cutlass.extract_mlir_values(obj) + values += v + self._lengths[name] = len(v) + else: + self._lengths[name] = 0 + return values + + def __new_from_mlir_values__(self, values): + new_fields = {} + offset = 0 + for name in ( + "tiled_mma", + "frag_A", + "frag_B", + "tiled_copy_s2r_A", + "tiled_copy_s2r_B", + "tiled_copy_r2s_A", + "tiled_copy_r2s_B", + ): + n = self._lengths[name] + if n > 0: + obj = getattr(self, name) + new_fields[name] = cutlass.new_from_mlir_values(obj, values[offset : offset + n]) + offset += n + else: + new_fields[name] = None + return replace(self, **new_fields) + + def acc(self, shape=None, dtype=Float32) -> cute.Tensor: + """Allocate an accumulator rmem tensor. `shape` defaults to logical (M, N). + Extra modes after (M, N) are appended to the partitioned C layout, e.g. + `(M, N, stage)` becomes `(MMA, MMA_M, MMA_N, stage)`. When swap_AB, the + physical wgmma C-side is (N, M) — we feed that to partition_shape_C; the + resulting rmem holds the transposed accumulator, but the user can treat + it as opaque (fill / pass to .fn / .r2s_C).""" + if shape is None: + shape = (self.M, self.N) + if self.swap_AB: + shape = (shape[1], shape[0], *shape[2:]) # physical (N, M, ...) + acc_shape = self.tiled_mma.partition_shape_C(shape[:2]) + for extra_mode in shape[2:]: + acc_shape = cute.append(acc_shape, extra_mode) + return cute.make_rmem_tensor(acc_shape, dtype) + + def clone_frag_A(self) -> cute.Tensor: + """Allocate another rmem tensor matching this MMA's frag_A — used for + multi-stage RS patterns where each stage needs its own A operand.""" + assert self.frag_A is not None, "no frag_A to clone (call bind_mma with source='RS')" + return cute.make_rmem_tensor(self.frag_A.layout, self.frag_A.element_type) + + def fn(self, acc, zero_init=False, frag_A=None, frag_B=None): + """Return a callable that captures `acc`/default frags — call per-iteration in a loop. + `frag_A`/`frag_B` override the bound fragments either here or at call time + (multi-stage RS pattern where the A fragment is produced in the loop). + When swap_AB, A_idx/B_idx are user-logical and get swapped internally.""" + default_A = frag_A if frag_A is not None else self.frag_A + default_B = frag_B if frag_B is not None else self.frag_B + + def _fn(A_idx=None, B_idx=None, wg_wait=-1, zero_init=zero_init, frag_A=None, frag_B=None): + fA = frag_A if frag_A is not None else default_A + fB = frag_B if frag_B is not None else default_B + if self.swap_AB: + return sm90_utils.gemm_w_idx( + self.tiled_mma, acc, fA, fB, zero_init, B_idx, A_idx, wg_wait + ) + return sm90_utils.gemm_w_idx( + self.tiled_mma, acc, fA, fB, zero_init, A_idx, B_idx, wg_wait + ) + + return _fn + + def fn_zero_init(self, shape=None, frag_A=None, frag_B=None): + """Return a partial for the zero-init gemm variant (allocates its own acc). + `shape` defaults to logical (M, N) — swapped to physical (N, M) when swap_AB. + A_idx/B_idx are user-logical and get swapped internally when swap_AB.""" + if shape is None: + shape = (self.M, self.N) + if self.swap_AB: + shape = (shape[1], shape[0]) + fA = frag_A if frag_A is not None else self.frag_A + fB = frag_B if frag_B is not None else self.frag_B + inner = partial(sm90_utils.gemm_zero_init, self.tiled_mma, shape, fA, fB) + if self.swap_AB: + + def _fn(A_idx=None, B_idx=None, **kw): + return inner(A_idx=B_idx, B_idx=A_idx, **kw) + + return _fn + return inner + + # SMEM<->RMEM helpers for the A and C operand positions of this MMA. + # Thin wrappers over `quack.copy_utils.get_smem_(load|store)_(A|C)` that + # bind `self.tiled_mma` so call sites read as `mma.r2s_C(sC, tidx)` instead + # of `copy_utils.get_smem_store_C(tiled_mma_pv, sC, tidx)`. Each returns the + # same `(copy_fn, thr_copy, partitioned_tensor)` tuple as the underlying + # helper. No B variants — WGMMA loads B via descriptor with no register + # staging path, mirroring `copy_utils`'s lack of `get_smem_(load|store)_B`. + def s2r_A(self, sA, thr, **kwargs): + return copy_utils.get_smem_load_A(self.tiled_mma, sA, thr, **kwargs) + + def r2s_A(self, sA, thr, **kwargs): + return copy_utils.get_smem_store_A(self.tiled_mma, sA, thr, **kwargs) + + def s2r_C(self, sC, thr, **kwargs): + # Mirror r2s_C's swap_AB handling so epilogue inputs staged in the + # logical (M, N) layout read back correctly from a swapped MMA. + if self.swap_AB: + kwargs["transpose"] = not kwargs.get("transpose", False) + sC = layout_utils.transpose_view(sC) + return copy_utils.get_smem_load_C(self.tiled_mma, sC, thr, **kwargs) + + def r2s_C(self, sC, thr, **kwargs): + # When swap_AB, the rmem accumulator is physically (N, M) but the user's + # `sC` is in logical (M, N) layout. Auto-fix: feed transpose_view(sC) so + # make_tiled_copy_C sees matching shape, and toggle the stmatrix transpose + # bit so the data lands in sC's underlying storage in logical orientation. + if self.swap_AB: + kwargs["transpose"] = not kwargs.get("transpose", False) + sC = layout_utils.transpose_view(sC) + return copy_utils.get_smem_store_C(self.tiled_mma, sC, thr, **kwargs) + + +@dataclass +class BoundMMASm100(BoundMMA): + """A tcgen05 tiled_mma plus its SMEM-descriptor operand fragments and full + MMA tile shape (M, N, K). The SM100 counterpart of `BoundMMA`, shaped by + the tcgen05 execution model: + + - `frag_A`/`frag_B` are multi-stage SMEM descriptor fragments + ((MMA, MMA_M/N, MMA_K, STAGE)); there is no register staging path. + - The accumulator lives in TMEM, not RMEM: `acc(tmem_ptr, stages=)` + builds the staged accumulator tensor at a retrieved TMEM pointer. + (Use `MatmulSpec.acc_layout_sm100` from warps that only read the + accumulator and never bind fragments, e.g. epilogue warps.) + - The MMA is issued from the MMA warp of the **leader CTA** only — for + cta_group=2 the peer CTA contributes operand SMEM but does not issue. + Gate `gemm()` with an `is_leader_cta` check; the single-thread election + within the warp is handled by `cute.gemm` itself. + - `gemm()` / `fn()` issue through a fresh `cute.make_mma_atom(tiled_mma.op)` + so ACCUMULATE/SFA/SFB mutations stay local to the helper and do not force + callers to loop-carry `tiled_mma` through dynamic regions. + + `M/N/K` are the user's logical full tile dims (the spec shapes). With + `swap_AB=True`, the physical tcgen05 C tile is `(N, M)`; accumulator + helpers derive that physical layout internally. The per-CTA storage shards + stay on `TensorSpec.storage_shape`, where TMA and SMEM views need them.""" + + A: Optional[TensorSpec] = None + B: Optional[TensorSpec] = None + + def _physical_acc_mn(self) -> Tuple[int, int]: + return (self.N, self.M) if self.swap_AB else (self.M, self.N) + + def _make_acc_frag(self, *, stages: Optional[int] = None) -> cute.Tensor: + shape = self.tiled_mma.partition_shape_C(self._physical_acc_mn()) + if stages is not None: + shape = cute.append(shape, stages) + return self.tiled_mma.make_fragment_C(shape) + + @property + def num_k_blocks(self) -> int: + return cute.size(self.frag_A, mode=[2]) + + def acc_layout(self, *, stages: Optional[int] = None): + """Staged TMEM accumulator layout: (MMA, MMA_M, MMA_N[, STAGE]).""" + return self._make_acc_frag(stages=stages).layout + + def acc(self, tmem_ptr, *, stages: Optional[int] = None) -> cute.Tensor: + """The staged accumulator tensor at a retrieved TMEM pointer.""" + return cute.make_tensor(tmem_ptr, self.acc_layout(stages=stages)) + + def t2r_C( + self, + acc: cute.Tensor, + tidx: Int32, + dst_dtype: Type[cutlass.Numeric], + *, + epi_tile: Optional[cute.Tile] = None, + num_cols: Optional[int] = None, + transpose: bool = False, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Build the tcgen05 accumulator TMEM->RMEM copy. + + Returns `(tiled_t2r, tTR_tAcc, tTR_rAcc)`, where `tTR_tAcc` is the + partitioned TMEM source and `tTR_rAcc` is the per-thread register + fragment to reuse for each epilogue subtile. Pass `epi_tile=None` for + non-epilogue spills that should partition the raw staged accumulator. + + With `swap_AB`, the accumulator is physically `(N, M)`. For epilogue + copies, `transpose` also presents the logical `(M, N)` accumulator view. + For raw `epi_tile=None` spills, `transpose` only selects the row/col + copy policy and the raw accumulator view is kept. The t2r atom is + selected from that source view shape, `transpose`-derived copy layout, + and destination dtype. + At least one of `epi_tile` or `num_cols` is required. `num_cols` uses + the same Blackwell atom-selection heuristic but substitutes a synthetic + `(tile_m, num_cols)` epilogue tile, useful for non-epilogue spills that + copy a narrower major-mode slice. + """ + assert epi_tile is not None or num_cols is not None, "pass epi_tile and/or num_cols" + tAcc = acc[(None, None), 0, 0, None] + copy_layout = LayoutEnum.COL_MAJOR if const_expr(transpose) else LayoutEnum.ROW_MAJOR + tile_m, tile_n = cute.size(tAcc.shape, mode=[0]), cute.size(tAcc.shape, mode=[1]) + is_2cta = cute.size(self.tiled_mma.thr_id.shape) == 2 + atom_tile = epi_tile if const_expr(epi_tile is not None) else (tile_m, num_cols) + copy_atom_t2r = sm100_utils.get_tmem_load_op( + (tile_m, tile_n, self.K), + copy_layout, + dst_dtype, + acc.element_type, + atom_tile, + is_2cta, + ) + if const_expr(epi_tile is None): + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc[None, None, 0]) + tTR_tAcc = tiled_copy_t2r.get_slice(tidx).partition_S(tAcc) + tTR_rAcc = copy_utils.tmem_reg_frag(tiled_copy_t2r, tTR_tAcc[..., 0]) + else: + tAcc_epi = cute.flat_divide(tAcc, epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[None, None, 0, 0, 0]) + tTR_tAcc = tiled_copy_t2r.get_slice(tidx).partition_S(tAcc_epi) + tTR_rAcc = copy_utils.tmem_reg_frag(tiled_copy_t2r, tTR_tAcc[..., 0, 0, 0]) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def r2s_C(self, tiled_t2r: cute.TiledCopy, sC, thr, **kwargs): + # SM100 epilogues chain SMEM stores from the TMEM-load tiled copy so + # register ownership matches the preceding t2r_C load. + if self.swap_AB: + kwargs["transpose"] = not kwargs.get("transpose", False) + sC = layout_utils.transpose_view(sC) + return copy_utils.get_smem_store_C(tiled_t2r, sC, thr, **kwargs) + + def s2r_C(self, tiled_t2r: cute.TiledCopy, sC, thr, r_layout, **kwargs): + """SMEM-load counterpart of `r2s_C`, for epilogue INPUTS that were + TMA-staged into the output buffer (gemm's C, ssd's z): the register + fragment is allocated at `r_layout` (pass the t2r fragment's layout) + so its linear element order matches the t2r/r2s fragments, and + swap_AB handling mirrors `r2s_C` (transpose_view + transposed copy). + + Unless a `copy_atom` is passed, this defaults to vectorized universal + loads rather than ldmatrix: the flat TensorSpec SMEM layouts' + per-thread partitions cannot statically prove ldmatrix's 16B source + alignment (the store side goes through CUTLASS's SM100 store-op + selector, which handles that; there is no load-side equivalent). + Scalar copies under a transposed view, where per-thread elements are + not contiguous. + + Returns `(tiled_copy, tRS_r, tSR_r, tSR_s)`; load via + `cute.copy(tiled_copy, tSR_s[..., idx], tSR_r)` then read `tRS_r`.""" + if self.swap_AB: + kwargs["transpose"] = not kwargs.get("transpose", False) + sC = layout_utils.transpose_view(sC) + if "copy_atom" not in kwargs: + dtype = sC.element_type + transpose = kwargs.get("transpose", False) + kwargs["copy_atom"] = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + dtype, + num_bits_per_copy=dtype.width * (1 if transpose else 2), + ) + return copy_utils.s2r_partition_from_t2r(tiled_t2r, sC, thr, r_layout, **kwargs) + + def fn(self, acc, zero_init=False, pre_kblock_fn=None): + """Return a per-iteration tcgen05 GEMM callable. + + This mirrors `BoundMMA.fn` for SM90 call sites. It is intentionally + stateless: callers only pass logical stage indices and do not thread + `tiled_mma` through the mainloop: + + `fn(A_idx=..., B_idx=..., acc_idx=...)` + + `acc` can be either a single TMEM accumulator tensor + `(MMA, MMA_M, MMA_N)` or a staged tensor + `(MMA, MMA_M, MMA_N, STAGE)`. When `acc_idx` is provided, that stage is + selected before issuing the MMA. + + For call sequences that accumulate across multiple dynamic calls, pass a + dynamic `zero_init` flag (true for the first call, false afterward) or + call the lower-level `gemm(...)` helper directly. + """ + + def _fn(A_idx, B_idx=None, acc_idx=None, zero_init=zero_init, pre_kblock_fn=pre_kblock_fn): + acc_cur = acc if acc_idx is None else acc[None, None, None, acc_idx] + self.gemm( + acc_cur, + A_idx, + stage_B=B_idx, + zero_init=zero_init, + pre_kblock_fn=pre_kblock_fn, + ) + + return _fn + + def gemm( + self, acc, stage, *, stage_B=None, tiled_mma=None, zero_init=False, pre_kblock_fn=None + ): + """Issue the unrolled k-block MMAs of one stage into `acc`. + + Call from the MMA warp of the leader CTA, between the pipeline + consumer_wait and consumer_release for the operand stages — the + pipeline choreography stays with the kernel. + + - `stage` indexes frag_A's stage mode; `stage_B` defaults to `stage` + and exists for kernels whose A and B operands come from different + pipelines (e.g. linear attention, where Q and K have different + stage counts). + - `zero_init=True` clears ACCUMULATE before the first k-block. In a + dynamic K-tile loop, pass `zero_init=(k_tile == 0)` so only the first + call starts a fresh accumulation. + - `pre_kblock_fn(mma_atom, k_blk)` runs before each k-block's gemm — + e.g. blockscaled SFA/SFB tmem pointer updates via `mma_atom.set`. + """ + if const_expr(self.swap_AB): + logical_A_stage = stage + logical_B_stage = stage if const_expr(stage_B is None) else stage_B + stage = logical_B_stage + stage_B = logical_A_stage + tm = tiled_mma if tiled_mma is not None else self.tiled_mma + spec_mma.gemm_sm100( + tm, + acc, + self.frag_A, + self.frag_B, + stage, + stage_B=stage_B, + zero_init=zero_init, + pre_kblock_fn=pre_kblock_fn, + ) + + +class MatmulSpec: + """Result of A @ B on two TensorSpecs. Deduces operand major modes from + storage layout + transposed flag; derives tiler_n from B's N dim. + + `M/N/K` are the FULL logical matmul dims (spec shapes are full tiles); + `cta_group` is read off the operand specs and must agree between A and B.""" + + def __init__(self, A: TensorSpec, B: TensorSpec): + assert A.shape[1] == B.shape[0], f"matmul shape mismatch: {A.shape} @ {B.shape}" + assert A.cta_group == B.cta_group, ( + f"matmul operand cta_group mismatch: {A.cta_group} vs {B.cta_group}" + ) + # A.dtype and B.dtype may differ for mixed-precision MMAs (e.g. fp8 ops + # with different a/b widths supported by the underlying tiled_mma). + self.A, self.B = A, B + self.M, self.K = A.shape + self.N = B.shape[1] + self.cta_group = A.cta_group + + def __extract_mlir_values__(self): + # MatmulSpec is a compile-time role view over two TensorSpecs, but the + # underlying specs may carry dynamic TMA/gmem values. Forward marshaling + # to A/B so MatmulSpec locals can remain live across DSL dynamic control + # flow (`if warp_idx == ...`, dynamic loops, etc.). + values = [] + a_values = cutlass.extract_mlir_values(self.A) + b_values = cutlass.extract_mlir_values(self.B) + values += a_values + values += b_values + self._n_A = len(a_values) + self._n_B = len(b_values) + return values + + def __new_from_mlir_values__(self, values): + offset = 0 + A = self.A + if self._n_A > 0: + A = cutlass.new_from_mlir_values(self.A, values[offset : offset + self._n_A]) + offset += self._n_A + B = self.B + if self._n_B > 0: + B = cutlass.new_from_mlir_values(self.B, values[offset : offset + self._n_B]) + offset += self._n_B + return MatmulSpec(A, B) + + def tiled_mma( + self, + source: Literal["SS", "RS", "TS"] = "SS", + *, + acc_dtype: Type[cutlass.Numeric] = Float32, + atom_layout_mnk: Tuple[int, int, int] = (1, 1, 1), + permutation_mnk: Optional[Tuple[int, int, int]] = None, + arch=None, + ) -> cute.TiledMma: + # cta_group is a storage property of the operand specs, not an MMA + # parameter — make_tiled_mma_for_arch reads it off `self.cta_group`. + return spec_mma.make_tiled_mma_for_arch( + self, + source=source, + atom_layout_mnk=atom_layout_mnk, + acc_dtype=acc_dtype, + permutation_mnk=permutation_mnk, + arch=arch, + ) + + def bind_mma( + self, + thr=None, + *, + sA: Optional[cute.Tensor] = None, + sB: Optional[cute.Tensor] = None, + source: Literal["SS", "RS", "TS"] = "SS", + tmem_A_ptr=None, + atom_layout_mnk: Tuple[int, int, int] = (1, 1, 1), + acc_dtype: Type[cutlass.Numeric] = Float32, + permutation_mnk: Optional[Tuple[int, int, int]] = None, + tiled_mma: Optional[cute.TiledMma] = None, + swap_AB: bool = False, + bind_operands: bool = True, + arch=None, + ) -> BoundMMA: + """Build an arch-appropriate BoundMMA. + + SM90/SM120 use the register-fragment path and accept `source="SS"` or + `"RS"`. SM100 uses the tcgen05/TMEM path and accepts `source="SS"` or + `"TS"`; passing `tmem_A_ptr` implies `"TS"`. + + - `thr` is the SM90/SM120 index (or layout) passed to + `tiled_mma.get_slice(...)`. + Pass `tidx` for single-thread frags or a wg-thread layout for warp- + group-partitioned frags. When omitted (`thr=None`), frag construction + is skipped (`frag_A`/`frag_B` are None) — useful when the caller only + needs the s2r/r2s tiled_copies and builds its own frags differently. + - `sA`/`sB` default to `spec.smem`, **auto-transposed when the operand's + major mode is "MN"**. + - For `source="RS"`, `frag_A` is auto-allocated as an rmem tensor with + shape `tiled_mma.partition_shape_A((M, K))`; on SM90 the physical A + operand is treated as K-major even if the logical TensorSpec also has + an MN-major SMEM backing. + - `tiled_mma`: pass a pre-built TiledMma to bypass the arch-dispatched + default (e.g. for warp-level MMA with custom permutation_mnk on SM120, + or any non-default MMA op selection). + - `bind_operands=False`: return a layout-only BoundMMA/BoundMMASm100 + with `frag_A`/`frag_B` left as None. This is useful for sizing a TMEM + accumulator before operand storage has been bound. + - `swap_AB=True`: physically compute `(B.T @ A.T)` instead of `(A @ B)` — + useful when logical M is too large but logical N would fit a single + wgmma instance (e.g. M=128, N=64 → 2 wgmma; swap to 64×128 → 1 wgmma). + The user keeps thinking in logical (A, B) terms; the returned BoundMMA's + `.acc()`/`.fn()`/`.r2s_C()` handle the physical swap automatically. + On SM100, `source="TS"` with `swap_AB=True` means logical B.T is the + physical TMEM A operand; `.fn(...)(A_idx=, B_idx=)` still accepts + logical indices and routes to physical stages. + """ + if arch is None: + arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum() + if arch.major in [10, 11]: + assert thr is None, "SM100 bind_mma does not use a thread slice; pass thr=None" + phys = MatmulSpec(self.B.T, self.A.T) if swap_AB else self + if phys.A.tmem is not None: + assert tmem_A_ptr is None, "pass either A.with_tmem(...) or tmem_A_ptr, not both" + source = "TS" + source = "TS" if tmem_A_ptr is not None else source + if tiled_mma is None: + tiled_mma = phys.tiled_mma( + acc_dtype=acc_dtype, + source=source, + permutation_mnk=permutation_mnk, + arch=arch, + ) + full_m, full_n, full_k = self.mma_tiler_mnk(tiled_mma) + if not bind_operands: + frag_A, frag_B = None, None + elif phys.A.tmem is not None: + assert sA is None, "TMEM A source cannot also pass sA" + frag_A = phys.tmem_view_A(tiled_mma) + elif tmem_A_ptr is not None: + assert sA is None, "pass either sA (SMEM A) or tmem_A_ptr (TMEM A), not both" + frag_A = phys.frag_A_tmem(tiled_mma, tmem_A_ptr) + else: + assert source != "TS", "SM100 TS MMA requires a bound TMEM physical A operand" + frag_A = tiled_mma.make_fragment_A(phys.smem_view_A(tiled_mma, sA)) + if bind_operands: + sB = phys.smem_view_B(tiled_mma, sB) + return BoundMMASm100( + tiled_mma=tiled_mma, + frag_A=frag_A, + frag_B=None if not bind_operands else tiled_mma.make_fragment_B(sB), + M=full_m, + N=full_n, + K=full_k, + A=self.A, + B=self.B, + swap_AB=swap_AB, + ) + + assert self.cta_group == 1, ( + f"{arch.name} bind_mma requires cta_group=1 operands, got {self.cta_group}" + ) + assert source in ("SS", "RS"), f"{arch.name} bind_mma source must be SS or RS, got {source}" + # `phys` is the *physical* MatmulSpec — what the wgmma is actually built + # against. When swap_AB, we flip operand roles: logical A becomes the + # physical B operand and vice versa (i.e., compute (B.T @ A.T) on hardware). + # Downstream code uses `phys.{A,B,M,N,K}` to construct tiled_mma + frags + + # smem partitions; the returned BoundMMA records the user's *logical* + # (M, N, K) and the swap_AB flag, so .acc/.fn/.r2s_C reconcile internally. + phys = MatmulSpec(self.B.T, self.A.T) if swap_AB else self + if tiled_mma is None: + tiled_mma = phys.tiled_mma( + source=source, + atom_layout_mnk=atom_layout_mnk, + acc_dtype=acc_dtype, + permutation_mnk=permutation_mnk, + arch=arch, + ) + if thr is not None: + thr_mma = tiled_mma.get_slice(thr) + # Lazily resolve sA/sB so register-only specs (e.g. source="RS" with + # an rmem-only A like dP) don't try to dereference a missing .smem. + sB_eff = sB if sB is not None else phys._smem_for(phys.B, is_A=False) + # Arch-dispatched frag construction. WGMMA (Hopper) frags are multi-stage + # descriptors that index into SMEM via A_idx/B_idx. Warp-level (SM120) + # frags are single-stage rmem tensors — the ldmatrix SMEM->RMEM step is + # the kernel's mainloop concern. + if arch.major in [8, 12]: # Warp-level — single-stage rmem frags + if source == "RS": + frag_A = cute.make_rmem_tensor( + tiled_mma.partition_shape_A((phys.M, phys.K)), phys.A.dtype + ) + else: + sA_eff = sA if sA is not None else phys._smem_for(phys.A, is_A=True) + frag_A = tiled_mma.make_fragment_A( + thr_mma.partition_A(sA_eff)[None, None, None, 0] + ) + frag_B = tiled_mma.make_fragment_B(thr_mma.partition_B(sB_eff)[None, None, None, 0]) + else: # WGMMA (Hopper) — multi-stage descriptors + if source == "RS": + frag_A = cute.make_rmem_tensor( + tiled_mma.partition_shape_A((phys.M, phys.K)), phys.A.dtype + ) + else: + sA_eff = sA if sA is not None else phys._smem_for(phys.A, is_A=True) + frag_A = tiled_mma.make_fragment_A(thr_mma.partition_A(sA_eff)) + frag_B = tiled_mma.make_fragment_B(thr_mma.partition_B(sB_eff)) + else: + frag_A, frag_B = None, None + # Ldmatrix/stmatrix copy atoms — generic (works on any arch). Caller derives + # the per-stage SMEM partition view at the use site: + # smem_view = mma.tiled_copy_s2r_A.get_slice(thr).partition_S(sA) + a_transpose = phys.A.layout.is_m_major_a() + b_transpose = phys.B.layout.is_n_major_b() + tiled_copy_s2r_A = cute.make_tiled_copy_A( + copy_utils.get_smem_load_atom(phys.A.dtype, transpose=a_transpose), tiled_mma + ) + tiled_copy_s2r_B = cute.make_tiled_copy_B( + copy_utils.get_smem_load_atom(phys.B.dtype, transpose=b_transpose), tiled_mma + ) + tiled_copy_r2s_A = cute.make_tiled_copy_A( + copy_utils.get_smem_store_atom(phys.A.dtype, transpose=a_transpose), tiled_mma + ) + tiled_copy_r2s_B = cute.make_tiled_copy_B( + copy_utils.get_smem_store_atom(phys.B.dtype, transpose=b_transpose), tiled_mma + ) + # M/N/K are LOGICAL (the user's matmul). swap_AB is the only flag the + # downstream BoundMMA needs to reconcile logical ↔ physical. + return BoundMMA( + tiled_mma=tiled_mma, + frag_A=frag_A, + frag_B=frag_B, + M=self.M, + N=self.N, + K=self.K, + tiled_copy_s2r_A=tiled_copy_s2r_A, + tiled_copy_s2r_B=tiled_copy_s2r_B, + tiled_copy_r2s_A=tiled_copy_r2s_A, + tiled_copy_r2s_B=tiled_copy_r2s_B, + swap_AB=swap_AB, + ) + + def _smem_for(self, t: TensorSpec, is_A: bool) -> cute.Tensor: + """Return the SMEM view to feed partition_{A,B}. + + partition_A wants storage in (M, K) order; partition_B wants (N, K). + The spec's logical `shape` respects `.T` — physical storage_shape is + `(shape[1], shape[0]) if transposed else shape`. So whether we need a + transpose-view is fully determined by `transposed`: + - A: transposed → storage is (K, M), need flip to (M, K). + - B: not transposed → storage is (K, N), need flip to (N, K). + The `layout` (ROW_MAJOR/COL_MAJOR) only affects the storage major mode + (which dim is contiguous), which is orthogonal to shape order and is + already handled in `_operand_major` for tiled_mma construction.""" + needs_transpose = t.transposed if is_A else not t.transposed + return layout_utils.transpose_view(t.smem) if needs_transpose else t.smem + + # ---- SM100 (tcgen05) ---------------------------------------------------- + # Role is an MMA concern, so it lives here rather than on the TensorSpec: + # the same spec (and the same SMEM bytes) can be the A operand of one MMA + # and the B operand of another. Spec shapes are full logical tiles; peer + # distribution (cta_group) is a storage property of the specs — the 2-CTA + # MMA splits A along M and B along N, each peer CTA holding half of both, + # and each spec's `storage_shape` is its per-CTA shard. + # The role-nested layouts below are byte-identical to the specs' flat + # `smem_layout()` (same swizzle, same addressing); they only differ in the + # mode nesting that `partition_A/B` / `make_fragment_A/B` expect. + + def mma_tiler_mnk(self, tiled_mma: cute.TiledMma) -> Tuple[int, int, int]: + """Full MMA tile (M, N, K). Spec shapes are full logical tiles, so this + is just (M, N, K); the tiled_mma's cta_group (read off `thr_id`) is + validated against the operands' storage distribution.""" + mma_cta_group = cute.size(tiled_mma.thr_id.shape) + assert mma_cta_group == self.cta_group, ( + f"tiled_mma cta_group {mma_cta_group} != operand spec cta_group {self.cta_group}" + ) + return (self.M, self.N, self.K) + + def smem_layout_A(self, tiled_mma: cute.TiledMma, *, stage: Optional[int] = None): + """Role-nested staged SMEM layout for the A operand — the + ((atom), rest_m, rest_k, stage) view that `partition_A` / + `make_fragment_A` expect.""" + stage = stage if stage is not None else self.A.stage + return sm100_utils.make_smem_layout_a( + tiled_mma, self.mma_tiler_mnk(tiled_mma), self.A.dtype, stage + ) + + def smem_layout_B(self, tiled_mma: cute.TiledMma, *, stage: Optional[int] = None): + """Role-nested staged SMEM layout for the B operand — the + ((atom), rest_n, rest_k, stage) view that `partition_B` / + `make_fragment_B` expect.""" + stage = stage if stage is not None else self.B.stage + return sm100_utils.make_smem_layout_b( + tiled_mma, self.mma_tiler_mnk(tiled_mma), self.B.dtype, stage + ) + + @staticmethod + def _view_smem_as(smem: cute.Tensor, layout) -> cute.Tensor: + return cute.make_tensor(smem.iterator, layout.outer if hasattr(layout, "outer") else layout) + + def smem_view_A( + self, + tiled_mma: cute.TiledMma, + smem: Optional[cute.Tensor] = None, + *, + stage: Optional[int] = None, + ) -> cute.Tensor: + """tcgen05 A-operand view over flat TensorSpec storage.""" + smem = smem if smem is not None else self.A.smem + assert smem is not None, "A smem not bound — call with_smem(...) or pass smem" + return self._view_smem_as(smem, self.smem_layout_A(tiled_mma, stage=stage)) + + def smem_view_B( + self, + tiled_mma: cute.TiledMma, + smem: Optional[cute.Tensor] = None, + *, + stage: Optional[int] = None, + ) -> cute.Tensor: + """tcgen05 B-operand view over flat TensorSpec storage.""" + smem = smem if smem is not None else self.B.smem + assert smem is not None, "B smem not bound — call with_smem(...) or pass smem" + return self._view_smem_as(smem, self.smem_layout_B(tiled_mma, stage=stage)) + + def tmem_view_A( + self, + tiled_mma: cute.TiledMma, + tmem: Optional[cute.Tensor] = None, + *, + stage: Optional[int] = None, + ) -> cute.Tensor: + """tcgen05 A-operand TMEM view over flat TensorSpec TMEM storage. + + The multi-stage stride comes from the bound TMEM tensor: aliased + storage (e.g. a bf16 P ring over a wider f32 accumulator ring, see + `alias_acc_as_tmem`) strides its stages by the ALIASED region's + footprint, not by this operand's own column footprint.""" + tmem = tmem if tmem is not None else self.A.tmem + assert tmem is not None, "A tmem not bound — call with_tmem(...) or pass tmem" + stage_stride = None + if cute.rank(tmem.layout) >= 3 and cute.size(tmem.layout, mode=[2]) > 1: + stage_stride = tmem.layout.stride[2] + return self.frag_A_tmem(tiled_mma, tmem.iterator, stage=stage, stage_stride=stage_stride) + + def acc_layout_sm100(self, tiled_mma: cute.TiledMma, *, stages: Optional[int] = None): + """Staged TMEM accumulator layout ((MMA, MMA_M, MMA_N[, STAGE])) for + this matmul. Standalone so warps that never bind operand fragments + (e.g. epilogue warps reading the accumulator) can build the tensor + with just the tiled_mma: `cute.make_tensor(tmem_ptr, layout)`.""" + shape = tiled_mma.partition_shape_C(self.mma_tiler_mnk(tiled_mma)[:2]) + if stages is not None: + shape = cute.append(shape, stages) + return tiled_mma.make_fragment_C(shape).layout + + def frag_A_tmem( + self, + tiled_mma: cute.TiledMma, + tmem_ptr, + *, + stage: Optional[int] = None, + stage_stride: Optional[int] = None, + ) -> cute.Tensor: + """TMEM-resident A-operand fragment ((MMA, MMA_M, MMA_K, STAGE)) at + `tmem_ptr`, for MMAs built with `source="TS"` (the A tile is produced + into TMEM by a previous stage, e.g. linear attention's masked Q@K^T fed + to P@V). Also usable standalone by the warp that *writes* the operand + into TMEM (tcgen05 store partitioning). + + `stage_stride` (elements of the operand dtype) overrides the trailing + stage-mode stride, for storage whose stages are NOT packed at this + operand's own footprint (e.g. a bf16 ring aliased over a wider f32 + accumulator ring).""" + layout = self.smem_layout_A(tiled_mma, stage=stage) + shape = layout.outer.shape if hasattr(layout, "outer") else layout.shape + fake = tiled_mma.make_fragment_A(shape) + frag_layout = fake.layout + if stage_stride is not None: + rank = cute.rank(frag_layout) + new_stride = tuple(frag_layout.stride[i] for i in range(rank - 1)) + (stage_stride,) + frag_layout = cute.make_layout(frag_layout.shape, stride=new_stride) + return cute.make_tensor(cute.recast_ptr(tmem_ptr, dtype=fake.element_type), frag_layout) + + @staticmethod + def _operand_major(t: TensorSpec, is_A: bool) -> Literal["K", "MN"]: + # Which logical matmul dim (K vs MN) is the fast dim in storage? + # Register-only A operands (in_rmem, no SMEM layout) follow CuTe's K-major + # fragment convention regardless of the spec's `transposed` flag. This is + # a WGMMA (SM90/SM120) concern only — the SM100 path uses + # `_storage_major` directly, both because tcgen05's A never lives in + # rmem and because `stage=None` there can simply mean "not yet known". + if is_A and t.in_rmem: + return "K" + return MatmulSpec._storage_major(t, is_A) + + @staticmethod + def _storage_major(t: TensorSpec, is_A: bool) -> Literal["K", "MN"]: + # ROW_MAJOR storage: storage[1] is fast; COL_MAJOR: storage[0] is fast. + # For A: matmul (M, K) maps to storage (0, 1) untransposed, (1, 0) transposed. + # For B: matmul (K, N) maps to storage (0, 1) untransposed, (1, 0) transposed. + is_row_major = t.layout == LayoutEnum.ROW_MAJOR + if is_A: + base = "K" if is_row_major else "MN" + else: + base = "MN" if is_row_major else "K" + return base if not t.transposed else ("MN" if base == "K" else "K") diff --git a/build/torch-cuda/quack/spec/tma.py b/build/torch-cuda/quack/spec/tma.py new file mode 100644 index 0000000000000000000000000000000000000000..28916f4f1f1f1051a389a3f9b77f3647ee576026 --- /dev/null +++ b/build/torch-cuda/quack/spec/tma.py @@ -0,0 +1,278 @@ +# Copyright (c) 2025-2026, Tri Dao. + +"""TMA helpers that sit below TensorSpec. + +TensorSpec-owned single-CTA TMA uses the same flat CTA-value map as CuTe's +generic tile helper. SM100 2-CTA dense loads need an explicit tcgen05 map +because a peer CTA owns instruction panels, not a contiguous half tile. +""" + +from typing import Any, Optional, Tuple, Type, Union, cast + +import cutlass +from cutlass import const_expr +from cutlass.cutlass_dsl import dsl_user_op +from cutlass._mlir import ir +import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir +import cutlass.cute as cute +import cutlass.cute.atom as cute_atom +import cutlass.cute.core as cute_core +from cutlass.cute.nvgpu import cpasync +from cutlass.cute.nvgpu.cpasync.copy import ( + CopyBulkTensorTileG2SNonExecTrait, + CopyBulkTensorTileG2SMulticastNonExecTrait, + CopyBulkTensorTileS2GNonExecTrait, + CopyReduceBulkTensorTileS2GNonExecTrait, +) +from cutlass.cute.nvgpu.cpasync.helpers import TmaInfo +from cutlass.cute.typing import NumericMeta + +from . import mma as spec_mma +from . import smem as spec_smem + + +@dsl_user_op +def _make_tiled_tma_atom_from_cta_v_map( + op: Union[ + cpasync.CopyBulkTensorTileG2SOp, + cpasync.CopyBulkTensorTileG2SMulticastOp, + cpasync.CopyBulkTensorTileS2GOp, + cpasync.CopyReduceBulkTensorTileS2GOp, + ], + gmem_tensor: cute.Tensor, + smem_layout, + cta_v_map, + num_multicast: int = 1, + *, + internal_type: Optional[Type[cutlass.Numeric]] = None, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> TmaInfo: + """Build a tiled TMA load atom from an explicit CTA-value map.""" + stored_smem_layout = smem_layout + smem_rank = cute.rank(smem_layout) + map_rank = cute.rank(cta_v_map) + if smem_rank == map_rank + 1: + smem_layout = cute.select(smem_layout, mode=list(range(map_rank))) + + smem_for_ir = smem_layout + if isinstance(smem_for_ir, cute_core._ComposedLayout): + smem_for_ir = smem_for_ir.value + + tma_format = None + if internal_type is not None: + itype: Any = internal_type + if not isinstance(internal_type, NumericMeta): + raise TypeError(f"internal_type must be a Numeric, but got {internal_type}") + + use_unpack = ( + itype.width == 8 + and isinstance(gmem_tensor.element_type, NumericMeta) + and gmem_tensor.element_type.width < 8 + ) + internal_mlir_type = gmem_tensor.element_type.mlir_type if use_unpack else itype.mlir_type + tma_format = _cute_nvgpu_ir.TmaDataFormat( + _cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack) + ) + + if isinstance(op, cpasync.CopyBulkTensorTileG2SOp): + if num_multicast != 1: + raise ValueError( + f"non-multicast G2S copies require num_multicast=1, got {num_multicast}" + ) + res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load( + cast(Any, gmem_tensor).value, + smem_for_ir, + cta_v_map, + op._to_ir(), + num_multicast=num_multicast, + tma_format=tma_format, + loc=loc, + ip=ip, + ) + return TmaInfo( + cute_atom.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), + res[1], + stored_smem_layout, + ) + + if isinstance(op, cpasync.CopyBulkTensorTileG2SMulticastOp): + if num_multicast < 1: + raise ValueError( + f"multicast G2S copies require num_multicast >= 1, got {num_multicast}" + ) + res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load( + cast(Any, gmem_tensor).value, + smem_for_ir, + cta_v_map, + op._to_ir(), + num_multicast=num_multicast, + tma_format=tma_format, + loc=loc, + ip=ip, + ) + return TmaInfo( + cute_atom.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])), + res[1], + stored_smem_layout, + ) + + if isinstance(op, cpasync.CopyBulkTensorTileS2GOp): + res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_store( + cast(Any, gmem_tensor).value, + smem_for_ir, + cta_v_map, + tma_format=tma_format, + loc=loc, + ip=ip, + ) + return TmaInfo( + cute_atom.CopyAtom(op, CopyBulkTensorTileS2GNonExecTrait(res[0])), + res[1], + stored_smem_layout, + ) + + if isinstance(op, cpasync.CopyReduceBulkTensorTileS2GOp): + res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_reduce( + cast(Any, gmem_tensor).value, + smem_for_ir, + cta_v_map, + op._to_ir(), + tma_format=tma_format, + loc=loc, + ip=ip, + ) + return TmaInfo( + cute_atom.CopyAtom(op, CopyReduceBulkTensorTileS2GNonExecTrait(res[0])), + res[1], + stored_smem_layout, + ) + + raise ValueError(f"expects a bulk tensor (TMA) Copy Op, but got {op}") + + +def _sm100_dense_tma_cta_v_map( + shape: Tuple[int, int], + smem_layout, +): + rows, cols = shape + tile_shape = spec_smem._sm100_smem_tile_shape(smem_layout) + atom_shape, rest_l, rest_k = tile_shape + atom_l_shape, atom_k = atom_shape + atom_l = cute.size(atom_l_shape) + assert rows == atom_l * rest_l, ( + f"SMEM layout leading shape {atom_l}*{rest_l} does not match operand rows {rows}" + ) + assert cols == atom_k * rest_k, ( + f"SMEM layout K shape {atom_k}*{rest_k} does not match operand cols {cols}" + ) + leading_panel_stride = rows * cute.E(0) if rest_l > 1 else 0 + + return cute.make_layout( + ((atom_l, atom_k), rest_l, rest_k), + stride=( + (cute.E(0), cute.E(1)), + leading_panel_stride, + atom_k * cute.E(1), + ), + ) + + +def _sm100_dense_tma_cta_v_map_from_shape( + dtype: Type[cutlass.Numeric], + shape: Tuple[int, int], + cta_group: int = 1, + mma_inst_k: Optional[int] = None, +): + """CTA-value map for a dense SM100 operand tile. + + `shape` is the local per-CTA tile. For cta_group=2 the TMA map must address + the full two-CTA tile in tcgen05 instruction panels, not a contiguous local + half. Example: local `(256, 64)` for full leading size 512 maps two 128-row + panels separated by a 256-row stride. + """ + rows, cols = shape + k_inst = spec_mma.resolve_mma_inst_k(dtype, mma_inst_k) + assert cols % k_inst == 0, f"TMA cols must be divisible by {k_inst}, got {cols}" + atom_rows, rest_rows = spec_mma.operand_leading_atom(rows, cta_group) + leading_panel_stride = rows * cute.E(0) if rest_rows > 1 else 0 + return cute.make_layout( + ((atom_rows, k_inst), rest_rows, cols // k_inst), + stride=( + (cute.E(0), cute.E(1)), + leading_panel_stride, + k_inst * cute.E(1), + ), + ) + + +def _sm100_dense_tma_flat_cta_v_map( + shape: Tuple[int, int], + cta_group: int = 1, +): + """CTA-value map for a flat role-free SM100 TMA storage view. + + `shape` is the local per-CTA tile. For a 2-CTA full leading tile of 512, + each CTA owns two 128-wide instruction panels: CTA0 maps rows + 0..127 and 256..383, while CTA1 maps the complementary panels. The gap + belongs in the CTA-value map; the SMEM descriptor can stay the normal flat + `(local_leading, k)` TensorSpec storage layout. + """ + rows, cols = shape + atom_rows, rest_rows = spec_mma.operand_leading_atom(rows, cta_group) + leading_panel_stride = rows * cute.E(0) if rest_rows > 1 else 0 + return cute.coalesce( + cute.make_layout( + ((atom_rows, rest_rows), cols), + stride=((cute.E(0), leading_panel_stride), cute.E(1)), + ), + target_profile=(1, 1), + ) + + +def slice_tma_tile_by_mma_cta( + g_tile: cute.Tensor, + rows_per_cta: int, + mma_tile_coord_v, + cta_group: int = 1, + mma_inst_k: Optional[int] = None, + *, + dtype: Optional[Type[cutlass.Numeric]] = None, + exact_layout: bool = False, +) -> cute.Tensor: + """Select this CTA's leading-mode slice from a full tcgen05 operand tile. + + TMA atom construction owns the CTA-value map for a per-CTA tile. If the + caller starts from a full MMA operand tile, it must first select the CTA + slice that `thr_mma.partition_A/B` would have selected before calling + `tma_partition`. By default this returns the flat per-CTA tile used by + TensorSpec TMA; `exact_layout=True` additionally reshapes static probes into + the nested layout produced by `partition_A/B`. + """ + if const_expr(cta_group == 1): + return g_tile + assert cta_group == 2, f"expected cta_group 1 or 2, got {cta_group}" + rank = cute.rank(g_tile) + sliced = cute.flat_divide(g_tile, (rows_per_cta,))[ + (None, mma_tile_coord_v, *([None] * (rank - 1))) + ] + if const_expr(not exact_layout): + return sliced + head_dtype = dtype if dtype is not None else g_tile.element_type + head_layout = _sm100_dense_tma_cta_v_map_from_shape( + head_dtype, + (rows_per_cta, cute.size(sliced.shape[1])), + cta_group=1, + mma_inst_k=mma_inst_k, + ) + if const_expr(rank == 2): + tiled_layout = head_layout + else: + tiled_layout = cute.make_layout( + (*head_layout.shape, *sliced.shape[2:]), + stride=( + *head_layout.stride, + *tuple(cute.E(i) for i in range(2, rank)), + ), + ) + return cute.composition(sliced, tiled_layout) diff --git a/build/torch-cuda/quack/spec/tmem.py b/build/torch-cuda/quack/spec/tmem.py new file mode 100644 index 0000000000000000000000000000000000000000..ac4929aa3f8e310b9c06b867f5c1af678dd7c3c6 --- /dev/null +++ b/build/torch-cuda/quack/spec/tmem.py @@ -0,0 +1,301 @@ +# Copyright (c) 2025-2026, Tri Dao. + +"""TMEM storage helpers for CuTe DSL kernels.""" + +from dataclasses import dataclass, replace +from types import SimpleNamespace +from typing import Literal, Optional, TYPE_CHECKING + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import tcgen05 +import cutlass.utils.blackwell_helpers as sm100_utils + +if TYPE_CHECKING: + from .tensor_spec import BoundMMASm100 + + +# The TMEM analogue of the SMEM SharedStorage @cute.struct, but +# column-addressed: TMEM is 128 lanes x 512 columns of 32-bit cells, a field's +# footprint is its column count (via `tcgen05.find_tmem_tensor_col_offset`), +# every field spans all 128 lanes, and offsets are added to the (32-bit-typed) +# TMEM base pointer. Field layouts come from a tiled_mma, not from +# (dtype, size), so fields are declared in MatmulSpec vocabulary, and dtype +# recasting (e.g. an f32 base pointer re-viewed as bf16 for a TMEM operand) +# happens inside the field. + + +def _tmem_dp_stride(dtype: type[cutlass.Numeric]) -> int: + assert dtype.width <= 32 and 32 % dtype.width == 0, ( + f"TMEM layout expects a sub-32b or 32b dtype, got width={dtype.width}" + ) + return (1 << 16) * (32 // dtype.width) + + +def m64_half_partition_offset( + dtype: type[cutlass.Numeric], partition: Literal["lower", "upper"] +) -> int: + """Base-pointer offset for the alternate M=64 half-subpartition.""" + assert partition in ("lower", "upper"), f"expected lower or upper, got {partition}" + return 0 if partition == "lower" else 16 * _tmem_dp_stride(dtype) + + +def make_tmem_layout( + dtype: type[cutlass.Numeric], + shape: tuple[int, int], + stage: int, + *, + interleaved: bool = False, +): + """Role-free logical TMEM layout for dense tcgen05 storage. + + SM100 TMEM addresses are `(dp << 16) | col` in 32-bit words. `tmem_ptr` + applies the sub-word scaling for `T`, so the DP-lane stride in element units + is `(1 << 16) * (32 / bits(T))`. Columns stay contiguous in element units. + + Only physical M=64 and M=128 are represented here. M=128 uses all DP lanes + linearly. M=64 uses half-subpartitions: rows are grouped as `(16, 4)` and + mapped to DPs `[0:16], [32:48], [64:80], [96:112]`. + + `interleaved=False` is the tcgen05 NonInterleaved layout used by TS-A and + TS C-fragments: every stage advances by the column footprint. For M=64, + `interleaved=True` packs stages in lower/upper half-subpartitions before + advancing columns, matching the 1SM SS accumulator C-fragment. Stage 3 is + rejected for interleaved layouts because the exact 3-stage pattern is not a + rectangular affine layout. + """ + assert len(shape) == 2, f"TMEM layout expects a 2D tile shape, got {shape}" + assert stage in (1, 2, 3, 4), f"TMEM layout supports stage 1, 2, 3, or 4, got {stage}" + rows, cols = shape + assert rows in (64, 128), f"tcgen05 A-source TMEM layout expects M=64 or 128, got {rows}" + elems_per_col = 32 // dtype.width + dp_stride = _tmem_dp_stride(dtype) + stage_stride = ((cols + elems_per_col - 1) // elems_per_col) * elems_per_col + if rows == 64: + if interleaved: + assert stage != 3, "interleaved M=64 TMEM layout does not support stage=3" + half_partition_stride = m64_half_partition_offset(dtype, "upper") + if stage == 1: + stage_shape = 1 + stage_stride = 0 + elif stage == 2: + stage_shape = 2 + stage_stride = half_partition_stride + else: + stage_shape = (2, 2) + stage_stride = (half_partition_stride, stage_stride) + return cute.make_layout( + ((16, 4), cols, stage_shape), + stride=((dp_stride, 32 * dp_stride), 1, stage_stride), + ) + return cute.make_layout( + ((16, 4), cols, stage), + stride=((dp_stride, 32 * dp_stride), 1, 0 if stage == 1 else stage_stride), + ) + return cute.make_layout( + cute.append(shape, stage), + stride=(dp_stride, 1, 0 if stage == 1 else stage_stride), + ) + + +@dataclass +class _TmemFieldBase: + """MLIR marshaling shared by TMEM field kinds. + + The field-owned MMA/copy objects are cute values; stage counts are static. + Threading a field, or the whole TmemStruct, across DSL region boundaries + re-binds those cute values. + """ + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.tiled_mma) + + def __new_from_mlir_values__(self, values): + return replace(self, tiled_mma=cutlass.new_from_mlir_values(self.tiled_mma, values)) + + +@dataclass +class TmemAcc(_TmemFieldBase): + """Accumulator region: staged (MMA, MMA_M, MMA_N[, STAGE]) TMEM tensor.""" + + mma: "BoundMMASm100" + stages: Optional[int] = None + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.mma) + + def __new_from_mlir_values__(self, values): + return replace(self, mma=cutlass.new_from_mlir_values(self.mma, values)) + + def _make_frag(self) -> cute.Tensor: + return self.mma._make_acc_frag(stages=self.stages) + + def num_cols(self) -> int: + return tcgen05.find_tmem_tensor_col_offset(self._make_frag()) + + def view(self, base_ptr, col_offset: int) -> cute.Tensor: + return cute.make_tensor(base_ptr + col_offset, self._make_frag().layout) + + +@dataclass +class TmemOperandA(_TmemFieldBase): + """TMEM-resident storage region later viewed as an MMA A operand.""" + + mma: "BoundMMASm100" + stage: Optional[int] = None + + def _physical_A(self): + # BoundMMASm100 keeps logical A/B plus swap_AB. TmemOperandA is about + # hardware operand A, so swapped MMAs source logical B.T from TMEM. + if self.mma.swap_AB: + assert self.mma.B is not None, "TmemOperandA requires BoundMMASm100.B when swapped" + return self.mma.B.T + assert self.mma.A is not None, "TmemOperandA requires BoundMMASm100.A" + return self.mma.A + + def _physical_mnk(self): + return ( + (self.mma.N, self.mma.M, self.mma.K) + if self.mma.swap_AB + else ( + self.mma.M, + self.mma.N, + self.mma.K, + ) + ) + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.mma) + + def __new_from_mlir_values__(self, values): + return replace(self, mma=cutlass.new_from_mlir_values(self.mma, values)) + + def num_cols(self) -> int: + A = self._physical_A() + assert A is not None, "TmemOperandA requires BoundMMASm100 physical A" + stage = self.stage if self.stage is not None else A.stage + layout = sm100_utils.make_smem_layout_a( + self.mma.tiled_mma, + self._physical_mnk(), + A.dtype, + stage, + ) + shape = layout.outer.shape if hasattr(layout, "outer") else layout.shape + return tcgen05.find_tmem_tensor_col_offset(self.mma.tiled_mma.make_fragment_A(shape)) + + def view(self, base_ptr, col_offset: int) -> cute.Tensor: + A = self._physical_A() + assert A is not None, "TmemOperandA requires BoundMMASm100 physical A" + cta_group = cute.size(self.mma.tiled_mma.thr_id.shape) + ptr = cute.recast_ptr(base_ptr + col_offset, dtype=A.dtype) + stage = self.stage if self.stage is not None else A.stage + rows, cols = A.shape + if cta_group == 2: + assert rows in (64, 128), ( + f"2CTA TS-A TMEM layout expects per-CTA M=64 or 128, got {rows}" + ) + rows = 128 + # Do not call TensorSpec.tmem_layout() here: that uses storage_shape so + # `.T` remains a view of the same backing storage. TmemOperandA allocates + # a fresh physical tcgen05 A tile, so `S.T` must materialize as shape + # `(D, N)`, not reuse `S`'s backing `(N, D)` storage shape. + return cute.make_tensor(ptr, make_tmem_layout(A.dtype, (rows, cols), stage)) + + +def alias_acc_as_tmem( + acc: cute.Tensor, + dtype: type[cutlass.Numeric], + shape: tuple[int, int], + *, + acc_cols: int, + stage: int, +) -> cute.Tensor: + """View the leading dtype tile of a TMEM accumulator allocation. + + This is for intentional subrange aliasing, e.g. linear attention's bf16 P + tile over the f32 QK accumulator. Recast the full accumulator stage to the + destination dtype first, then compose out the logical tile. This preserves + dtype-scaled DP strides and the accumulator's physical stage stride, so if + QK uses physical columns `[0, 128)` and `[128, 256)`, bf16 P uses logical + columns `[0, 128)` backed by physical columns `[0, 64)` and `[128, 192)`. + + This currently covers the 1CTA physical-M=128 use case. If we reuse it for + 1CTA M=64 or 2CTA TS-A aliasing, it should grow the same half-subpartition / + duplicated-local-view handling as `TensorSpec.tmem_layout(cta_group=...)`. + """ + rows, cols = shape + assert rows in (64, 128), f"TMEM alias expects M=64 or 128, got {rows}" + assert dtype.width <= 32, f"TMEM alias expects <=32-bit dtype, got width={dtype.width}" + elems_per_col = 32 // dtype.width + assert cols <= acc_cols * elems_per_col, ( + f"TMEM alias shape {shape} does not fit in {acc_cols} physical accumulator cols" + ) + dp_stride = _tmem_dp_stride(dtype) + stage_stride = 0 if stage == 1 else acc_cols * elems_per_col + if rows == 64: + layout = cute.make_layout( + ((16, 4), cols, stage), + stride=((dp_stride, 32 * dp_stride), 1, stage_stride), + ) + else: + layout = cute.make_layout(cute.append(shape, stage), stride=(dp_stride, 1, stage_stride)) + return cute.make_tensor(cute.recast_ptr(acc.iterator, dtype=dtype), layout) + + +class TmemStruct: + """Named TMEM regions for a kernel, packed back-to-back in declaration order.""" + + def __init__(self, **fields): + self._fields = fields + self._offsets = {} + field_cols = { + name: 0 if field is None else field.num_cols() for name, field in fields.items() + } + offset = 0 + for name, num_cols in field_cols.items(): + self._offsets[name] = offset + offset += num_cols + num_cols = 32 + while num_cols < offset: + num_cols *= 2 + max_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + assert num_cols <= max_cols, ( + f"TMEM plan needs {offset} cols ({field_cols}); " + f"power-of-2 allocation {num_cols} exceeds the {max_cols}-col capacity" + ) + self.num_cols = num_cols + + def col_offset(self, name: str) -> int: + return self._offsets[name] + + def bind(self, base_ptr) -> SimpleNamespace: + """Materialize all field views at the retrieved TMEM base pointer.""" + return SimpleNamespace( + **{ + name: None if f is None else f.view(base_ptr, self._offsets[name]) + for name, f in self._fields.items() + } + ) + + def __extract_mlir_values__(self): + values = [] + self._field_lengths = [] + for f in self._fields.values(): + v = [] if f is None else cutlass.extract_mlir_values(f) + values += v + self._field_lengths.append(len(v)) + return values + + def __new_from_mlir_values__(self, values): + new = object.__new__(TmemStruct) + new_fields = {} + offset = 0 + for (name, f), n in zip(self._fields.items(), self._field_lengths): + new_fields[name] = ( + None if f is None else cutlass.new_from_mlir_values(f, values[offset : offset + n]) + ) + offset += n + new._fields = new_fields + new._offsets = dict(self._offsets) + new.num_cols = self.num_cols + return new diff --git a/build/torch-cuda/quack/sync/__init__.py b/build/torch-cuda/quack/sync/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..55a7a0f77f6cf77df3977de06127c4607edbda9c --- /dev/null +++ b/build/torch-cuda/quack/sync/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026, Tri Dao. + +"""Synchronization helpers for CuTe DSL kernels.""" + +from .barrier import GlobalSemaphore, Semaphore, arrive_inc, wait_eq + +__all__ = ["Semaphore", "GlobalSemaphore", "wait_eq", "arrive_inc"] diff --git a/build/torch-cuda/quack/sync/barrier.py b/build/torch-cuda/quack/sync/barrier.py new file mode 100644 index 0000000000000000000000000000000000000000..69375b46a06ccb454879b9debef53047dde18e81 --- /dev/null +++ b/build/torch-cuda/quack/sync/barrier.py @@ -0,0 +1,130 @@ +# Copyright (c) 2026, Tri Dao. + +"""Global-memory synchronization helpers for CuTe DSL kernels. + +These mirror the small counter-semaphore pattern used by CUTLASS C++ +(`cutlass/barrier.h` / `cutlass/semaphore.h`): one elected thread spins on an +acquire load of a global flag, and one elected thread publishes progress with a +release global reduction. They intentionally do not perform a CTA/warp sync; +callers should pair them with the appropriate warp, named-barrier, or pipeline +synchronization for their producer/consumer protocol. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal, Optional + +import cutlass.cute as cute +from cutlass import Int32, Int64, const_expr + + +@cute.jit +def wait_eq( + lock_ptr: cute.Pointer, + thread_idx: Int32, + flag_offset: Int32 | Int64, + val: Int32, + skip_zero: bool = False, + sync: Literal["none", "warp", "cta"] = "none", +) -> None: + """Wait until ``lock_ptr[flag_offset] == val`` using ``thread_idx == 0``.""" + if const_expr(skip_zero): + if val != 0: + flag_ptr = lock_ptr + flag_offset + if thread_idx == 0: + read_val = Int32(0) + while read_val != val: + read_val = cute.arch.load(flag_ptr, Int32, sem="acquire", scope="gpu") + if const_expr(sync == "warp"): + cute.arch.sync_warp() + elif const_expr(sync == "cta"): + cute.arch.sync_threads() + else: + flag_ptr = lock_ptr + flag_offset + if thread_idx == 0: + read_val = Int32(0) + while read_val != val: + read_val = cute.arch.load(flag_ptr, Int32, sem="acquire", scope="gpu") + if const_expr(sync == "warp"): + cute.arch.sync_warp() + elif const_expr(sync == "cta"): + cute.arch.sync_threads() + + +@cute.jit +def arrive_inc( + lock_ptr: cute.Pointer, + thread_idx: Int32, + flag_offset: Int32 | Int64, + val: Int32 = 1, +) -> None: + """Increment ``lock_ptr[flag_offset]`` by ``val`` using ``thread_idx == 0``.""" + flag_ptr = lock_ptr + flag_offset + if thread_idx == 0: + cute.arch.red(flag_ptr, Int32(val), op="add", dtype="s32", sem="release", scope="gpu") + + +@dataclass(frozen=True) +class Semaphore: + """Global-memory counter semaphore. + + This keeps the pointer, participating thread index, and flag offset together + so call sites can read like the CUTLASS C++ semaphore helpers: + + .. code-block:: python + + sem = Semaphore(ptr, tidx, flag_offset) + sem.wait_eq(expected) + sem.arrive_inc() + + ``sync`` optionally mirrors CUTLASS's ``GenericBarrier`` behavior: + wait synchronizes after the acquire loop, arrive synchronizes before the + release increment. Leave it as ``"none"`` when the call site already has a + surrounding named barrier / pipeline sync or only needs the raw counter op. + """ + + lock_ptr: cute.Pointer + thread_idx: int | Int32 + sync: Literal["none", "warp", "cta"] = "none" + + def _sync(self) -> None: + if self.sync == "warp": + cute.arch.sync_warp() + elif self.sync == "cta": + cute.arch.sync_threads() + + def wait_eq( + self, + val: int | Int32, + flag_offset: Optional[int | Int32 | Int64] = None, + skip_zero: bool = False, + ) -> None: + flag_offset = 0 if flag_offset is None else flag_offset + wait_eq( + self.lock_ptr, + self.thread_idx, + flag_offset, + val, + skip_zero=skip_zero, + sync=self.sync, + ) + + def arrive_inc( + self, val: int | Int32 = 1, flag_offset: Optional[int | Int32 | Int64] = None + ) -> None: + flag_offset = 0 if flag_offset is None else flag_offset + self._sync() + arrive_inc(self.lock_ptr, self.thread_idx, flag_offset, val) + + +# More explicit alias for call sites where plain ``Semaphore`` is ambiguous. +GlobalSemaphore = Semaphore + + +__all__ = [ + "Semaphore", + "GlobalSemaphore", + "wait_eq", + "arrive_inc", +] diff --git a/build/torch-cuda/quack/tensormap_manager.py b/build/torch-cuda/quack/tensormap_manager.py index a25e68c14798efcfa0d1b95adbaedab83004c2fb..462c8edaaa29c8a9547d8b6d035afacee5101c0d 100644 --- a/build/torch-cuda/quack/tensormap_manager.py +++ b/build/torch-cuda/quack/tensormap_manager.py @@ -7,7 +7,6 @@ import cutlass import cutlass.cute as cute from cutlass.cutlass_dsl import Boolean, const_expr, Int32 from cutlass.utils import TensorMapUpdateMode, TensorMapManager -from cutlass._mlir.dialects import llvm @dataclass(frozen=True) @@ -75,18 +74,15 @@ class TensorMapManagerSm90(TensorMapManager): if is_manager_warp: if const_expr(self.tensormap_update_mode == TensorMapUpdateMode.SMEM): for smem_ptr, shape, order in zip(tensormap_smem_ptr, shapes, orders): - smem_ptr_i32 = smem_ptr.toint().ir_value() - llvm.inline_asm( - None, - [smem_ptr_i32, Int32(shape).ir_value()], + smem_ptr_i32 = smem_ptr.toint() + cute.arch.inline_ptx( "{\n\t" ".reg .b64 smem_ptr_i64;\n\t" - "cvt.u64.u32 smem_ptr_i64, $0;\n\t" - f"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [smem_ptr_i64], {order}, $1;\n\t" + "cvt.u64.u32 smem_ptr_i64, {$r0};\n\t" + f"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 " + f"[smem_ptr_i64], {order}, {{$r1}};\n\t" "}\n", - "r,r", - has_side_effects=True, - is_align_stack=False, + read_only_args=[smem_ptr_i32, Int32(shape)], ) # wait until it's safe to update tensormap in global memory with cute.arch.elect_one(): @@ -100,14 +96,11 @@ class TensorMapManagerSm90(TensorMapManager): else: assert len(shapes) == len(orders) == len(tensormap_gmem_ptr) for gmem_ptr, shape, order in zip(tensormap_gmem_ptr, shapes, orders): - gmem_ptr_i64 = gmem_ptr.toint().ir_value() - llvm.inline_asm( - None, - [gmem_ptr_i64, Int32(shape).ir_value()], - f"tensormap.replace.tile.global_dim.global.b1024.b32 [$0], {order}, $1;", - "l,r", - has_side_effects=True, - is_align_stack=False, + gmem_ptr_i64 = gmem_ptr.toint() + cute.arch.inline_ptx( + f"tensormap.replace.tile.global_dim.global.b1024.b32 " + f"[{{$r0}}], {order}, {{$r1}};", + read_only_args=[gmem_ptr_i64, Int32(shape)], ) cute.arch.sync_warp() cute.nvgpu.cpasync.fence_tma_desc_release() diff --git a/build/torch-cuda/quack/testing/__init__.py b/build/torch-cuda/quack/testing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..29483b79da381a0e4f185c0ba810aad205ab1b3b --- /dev/null +++ b/build/torch-cuda/quack/testing/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, Tri Dao. +"""Reusable pytest plugin for QuACK test workflows. + +This subpackage's only contents is the reusable pytest plugin in +:mod:`quack.testing.pytest_plugin`, which wires the ``--async-compile`` pool +(defer-and-retry kernel compilation) into a pytest run:: + + # In a downstream project's conftest.py: + pytest_plugins = ["quack.testing.pytest_plugin"] +""" diff --git a/build/torch-cuda/quack/testing/pytest_plugin.py b/build/torch-cuda/quack/testing/pytest_plugin.py new file mode 100644 index 0000000000000000000000000000000000000000..f6192826c276ed1d35d728907eae8fe7e3d305cb --- /dev/null +++ b/build/torch-cuda/quack/testing/pytest_plugin.py @@ -0,0 +1,595 @@ +# Copyright (c) 2026, Tri Dao. +"""Reusable pytest plugin for QuACK async kernel compilation. + +Adds the ``--async-compile[=N]`` CLI flag: on a kernel-compile ``.o``-cache +miss, the compile is shipped to a pool of N CPU workers (forkserver sidecar, +see :mod:`quack.cache.async_compile`), the test is deferred while other tests +run, and it is retried once the ``.o`` is exported. Zero overhead when the +cache is warm. Works single-process and under xdist (both ``load`` and +``worksteal`` schedulers). + +To opt in, add this line to your ``conftest.py``:: + + pytest_plugins = ["quack.testing.pytest_plugin"] +""" + +from __future__ import annotations + +import pytest + + +# Saved originals so ``pytest_unconfigure`` can restore pytest internals we +# monkey-patched in ``pytest_configure``. Set to ``None`` when the +# corresponding patch was skipped (e.g. pytest internals didn't match what +# we expected, or the env-var opt-out was set). +_orig_compat_getfuncargnames = None +_orig_fixtures_getfuncargnames = None + + +def pytest_addoption(parser): + parser.addoption( + "--async-compile", + type=int, + nargs="?", + const=32, + default=None, + metavar="N", + help=( + "On a kernel-compile cache miss, submit the compile to a pool of " + "N CPU workers, defer the test, and retry it once the .o is " + "ready. Zero overhead when the cache is warm." + ), + ) + + +def _install_getfuncargnames_cache() -> None: + """Cache ``_pytest.compat.getfuncargnames`` by stable function identity. + + Performance background + ---------------------- + Pytest's fixture resolution calls ``getfuncargnames`` ~20 times per test + (once per active fixture / parametrize axis). Each call runs + ``inspect.signature(function)`` from scratch — a deep AST/Signature build + with no caching upstream. For a 2k-test file this is ~40k + ``inspect.signature`` invocations and ~1 s of pure CPU per worker. See + pytest-dev/pytest#11284 (open since 2023) for the underlying root cause: + pytest probes the entire fixture closure for every test instead of just + the test's ``initialnames``. + + Caching by ``id(function)`` does NOT work: the ``request`` fixture is + rebuilt per test, producing thousands of distinct function objects with + identical signatures. Cache by ``(qualname, co_filename, co_firstlineno)`` + instead — that triple is stable across the per-test wrappers and uniquely + identifies a Python function definition. + + Deliberate tradeoffs + -------------------- + This is a monkey-patch of pytest's private API (``_pytest.compat`` and + ``_pytest.fixtures``). We accept this because (a) the upstream fix + (#11284) has been stuck in deprecation cycles for years, and (b) the + speedup is ~1 s wall per file for parametrize-heavy suites. To minimize + risk we: + + * Validate that ``getfuncargnames`` has the expected ``(function, *, name, + cls)`` signature before patching, and skip silently with a warning if + pytest's internals have changed. + * Stash the originals so ``pytest_unconfigure`` restores them. + * Honor ``QUACK_PYTEST_NO_GETFUNCARGNAMES_CACHE=1`` to opt out at runtime. + + Subtlety: ``_pytest.fixtures`` does ``from .compat import getfuncargnames`` + at import time, so we must rebind the name on both modules. + """ + global _orig_compat_getfuncargnames, _orig_fixtures_getfuncargnames + + import os + import warnings + + if os.environ.get("QUACK_PYTEST_NO_GETFUNCARGNAMES_CACHE"): + return + + try: + import _pytest.compat as _compat + import _pytest.fixtures as _fixtures + except ImportError: + return # pytest internals not where we expect them + + orig = getattr(_compat, "getfuncargnames", None) + if orig is None or getattr(_fixtures, "getfuncargnames", None) is not orig: + warnings.warn( + "quack.testing.pytest_plugin: skipping getfuncargnames cache; " + "pytest internals (_pytest.compat / _pytest.fixtures) do not " + "match expected shape. Tests still run, just without the " + "~1 s/file fixture-resolution speedup.", + stacklevel=2, + ) + return + + # Verify the signature is still `(function, *, name, cls)`. If pytest bumps + # the API we'd rather skip than silently miscache. + import inspect + + try: + sig = inspect.signature(orig) + params = sig.parameters + expected = ("function", "name", "cls") + if not all(p in params for p in expected): + raise ValueError(f"unexpected params {tuple(params)!r}") + except (TypeError, ValueError) as e: + warnings.warn( + f"quack.testing.pytest_plugin: skipping getfuncargnames cache; " + f"signature check failed ({e!r}). Tests still run normally.", + stacklevel=2, + ) + return + + cache: dict = {} + + def _identity_key(function): + code = getattr(function, "__code__", None) + if code is None: + return ("__obj__", function) + return ( + getattr(function, "__qualname__", None) or function.__name__, + code.co_filename, + code.co_firstlineno, + ) + + def _patched(function, *, name="", cls=None): + try: + key = (_identity_key(function), name, cls) + except (AttributeError, TypeError): + return orig(function, name=name, cls=cls) + cached = cache.get(key) + if cached is not None: + return cached + result = orig(function, name=name, cls=cls) + cache[key] = result + return result + + _orig_compat_getfuncargnames = orig + _orig_fixtures_getfuncargnames = _fixtures.getfuncargnames # == orig + _compat.getfuncargnames = _patched + # Captured via `from .compat import getfuncargnames` at import time. + _fixtures.getfuncargnames = _patched + + +def _restore_getfuncargnames_cache() -> None: + """Undo ``_install_getfuncargnames_cache``. No-op if not installed.""" + global _orig_compat_getfuncargnames, _orig_fixtures_getfuncargnames + if _orig_compat_getfuncargnames is None: + return + try: + import _pytest.compat as _compat + import _pytest.fixtures as _fixtures + except ImportError: + return + _compat.getfuncargnames = _orig_compat_getfuncargnames + _fixtures.getfuncargnames = _orig_fixtures_getfuncargnames + _orig_compat_getfuncargnames = None + _orig_fixtures_getfuncargnames = None + + +def _disable_unused_accelerator_lazy_call() -> None: + """No-op ``_lazy_call`` for accelerators that aren't available. + + Every ``torch.manual_seed(seed)`` fans out across CUDA, MPS, XPU, MTIA, and + any custom device. For each *uninitialized* backend, ``_lazy_call`` takes + a slow path that calls ``traceback.format_stack()`` to record where the + seed was queued from. Under pytest the call stack is deep, so each + ``format_stack`` costs ~1 ms; across thousands of tests this is several + seconds of pure CPU overhead per worker, even though no XPU/MTIA work is + ever submitted. + + Subtlety: ``torch.xpu/random.py`` and ``torch.cuda/random.py`` do + ``from . import _lazy_call`` at import time, so we must replace the name + on the submodule too — patching only the package attribute is not enough. + """ + import torch + + nop = lambda callable, **kwargs: None # noqa: E731 + if not torch.xpu.is_available(): + torch.xpu._lazy_call = nop + torch.xpu.random._lazy_call = nop # captured via `from . import _lazy_call` + if not torch.mtia.is_available(): + torch.mtia._lazy_call = nop + + +def pytest_configure(config): + """Set up the async compile pool and pytest-internal speedup patches.""" + jobs = config.getoption("--async-compile", default=None) + if jobs is not None: + import os as _os + + worker = _os.environ.get("PYTEST_XDIST_WORKER") + is_xdist_master = worker is None and getattr(config.option, "numprocesses", None) + if not is_xdist_master: + # Real test-running process (single-proc pytest or xdist worker). + from ..cache.async_compile import activate + + n_workers = int(_os.environ.get("PYTEST_XDIST_WORKER_COUNT", "1")) + pool = activate(max(2, jobs // n_workers)) + pool.prewarm() # sidecar import overlaps collection, not the first miss + if worker is not None: + config.pluginmanager.register(_XdistWorkerDefer(pool), "quack-xdist-defer") + else: + config.pluginmanager.register(_SingleProcDeferLoop(pool), "quack-defer-loop") + + # Speed up manual_seed by short-circuiting the queued-seed path on + # unavailable accelerators. + _disable_unused_accelerator_lazy_call() + + # Cache inspect.signature() lookups behind _pytest.compat.getfuncargnames. + # Pytest re-builds signatures ~20x per test during fixture resolution; on a + # 2k-test file this is several seconds of pure CPU we can avoid. + _install_getfuncargnames_cache() + + +def pytest_unconfigure(config): + """Tear down the compile pool and undo any pytest-internal patches.""" + from ..cache.async_compile import get_active_pool, deactivate + + pool = get_active_pool() + if pool is not None: + stats = pool.stats() + defer_plugin = config.pluginmanager.get_plugin( + "quack-defer-loop" + ) or config.pluginmanager.get_plugin("quack-xdist-defer") + defers = defer_plugin.defer_count if defer_plugin else 0 + print( + f"\nasync-compile: {stats['submitted']} keys submitted, " + f"{stats['failed']} failed, {defers} test deferrals" + ) + for sha, err in stats["errors"][:5]: + print(f" pool compile error [{sha[:12]}]: {err}") + deactivate() + + # Always undo the global monkey-patches we installed. This keeps the + # process clean for downstream callers (e.g. notebook hosts that + # pytest-main multiple times in the same interpreter). + _restore_getfuncargnames_cache() + + +# --- CompilePending deferral hooks ------------------------------------------ + + +def _defer_if_compile_pending(item, outcome, force_pass: bool) -> bool: + """If the phase raised CompilePending, flag the item for deferral. + + ``force_pass=True`` (call phase): convert the outcome to a pass. The + defer loop discards all reports of a deferred attempt anyway, and + letting the exception stand would make pytest build a full failure + ``longrepr`` (source-loading traceback format, ~100 ms) for every + deferral — measured to nearly double a cold run's in-session time. + + ``force_pass=False`` (setup phase): leave the exception so the phase + reports as failed and pytest skips the call phase — half-built fixtures + must not run the test body. Setup-phase compiles are rare, so the + longrepr cost is negligible there. + """ + if outcome.excinfo is None: + return False + from ..cache.async_compile import CompilePending + + if not issubclass(outcome.excinfo[0], CompilePending): + return False + item._quack_pending_sha = outcome.excinfo[1].sha + if force_pass: + outcome.force_result(None) + return True + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_setup(item): + outcome = yield + _defer_if_compile_pending(item, outcome, force_pass=False) + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_call(item): + outcome = yield + _defer_if_compile_pending(item, outcome, force_pass=True) + + +# --- Async compile pool: defer-and-retry run loop --------------------------- +# +# With --async-compile, jit_cache misses raise CompilePending after shipping +# the key to a CPU subprocess pool. A deferred test's reports are discarded +# (no logging) and the test is retried once its pending sha completes; +# everything else runs and reports normally. Tests that keep deferring past +# _MAX_ATTEMPTS run one final attempt with the pool suppressed (in-process +# compile) so persistent failures surface as ordinary test failures. +# +# Two execution modes: +# +# * Single-process pytest: ``_SingleProcDeferLoop`` replaces the default +# ``pytest_runtestloop`` with a deque that rotates deferred items to the +# back. +# +# * xdist worker: xdist's ``WorkerInteractor`` owns the runtestloop, but it +# invokes ``pytest_runtest_protocol`` as a hook, so ``_XdistWorkerDefer`` +# takes over the protocol (tryfirst + firstresult). Deferred items are +# stashed; the master receives ``runtest_protocol_complete`` immediately +# (its ``mark_test_complete`` is a plain order-independent ``.remove()``) +# and keeps streaming new items to the worker -- skip-ahead for free. +# Stashed items are opportunistically re-run between incoming items as +# their compiles finish, and fully drained in a ``pytest_runtestloop`` +# hookwrapper after xdist's inner loop exits (channel still open, so late +# reports flow to the master normally). + +# A test needing K distinct cold kernels defers K times (keys are discovered +# serially: each attempt stops at the first missing kernel). Attempts are +# cheap — an item is only re-run once its awaited sha resolves — so the cap +# just needs to exceed the realistic kernels-per-test count; past it, the +# final attempt compiles in-process (pool suppressed) so a wedged pool can't +# defer a test forever. +_MAX_ATTEMPTS = 20 + +# The item that the most recent runtestprotocol call *predicted* would run +# next (its ``nextitem``). Fixture teardown is scoped to that prediction +# (``SetupState.teardown_exact``), and pytest asserts the prediction was +# right at the next setup ("previous item was not torn down properly"). +# Deferral breaks the chain — a drained item can run while the fixture stack +# is scoped for a different module — so ``_run_protocol`` detects the +# misprediction and forces a full teardown first. +_LAST_PREDICTED_NEXT = [None] + + +def _run_protocol(item, nextitem, force_sync: bool): + """Run one test protocol without logging. Returns (pending_sha, reports).""" + from _pytest.runner import runtestprotocol + + from ..cache.async_compile import suppress_pool + + if _LAST_PREDICTED_NEXT[0] is not item: + ss = item.session._setupstate + if ss.stack: + ss.teardown_exact(None) + _LAST_PREDICTED_NEXT[0] = nextitem + + item._quack_pending_sha = None + if force_sync: + with suppress_pool(): + reports = runtestprotocol(item, nextitem=nextitem, log=False) + else: + reports = runtestprotocol(item, nextitem=nextitem, log=False) + return getattr(item, "_quack_pending_sha", None), reports + + +def _log_reports(item, reports): + ihook = item.ihook + ihook.pytest_runtest_logstart(nodeid=item.nodeid, location=item.location) + for rep in reports: + ihook.pytest_runtest_logreport(report=rep) + ihook.pytest_runtest_logfinish(nodeid=item.nodeid, location=item.location) + + +def _log_deferred_teardown_errors(item, reports, log_fn): + """Surface real teardown failures from an otherwise-discarded attempt. + + A deferred attempt's reports are dropped (the retry produces the ones + that count), but the attempt's teardown *did run* — an error there is a + genuine bug (leaked fixture state) that the retry may not reproduce. + Forward just the failed teardown report so it isn't silently hidden. + """ + for rep in reports: + if rep.when == "teardown" and rep.failed: + log_fn(item, [rep]) + + +class _SingleProcDeferLoop: + """Custom runtestloop for non-xdist runs: defer = rotate to back of deque.""" + + def __init__(self, pool): + self.pool = pool + self.defer_count = 0 + + #: How long a deferred test may wait on its pool compile before the loop + #: stops trusting the pool and re-runs it with the pool suppressed + #: (in-process compile). Guards against a permanently-"pending" sha + #: (wedged worker, foreign flock holder that never produces the .o) — + #: without it a rotation-only item would spin forever, since attempts + #: increment only on actual runs. + _WEDGE_TIMEOUT_S = 600.0 + + @pytest.hookimpl(tryfirst=True) + def pytest_runtestloop(self, session): + if session.config.option.collectonly: + return None + import time + from collections import Counter, deque + + queue = deque(session.items) + attempts: Counter = Counter() + deadline: dict = {} # nodeid -> wedge deadline for the awaited sha + spins = 0 # consecutive queue rotations without running anything + while queue: + item = queue.popleft() + awaiting = getattr(item, "_quack_await_sha", None) + wedged = awaiting is not None and time.monotonic() > deadline.get( + item.nodeid, float("inf") + ) + if awaiting is not None and not wedged: + state, _ = self.pool.poll(awaiting) + if state == "pending": + queue.append(item) + spins += 1 + if spins >= len(queue): + time.sleep(0.2) # everything is blocked on the pool + spins = 0 + continue + spins = 0 + attempts[item.nodeid] += 1 + nextitem = queue[0] if queue else None + pending_sha, reports = _run_protocol( + item, nextitem, force_sync=wedged or attempts[item.nodeid] > _MAX_ATTEMPTS + ) + if pending_sha: + item._quack_await_sha = pending_sha + deadline.setdefault(item.nodeid, time.monotonic() + self._WEDGE_TIMEOUT_S) + _log_deferred_teardown_errors(item, reports, _log_reports) + queue.append(item) + self.defer_count += 1 + continue + item._quack_await_sha = None + _log_reports(item, reports) + if session.shouldstop: + raise session.Interrupted(session.shouldstop) + if session.shouldfail: + raise session.Failed(session.shouldfail) + return True + + +class _XdistWorkerDefer: + """Defer-and-retry inside an xdist worker. + + The worker's ``run_one_test`` sends ``runtest_protocol_complete`` to the + master right after our protocol hook returns -- including for deferred + attempts. That is what we want: the master stops tracking the item and + keeps the worker's queue full. The item is then entirely this worker's + responsibility; we re-run it once its compile lands and forward the + reports late (the master's report handling is order-independent). + """ + + def __init__(self, pool): + self.pool = pool + self.deferred = [] # list of (item, awaited_sha) + self.attempts = {} # nodeid -> int + self.defer_count = 0 + self._interactor = None + + def _get_interactor(self, item): + """Find xdist's WorkerInteractor (it registers itself as a plugin).""" + if self._interactor is None: + for p in item.config.pluginmanager.get_plugins(): + if type(p).__name__ == "WorkerInteractor": + self._interactor = p + break + return self._interactor + + def _log_reports_as(self, item, reports) -> None: + """Forward reports with the interactor's item_index pointing at *item*. + + ``WorkerInteractor.pytest_runtest_logreport`` asserts (and serializes) + ``session.items[self.item_index].nodeid == report.nodeid``. For the + incoming item that index is already correct, but reports for a + drained deferred item are forwarded while a *different* item is + current — so temporarily repoint the index. + """ + interactor = self._get_interactor(item) + if interactor is None: + _log_reports(item, reports) + return + saved = interactor.item_index + try: + interactor.item_index = item.session.items.index(item) + _log_reports(item, reports) + finally: + interactor.item_index = saved + + def _attempt(self, item, nextitem) -> None: + """Run item; either log its reports or stash it as deferred.""" + n = self.attempts.get(item.nodeid, 0) + 1 + self.attempts[item.nodeid] = n + pending_sha, reports = _run_protocol(item, nextitem, force_sync=n > _MAX_ATTEMPTS) + if pending_sha: + self.deferred.append((item, pending_sha)) + self.defer_count += 1 + _log_deferred_teardown_errors(item, reports, self._log_reports_as) + else: + self._log_reports_as(item, reports) + + def _drain_ready(self, nextitem) -> None: + """Re-run any deferred items whose compile finished (or failed).""" + pending = self.deferred + self.deferred = [] + for item, sha in pending: + state, _ = self.pool.poll(sha) + if state == "pending": + self.deferred.append((item, sha)) + else: + self._attempt(item, nextitem) # may re-append to self.deferred + + @pytest.hookimpl(tryfirst=True) + def pytest_runtest_protocol(self, item, nextitem): + # Retry any ready deferred items first. nextitem correctness is + # handled centrally by _run_protocol's misprediction guard. + self._drain_ready(nextitem=item) + self._attempt(item, nextitem) + return True # protocol handled; suppress the default impl + + @pytest.hookimpl(hookwrapper=True) + def pytest_runtestloop(self, session): + outcome = yield # xdist WorkerInteractor loop: all assigned items + if outcome.excinfo is not None: + return + if session.shouldstop or session.shouldfail: + # -x / --maxfail / master-initiated stop: the session is being + # aborted; running deferred tests now would both delay shutdown + # and report tests "after" the stop point. Drop them (the master + # already considers them complete; the run is failing anyway). + self.deferred.clear() + return + import time + + # Drain remaining deferred items, blocking on the pool. The channel + # to the master is still open, so reports flow normally. + deadline = time.monotonic() + 600 + while self.deferred: + before = len(self.deferred) + self._drain_ready(nextitem=None) + if self.deferred and len(self.deferred) == before: + if time.monotonic() > deadline: + # Pool wedged: force remaining items through in-process. + for item, _ in self.deferred: + self.attempts[item.nodeid] = _MAX_ATTEMPTS + 1 + deadline = float("inf") + time.sleep(0.2) + + +# --- Session-end report-integrity check -------------------------------------- +# +# The defer machinery reports an xdist item's protocol as complete *before* +# its reports are sent (that's what keeps the master streaming new items to +# the worker). The one integrity hole this opens: a worker that crashes with +# deferred items still stashed loses them SILENTLY — the master already +# counted them complete, so nothing else will notice the missing reports. +# This check closes the hole: at session end, any collected-but-unreported, +# non-deselected test flips the exit status to failure. + +_reported_nodeids: set = set() +_deselected_nodeids: set = set() + + +def pytest_deselected(items): + for item in items: + _deselected_nodeids.add(item.nodeid) + + +def pytest_runtest_logreport(report): + _reported_nodeids.add(report.nodeid) + + +@pytest.hookimpl(trylast=True) +def pytest_sessionfinish(session, exitstatus): + import os + + config = session.config + if config.getoption("--async-compile", default=None) is None: + return # integrity risk only exists with the defer machinery + if os.environ.get("PYTEST_XDIST_WORKER"): + return # workers see the full collection but run a subset + if getattr(config.option, "collectonly", False): + return + if exitstatus not in (0, 1): + return # interrupted (-x, ^C, internal error): missing reports expected + missing = {item.nodeid for item in session.items} - _reported_nodeids - _deselected_nodeids + if not missing: + return + tr = config.pluginmanager.get_plugin("terminalreporter") + lines = [ + f"async-compile INTEGRITY ERROR: {len(missing)} collected test(s) produced no " + "report (deferred tests lost to a worker crash?):" + ] + [f" {nodeid}" for nodeid in sorted(missing)[:20]] + for line in lines: + (tr.write_line if tr else print)(line) + session.exitstatus = 1 diff --git a/build/torch-cuda/quack/testing/trace.py b/build/torch-cuda/quack/testing/trace.py new file mode 100644 index 0000000000000000000000000000000000000000..5434324abdfa9c398421feb76eb06b03d5c0bd17 --- /dev/null +++ b/build/torch-cuda/quack/testing/trace.py @@ -0,0 +1,37 @@ +# Copyright (c) 2026, Tri Dao. +"""Run trace-time checks under the DSL-managed MLIR context. + +Tests must NOT create their own ``with ir.Context():`` to build cute IR +(identity tensors, layout algebra, ``mlir_type`` queries, ...). Building cute +IR inside a user-created raw ``ir.Context`` leaves dangling references in the +DSL's process-global state; the next sizeable ``cute.compile`` then writes +through them and corrupts the heap — glibc aborts with ``malloc(): unaligned +tcache chunk detected`` (or SIGSEGV) in whatever kernel happens to compile +next. This was the cause of the intermittent xdist worker crashes in CI +(worker dies during the cross-entropy backward compile). Upstream +nvidia-cutlass-dsl bug, reproduced with pure-cutlass code on 4.6.0.dev0. + +Use :func:`run_traced` instead: the check runs at ``cute.jit`` trace time, +inside a context whose lifecycle the DSL owns. +""" + +import cutlass +import cutlass.cute as cute + +__all__ = ["run_traced"] + + +@cute.jit +def _traced_runner(fn: cutlass.Constexpr): + fn() + + +def run_traced(fn) -> None: + """Call ``fn()`` at ``cute.jit`` trace time under the DSL-managed context. + + ``fn`` is invoked as a compile-time constexpr while tracing a trivial + host-only jit function, so a live MLIR context (with all cute dialects + registered) is current for the duration of the call. Assertion failures + inside ``fn`` propagate to the caller like any Python exception. + """ + cute.compile(_traced_runner, fn) diff --git a/build/torch-cuda/quack/tile_scheduler.py b/build/torch-cuda/quack/tile_scheduler.py index a42a28f2858d9be8a30093ab8153fc23ccc04a12..acdf945789467330f60356ca0fc968994365b286 100644 --- a/build/torch-cuda/quack/tile_scheduler.py +++ b/build/torch-cuda/quack/tile_scheduler.py @@ -6,7 +6,10 @@ from enum import IntEnum import cutlass import cutlass.cute as cute -from cutlass import Int32, Float32, Boolean, const_expr +from cutlass import Int32, Uint32, Float32, Boolean, const_expr +from cutlass._mlir.dialects import nvvm +from cutlass.cute.experimental import iket + from . import utils as utils from .fast_math import FastDivmod @@ -29,9 +32,63 @@ class PersistenceMode(IntEnum): NONE = 0 STATIC = 1 DYNAMIC = 2 + # Cluster-launch-control work stealing, with the try_cancel response multicast + # by hardware into every CTA's smem; each consumer warp decodes + swizzles + # locally. The work idx comes from the canceled cluster's x coordinate rather + # than a persistent linear counter in the z coordinate. CLC = 3 +# Bytes per sched_smem stage slot: 4 Int32 — either the STAS-broadcast +# (pid_m, pid_n, batch_idx, is_valid) or the CLC try_cancel response. Also the +# expect_tx count both producers arm on the full barrier. +SCHED_SLOT_BYTES = 16 + +# Cap on fire-and-forget try_cancels a retiring cluster sprays to drain the pending +# pool tail (see TileScheduler.cancel_pending_tail). There is no loop: looping +# "until the pool is empty" requires observing responses, which is the synchronous +# drain this design replaces — a fixed blind count is the only non-observing option, +# and the launched-straggler cascade acts as the loop across generations. +# +# Sizing rule: CLC_DRAIN_CANCELS * num_resident_clusters >= typical padding, so the +# residents' first volley covers the pool in generation zero. Varlen padding is +# bounded by (L - 1) M-slots * ncluster_n (~2k for L=512, N-tiles=8); 32 * ~74 +# residents = ~2.4k covers it (traced: only ~19 cancel/launch-race stragglers +# launch, no generational waves). Bounded from above by three costs: (1) on +# exact-grid kernels (dense/symmetric) the pool is already empty at retirement, so +# ALL sprays fail — a per-kernel tax that must stay cheap (measured free at 32); +# (2) the issuer stalls on async-proxy backpressure while enqueueing, holding its +# SM slot and delaying kernel end when there is nothing left to cancel; (3) past +# the pool size, extra cancels buy nothing. +# +# The budget is dynamic between these bounds (blog-style tiering, see +# cancel_pending_tail): a retiring cluster estimates the remaining tail from the +# phantom index it just decoded (tail <= grid_total - w) and sprays +# ceil(tail / max_active_clusters), so block-aligned seqlens (maximal padding, +# ~2x the random-length average) still drain in generation zero. The MIN keeps +# the estimate-free fallback; the MAX bounds the enqueue-backpressure stall a +# retiring cluster's SM slot endures (~256 * ~8ns = ~2us) — beyond it, extra +# generations (~20us empty-cluster waves) are cheaper than deeper stalls, and the +# batched-spray-plus-one-peek design is the real upgrade path. +CLC_DRAIN_CANCELS_MIN = 32 +CLC_DRAIN_CANCELS_MAX = 256 + + +@cute.jit +def cluster_idx_from_block_idx( + cluster_shape_mnk: cutlass.Constexpr[cute.Shape], *, loc=None, ip=None +) -> Tuple[Int32, Int32, Int32]: + """blockIdx // cluster_shape with the cluster shape as a compile-time constant. + cute.arch.cluster_idx() divides by the *runtime* cluster dims from special + registers, which lowers to an I2F/FMUL/F2I float-reciprocal chain per component; + the constexpr division here is a shift (or compile-time magic) instead.""" + bidx = cute.arch.block_idx() + return tuple( + Int32(b) if const_expr(s == 1) else Int32(Uint32(b) // s) + for b, s in zip(bidx, cluster_shape_mnk) + ) + + @cute.jit def get_raster_order_from_option( raster_order_option: RasterOrderOption, problem_shape_ncluster_mn: cute.Shape, group_size: Int32 @@ -72,6 +129,13 @@ class TileSchedulerArguments: class TileScheduler: + # Whether the launched grid can exceed the real work, i.e. whether padding work + # indices exist. Exact-grid schedulers retire only on pool-empty (every granted + # steal is a real tile), so the retirement cancel spray is dead code for them; + # the varlen scheduler over-provisions (worst-case per-batch padding, see its + # get_grid_shape) and overrides this. + grid_may_exceed_work: bool = False + @dataclass class Params: problem_shape_ncluster_mnl: cute.Shape @@ -83,16 +147,17 @@ class TileScheduler: num_clusters_in_group_fdd: FastDivmod tile_count_semaphore: Optional[cute.Pointer] batch_idx_permute: Optional[cute.Tensor] - cluster_shape_mn: cutlass.Constexpr[cute.Shape] + cluster_shape_mnk: cutlass.Constexpr[cute.Shape] persistence_mode: cutlass.Constexpr[PersistenceMode] @staticmethod @cute.jit def create(args: TileSchedulerArguments, *, loc=None, ip=None) -> "TileScheduler.Params": - assert args.cluster_shape_mnk[2] == 1 - cluster_shape_mn = const_expr(cute.select(args.cluster_shape_mnk, mode=[0, 1])) problem_shape_ntile_mn = cute.select(args.problem_shape_ntile_mnl, mode=[0, 1]) - problem_shape_ncluster_mn = cute.ceil_div(problem_shape_ntile_mn, cluster_shape_mn) + problem_shape_ncluster_mn = ( + cute.ceil_div(problem_shape_ntile_mn[0], args.cluster_shape_mnk[0]), + cute.ceil_div(problem_shape_ntile_mn[1], args.cluster_shape_mnk[1]), + ) problem_shape_ncluster_mnl = problem_shape_ncluster_mn + ( args.problem_shape_ntile_mnl[2], ) @@ -129,7 +194,7 @@ class TileScheduler: if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC) else None, args.batch_idx_permute, - cluster_shape_mn, + args.cluster_shape_mnk, args.persistence_mode, ) @@ -142,6 +207,7 @@ class TileScheduler: sched_smem: Optional[cute.Tensor], scheduler_pipeline: Optional[cutlass.pipeline.PipelineAsync], pipeline_state: PipelineStateWAdvance, + throttle_barrier: Optional[cutlass.pipeline.NamedBarrier], params: Params, *, loc=None, @@ -154,6 +220,7 @@ class TileScheduler: self._sched_smem = sched_smem self._scheduler_pipeline = scheduler_pipeline self._pipeline_state = pipeline_state + self._throttle_barrier = throttle_barrier self.params = params self._loc = loc self._ip = ip @@ -162,20 +229,16 @@ class TileScheduler: def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params: return TileScheduler.Params.create(args, loc=loc, ip=ip) - @staticmethod - @cute.jit - def _init_clc_mbarrier(sched_smem: Optional[cute.Tensor] = None, *, loc=None, ip=None) -> None: - # We use 4 ints to store (pid_m, pid_n, batch_idx, is_valid), - # another 4 ints to store clc response, and 2 ints to store the mbarrier for CLC - # Since only the scheduler warp will touch the mbarrier (we don't use multicast when trying - # to cancel workID), we only need the scheduler warp to initialize and sync. - # If we use multicast when canceling workID, we would need all threads to sync. - assert cute.size(sched_smem, mode=[0]) >= 12 - clc_mbar_ptr = sched_smem[None, 0].iterator + 8 - with cute.arch.elect_one(): - cute.arch.mbarrier_init(clc_mbar_ptr, 1) - cute.arch.mbarrier_init_fence() - cute.arch.sync_warp() + def _producer_state(self) -> PipelineStateWAdvance: + """Producer-side view of this warp's consumer pipeline state: same stage + index/count, phase flipped — the producer's phase is always the consumer's + phase ^ 1, since each slot is filled exactly once per consume cycle.""" + return PipelineStateWAdvance( + self._pipeline_state.stages, + self._pipeline_state.count, + self._pipeline_state.index, + self._pipeline_state.phase ^ 1, + ) @staticmethod @cute.jit @@ -191,33 +254,32 @@ class TileScheduler: batch_idx = None return current_work_idx, batch_idx - @staticmethod + @classmethod @cute.jit def create( + cls, params: Params, sched_smem: Optional[cute.Tensor] = None, scheduler_pipeline: Optional[cutlass.pipeline.PipelineAsync] = None, is_scheduler_warp: bool | Boolean = False, + throttle_barrier: Optional[cutlass.pipeline.NamedBarrier] = None, *, loc=None, ip=None, ) -> "TileScheduler": - """is_scheduler_warp should only be true for one warp in the whole cluster""" - current_work_idx, _ = TileScheduler._cluster_idx_to_work_idx_batch( - params, cute.arch.cluster_idx(), loc=loc, ip=ip + """Shared by all scheduler subclasses (cls dispatches Params and + _cluster_idx_to_work_idx_batch overrides). is_scheduler_warp should only be + true for one warp in the whole cluster.""" + cluster_idx = cluster_idx_from_block_idx(params.cluster_shape_mnk, loc=loc, ip=ip) + current_work_idx, _ = cls._cluster_idx_to_work_idx_batch( + params, cluster_idx, loc=loc, ip=ip ) stages = 0 - if const_expr( - params.persistence_mode - in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC] - ): + if const_expr(params.persistence_mode != PersistenceMode.NONE): assert sched_smem is not None assert scheduler_pipeline is not None stages = const_expr(cute.size(sched_smem, mode=[1])) - if const_expr(params.persistence_mode == PersistenceMode.CLC): - if is_scheduler_warp: - TileScheduler._init_clc_mbarrier(sched_smem, loc=loc, ip=ip) - return TileScheduler( + return cls( current_work_idx, Int32(0), # num_tiles_executed Int32(0), # current_batch_idx @@ -225,6 +287,7 @@ class TileScheduler: sched_smem, scheduler_pipeline, PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(0)), + throttle_barrier, params, loc=loc, ip=ip, @@ -241,20 +304,24 @@ class TileScheduler: ) -> Tuple[Int32, Int32, Int32]: if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]): return ( - params.cluster_shape_mn[0] * cute.size(params.problem_shape_ncluster_mnl[:2]), - params.cluster_shape_mn[1], - params.problem_shape_ncluster_mnl[2], + params.cluster_shape_mnk[0] * cute.size(params.problem_shape_ncluster_mnl[:2]), + params.cluster_shape_mnk[1], + params.cluster_shape_mnk[2] * params.problem_shape_ncluster_mnl[2], ) else: num_ctas_in_problem = cute.size( params.problem_shape_ncluster_mnl, loc=loc, ip=ip - ) * cute.size(params.cluster_shape_mn) - num_ctas_per_cluster = cute.size(params.cluster_shape_mn, loc=loc, ip=ip) + ) * cute.size(params.cluster_shape_mnk) + num_ctas_per_cluster = cute.size(params.cluster_shape_mnk, loc=loc, ip=ip) # Total ctas that can run in one wave num_ctas_per_wave = max_active_clusters * num_ctas_per_cluster num_persistent_ctas = cutlass.min(num_ctas_in_problem, num_ctas_per_wave) num_persistent_clusters = num_persistent_ctas // num_ctas_per_cluster - return (*params.cluster_shape_mn, num_persistent_clusters) + return ( + params.cluster_shape_mnk[0], + params.cluster_shape_mnk[1], + params.cluster_shape_mnk[2] * num_persistent_clusters, + ) @cute.jit def _swizzle_cta( @@ -266,12 +333,8 @@ class TileScheduler: cid_fast_in_group, cid_slow = Int32(0), Int32(0) if group_id < params.num_groups_regular: cid_slow, cid_fast_in_group = divmod(id_in_group, params.group_size_fdd) - # if cid_slow % 2 == 1: # inner serpentine - # cid_fast_in_group = params.group_size_fdd.divisor - 1 - cid_fast_in_group else: # tail part cid_slow, cid_fast_in_group = divmod(id_in_group, params.group_size_tail_fdd) - # if cid_slow % 2 == 1: # inner serpentine - # cid_fast_in_group = params.group_size_tail_fdd.divisor - 1 - cid_fast_in_group if group_id % 2 == 1: # serpentine order ncluster_slow = ( params.problem_shape_ncluster_mnl[1] @@ -289,13 +352,15 @@ class TileScheduler: def _cluster_id_to_cta_id( self, cid_m: Int32, cid_n: Int32, *, block_zero_only: bool = False, loc=None, ip=None ) -> Tuple[Int32, Int32]: - if const_expr(block_zero_only): + if const_expr( + block_zero_only or cute.size(self.params.cluster_shape_mnk, loc=loc, ip=ip) == 1 + ): bidx_in_cluster = (Int32(0), Int32(0)) else: # Get the pid from cluster id bidx_in_cluster = cute.arch.block_in_cluster_idx() - pid_m = cid_m * self.params.cluster_shape_mn[0] + bidx_in_cluster[0] - pid_n = cid_n * self.params.cluster_shape_mn[1] + bidx_in_cluster[1] + pid_m = cid_m * self.params.cluster_shape_mnk[0] + bidx_in_cluster[0] + pid_n = cid_n * self.params.cluster_shape_mnk[1] + bidx_in_cluster[1] return pid_m, pid_n @cute.jit @@ -321,11 +386,13 @@ class TileScheduler: if is_valid: if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]): cluster_id_in_problem = work_idx - _, _, bidz_ = cute.arch.block_idx() + bidz_ = ( + bidz + if const_expr(bidz is not None) + else cluster_idx_from_block_idx(params.cluster_shape_mnk, loc=loc, ip=ip)[2] + ) else: bidz_, cluster_id_in_problem = divmod(work_idx, params.num_clusters_per_problem_fdd) - if const_expr(bidz is not None): - bidz_ = bidz cid_m, cid_n = self._swizzle_cta(cluster_id_in_problem, loc=loc, ip=ip) pid_m, pid_n = self._cluster_id_to_cta_id( cid_m, cid_n, block_zero_only=block_zero_only, loc=loc, ip=ip @@ -341,90 +408,184 @@ class TileScheduler: @cute.jit def get_current_work(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo: params = self.params + if const_expr(params.persistence_mode == PersistenceMode.CLC): + return self._get_current_work_clc(loc=loc, ip=ip) pid_m, pid_n, batch_idx, is_valid = Int32(0), Int32(0), Int32(0), Boolean(False) if const_expr(params.persistence_mode == PersistenceMode.NONE): pass - # elif const_expr(params.persistence_mode == PersistenceMode.STATIC): - # return self._delinearize_work_idx(loc=loc, ip=ip) else: + iket.range_push("fetch_wait") self._scheduler_pipeline.consumer_wait(self._pipeline_state) + iket.range_pop() + iket.range_push("fetch_decode") pid_m, pid_n, batch_idx, is_valid_i32 = [ self._sched_smem[i, self._pipeline_state.index] for i in range(4) ] # Need this fence since the STAS from the producer is using the async proxy. # Without this, we get race condition / deadlock. - if const_expr(cute.size(params.cluster_shape_mn) > 1): + if const_expr(cute.size(params.cluster_shape_mnk) > 1): cute.arch.fence_view_async_shared() - cute.arch.sync_warp() - with cute.arch.elect_one(): - self._scheduler_pipeline.consumer_release(self._pipeline_state) + self._scheduler_pipeline.consumer_release(self._pipeline_state) self._pipeline_state.advance() is_valid = Boolean(is_valid_i32) + iket.range_pop() tile_coord_mnkl = (pid_m, pid_n, None, batch_idx) return cutlass.utils.WorkTileInfo(tile_coord_mnkl, Boolean(is_valid)) - # @cute.jit + @cute.jit + def _get_current_work_clc(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo: + """Consumer side of the multicast CLC pipeline, called by every consumer warp + in every CTA of the cluster. The hardware has multicast the 16-byte CLC response + into this CTA's smem slot (completing the local full barrier), so each warp + decodes the response and computes the swizzle itself instead of reading + coordinates decoded once by the scheduler warp.""" + params = self.params + iket.range_push("fetch_wait") + self._scheduler_pipeline.consumer_wait(self._pipeline_state) + iket.range_pop() + iket.range_push("fetch_decode") + clc_response_ptr = self._sched_smem[None, self._pipeline_state.index].iterator + bidx, bidy, bidz, valid = cute.arch.clc_response(clc_response_ptr, loc=loc, ip=ip) + # The CLC response is written by the async proxy; fence so our generic-proxy + # read is ordered before the release below lets the producer's next CLC + # query overwrite the slot. + cute.arch.fence_view_async_shared() + self._scheduler_pipeline.consumer_release(self._pipeline_state) + self._pipeline_state.advance() + # Deliberately decode/swizzle AFTER the release: only the b128 response load + # needs the slot; freeing it here lets the scheduler warp recycle the stage + # for the next query while this warp runs the (possibly expensive, e.g. + # varlen scan) delinearization. + cluster_idx = ( + Int32(Uint32(bidx) // params.cluster_shape_mnk[0]), + Int32(Uint32(bidy) // params.cluster_shape_mnk[1]), + Int32(Uint32(bidz) // params.cluster_shape_mnk[2]), + ) + work_idx, batch_idx = self._cluster_idx_to_work_idx_batch(params, cluster_idx) + # Remember the last decoded work index: at retirement it is the first phantom + # this cluster saw, giving cancel_pending_tail its remaining-tail estimate. + self._current_work_idx = work_idx + ret = self._delinearize_work_idx(work_idx, batch_idx, Boolean(valid), loc=loc, ip=ip) + iket.range_pop() + return ret + + @cute.jit + def _issue_clc_query_multicast(self, *, loc=None, ip=None) -> None: + """Producer side of the multicast CLC pipeline; called only by the scheduler + warp of CTA 0 in the cluster. Waits for all consumers (cluster-wide) to have + released the slot, arms every CTA's full barrier with a 16-byte transaction, + then issues one multicast CLC query. No STAS re-broadcast: the response lands + in all CTAs' smem directly from the hardware.""" + params = self.params + pipeline_state_producer = self._producer_state() + self._scheduler_pipeline.producer_acquire(pipeline_state_producer) + mbar_ptr = self._scheduler_pipeline.producer_get_barrier(pipeline_state_producer) + lane_idx = cute.arch.lane_idx() + if lane_idx < cute.size(params.cluster_shape_mnk): + # Arm each CTA's full barrier: fused arrive (count 1, matching the + # producer group) + expect_tx(16) for the multicast response. + cute.arch.mbarrier_arrive_and_expect_tx(mbar_ptr, SCHED_SLOT_BYTES, lane_idx) + clc_response_ptr = self._sched_smem[None, self._pipeline_state.index].iterator + with cute.arch.elect_one(): + cute.arch.issue_clc_query(mbar_ptr, clc_response_ptr, multicast=True, loc=loc, ip=ip) + + @cute.jit + def throttle_producer_commit( + self, is_producer_warp: bool | Boolean = True, *, loc=None, ip=None + ) -> None: + """Called once per work tile by the main load warp (CTA 0 of the cluster only), + before it starts issuing the tile's loads. Signals the scheduler warp that one + more multicast CLC query may be issued.""" + if const_expr(self._throttle_barrier is not None): + if is_producer_warp: + self._throttle_barrier.arrive() + + @cute.jit + def cancel_pending_tail(self, *, loc=None, ip=None) -> None: + """Fire-and-forget drain of the pending-cluster tail, called by the scheduler + warp when its persistent loop exits (i.e. a steal decoded to an invalid tile). + + CORRECTNESS ASSUMPTION (grant monotonicity): once any fetch decodes into the + invalid/padding region, no pending cluster maps to real work — so canceling + arbitrary pending clusters without inspecting them is safe. PTX does not + document try_cancel grant order; this holds for the observed FIFO-ish drain + and is the same assumption made by the capped spray-and-pray drain in + https://drisspg.github.io/nuggets/A-Tale-of-Two-Schedulers (which hits this + problem at up to 64x padding in capacity-sized grouped GEMM). If it were + violated, a real tile could be canceled unprocessed. + + Fires CLC_DRAIN_CANCELS non-multicast try_cancels at issue rate with no + response waits (responses land in the dead stage-0 slot, tx pre-armed so the + barrier stays balanced; nobody observes either again). The cancel requests + outlive this cluster: their pool-removal effect happens at the work + distributor whether or not the issuer is still resident; only the (unread) + response write-back is orphaned by the exit. + + Pending clusters that launch anyway (cancel/launch races at retirement, or + padding beyond the residents' first volley) see an invalid initial tile, + skip their loop, and spray again on exit. Launches are gated by SM capacity + and the sprayers die near-simultaneously (shared CWD backlog), so + stragglers arrive in machine-width waves, each min(num_residents, + remaining pool) clusters and costing ~one empty-cluster lifetime — a + decaying cascade instead of the full launch stampede. See + CLC_DRAIN_CANCELS for the cap sizing and why there is no drain loop.""" + if const_expr( + self.params.persistence_mode == PersistenceMode.CLC and self.grid_may_exceed_work + ): + params = self.params + # Remaining tail <= total work indices - the phantom index we just drew; + # split it across the resident clusters, which all retire around now. + grid_total = Int32(Uint32(cute.arch.grid_dim()[0]) // params.cluster_shape_mnk[0]) + tail = grid_total - self._current_work_idx + budget = cutlass.min( + Int32(CLC_DRAIN_CANCELS_MAX), + cutlass.max( + Int32(CLC_DRAIN_CANCELS_MIN), + (tail + params.max_active_clusters - 1) // params.max_active_clusters, + ), + ) + state0 = PipelineStateWAdvance( + self._pipeline_state.stages, Int32(0), Int32(0), Int32(0) + ) + mbar_ptr = self._scheduler_pipeline.producer_get_barrier(state0) + resp_ptr = self._sched_smem[None, 0].iterator + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_ptr, SCHED_SLOT_BYTES * budget) + for _ in cutlass.range(budget): + cute.arch.issue_clc_query(mbar_ptr, resp_ptr, multicast=False, loc=loc, ip=ip) + def initial_work_tile_info(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo: return self._delinearize_work_idx(self._current_work_idx, loc=loc, ip=ip) - # if is_scheduler_warp: - # work_tile_info = self._delinearize_work_idx(block_zero_only=True, loc=loc, ip=ip) - # self.write_work_tile_to_smem(work_tile_info, loc=loc, ip=ip) - # self.write_work_tile_to_smem(self._delinearize_work_idx(block_zero_only=True, loc=loc, ip=ip), loc=loc, ip=ip) @cute.jit - def _fetch_next_work_idx(self, *, loc=None, ip=None) -> Int32 | Tuple[Int32, Int32, Boolean]: + def _fetch_next_work_idx(self, *, loc=None, ip=None) -> Int32: """should only be called by the scheduler warp""" params = self.params - num_persistent_clusters = Int32(cute.arch.grid_dim()[2]) + num_persistent_clusters = Int32( + Uint32(cute.arch.grid_dim()[2]) // params.cluster_shape_mnk[2] + ) if const_expr(params.persistence_mode == PersistenceMode.STATIC): return self._current_work_idx + num_persistent_clusters - # Serpentine: alternate wave direction for a bit better load balancing - # But currently seems a tiny bit slower, disabling for now. - # c = Int32(cute.arch.cluster_idx()[2]) - # next_work_idx = self._current_work_idx + 2 * c + 1 - # if self.num_tiles_executed % 2 == 1: - # next_work_idx = self._current_work_idx + 2 * (num_persistent_clusters - 1 - c) + 1 - # return next_work_idx elif const_expr(params.persistence_mode == PersistenceMode.DYNAMIC): next_work_linear_idx = Int32(0) if cute.arch.lane_idx() == 0: # If varlen_m, problem_shape_ncluster_mnl[0] is None, so we use atomic_add # instead of atomic_inc, and at the end of the kernel must reset the semaphore to 0. - # # cute.printf("before atomicadd, tidx = {}, bidz = {}, idx = {}", cute.arch.thread_idx()[0], cute.arch.block_idx()[2], current_work_idx) if const_expr(params.problem_shape_ncluster_mnl[0] is not None): - next_work_linear_idx = num_persistent_clusters + utils.atomic_inc_i32( - cute.size(params.problem_shape_ncluster_mnl) - 1, - params.tile_count_semaphore, + next_work_linear_idx = num_persistent_clusters + Int32( + nvvm.atomicrmw( + op=nvvm.AtomicOpKind.INC, + ptr=params.tile_count_semaphore.llvm_ptr, + a=Int32(cute.size(params.problem_shape_ncluster_mnl) - 1).ir_value(), + loc=loc, + ip=ip, + ) ) else: # varlen_m - next_work_linear_idx = num_persistent_clusters + utils.atomic_add_i32( - 1, params.tile_count_semaphore + next_work_linear_idx = num_persistent_clusters + cute.arch.atomic_add( + params.tile_count_semaphore, Int32(1), loc=loc, ip=ip ) - # cute.printf("after atomicadd, tidx = {}, bidz = {}, idx = {}", cute.arch.thread_idx()[0], cute.arch.block_idx()[2], current_work_idx) return cute.arch.shuffle_sync(next_work_linear_idx, 0) - elif const_expr(params.persistence_mode == PersistenceMode.CLC): - clc_response_ptr = self._sched_smem[None, self._pipeline_state.index].iterator + 4 - mbarrier_addr = self._sched_smem[None, 0].iterator + 8 - cute.arch.sync_warp() - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx(mbarrier_addr, 16, loc=loc, ip=ip) - # cute.arch.issue_clc_query(mbarrier_addr, clc_response_ptr, loc=loc, ip=ip) - utils.issue_clc_query_nomulticast(mbarrier_addr, clc_response_ptr, loc=loc, ip=ip) - cute.arch.sync_warp() - cute.arch.mbarrier_wait(mbarrier_addr, self._pipeline_state.phase, loc=loc, ip=ip) - bidx, bidy, bidz, valid = cute.arch.clc_response(clc_response_ptr, loc=loc, ip=ip) - cute.arch.fence_view_async_shared() - cluster_idx = ( - bidx // params.cluster_shape_mn[0], - bidy // params.cluster_shape_mn[1], - bidz, - ) - cluster_idx, batch_idx = type(self)._cluster_idx_to_work_idx_batch( - params, cluster_idx, loc=loc, ip=ip - ) - return cluster_idx, batch_idx, Boolean(valid) - else: - return Int32(0) @cute.jit def write_work_tile_to_smem( @@ -432,13 +593,7 @@ class TileScheduler: ): params = self.params if const_expr(self._sched_smem is not None): - # producer phase is always consumer_phase ^ 1 - pipeline_state_producer = PipelineStateWAdvance( - self._pipeline_state.stages, - self._pipeline_state.count, - self._pipeline_state.index, - self._pipeline_state.phase ^ 1, - ) + pipeline_state_producer = self._producer_state() self._scheduler_pipeline.producer_acquire(pipeline_state_producer) sched_data = [ work_tile_info.tile_idx[0], @@ -447,21 +602,24 @@ class TileScheduler: Int32(work_tile_info.is_valid_tile), ] lane_idx = cute.arch.lane_idx() - if lane_idx < cute.size(params.cluster_shape_mn): - # cute.printf("Producer pid_m = {}, pid_n = {}, batch_idx = {}, is_valid = {}, after empty wait, idx = {}", sched_data[0], sched_data[1], sched_data[2], sched_data[3], self._current_work_idx) + if lane_idx < cute.size(params.cluster_shape_mnk): pipeline_idx = self._pipeline_state.index - if const_expr(cute.size(params.cluster_shape_mn) == 1): + if const_expr(cute.size(params.cluster_shape_mnk) == 1): for i in cutlass.range_constexpr(4): self._sched_smem[i, pipeline_idx] = sched_data[i] self._scheduler_pipeline.producer_commit(self._pipeline_state) else: peer_cta_rank_in_cluster = lane_idx # Here we assume that the block idx in cluster is linearized such that - # x is the fastest moving direction. - bidx_in_cluster = peer_cta_rank_in_cluster % params.cluster_shape_mn[0] - bidy_in_cluster = peer_cta_rank_in_cluster // params.cluster_shape_mn[0] + # x is the fastest moving direction, followed by y, then z. + bidx_in_cluster = peer_cta_rank_in_cluster % params.cluster_shape_mnk[0] + bidy_in_cluster = ( + peer_cta_rank_in_cluster // params.cluster_shape_mnk[0] + ) % params.cluster_shape_mnk[1] mbar_ptr = self._scheduler_pipeline.producer_get_barrier(self._pipeline_state) - cute.arch.mbarrier_arrive_and_expect_tx(mbar_ptr, 16, peer_cta_rank_in_cluster) + cute.arch.mbarrier_arrive_and_expect_tx( + mbar_ptr, SCHED_SLOT_BYTES, peer_cta_rank_in_cluster + ) utils.store_shared_remote_x4( sched_data[0] + bidx_in_cluster, sched_data[1] + bidy_in_cluster, @@ -481,10 +639,10 @@ class TileScheduler: loc=None, ip=None, ): - """is_scheduler_warp should only be true for one warp in the whole cluster. - Moreover, we assume that only block zero in the cluster is calling this function. - If calling with is_scheduler_warp = True, advance_count must be 1. - """ + """Called by every consumer warp; only the producer work (fetch/query) is + gated on is_scheduler_warp, which must be true for exactly one warp in the + whole cluster (CTA 0's scheduler warp). If calling with + is_scheduler_warp=True, advance_count must be 1.""" params = self.params self.num_tiles_executed += Int32(advance_count) if const_expr(self._pipeline_state is not None and advance_count > 1): @@ -500,21 +658,24 @@ class TileScheduler: elif const_expr(params.persistence_mode == PersistenceMode.CLC): # We assume here that advance_count is 1 for scheduler_warp if is_scheduler_warp: - self._current_work_idx, batch, is_valid = self._fetch_next_work_idx(loc=loc, ip=ip) - work_tile_info = self._delinearize_work_idx( - self._current_work_idx, batch, is_valid, block_zero_only=True, loc=loc, ip=ip - ) - self.write_work_tile_to_smem(work_tile_info, loc=loc, ip=ip) + if const_expr(self._throttle_barrier is not None): + # Throttle: pace queries to tiles actually started by the load warp. + # Without this, the multi-stage lookahead lets a cluster issue queries + # at CLC-round-trip cadence (~1us) instead of tile cadence, + # over-canceling pending clusters and starving other persistent + # workers of steals (cutlass's CLCThrottlePipeline serves this purpose + # with an mbarrier pipeline). A single named barrier suffices: the + # dependency chain (commit k+1 needs fetch k+1 needs query k+1 needs + # this sync k) guarantees producer/consumer arrivals strictly + # alternate, so at most one credit is ever outstanding. bar.sync also + # gives a hardware-scheduled wakeup instead of mbarrier + # PHASECHK+NANOSLEEP polling. + self._throttle_barrier.arrive_and_wait() + self._issue_clc_query_multicast(loc=loc, ip=ip) def producer_tail(self): if const_expr(self._scheduler_pipeline is not None): - pipeline_state_producer = PipelineStateWAdvance( - self._pipeline_state.stages, - self._pipeline_state.count, - self._pipeline_state.index, - self._pipeline_state.phase ^ 1, - ) - self._scheduler_pipeline.producer_tail(pipeline_state_producer) + self._scheduler_pipeline.producer_tail(self._producer_state()) def __extract_mlir_values__(self): values, self._values_pos = [], [] @@ -526,6 +687,7 @@ class TileScheduler: self._sched_smem, self._scheduler_pipeline, self._pipeline_state, + self._throttle_barrier, self.params, ]: obj_values = cutlass.extract_mlir_values(obj) @@ -544,6 +706,7 @@ class TileScheduler: self._sched_smem, self._scheduler_pipeline, self._pipeline_state, + self._throttle_barrier, self.params, ], self._values_pos, @@ -559,7 +722,7 @@ def triangular_idx_to_coord(idx: Int32) -> Tuple[Int32, Int32]: Convert a triangular index to 2D coordinates. This is used to convert the linear index to 2D coordinates for triangular matrices. """ - row = utils.ceil((utils.sqrt(2 * idx + 2.25) - 0.5)) - 1 + row = Int32(cute.math.ceil(cute.math.sqrt(2 * idx + 2.25, approx=True) - 0.5)) - 1 col = idx - (row * (row + 1)) // 2 return row, col @@ -578,7 +741,7 @@ class TriangularTileScheduler(TileScheduler): group_size_mul_group_size_fdd: FastDivmod group_size_tail_mul_group_size_fdd: FastDivmod tile_count_semaphore: Optional[cute.Pointer] - cluster_shape_mn: cutlass.Constexpr[cute.Shape] + cluster_shape_mnk: cutlass.Constexpr[cute.Shape] persistence_mode: cutlass.Constexpr[PersistenceMode] @staticmethod @@ -587,9 +750,11 @@ class TriangularTileScheduler(TileScheduler): args: TileSchedulerArguments, *, loc=None, ip=None ) -> "TriangularTileScheduler.Params": assert args.cluster_shape_mnk[2] == 1 - cluster_shape_mn = const_expr(cute.select(args.cluster_shape_mnk, mode=[0, 1])) problem_shape_ntile_mn = cute.select(args.problem_shape_ntile_mnl, mode=[0, 1]) - problem_shape_ncluster_mn = cute.ceil_div(problem_shape_ntile_mn, cluster_shape_mn) + problem_shape_ncluster_mn = ( + cute.ceil_div(problem_shape_ntile_mn[0], args.cluster_shape_mnk[0]), + cute.ceil_div(problem_shape_ntile_mn[1], args.cluster_shape_mnk[1]), + ) problem_shape_ncluster_mnl = problem_shape_ncluster_mn + ( args.problem_shape_ntile_mnl[2], ) @@ -614,7 +779,7 @@ class TriangularTileScheduler(TileScheduler): args.tile_count_semaphore if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC) else None, - cluster_shape_mn, + args.cluster_shape_mnk, args.persistence_mode, ) @@ -622,44 +787,6 @@ class TriangularTileScheduler(TileScheduler): def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params: return TriangularTileScheduler.Params.create(args, loc=loc, ip=ip) - @staticmethod - @cute.jit - def create( - params: Params, - sched_smem: Optional[cute.Tensor] = None, - scheduler_pipeline: Optional[cutlass.pipeline.PipelineAsync] = None, - is_scheduler_warp: bool | Boolean = False, - *, - loc=None, - ip=None, - ) -> "TriangularTileScheduler": - current_work_idx, _ = TileScheduler._cluster_idx_to_work_idx_batch( - params, cute.arch.cluster_idx(), loc=loc, ip=ip - ) - stages = 0 - if const_expr( - params.persistence_mode - in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC] - ): - assert sched_smem is not None - assert scheduler_pipeline is not None - stages = const_expr(cute.size(sched_smem, mode=[1])) - if const_expr(params.persistence_mode == PersistenceMode.CLC): - if is_scheduler_warp: - TileScheduler._init_clc_mbarrier(sched_smem, loc=loc, ip=ip) - return TriangularTileScheduler( - current_work_idx, - Int32(0), # num_tiles_executed - Int32(0), # current_batch_idx - Int32(0), # num_work_idx_before_cur_batch - sched_smem, - scheduler_pipeline, - PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(0)), - params, - loc=loc, - ip=ip, - ) - # called by host @staticmethod def get_grid_shape( @@ -670,19 +797,25 @@ class TriangularTileScheduler(TileScheduler): ip=None, ) -> Tuple[Int32, Int32, Int32]: clusters = (params.num_clusters_per_problem_fdd.divisor, 1) - num_ctas_mnl = tuple(x * y for x, y in zip(clusters, params.cluster_shape_mn)) + ( - params.problem_shape_ncluster_mnl[2], + num_ctas_mnl = ( + clusters[0] * params.cluster_shape_mnk[0], + clusters[1] * params.cluster_shape_mnk[1], + params.cluster_shape_mnk[2] * params.problem_shape_ncluster_mnl[2], ) if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]): return num_ctas_mnl else: num_ctas_in_problem = cute.size(num_ctas_mnl, loc=loc, ip=ip) - num_ctas_per_cluster = cute.size(params.cluster_shape_mn, loc=loc, ip=ip) + num_ctas_per_cluster = cute.size(params.cluster_shape_mnk, loc=loc, ip=ip) # Total ctas that can run in one wave num_ctas_per_wave = max_active_clusters * num_ctas_per_cluster num_persistent_ctas = cutlass.min(num_ctas_in_problem, num_ctas_per_wave) num_persistent_clusters = num_persistent_ctas // num_ctas_per_cluster - return (*params.cluster_shape_mn, num_persistent_clusters) + return ( + params.cluster_shape_mnk[0], + params.cluster_shape_mnk[1], + params.cluster_shape_mnk[2] * num_persistent_clusters, + ) @cute.jit def _swizzle_cta( @@ -692,8 +825,11 @@ class TriangularTileScheduler(TileScheduler): params = self.params group_size = params.group_size_fdd.divisor group_id = ( - utils.ceil( - (utils.sqrt(2 * cluster_id_in_problem + 2.25) - 0.5) * params.group_size_inv_f32 + Int32( + cute.math.ceil( + (cute.math.sqrt(2 * cluster_id_in_problem + 2.25, approx=True) - 0.5) + * params.group_size_inv_f32 + ) ) - 1 ) @@ -748,22 +884,20 @@ class TriangularTileScheduler(TileScheduler): if is_valid: if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]): cluster_id_in_problem = work_idx - _, _, bidz_ = cute.arch.block_idx() + bidz_ = ( + bidz + if const_expr(bidz is not None) + else cluster_idx_from_block_idx(params.cluster_shape_mnk, loc=loc, ip=ip)[2] + ) else: bidz_, cluster_id_in_problem = divmod(work_idx, params.num_clusters_per_problem_fdd) cluster_id_in_problem = Int32(cluster_id_in_problem) # divmod returns IntValue - if const_expr(bidz is not None): - bidz_ = bidz cid_m, cid_n = self._swizzle_cta(cluster_id_in_problem, loc=loc, ip=ip) pid_m, pid_n = self._cluster_id_to_cta_id( cid_m, cid_n, block_zero_only=block_zero_only, loc=loc, ip=ip ) batch_idx = bidz_ tile_coord_mnkl = (pid_m, pid_n, None, batch_idx) - # tidx, _, _ = cute.arch.thread_idx() - # if tidx == 0: - # cute.printf("bidx = {}, bidy = {}, group_id = {}, id_in_group = {}, group_size_actual = {}, group_col = {}, group_remainder = {}, cid_n_in_group = {}, cid_m_in_group = {}, cid_m = {}, cid_n = {}, is_valid = {}", - # bidx, bidy, group_id, id_in_group, group_size_actual, group_col, group_remainder, cid_n_in_group, cid_m_in_group, cid_m, cid_n, is_valid) return cutlass.utils.WorkTileInfo(tile_coord_mnkl, is_valid) @@ -772,6 +906,7 @@ class VarlenMTileSchedulerArguments: problem_shape_ntile_mnl: cute.Shape total_m: Int32 cu_seqlens_m: cute.Tensor + max_active_clusters: Int32 raster_order: cutlass.Constexpr[RasterOrderOption] group_size: Int32 tile_shape_mn: cutlass.Constexpr[cute.Shape] @@ -781,11 +916,14 @@ class VarlenMTileSchedulerArguments: class VarlenMTileScheduler(TileScheduler): + grid_may_exceed_work: bool = True + @dataclass class Params: problem_shape_ncluster_mnl: cute.Shape total_m: Int32 cu_seqlens_m: cute.Tensor + max_active_clusters: Int32 raster_order: cutlass.Constexpr[RasterOrder] group_size: Int32 group_size_fdd: Optional[FastDivmod] @@ -793,7 +931,7 @@ class VarlenMTileScheduler(TileScheduler): num_clusters_in_group_fdd: FastDivmod tile_shape_mn: cutlass.Constexpr[cute.Shape] tile_count_semaphore: Optional[cute.Pointer] - cluster_shape_mn: cutlass.Constexpr[cute.Shape] + cluster_shape_mnk: cutlass.Constexpr[cute.Shape] persistence_mode: cutlass.Constexpr[PersistenceMode] @staticmethod @@ -801,13 +939,11 @@ class VarlenMTileScheduler(TileScheduler): def create( args: TileSchedulerArguments, *, loc=None, ip=None ) -> "VarlenMTileScheduler.Params": - assert args.cluster_shape_mnk[2] == 1 - cluster_shape_mn = const_expr(cute.select(args.cluster_shape_mnk, mode=[0, 1])) # problem_shape_ntile_mnl[0] will be None for VarlenM problem_shape_ntile_mn = cute.select(args.problem_shape_ntile_mnl, mode=[0, 1]) problem_shape_ncluster_mn = ( None, - cute.ceil_div(problem_shape_ntile_mn[1], cluster_shape_mn[1]), + cute.ceil_div(problem_shape_ntile_mn[1], args.cluster_shape_mnk[1]), ) problem_shape_ncluster_mnl = problem_shape_ncluster_mn + ( args.problem_shape_ntile_mnl[2], @@ -837,6 +973,7 @@ class VarlenMTileScheduler(TileScheduler): problem_shape_ncluster_mnl, args.total_m, args.cu_seqlens_m, + args.max_active_clusters, raster_order, group_size, FastDivmod(group_size) if ncluster_fast is not None else None, @@ -849,35 +986,10 @@ class VarlenMTileScheduler(TileScheduler): args.tile_count_semaphore if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC) else None, - cluster_shape_mn, + args.cluster_shape_mnk, args.persistence_mode, ) - def __init__( - self, - current_work_idx: Int32, - num_tiles_executed: Int32, - current_batch_idx: Int32, - num_work_idx_before_cur_batch: Int32, - sched_smem: Optional[cute.Tensor], - scheduler_pipeline: Optional[cutlass.pipeline.PipelineAsync], - pipeline_state: PipelineStateWAdvance, - params: Params, - *, - loc=None, - ip=None, - ): - self._current_work_idx = current_work_idx - self.num_tiles_executed = num_tiles_executed - self._current_batch_idx = current_batch_idx - self._num_work_idx_before_cur_batch = num_work_idx_before_cur_batch - self._sched_smem = sched_smem - self._scheduler_pipeline = scheduler_pipeline - self._pipeline_state = pipeline_state - self.params = params - self._loc = loc - self._ip = ip - @staticmethod def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params: return VarlenMTileScheduler.Params.create(args, loc=loc, ip=ip) @@ -894,44 +1006,6 @@ class VarlenMTileScheduler(TileScheduler): batch_idx = None return current_work_idx, batch_idx - @staticmethod - @cute.jit - def create( - params: Params, - sched_smem: Optional[cute.Tensor] = None, - scheduler_pipeline: Optional[cutlass.pipeline.PipelineAsync] = None, - is_scheduler_warp: bool | Boolean = False, - *, - loc=None, - ip=None, - ) -> "VarlenMTileScheduler": - current_work_idx, _ = VarlenMTileScheduler._cluster_idx_to_work_idx_batch( - params, cute.arch.cluster_idx(), loc=loc, ip=ip - ) - stages = 0 - if const_expr( - params.persistence_mode - in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC] - ): - assert sched_smem is not None - assert scheduler_pipeline is not None - stages = const_expr(cute.size(sched_smem, mode=[1])) - if const_expr(params.persistence_mode == PersistenceMode.CLC): - if is_scheduler_warp: - TileScheduler._init_clc_mbarrier(sched_smem, loc=loc, ip=ip) - return VarlenMTileScheduler( - current_work_idx, - Int32(0), # num_tiles_executed - Int32(0), # current_batch_idx - Int32(0), # num_work_idx_before_cur_batch - sched_smem, - scheduler_pipeline, - PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(0)), - params, - loc=loc, - ip=ip, - ) - # called by host @staticmethod def get_grid_shape( @@ -941,15 +1015,27 @@ class VarlenMTileScheduler(TileScheduler): loc=None, ip=None, ) -> Tuple[Int32, Int32, Int32]: - block_size = params.tile_shape_mn[0] * params.cluster_shape_mn[0] + block_size = params.tile_shape_mn[0] * params.cluster_shape_mnk[0] num_batch = params.problem_shape_ncluster_mnl[2] + # Tight upper bound on sum(ceil(len_i / block)) given only (total_m, L): + # achieved by adversarial lengths ≡ 1 (mod block), so no smaller grid is safe + # without per-batch seqlens (a too-small grid = tiles with no work index = + # wrong results under CLC). cancel_pending_tail makes the padding slots cheap. total_clusters_m_max = (params.total_m + num_batch * (block_size - 1)) // block_size total_clusters_max = total_clusters_m_max * params.problem_shape_ncluster_mnl[1] if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]): - return (params.cluster_shape_mn[0] * total_clusters_max, params.cluster_shape_mn[1], 1) + return ( + params.cluster_shape_mnk[0] * total_clusters_max, + params.cluster_shape_mnk[1], + params.cluster_shape_mnk[2], + ) else: num_persistent_clusters = cutlass.min(max_active_clusters, total_clusters_max) - return (*params.cluster_shape_mn, num_persistent_clusters) + return ( + params.cluster_shape_mnk[0], + params.cluster_shape_mnk[1], + params.cluster_shape_mnk[2] * num_persistent_clusters, + ) @cute.jit def _swizzle_cta( @@ -970,12 +1056,8 @@ class VarlenMTileScheduler(TileScheduler): num_clusters = num_clusters_m * params.problem_shape_ncluster_mnl[1] if (group_id + 1) * num_clusters_in_group <= num_clusters: cid_slow, cid_fast_in_group = divmod(id_in_group, params.group_size_fdd) - # if cid_slow % 2 == 1: # inner serpentine - # cid_fast_in_group = params.group_size_fdd.divisor - 1 - cid_fast_in_group else: # tail part cid_slow, cid_fast_in_group = divmod(id_in_group, params.group_size_tail_fdd) - # if cid_slow % 2 == 1: # inner serpentine - # cid_fast_in_group = params.group_size_tail_fdd.divisor - 1 - cid_fast_in_group else: assert params.raster_order == RasterOrder.AlongM group_size_actual = cutlass.min( @@ -983,8 +1065,6 @@ class VarlenMTileScheduler(TileScheduler): ) cid_slow = id_in_group // group_size_actual cid_fast_in_group = id_in_group - cid_slow * group_size_actual - # if cid_slow % 2 == 1: # inner serpentine - # cid_fast_in_group = group_size_actual - 1 - cid_fast_in_group if group_id % 2 == 1: # serpentine order ncluster_slow = ( params.problem_shape_ncluster_mnl[1] @@ -1030,15 +1110,16 @@ class VarlenMTileScheduler(TileScheduler): params = self.params lane_idx = cute.arch.lane_idx() num_batch = self.params.problem_shape_ncluster_mnl[2] - block_size = params.tile_shape_mn[0] * params.cluster_shape_mn[0] + block_size = params.tile_shape_mn[0] * params.cluster_shape_mnk[0] batch_idx = self._current_batch_idx next_tile_idx = work_idx problems_end_tile = self._num_work_idx_before_cur_batch + # Pre-init: assigned under a dynamic `if` below, but read outside it (DSL + # scoping requires the outer definition). + num_work_idx_before_cur_batch = self._num_work_idx_before_cur_batch num_clusters_m, num_clusters_cumulative, clusters_in_problems = Int32(0), Int32(0), Int32(0) - is_valid = True - if const_expr(is_valid_ is not None): - is_valid = is_valid_ + is_valid = True if const_expr(is_valid_ is None) else is_valid_ if is_valid: while problems_end_tile <= next_tile_idx: num_clusters_m = self._get_num_m_blocks( @@ -1058,17 +1139,10 @@ class VarlenMTileScheduler(TileScheduler): problems_end_tile = next_tile_idx + 1 else: batch_idx = Int32(num_batch) - - is_valid = batch_idx < num_batch - if const_expr(params.persistence_mode == PersistenceMode.NONE): - is_valid &= self.num_tiles_executed == 0 - cid_m, cid_n = Int32(0), Int32(0) - num_work_idx_before_cur_batch = self._num_work_idx_before_cur_batch - if is_valid: + if batch_idx < num_batch: problems_start_tile = problems_end_tile - clusters_in_problems - # if cute.arch.thread_idx()[0] == 128 + 31: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, problems_end_tile = %d, num_clusters_m=%d, batch_idx = %d", self._tile_idx, problems_end_tile, num_clusters_m, batch_idx) - # The next problem to process is the first one that does not have ending tile position - # that is greater than or equal to tile index. + # The next problem to process is the first one that does not have ending tile + # position that is greater than or equal to tile index. batch_idx_in_problems = cute.arch.popc( cute.arch.vote_ballot_sync( problems_start_tile + num_clusters_cumulative <= next_tile_idx @@ -1082,8 +1156,13 @@ class VarlenMTileScheduler(TileScheduler): ) num_clusters_m = cute.arch.shuffle_sync(num_clusters_m, batch_idx_in_problems) num_work_idx_before_cur_batch = problems_start_tile + num_clusters_prev_lane + + is_valid = batch_idx < num_batch + if const_expr(params.persistence_mode == PersistenceMode.NONE): + is_valid &= self.num_tiles_executed == 0 + cid_m, cid_n = Int32(0), Int32(0) + if is_valid: cluster_id_in_problem = next_tile_idx - num_work_idx_before_cur_batch - # if cute.arch.thread_idx()[0] == 128: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, batch_idx=%d, cid_n=%d, cid_m=%d, is_valid = %d", self._tile_idx, batch_idx, cid_n, cid_m, is_valid) cid_m, cid_n = self._swizzle_cta(cluster_id_in_problem, num_clusters_m, loc=loc, ip=ip) pid_m, pid_n = self._cluster_id_to_cta_id( cid_m, cid_n, block_zero_only=block_zero_only, loc=loc, ip=ip diff --git a/build/torch-cuda/quack/topk.py b/build/torch-cuda/quack/topk.py index 7cd089b38f581e772ae9ba9fcc503ee7ce1cc8a4..6aa2add24180fefa9ca97d3ac89200ec953ec984 100644 --- a/build/torch-cuda/quack/topk.py +++ b/build/torch-cuda/quack/topk.py @@ -6,7 +6,7 @@ from typing import Type, Optional import torch -from ._ops_compat import add_quack_op_namespace_prefix +from ._ops_compat import add_op_namespace_prefix import cuda.bindings.driver as cuda import cutlass @@ -18,8 +18,9 @@ from . import copy_utils as copy_utils from .compile_utils import make_fake_tensor as fake_tensor from .reduction_base import ReductionBase from .reduce import row_reduce -from .cache_utils import jit_cache +from .cache import jit_cache from .cute_dsl_utils import torch2cute_dtype_map +from .dsl import cute_op from .sort.bitonic_sort import bitonic_topk @@ -107,8 +108,7 @@ class TopK: if tXcX[0][0] < shape[0]: copy(tXgX, tXrX) - tXrX_f32 = cute.make_rmem_tensor(tXrX.shape, Float32) - tXrX_f32.store(tXrX.load().to(Float32)) + tXrX_f32 = tXrX.to(Float32) # Encode the indices into the bottom bits of values. log_N = int(math.log2(self.N)) @@ -187,8 +187,7 @@ class TopK: topk_vals_split.store(exp_x * cute.arch.rcp_approx(denom)) # Convert cleaned values to output type - topk_vals_out = cute.make_rmem_tensor_like(topk_vals_split, mValues.element_type) - topk_vals_out.store(topk_vals_split.load().to(mValues.element_type)) + topk_vals_out = topk_vals_split.to(mValues.element_type) row = tXcX[0][0] # # Only the 1st thread in this row writes the top-k values and indices @@ -215,8 +214,25 @@ class TopK: cute.autovec_copy(topk_vals_out[None, i], mValues_store[None, col]) cute.autovec_copy(topk_indices[None, i], mIndices_store[None, col]) + @staticmethod + @jit_cache + def compile(dtype, N, k, softmax): + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + x_cute = fake_tensor(dtype, (batch_sym, N), div) + values_cute = fake_tensor(dtype, (batch_sym, k), div) + indices_cute = fake_tensor(Int32, (batch_sym, k), div) + return cute.compile( + TopK(dtype, N, k, softmax=softmax), + x_cute, + values_cute, + indices_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + -@torch.library.custom_op(add_quack_op_namespace_prefix("_topk_fwd"), mutates_args={"values", "indices"}) +@cute_op(add_op_namespace_prefix("_topk_fwd"), mutates_args={"values", "indices"}) def _topk_fwd( x: torch.Tensor, k: int, softmax: bool, values: torch.Tensor, indices: torch.Tensor ) -> None: @@ -229,47 +245,14 @@ def _topk_fwd( Tuple of (values tensor of shape (M, k), indices tensor of shape (M, k)) """ assert x.dim() == 2, "Input must be 2D" - assert x.is_cuda, "Tensor must be on CUDA device" assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype" assert k > 0 and k <= x.shape[1], "k must be positive and <= N" + if x.numel() == 0: + return N = x.size(1) dtype = torch2cute_dtype_map[x.dtype] - _compile_topk_fwd(dtype, N, k, softmax)(x, values, indices) - - -@_topk_fwd.register_fake -def _topk_fwd_fake( - x: torch.Tensor, k: int, softmax: bool, values: torch.Tensor, indices: torch.Tensor -) -> None: - # See softmax.py _softmax_fwd_fake for why register_fake is needed. - from .cache_utils import COMPILE_ONLY - - has_symint = isinstance(x.size(1), torch.SymInt) or isinstance(k, torch.SymInt) - if COMPILE_ONLY and not has_symint: - N = x.size(1) - dtype = torch2cute_dtype_map[x.dtype] - dx_dtype = torch2cute_dtype_map[x.dtype] - _compile_topk_fwd(dtype, N, k, softmax) - _compile_topk_bwd(dtype, dtype, dx_dtype, N, k, softmax) - - -@jit_cache -def _compile_topk_fwd(dtype, N, k, softmax): - batch_sym = cute.sym_int() - div = math.gcd(128 // dtype.width, N) - x_cute = fake_tensor(dtype, (batch_sym, N), div) - values_cute = fake_tensor(dtype, (batch_sym, k), div) - indices_cute = fake_tensor(Int32, (batch_sym, k), div) - topk_op = TopK(dtype, N, k, softmax=softmax) - return cute.compile( - topk_op, - x_cute, - values_cute, - indices_cute, - cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), - options="--enable-tvm-ffi", - ) + TopK.compile(dtype, N, k, softmax)(x, values, indices) def topk_fwd(x: torch.Tensor, k: int, softmax: bool = False): @@ -456,8 +439,27 @@ class TopKBackward(ReductionBase): if row < shape[0]: copy_dx(tXrdX, tXgdX) + @staticmethod + @jit_cache + def compile(dtype, val_dtype, dx_dtype, N, k, softmax): + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + dvalues_cute = fake_tensor(dtype, (batch_sym, k), div) + values_cute = fake_tensor(val_dtype, (batch_sym, k), div) if val_dtype is not None else None + indices_cute = fake_tensor(Int32, (batch_sym, k), div) + dx_cute = fake_tensor(dx_dtype, (batch_sym, N), div) + return cute.compile( + TopKBackward(dtype, N, k, softmax=softmax), + dvalues_cute, + values_cute, + indices_cute, + dx_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + -@torch.library.custom_op(add_quack_op_namespace_prefix("_topk_bwd"), mutates_args={"dx"}) +@cute_op(add_op_namespace_prefix("_topk_bwd"), mutates_args={"dx"}) def _topk_bwd( dvalues: torch.Tensor, values: Optional[torch.Tensor], @@ -479,54 +481,15 @@ def _topk_bwd( if values is not None: assert values.dim() == 2, "values must be 2D" assert indices.dim() == 2, "indices must be 2D" - assert dvalues.is_cuda and indices.is_cuda, "Tensors must be on CUDA device" assert dvalues.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype" + if dvalues.numel() == 0: + return N = dx.size(1) dtype = torch2cute_dtype_map[dvalues.dtype] val_dtype = torch2cute_dtype_map[values.dtype] if values is not None else None dx_dtype = torch2cute_dtype_map[dx.dtype] - _compile_topk_bwd(dtype, val_dtype, dx_dtype, N, k, softmax)(dvalues, values, indices, dx) - - -@_topk_bwd.register_fake -def _topk_bwd_fake( - dvalues: torch.Tensor, - values: Optional[torch.Tensor], - indices: torch.Tensor, - k: int, - softmax: bool, - dx: torch.Tensor, -) -> None: - # See softmax.py _softmax_fwd_fake for why register_fake is needed. - from .cache_utils import COMPILE_ONLY - - if COMPILE_ONLY and not isinstance(dx.size(1), torch.SymInt): - N = dx.size(1) - dtype = torch2cute_dtype_map[dvalues.dtype] - val_dtype = torch2cute_dtype_map[values.dtype] if values is not None else None - dx_dtype = torch2cute_dtype_map[dx.dtype] - _compile_topk_bwd(dtype, val_dtype, dx_dtype, N, k, softmax) - - -@jit_cache -def _compile_topk_bwd(dtype, val_dtype, dx_dtype, N, k, softmax): - batch_sym = cute.sym_int() - div = math.gcd(128 // dtype.width, N) - dvalues_cute = fake_tensor(dtype, (batch_sym, k), div) - values_cute = fake_tensor(val_dtype, (batch_sym, k), div) if val_dtype is not None else None - indices_cute = fake_tensor(Int32, (batch_sym, k), div) - dx_cute = fake_tensor(dx_dtype, (batch_sym, N), div) - topk_bwd_op = TopKBackward(dtype, N, k, softmax=softmax) - return cute.compile( - topk_bwd_op, - dvalues_cute, - values_cute, - indices_cute, - dx_cute, - cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), - options="--enable-tvm-ffi", - ) + TopKBackward.compile(dtype, val_dtype, dx_dtype, N, k, softmax)(dvalues, values, indices, dx) def topk_bwd( diff --git a/build/torch-cuda/quack/trace.py b/build/torch-cuda/quack/trace.py index 7ec4d7dbec830ee1be630a01e616f02c5cb5e493..8499cfb31deaa224e30f5e7f7261171f8030e119 100644 --- a/build/torch-cuda/quack/trace.py +++ b/build/torch-cuda/quack/trace.py @@ -1,820 +1,9 @@ -# Copyright (c) 2025-2026, Tri Dao. -"""Intra-kernel trace profiler for CuTe-DSL kernels. - -Emits Chrome Trace JSON (compatible with Perfetto / chrome://tracing) from -per-warp instrumentation inserted directly into CuTe-DSL kernels. - -Toggle with QUACK_TRACE=1 env var. When disabled (the default) every trace -call is a compile-time no-op — the JIT never emits any profiling PTX. - -Design decisions ----------------- -**Two-timer approach (inspired by Triton Proton).** -NVIDIA GPUs expose two timers accessible from PTX: - - %globaltimer — device-wide, ~1 GHz, synchronized across all SMs. - - %clock64 — per-SM cycle counter, ~2.1 GHz on H100, *not* synchronized - across SMs (confirmed empirically: cross-SM spread > 400M ticks - vs ~200 ticks for globaltimer on the same launch). -We read %globaltimer once at init and once at flush (per warp) to anchor each -warp's timeline to a device-wide epoch, then read %clock64 for every event. -This gives us low-overhead per-event timestamps (local SM register read) while -retaining cross-SM comparability. During post-processing the per-slot pair - (init_globaltimer, init_clock64) and (final_globaltimer, final_clock64) -auto-calibrates the clock64-to-nanosecond conversion: - ratio = (final_gt - init_gt) / (final_clk - init_clk) - event_ns = init_gt + (event_clk - init_clk) * ratio - -**Compact events (inspired by ThunderKittens).** -Each event is 8 bytes: a raw 32-bit %clock value and a packed (region_id, -event_type) tag, stored with a single v2.u32 streaming store. The device -writes the raw clock — no subtraction needed. The host computes deltas -during post-processing using init_clock from metadata with proper u32 -wraparound. Block and warp identity (constant per slot) are stored once -in per-slot metadata instead of per event. - -**Minimal live registers.** -The TraceContext dataclass carries only 3 DSL values across loop iterations: - - slot_ptr (64-bit) — base of this warp's interleaved [metadata|events] - - cnt (32-bit) — circular buffer write index - - is_active (1-bit) — predicate for stores (warp leader AND warp sampling) -init_clk is NOT stored — the device writes raw clock values and the host -subtracts init_clk during post-processing. This saves one register vs -computing deltas on device. - -**Interleaved per-slot layout.** -Each warp's metadata and events are contiguous in memory: - [meta₀ events₀ | meta₁ events₁ | ...] -This means the device needs only ONE pointer (slot_ptr) instead of separate -metadata and event pointers, saving another register. - -**Warp sampling.** -An optional warp_ids parameter restricts profiling to specific warps. -Non-selected warps execute predicated stores that the GPU evaluates to -hardware no-ops — zero store bandwidth and no branch divergence. - -**Auto-interned region names.** -ctx.b("mma") / ctx.e("mma") auto-assign integer IDs via a module-level -registry at JIT time. The host reads the same registry at write_trace time. -No region_names parameter needed on either side. - -Usage ------ -Host: - with TraceSession("trace.json", grid_size=G, block_size=B) as sess: - my_kernel[grid, block](..., sess.ptr) - -Device (safe to call from all lanes): - ctx = TraceContext.create(trace_ptr) - ctx.b("load"); ctx.e("load") - ctx.flush() -""" - -from __future__ import annotations - -import json -import math -import os -import struct -from typing import Optional -from collections import defaultdict -from dataclasses import dataclass - -import torch - -import cutlass -import cutlass.cute as cute -from cutlass import Int32, Int64, const_expr -from cutlass.base_dsl.arch import Arch -from cutlass._mlir.dialects import nvvm -from cutlass.cutlass_dsl import T - -from .copy_utils import store, store_v2 -from .cute_dsl_utils import ParamsBase - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- -QUACK_TRACE_ENV = "QUACK_TRACE" - -EVENT_BEGIN = 0 -EVENT_END = 1 -EVENT_MARK = 2 - -# Per-event record: (u32 raw_clock, u16 region_id, u16 event_type) -# raw_clock is the 32-bit %clock value at event time; the host converts to -# nanoseconds by subtracting init_clock and applying the calibration ratio. -EVENT_SIZE = 8 -EVENT_STRUCT = struct.Struct(" bool: - """Check QUACK_TRACE=1. Evaluated at JIT time so disabled = no codegen.""" - return os.environ.get(QUACK_TRACE_ENV, "") == "1" - - -# Module-level registry: auto-populated by TraceContext at JIT time, -# read by TraceSession at write_trace time. No need for the user to -# pass region_names to both sides. -_REGION_REGISTRY: dict[int, str] = {} - - -def _intern_region(name: str) -> int: - """Assign a stable integer ID to a region name. JIT-time only.""" - for rid, n in _REGION_REGISTRY.items(): - if n == name: - return rid - rid = len(_REGION_REGISTRY) - _REGION_REGISTRY[rid] = name - return rid - - -def _reset_region_registry(): - """Clear the registry. Called by TraceContext.create so each kernel starts fresh.""" - _REGION_REGISTRY.clear() - - -# --------------------------------------------------------------------------- -# Contiguous buffer layout (shared between host and device) -# --------------------------------------------------------------------------- -# Per-slot data is interleaved: metadata followed by events for each slot. -# This means the device only needs ONE pointer per warp (the slot base). -# -# ┌──────────────────────────────────────┐ slot 0 -# │ metadata (40B) │ events (8B × cap) │ -# ├──────────────────────────────────────┤ slot 1 -# │ metadata (40B) │ events (8B × cap) │ -# ├──────────────────────────────────────┤ ... -# │ ... │ -# └──────────────────────────────────────┘ -# -# In u32 elements: slot_size = META_ELEMS + per_warp_cap * EVENT_ELEMS - -META_ELEMS = METADATA_SIZE // 4 # 10 u32 elements per slot's metadata -EVENT_ELEMS = EVENT_SIZE // 4 # 2 u32 elements per event - - -def _slot_size(per_warp_cap: int) -> int: - """Per-slot size in bytes (metadata + events).""" - return METADATA_SIZE + per_warp_cap * EVENT_SIZE - - -def _buf_total_bytes(total_slots: int, per_warp_cap: int) -> int: - return total_slots * _slot_size(per_warp_cap) - - -# --------------------------------------------------------------------------- -# Device-side helpers -# --------------------------------------------------------------------------- -# Special register reads use NVVM intrinsics. Unpredicated stores use -# cute.arch.store (which wraps nvvm.store_ext with nice pointer/Numeric -# handling). Predicated stores still need inline asm since the NVVM store -# op doesn't support PTX predication. - - -def _read_globaltimer(): - return Int64(nvvm.read_ptx_sreg_globaltimer(T.i64())) - - -def _read_clock64(): - return Int64(nvvm.read_ptx_sreg_clock64(T.i64())) - - -def _read_clock(): - """Read %clock (32-bit, low half of clock64). Used in the hot path for delta encoding.""" - return cutlass.Int32(nvvm.read_ptx_sreg_clock(T.i32())) - - -def _read_smid(): - return cutlass.Int32(nvvm.read_ptx_sreg_smid(T.i32())) - - -def _gmem_ptr(dtype, addr): - """Create a cute global-memory pointer from an Int64 address.""" - return cute.make_ptr(dtype, Int64(addr), cute.AddressSpace.gmem) - - -def _is_warp_leader(): - """Return a DSL predicate for the warp leader thread. - - Uses nvvm.elect_sync() on SM90+ (hardware single-thread election), - falls back to lane_idx() == 0 on older architectures. - """ - if cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum() >= Arch.sm_90: - if cutlass.const_expr(cutlass.CUDA_VERSION.major) == 12: - return cutlass.Boolean(nvvm.elect_sync(T.bool())) - elif cutlass.const_expr(cutlass.CUDA_VERSION.major) == 13: - return cutlass.Boolean(nvvm.elect_sync()) - else: - raise ValueError(f"CUDA_VERSION.major must be >= 12, got {cutlass.CUDA_VERSION.major}") - return cute.arch.lane_idx() == 0 - - -# --------------------------------------------------------------------------- -# Device-side: TraceContext -# --------------------------------------------------------------------------- - - -@dataclass -class TraceContext(ParamsBase): - """Per-warp trace recorder for use inside CuTe-DSL kernels. - - Use the ``create`` classmethod (not ``__init__``) to construct. Named - regions (ctx.b("mma") / ctx.e("mma")) are resolved to integer IDs at JIT - time. Optional warp_ids restricts profiling to specific warps. - - Usage:: - - ctx = TraceContext.create(trace_ptr) - ctx.b("load"); ctx.e("load") - ctx.flush() - """ - - # Compile-time constants (auto-detected as static by ParamsBase) - per_warp_cap: int = 0 - warp_ids: tuple | None = None - - # DSL values (auto-serialized by ParamsBase across cutlass.range loops). - # slot_ptr points to this warp's interleaved [metadata | events] region. - # Metadata at slot_ptr+0, events at slot_ptr+META_ELEMS. - slot_ptr: cute.Pointer = None - cnt: cutlass.Int32 = None - is_active: cutlass.Boolean = None - - # ── Public factory ────────────────────────────────────────────────────── - - @classmethod - def create( - cls, - buf_ptr: Optional[Int64], - per_warp_cap: int = 4096, - warp_ids: tuple[int, ...] | list[int] | None = None, - ): - """Create and initialize a TraceContext. Safe to call from all lanes. - - Only lane 0 (warp leader) performs stores; all other lanes execute - the arithmetic but skip the writes via predication. The caller does - NOT need an ``if is_warp_leader():`` guard. - - Region names are auto-interned by ctx.b("name") / ctx.e("name") via a - module-level registry — no explicit region_names list needed. - """ - assert (per_warp_cap & (per_warp_cap - 1)) == 0, "per_warp_cap must be power of 2" - _reset_region_registry() - warp_ids = tuple(warp_ids) if warp_ids is not None else None - - if not enabled() or const_expr(buf_ptr is None): - return cls( - per_warp_cap=per_warp_cap, - warp_ids=warp_ids, - slot_ptr=None, - cnt=None, - is_active=None, - ) - - SLOT_ELEMS = META_ELEMS + per_warp_cap * EVENT_ELEMS # u32 elements per slot - - bdx, bdy, bdz = cute.arch.block_dim() - warps_per_block = (bdx * bdy * bdz + cute.arch.WARP_SIZE - 1) // cute.arch.WARP_SIZE - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - smid = _read_smid() - - # Linearize blockIdx across all grid dimensions. - bidx, bidy, bidz = cute.arch.block_idx() - gdx, gdy, gdz = cute.arch.grid_dim() - linear_block = bidx + bidy * gdx + bidz * gdx * gdy - slot = linear_block * warps_per_block + warp_idx - - # Single pointer to this warp's interleaved [metadata | events] region. - buf = _gmem_ptr(Int32, Int64(buf_ptr)) - slot_ptr = buf + slot * SLOT_ELEMS - - # is_active gates all stores: warp leader only, AND warp sampling if set. - is_leader = _is_warp_leader() - if warp_ids is not None: - is_active = cutlass.Boolean(False) - for wid in warp_ids: - is_active = is_active | (warp_idx == wid) - is_active = is_active & is_leader - else: - is_active = is_leader - - # Pack warp + smid into 16 bits: warp[5:0] | smid[15:6] - packed = (warp_idx & 0x3F) | ((smid & 0x3FF) << 6) - info = linear_block | (packed << 16) - - # Read timers for metadata (host-side calibration). - gt = _read_globaltimer() - clk64 = _read_clock64() - - # Write init metadata at slot_ptr. cnt is written by flush(). - store(slot_ptr, gt, is_active, cop="cs") # offset 0: init_gt - store(slot_ptr + 2, clk64, is_active, cop="cs") # offset 8: init_clk64 - store(slot_ptr + 8, info, is_active, cop="cs") # offset 32: info - - return cls( - per_warp_cap=per_warp_cap, - warp_ids=warp_ids, - slot_ptr=slot_ptr, - cnt=Int32(0), - is_active=is_active, - ) - - def flush(self): - """Write final timer pair and event count. Safe to call from all lanes.""" - if self.slot_ptr is None: - return - gt = _read_globaltimer() - clk = _read_clock64() - store(self.slot_ptr + 4, gt, self.is_active, cop="cs") # final_gt - store(self.slot_ptr + 6, clk, self.is_active, cop="cs") # final_clk - store(self.slot_ptr + 9, self.cnt, self.is_active, cop="cs") # cnt - - # ── Recording ─────────────────────────────────────────────────────────── - - def _record(self, region_id: int, event_type: int): - if self.slot_ptr is None: - return - clk = _read_clock() # raw 32-bit clock; host subtracts init_clk - evt_off = META_ELEMS + (self.cnt & (self.per_warp_cap - 1)) * EVENT_ELEMS - tag = Int32(region_id) | (Int32(event_type) << 16) - store_v2(self.slot_ptr + evt_off, clk, tag, self.is_active, cop="cs") - self.cnt += 1 - - # Integer-ID API - def record_b(self, region_id: int): - self._record(region_id, EVENT_BEGIN) - - def record_e(self, region_id: int): - self._record(region_id, EVENT_END) - - def record_m(self, region_id: int): - self._record(region_id, EVENT_MARK) - - # Named-region API (string → int resolved at JIT time via module registry) - def b(self, name: str): - self._record(_intern_region(name), EVENT_BEGIN) - - def e(self, name: str): - self._record(_intern_region(name), EVENT_END) - - def m(self, name: str): - self._record(_intern_region(name), EVENT_MARK) - - -# ═══════════════════════════════════════════════════════════════════════════ -# Host-side -# ═══════════════════════════════════════════════════════════════════════════ - - -def _unpack_warp(packed: int) -> int: - return packed & 0x3F - - -def _unpack_smid(packed: int) -> int: - return packed >> 6 - - -@dataclass -class _Event: - """Reconstructed event with absolute timestamp (nanoseconds).""" - - ts: int - id: int - type: int - block: int - warp_smid: int - - @property - def warp(self) -> int: - return _unpack_warp(self.warp_smid) - - @property - def smid(self) -> int: - return _unpack_smid(self.warp_smid) - - -@dataclass -class _SlotMeta: - """Per-slot metadata read back from device.""" - - init_gt: int - init_clk: int - final_gt: int - final_clk: int - info: int - cnt: int - - @property - def block(self) -> int: - return self.info & 0xFFFF - - @property - def warp_smid(self) -> int: - return (self.info >> 16) & 0xFFFF - - @property - def init_clk32(self) -> int: - """Low 32 bits of init_clock64 (%clock at init time).""" - return self.init_clk & 0xFFFFFFFF - - @property - def ratio(self) -> float: - """clock64 ticks → nanoseconds conversion factor for this slot.""" - dclk = self.final_clk - self.init_clk - return (self.final_gt - self.init_gt) / dclk if dclk > 0 else 1.0 - - def clock_to_ns(self, raw_clock32: int) -> float: - """Convert a raw 32-bit clock value to absolute nanoseconds.""" - delta = (raw_clock32 - self.init_clk32) & 0xFFFFFFFF # u32 wraparound - return self.init_gt + delta * self.ratio - - -@dataclass -class TraceWriteOptions: - scale: float = 1e-3 # globaltimer is ns; Chrome trace displayTimeUnit is also "ns" - emit_complete_events: bool = True # pair B/E into ph:"X" (more robust in viewers) - group_by_smid: bool = False # pid = block id (ordered); True = pid = SM id - emit_summary_json: bool = True - summary_hist_bins: int = 128 - - -class TraceSession: - """Host-side profiling session. - - Allocates a single contiguous device buffer, provides one pointer (sess.ptr) - to pass to the kernel, and writes Chrome Trace JSON on exit. - - Can be used as a context manager for automatic sync + write: - - with TraceSession("trace.json", grid_size=G, block_size=B) as sess: - my_kernel[grid, block](..., sess.ptr) - # trace.json written here - """ - - def __init__( - self, - path: str | None = None, - *, - per_warp_cap: int = 4096, - grid_size: int = 1, - block_size: int = 128, - warp_ids: list[int] | tuple[int, ...] | None = None, - device: str | torch.device = "cuda", - ): - assert (per_warp_cap & (per_warp_cap - 1)) == 0, "per_warp_cap must be power of 2" - self.path = path - self.per_warp_cap = per_warp_cap - self.total_blocks = grid_size - self.warps_per_block = (block_size + 31) // 32 - self.warp_ids = tuple(warp_ids) if warp_ids is not None else None - self.device = device - - if not enabled(): - self.d_buf = None - return - - total_slots = self.total_blocks * self.warps_per_block - self.d_buf = torch.zeros( - _buf_total_bytes(total_slots, per_warp_cap), - dtype=torch.uint8, - device=device, - ) - - @property - def ptr(self): - """Device pointer as Int64, or None when tracing is disabled. - Pass directly as an Optional[Int64] kernel argument.""" - from cutlass.cutlass_dsl import Int64 - - return Int64(self.d_buf.data_ptr()) if self.d_buf is not None else None - - def reset(self): - if self.d_buf is not None: - self.d_buf.zero_() - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - if self.path and enabled(): - torch.cuda.synchronize() - self.write_trace(self.path) - return False - - # ── Read helpers ──────────────────────────────────────────────────────── - - def _raw_bytes(self): - return self.d_buf.cpu().numpy() - - def _read_metadata(self, raw) -> list[_SlotMeta]: - total_slots = self.total_blocks * self.warps_per_block - slot_bytes = _slot_size(self.per_warp_cap) - return [ - _SlotMeta(*METADATA_STRUCT.unpack_from(raw, s * slot_bytes)) for s in range(total_slots) - ] - - def _read_events(self, raw, metas) -> list[_Event]: - total_slots = self.total_blocks * self.warps_per_block - slot_bytes = _slot_size(self.per_warp_cap) - events = [] - for s in range(total_slots): - cnt = metas[s].cnt - n = min(cnt, self.per_warp_cap) - start = (cnt & (self.per_warp_cap - 1)) if cnt > self.per_warp_cap else 0 - # Events start after metadata within this slot. - slot_events_off = s * slot_bytes + METADATA_SIZE - meta = metas[s] - for i in range(n): - idx = (start + i) & (self.per_warp_cap - 1) - raw_clk, eid, etype = EVENT_STRUCT.unpack_from( - raw, - slot_events_off + idx * EVENT_SIZE, - ) - events.append( - _Event( - ts=int(meta.clock_to_ns(raw_clk)), - id=eid, - type=etype, - block=meta.block, - warp_smid=meta.warp_smid, - ) - ) - events.sort( - key=lambda ev: ( - ev.ts, - _unpack_smid(ev.warp_smid), - ev.block, - _unpack_warp(ev.warp_smid), - 0 if ev.type == 0 else (1 if ev.type == 2 else 2), - ev.id, - ) - ) - return events - - def _region_name(self, rid: int) -> str: - return _REGION_REGISTRY.get(rid, str(rid)) - - # ── Chrome Trace JSON output ──────────────────────────────────────────── - - def write_trace(self, path: str, opt: TraceWriteOptions | None = None): - if not enabled(): - return - opt = opt or TraceWriteOptions() - - raw = self._raw_bytes() - metas = self._read_metadata(raw) - events = self._read_events(raw, metas) - if not events: - print("intra_kernel_profiler::trace: 0 events") - return - - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - min_ts = events[0].ts - trace_events: list[dict] = [] - - # Collect unique pids/tids and build name metadata entries. - used_pids: set[int] = set() - used_threads: set[tuple[int, int]] = set() - block_to_smid: dict[int, int] = {} - for e in events: - sm, b, w = e.smid, e.block, e.warp - block_to_smid.setdefault(b, sm) - pid = sm if opt.group_by_smid else b - tid = ((b << 6) | w) if opt.group_by_smid else (w * 32) - used_pids.add(pid) - used_threads.add((pid, tid)) - - for pid in sorted(used_pids): - pname = ( - f"SM {pid:03d}" - if opt.group_by_smid - else f"SM {block_to_smid.get(pid, 0):03d} Block {pid:04d}" - ) - trace_events.append( - {"ph": "M", "name": "process_name", "pid": pid, "tid": 0, "args": {"name": pname}} - ) - trace_events.append( - { - "ph": "M", - "name": "process_sort_index", - "pid": pid, - "tid": 0, - "args": {"sort_index": pid}, - } - ) - for pid, tid in sorted(used_threads): - if opt.group_by_smid: - tname = f"Block {tid >> 6:04d} Warp {tid & 0x3F:02d}" - else: - tname = f"Warp {tid // 32:02d}" - trace_events.append( - {"ph": "M", "name": "thread_name", "pid": pid, "tid": tid, "args": {"name": tname}} - ) - trace_events.append( - { - "ph": "M", - "name": "thread_sort_index", - "pid": pid, - "tid": tid, - "args": {"sort_index": tid}, - } - ) - - # Convert events to Chrome Trace format. - if opt.emit_complete_events: - out_events = self._pair_begin_end(events, opt) - for ts, dur, pid, tid, rid, kind, b, w, sm in out_events: - ev = { - "name": self._region_name(rid), - "pid": pid, - "tid": tid, - "cname": _CNAME_LIST[rid % len(_CNAME_LIST)], - "args": {"sm": sm, "block": b, "warp": w}, - } - if kind == 0: - ev.update(ph="X", ts=(ts - min_ts) * opt.scale, dur=dur * opt.scale) - else: - ev.update(ph="i", s="t", ts=(ts - min_ts) * opt.scale) - trace_events.append(ev) - else: - out_events = [] - for e in events: - sm, b, w = e.smid, e.block, e.warp - pid = sm if opt.group_by_smid else b - tid = ((b << 6) | w) if opt.group_by_smid else (w * 32) - ph = "B" if e.type == 0 else ("E" if e.type == 1 else "i") - ev = { - "name": self._region_name(e.id), - "ph": ph, - "ts": (e.ts - min_ts) * opt.scale, - "pid": pid, - "tid": tid, - "cname": _CNAME_LIST[e.id % len(_CNAME_LIST)], - "args": {"sm": sm, "block": b, "warp": w}, - } - if e.type == EVENT_MARK: - ev["s"] = "t" - trace_events.append(ev) - - with open(path, "w") as f: - json.dump({"displayTimeUnit": "ns", "traceEvents": trace_events}, f) - print(f"intra_kernel_profiler::trace: {len(events)} events -> {path}") - - if opt.emit_complete_events and opt.emit_summary_json: - self._write_summary(path, out_events, opt) - - @staticmethod - def _pair_begin_end(events, opt): - """Match B/E events into (ts, dur, pid, tid, rid, kind, block, warp, sm) tuples.""" - thread_states: dict[tuple, dict[int, list[int]]] = defaultdict(lambda: defaultdict(list)) - out = [] - for e in events: - sm, b, w = e.smid, e.block, e.warp - pid = sm if opt.group_by_smid else b - tid = ((b << 6) | w) if opt.group_by_smid else (w * 32) - key = (pid, tid) - if e.type == EVENT_BEGIN: - thread_states[key][e.id].append(e.ts) - elif e.type == EVENT_END: - stack = thread_states[key][e.id] - if stack: - t0 = stack.pop() - if e.ts >= t0: - out.append((t0, e.ts - t0, pid, tid, e.id, 0, b, w, sm)) - else: - out.append((e.ts, 0, pid, tid, e.id, 1, b, w, sm)) - return out - - # ── Summary JSON ──────────────────────────────────────────────────────── - - def _write_summary(self, trace_path, out_events, opt): - base = trace_path.rsplit(".json", 1)[0] if trace_path.endswith(".json") else trace_path - summary_path = base + "_summary.json" - - region_stats: dict[int, list[float]] = defaultdict(list) - for ts, dur, pid, tid, rid, kind, b, w, sm in out_events: - if kind == 0: - region_stats[rid].append(dur * opt.scale) - - regions = [] - for rid in sorted(region_stats): - durs = region_stats[rid] - n = len(durs) - if n == 0: - continue - mean = sum(durs) / n - min_d, max_d = min(durs), max(durs) - var_pop = sum((d - mean) ** 2 for d in durs) / n - var_sample = sum((d - mean) ** 2 for d in durs) / (n - 1) if n > 1 else 0 - cv = math.sqrt(var_sample) / abs(mean) if abs(mean) > 0 and n > 1 else None - - bins = opt.summary_hist_bins or 128 - hist = [0] * bins - if max_d > min_d: - for d in durs: - hist[ - min(int(max(0.0, min(1.0, (d - min_d) / (max_d - min_d))) * bins), bins - 1) - ] += 1 - else: - hist[0] = n - - w_bin = (max_d - min_d) / bins if max_d > min_d else 0 - pcts = {} - for p in (5, 10, 25, 50, 75, 90, 95, 99): - q, cum, val = p / 100.0, 0.0, min_d - for i, c in enumerate(hist): - prev = cum - cum += c / n - if cum >= q: - prob = c / n - frac = max(0.0, min(1.0, (q - prev) / prob)) if prob > 0 else 0 - val = min_d + w_bin * i + frac * w_bin - break - pcts[f"p{p}"] = val - - regions.append( - { - "region": rid, - "name": self._region_name(rid), - "count": n, - "mean_dur": mean, - "cv_dur": cv, - "min_dur": min_d, - "max_dur": max_d, - "var_dur_pop": var_pop, - "var_dur_sample": var_sample, - "percentiles": pcts, - "hist": { - "bins": bins, - "min": min_d, - "max": max_d, - "prob": [c / n for c in hist], - }, - } - ) - - with open(summary_path, "w") as f: - json.dump( - { - "trace": trace_path, - "displayTimeUnit": "ns", - "scale": opt.scale, - "blocks": self.total_blocks, - "warps_per_block": self.warps_per_block, - "per_warp_cap": self.per_warp_cap, - "regions": regions, - }, - f, - indent=2, - ) - print(f"intra_kernel_profiler::trace: summary -> {summary_path}") +"""Compatibility tombstone for the removed QuACK trace API.""" + +raise ImportError( + "quack.trace has been removed. QuACK now uses NVIDIA IKET directly; import " + "`cutlass.cute.experimental.iket` and run workloads under " + "`python -m iket.cli.main ... profile -- ...`. See " + "`examples/example_iket_trace.py` for a minimal marker workload and " + "`examples/example_gemm_trace.py` for a real GEMM trace." +) diff --git a/build/torch-cuda/quack/transform/__init__.py b/build/torch-cuda/quack/transform/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..751ff23defb537558f9446734328a15624d054cb --- /dev/null +++ b/build/torch-cuda/quack/transform/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2026, QuACK team. +"""Transform kernels.""" + +from .hadamard import ( + hadamard_transform, + hadamard_transform_fwd, + hadamard_transform_ref, +) + +__all__ = [ + "hadamard_transform", + "hadamard_transform_fwd", + "hadamard_transform_ref", +] diff --git a/build/torch-cuda/quack/transform/hadamard.py b/build/torch-cuda/quack/transform/hadamard.py new file mode 100644 index 0000000000000000000000000000000000000000..3c39d5b8187ee162ae4935279b96ffd0cc06948f --- /dev/null +++ b/build/torch-cuda/quack/transform/hadamard.py @@ -0,0 +1,897 @@ +# Copyright (c) 2026, QuACK team. + +import math +import os +from functools import partial +from typing import Literal, NamedTuple, Type + +import torch +from .._ops_compat import add_op_namespace_prefix +from torch import Tensor + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Boolean, Float32, Int32, const_expr + +from . import copy_utils +from . import layout_utils +from ..cache import jit_cache +from ..compile_utils import make_fake_tensor as fake_tensor +from ..cute_dsl_utils import get_device_multiprocessor_count, torch2cute_dtype_map +from ..dsl import cute_op + + +# Largest supported transform length. Tied to `_EPT_BY_N_PADDED` below: every +# padded length up to `_MAX_N` must have an entry there. +_MAX_N = 32768 + + +# Default elements-per-thread for each padded N. Tuned empirically. +_EPT_BY_N_PADDED = { + 2: 2, + 4: 4, + 8: 4, + 16: 4, + 32: 8, + 64: 8, + 128: 8, + 256: 8, + 512: 16, + 1024: 32, + 2048: 32, + 4096: 32, + 8192: 32, + 16384: 64, + 32768: 32, +} +assert max(_EPT_BY_N_PADDED) == _MAX_N + + +def _next_power_of_2(n: int) -> int: + return 1 << math.ceil(math.log2(n)) + + +def _log2_exact(n: int, name: str) -> int: + assert n >= 1 and (n & (n - 1) == 0), f"{name} must be a power of 2" + return int(math.log2(n)) + + +def _ensure_last_dim_contiguous(t: Tensor) -> Tensor: + if torch.compiler.is_compiling(): + return t.contiguous() + return t if t.stride(-1) == 1 else t.contiguous() + + +def _get_num_sms(device: torch.device) -> int: + if not torch.cuda.is_available(): + return 1 + device_id = torch.cuda.current_device() if device.index is None else device.index + return max(1, get_device_multiprocessor_count(device_id)) + + +def _should_use_persistent(x: Tensor, N: int) -> bool: + # Persistent mode is intended only for the known occupancy-1 case: 32K rows + # with 16-bit elements. Other sizes keep the regular CTA-per-row launch. + return N == _MAX_N and x.dtype in (torch.float16, torch.bfloat16) + + +# ─── Compile-time bit ownership ───────────────────────────────────────────── + + +class BitOrder(NamedTuple): + """Compile-time ownership of logical Hadamard bits. + + `t` is the bit order owned by thread id bits, and `v` is the bit order owned + by values/registers inside each thread. The class is immutable and still + tuple-like, but named fields make layout reasoning less error-prone. + """ + + t: tuple[int, ...] + v: tuple[int, ...] + + @classmethod + def initial( + cls, + log_n: int, + log_threads_per_transform: int, + log_copy_vecsize: int, + ) -> "BitOrder": + """Initial ownership order induced by `copy_utils.tiled_copy_2d`. + + Low vector bits are per-thread values, then transform-thread bits, then + any remaining high bits return to per-thread values. For + N=2048/ept=64/vec=8 this is: + t=(3, 4, 5, 6, 7), v=(0, 1, 2, 8, 9, 10) + """ + log_ept = log_n - log_threads_per_transform + assert 0 <= log_copy_vecsize <= log_ept + thread_start = log_copy_vecsize + thread_stop = thread_start + log_threads_per_transform + thread_bits = tuple(range(thread_start, thread_stop)) + value_bits = tuple(range(log_copy_vecsize)) + tuple(range(thread_stop, log_n)) + assert len(thread_bits) == log_threads_per_transform + assert len(value_bits) == log_ept + return cls(thread_bits, value_bits) + + def __str__(self) -> str: + return f"t={self.t} v={self.v}" + + def __repr__(self) -> str: + return f"BitOrder(t={self.t}, v={self.v})" + + def hadamard_bits(self, bit_shift: int | None = None) -> tuple[int, ...]: + if bit_shift is None: + bit_shift = len(self.v) + return self.v[:bit_shift] + + def exchange( + self, bit_shift: int | None = None, tail_register_permute: bool = False + ) -> "BitOrder": + """Ownership after `exchange(..., bit_shift=bit_shift)`. + + Store addresses are old thread bits followed by old value bits. The load + view uses the low `bit_shift` address bits as the exchanged value group + and the next thread bits as the new owner thread. If omitted, + `bit_shift=len(v)`, i.e. exchange all current value bits. + """ + if bit_shift is None: + bit_shift = len(self.v) + log_threads = len(self.t) + address_bits = self.t + self.v + assert 0 <= bit_shift <= len(self.v) + new_thread_bits = address_bits[bit_shift : bit_shift + log_threads] + l, r = address_bits[:bit_shift], address_bits[bit_shift + log_threads :] + new_value_bits = l + r if not tail_register_permute else r + l + assert len(new_thread_bits) == len(self.t) + assert len(new_value_bits) == len(self.v) + return BitOrder(new_thread_bits, new_value_bits) + + def tail_direct_store(self, bit_shift: int, log_copy_vecsize: int) -> "BitOrder": + """Ownership after the skipped-tail-exchange register permutation.""" + split = bit_shift + log_copy_vecsize + prefix = self.v[:split] + value_bits = prefix[bit_shift:] + prefix[:bit_shift] + self.v[split:] + return BitOrder(self.t, value_bits) + + def owner_runs(self, log_n: int) -> list[tuple[bool, int, int]]: + """Group physical output bits 0..log_n-1 into maximal value/thread runs. + + Consecutive runs alternate by construction because ownership is binary. + Returns `(is_value, start, width)` tuples. + """ + value_set = set(self.v) + runs: list[tuple[bool, int, int]] = [] + bit = 0 + while bit < log_n: + is_value = bit in value_set + start = bit + while bit < log_n and (bit in value_set) == is_value: + bit += 1 + runs.append((is_value, start, bit - start)) + return runs + + def run_offset(self, start: int, width: int, on: Literal["t", "v"]) -> int | None: + """Where physical run `start:start+width` appears in this t- or v-order. + + Returns the offset, or `None` if the run is not present contiguously. + """ + order = self.t if on == "t" else self.v + try: + offset = order.index(start) + except ValueError: + return None + target = tuple(range(start, start + width)) + return offset if order[offset : offset + width] == target else None + + +# ─── Tail direct store layout ─────────────────────────────────────────────── + + +class TailDirectStorePlan(NamedTuple): + """Custom TiledCopy + gmem layout for skipping the final smem exchange. + + Think of the destination N dimension as physical output bits 0..log_n-1. + After the tail local Hadamard, the direct-store path first permutes rmem so + the vectorized gmem bits are leading in value order. We can then skip the + final smem exchange iff: + * value order starts with vector bits, and + * the warp-lane portion of thread order owns the remaining low address bits + needed for a 128B coalesced segment. + + Running example, N=4096/ept=32/bf16 just before the skipped tail exchange: + order = BitOrder(t=(10,11,3,4,5,6,7), v=(8,9,0,1,2)), bit_shift=2. + The final local Hadamard consumes value bits (8,9), then the rmem permute + changes value order to (0,1,2,8,9). Physical output bits then group into: + value run 0..2, thread run 3..7, value run 8..9, thread run 10..11. + This produces layouts equivalent to: + gmem = (row, v012, t34567, v89, t1011) + thr = (row, 1, 32, 1, 4) with compact tid strides (row,0,4,0,1) + val = ( 1, 8, 1, 4, 1) with compact val strides (0,1,0,8,0). + """ + + gmem_shape: tuple[int, ...] + gmem_stride: tuple[int, ...] + thr_shape: tuple[int, ...] + thr_stride: tuple[int, ...] + val_shape: tuple[int, ...] + val_stride: tuple[int, ...] + value_store_bits: tuple[int, ...] + + @classmethod + def is_feasible( + cls, + log_n: int, + order: BitOrder, + bit_shift: int, + log_copy_vecsize: int, + dtype_width: int, + ) -> bool: + if log_copy_vecsize == 0: + return False + value_store_bits = order.tail_direct_store(bit_shift, log_copy_vecsize).v + vector_bits = tuple(range(log_copy_vecsize)) + # Physical bits are element-index bits. A 128B segment spans 64 bf16/fp16 + # elements (bits 0..5) or 32 fp32 elements (bits 0..4). The vector bits + # come from registers; the remaining low bits should come from warp lanes. + log_elems_per_128b = _log2_exact(1024 // dtype_width, "elements_per_128B") + warp_thread_bits = tuple(range(log_copy_vecsize, min(log_elems_per_128b, log_n))) + lane_thread_bits = order.t[: min(5, len(order.t))] + if value_store_bits[:log_copy_vecsize] != vector_bits: + return False + if any(bit not in lane_thread_bits for bit in warp_thread_bits): + return False + # Each run must be contiguous in either the value or thread bit order. + store_order = BitOrder(order.t, value_store_bits) + return all( + store_order.run_offset(start, width, on="v" if is_value else "t") is not None + for is_value, start, width in store_order.owner_runs(log_n) + ) + + @classmethod + def build( + cls, + log_n: int, + order: BitOrder, + bit_shift: int, + log_copy_vecsize: int, + dtype_width: int, + ) -> "TailDirectStorePlan": + """Construct the gmem/thr/val layout specs for the skipped-tail store.""" + assert cls.is_feasible(log_n, order, bit_shift, log_copy_vecsize, dtype_width) + value_store_bits = order.tail_direct_store(bit_shift, log_copy_vecsize).v + store_order = BitOrder(order.t, value_store_bits) + + gmem_shape = [] + gmem_stride = [] + thr_shape = [] + thr_stride = [] + val_shape = [] + val_stride = [] + runs = store_order.owner_runs(log_n) + assert all(is_value == (i % 2 == 0) for i, (is_value, _, _) in enumerate(runs)) + for is_value, start, width in runs: + order_offset = store_order.run_offset(start, width, on="v" if is_value else "t") + assert order_offset is not None + run_shape = 1 << width + run_stride = 1 << order_offset + gmem_shape.append(run_shape) + gmem_stride.append(1 << start) + thr_shape.append(1 if is_value else run_shape) + thr_stride.append(0 if is_value else run_stride) + val_shape.append(run_shape if is_value else 1) + val_stride.append(run_stride if is_value else 0) + + return cls( + gmem_shape=tuple(gmem_shape), + gmem_stride=tuple(gmem_stride), + thr_shape=tuple(thr_shape), + thr_stride=tuple(thr_stride), + val_shape=tuple(val_shape), + val_stride=tuple(val_stride), + value_store_bits=value_store_bits, + ) + + +# ─── Host-side schedule ───────────────────────────────────────────────────── + + +class HadamardTransformPlan: + """Single source of truth for all host-time Hadamard configuration. + + Owns: + * Sizes: `dtype`, `N`, `N_padded`, `ept`, `vecsize`, `copy_vecsize`, + `threads_per_transform`, `rows_per_block`, `num_threads` + * Log variants: `log_n`, `log_ept`, `log_threads_per_transform`, + `log_copy_vecsize` + * Stage schedule: `stages`, `bit_shifts`, `tail_bit_shift`, `bit_orders` + * Tail-store decision: `tail_store: TailDirectStorePlan | None` + """ + + @staticmethod + def default_ept_for(N_padded: int) -> int: + return _EPT_BY_N_PADDED[N_padded] + + @staticmethod + def default_rows_per_block(threads_per_transform: int, N_padded: int) -> int: + """Pick a default rows-per-block. + + Small transforms pack multiple rows into one CTA to improve occupancy + and amortize launch overhead. Larger transforms default to one row per + CTA because the extra shared-memory exchange/barrier work outweighs the + occupancy gain. + """ + if threads_per_transform > cute.arch.WARP_SIZE: + return 1 + thread_cap = max(1, 1024 // threads_per_transform) + reg_cap = max(1, 8192 // N_padded) + return min(thread_cap, reg_cap) + + def __init__( + self, + dtype: Type[cutlass.Numeric], + N: int, + ept: int | None = None, + rows_per_block: int | None = None, + tail_direct_store_enabled: bool | None = None, + ): + # ── Sizes + assert 2 <= N <= _MAX_N, f"Hadamard requires 2 <= last dim <= {_MAX_N}" + self.dtype = dtype + self.N = N + self.N_padded = _next_power_of_2(N) + assert self.N_padded <= _MAX_N, f"Padded Hadamard dim must be <= {_MAX_N}" + + if ept is None: + ept = self.default_ept_for(self.N_padded) + assert ept >= 1 and (ept & (ept - 1) == 0), "ept must be a power of 2" + assert self.N_padded % ept == 0, "Padded Hadamard dimension must be divisible by ept" + self.ept = ept + self.threads_per_transform = self.N_padded // ept + assert self.threads_per_transform >= 1 and ( + self.threads_per_transform & (self.threads_per_transform - 1) == 0 + ), "N_padded / ept must be a power of 2" + + max_vecsize = 4 if dtype.width == 32 else 8 + self.vecsize = min(max_vecsize, self.ept) + # Row stride is the original N, so padded/non-power-of-two rows may not + # be aligned enough for the full compute vector width in global memory. + self.copy_vecsize = math.gcd(self.vecsize, self.N) + assert self.vecsize & (self.vecsize - 1) == 0 + assert self.copy_vecsize & (self.copy_vecsize - 1) == 0 + assert self.ept % self.vecsize == 0 + + if rows_per_block is None: + rows_per_block = self.default_rows_per_block(self.threads_per_transform, self.N_padded) + assert rows_per_block >= 1, "rows_per_block must be positive" + assert rows_per_block * self.threads_per_transform <= 1024, ( + "block size exceeds 1024 threads" + ) + self.rows_per_block = rows_per_block + self.num_threads = self.threads_per_transform * rows_per_block + + # ── Log variants + self.log_n = _log2_exact(self.N_padded, "N_padded") + self.log_threads_per_transform = _log2_exact( + self.threads_per_transform, "threads_per_transform" + ) + self.log_copy_vecsize = _log2_exact(self.copy_vecsize, "copy_vecsize") + self.log_ept = _log2_exact(self.ept, "ept") + assert self.log_ept == self.log_n - self.log_threads_per_transform + + # ── Stage schedule + self.stages = (self.log_n + self.log_ept - 1) // self.log_ept + self.tail_bit_shift = self.log_n % self.log_ept or self.log_ept + self.bit_shifts = tuple( + self.tail_bit_shift if stage == self.stages - 1 else self.log_ept + for stage in range(self.stages) + ) + + # ── Tail direct store decision + # + # The final stage can be a "tail" when log_N is not an exact multiple of + # log_ept. Normally we would run the final local Hadamard, do one more + # full fp32 smem exchange to restore the original gmem ownership order, + # then use the regular tiled gmem store. In some schedules, however, + # the state *before* that final exchange already has: + # - the vectorized gmem bits in per-thread values after a small rmem + # permute, and + # - the remaining low gmem bits for a 128B segment owned by the warp-lane + # thread bits (e.g. bf16 copy_vecsize=8 needs 3,4,5; bf16 + # copy_vecsize=4 needs 2,3,4,5; fp32 copy_vecsize=4 needs 2,3,4). + # When `TailDirectStorePlan.is_feasible` says yes we skip that last smem + # round trip and store directly to gmem with a custom TiledCopy while + # preserving coalesced vector stores. + if tail_direct_store_enabled is None: + tail_direct_store_enabled = os.getenv("QUACK_HADAMARD_TAIL_DIRECT_STORE", "1") != "0" + + # Compute bit orders without tail direct store first, so feasibility can + # be checked on `bit_orders[-2]` (the order before the final stage, + # which is identical in both schedules). + self.bit_orders = self._compute_bit_orders(tail_direct_store=False) + use_tail_direct_store = ( + tail_direct_store_enabled + and self.tail_bit_shift < self.log_ept + and TailDirectStorePlan.is_feasible( + self.log_n, + self.bit_orders[-2], + self.tail_bit_shift, + self.log_copy_vecsize, + self.dtype.width, + ) + ) + if use_tail_direct_store: + self.bit_orders = self._compute_bit_orders(tail_direct_store=True) + self.tail_store: TailDirectStorePlan | None = TailDirectStorePlan.build( + self.log_n, + self.bit_orders[-2], + self.tail_bit_shift, + self.log_copy_vecsize, + self.dtype.width, + ) + else: + self.tail_store = None + + # If one transform fits in a warp, every smem exchange stays within that + # warp. The rows_per_block packing uses a power-of-two x dimension, so + # each transform is warp-contained and row-local smem slices do not + # require a CTA-wide barrier. + self.exchange_uses_syncwarp = self.threads_per_transform <= cute.arch.WARP_SIZE + + @property + def use_tail_direct_store(self) -> bool: + return self.tail_store is not None + + def __str__(self) -> str: + lines = [ + "Hadamard bit order: " + f"N={self.N} N_padded={self.N_padded} ept={self.ept} " + f"threads={self.threads_per_transform} copy_vecsize={self.copy_vecsize}" + ] + for stage, bit_shift in enumerate(self.bit_shifts): + before = self.bit_orders[stage] + after = self.bit_orders[stage + 1] + lines.append( + f" stage {stage} bit_shift={bit_shift} " + f"hadamard_bits={before.hadamard_bits(bit_shift)} " + f"before {before} -> after {after}" + ) + return "\n".join(lines) + + def __repr__(self) -> str: + return ( + "HadamardTransformPlan(" + f"N={self.N}, N_padded={self.N_padded}, ept={self.ept}, " + f"threads_per_transform={self.threads_per_transform}, " + f"rows_per_block={self.rows_per_block}, " + f"copy_vecsize={self.copy_vecsize}, " + f"use_tail_direct_store={self.use_tail_direct_store}, " + f"exchange_uses_syncwarp={self.exchange_uses_syncwarp})" + ) + + def _compute_bit_orders(self, tail_direct_store: bool) -> tuple[BitOrder, ...]: + """Compute ownership order before each stage, plus the final store order.""" + order = BitOrder.initial(self.log_n, self.log_threads_per_transform, self.log_copy_vecsize) + bit_orders = [order] + for stage, bit_shift in enumerate(self.bit_shifts): + assert bit_shift <= len(order.v) + is_tail_stage = stage == len(self.bit_shifts) - 1 + if tail_direct_store and is_tail_stage: + # The skipped-exchange store first permutes local values so the + # contiguous vector bits are leading: v=(low-vector, tail, rest). + order = order.tail_direct_store(bit_shift, self.log_copy_vecsize) + else: + tail_register_permute = is_tail_stage and bit_shift < len(order.v) + order = order.exchange(bit_shift, tail_register_permute=tail_register_permute) + bit_orders.append(order) + + consumed_bits = tuple( + bit + for order, bit_shift in zip(bit_orders, self.bit_shifts) + for bit in order.hadamard_bits(bit_shift) + ) + assert sorted(consumed_bits) == list(range(self.log_n)), ( + "Hadamard stage bit order should consume each logical bit exactly once" + ) + return tuple(bit_orders) + + +# ─── CuTe DSL helpers ─────────────────────────────────────────────────────── +# +# Kept as module-level free functions because CuTe DSL relies on +# `inspect.getsourcelines()`, which makes class-method `@cute.jit` definitions +# fragile (see AGENTS.md). + + +@cute.jit +def _tail_store_pred(coords: cute.Tensor, layout: cute.Layout, row_stride: Int32) -> cute.Tensor: + # `coords` has a vectorized first mode `((copy_elem, packet), ...)` while + # cute.copy wants one predicate per vector packet: `(packet, ...)`. + pred_shape = (coords.shape[0][1], *coords.shape[1:]) + pred = cute.make_rmem_tensor(pred_shape, Boolean) + pred_id = cute.make_identity_tensor(pred_shape) + flat_pred = cute.coalesce(pred) + flat_id = cute.coalesce(pred_id) + for i in cutlass.range_constexpr(cute.size(flat_pred)): + pred_coord = flat_id[i] + coord = ((0, pred_coord[0]), *pred_coord[1:]) + output_coord = coords[coord] + offset = cute.crd2idx(output_coord, layout) - output_coord[0] * row_stride + flat_pred[i] = offset < row_stride + return pred + + +@cute.jit +def _hadamard_thread_col(vals: cute.Tensor) -> None: # (N, col) + n = cute.size(vals, mode=[0]) + log_n = int(math.log2(n)) + assert n == 1 << log_n, "hadamard_thread_col requires power-of-two size" + for step in cutlass.range_constexpr(log_n): + stride = 1 << step + for j in cutlass.range(1 << (log_n - 1), unroll_full=True): + lo = j & (stride - 1) + idx = (j - lo) * 2 + lo + for col in cutlass.range(cute.size(vals, mode=[1]), unroll_full=True): + a, b = vals[idx, col], vals[idx + stride, col] + vals[idx, col] = a + b + vals[idx + stride, col] = a - b + + +@cute.jit +def _hadamard_warp( + vals: cute.Tensor, + tidx: Int32, + log_width: cutlass.Constexpr[int], +) -> None: + for step in cutlass.range_constexpr(log_width): + offset = const_expr(1 << step) + sign_bit = tidx & offset + sign = Float32(1.0) if sign_bit == 0 else Float32(-1.0) + for i in cutlass.range(cute.size(vals), unroll_full=True): + vals[i] = sign * vals[i] + cute.arch.shuffle_sync_bfly(vals[i], offset=offset) + + +@cute.jit +def exchange( + vals: cute.Tensor, # size ept + smem: cute.Tensor, # compact backing smem, size ept * threads_per_transform + tidx: Int32, + sync_fn, + bit_shift: int | None = None, # if None, shift all value bits, i.e. log2(ept) +) -> cute.Tensor: # (ept,) + ept = cute.size(vals) + log_ept = int(math.log2(ept)) + if const_expr(bit_shift is None): + bit_shift = log_ept + assert 0 <= bit_shift <= log_ept, "bit_shift must be in [0, log2(ept)]" + radix = const_expr(1 << bit_shift) + threads_per_transform = cute.size(smem) // ept + log_threads_per_transform = int(math.log2(threads_per_transform)) + # Store address bits are: old thread bits, then old value bits. + smem_store_base = cute.make_layout((threads_per_transform, ept)) + # Load address bits are: exchanged value bits, new thread bits, leftover value bits. + # Tail exchanges are register-permuted below, not loaded in a different smem order. + smem_load_base = cute.make_layout( + (threads_per_transform, (radix, ept // radix)), + stride=(radix, (1, threads_per_transform * radix)), + ) + + log_vecsize = min(2, bit_shift) # vectorize 4 elements during load, unless bit_shift < 2 + log_threads_in_phase = min(5, bit_shift, log_threads_per_transform) - log_vecsize + swizzle = cute.make_swizzle(log_threads_in_phase, log_vecsize, bit_shift - log_vecsize) + smem_store_layout = cute.make_composed_layout(swizzle, 0, smem_store_base) + smem_load_layout = cute.make_composed_layout(swizzle, 0, smem_load_base) + smem_store = cute.make_tensor(smem.iterator, smem_store_layout) + smem_load = cute.make_tensor(smem.iterator, smem_load_layout) + sX_store = smem_store[tidx, None] # (ept) + sX_load = smem_load[tidx, None] + cute.autovec_copy(cute.composition(vals, (ept,)), sX_store) + sync_fn() + vals_exchanged = copy_utils.load_s2r(sX_load) + # `load_s2r` preserves a layout like the swizzled smem source; the values are correct, but + # that swizzled rmem layout is not contiguous, e.g, ((1,(4,2,2,2,2))):((0,(1,4,8,16,32))). + # We call contiguous to materialize into a compact rmem layout. + vals_exchanged = vals_exchanged.contiguous() + if const_expr(bit_shift < log_ept): + # Tail exchanges leave the loaded value bits in `(exchanged, leftover)` order. + # Before the N=2048/ept=64 tail exchange, for example, ownership is + # t=(1,2,8,9,10), v=(3,4,5,6,7,0), bit_shift=5. + # The raw load produces + # t=(3,4,5,6,7), v=(1,2,8,9,10,0), + # but the original gmem store layout expects local v=(0,1,2,8,9,10). View the + # 64 registers as (32,2) = (exchanged bits, leftover bit), transpose to (2,32), + # and compact so the leftover bit 0 becomes the leading local value bit again. + vals_exchanged = cute.composition(vals_exchanged, cute.make_layout((radix, ept // radix))) + vals_exchanged = layout_utils.select(vals_exchanged, [1, 0]).contiguous() + return cute.composition(vals_exchanged, cute.make_layout((ept, 1))) + + +# ─── Kernel object ────────────────────────────────────────────────────────── + + +class HadamardTransform: + """Hadamard kernel wrapper. Holds a `HadamardTransformPlan` and launches it.""" + + def __init__( + self, + dtype: Type[cutlass.Numeric], + N: int, + ept: int | None = None, + rows_per_block: int | None = None, + persistent: bool = False, + ): + self.plan = HadamardTransformPlan(dtype, N, ept=ept, rows_per_block=rows_per_block) + self.persistent = persistent + self.async_load = persistent + self.use_shuffle = N <= 128 # slightly faster to use shuffles than smem for small N + # print(self.plan) # Uncomment to see the generated schedule and bit orders. + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO: cute.Tensor, + scale: Float32, + num_sms: Int32, + stream: cuda.CUstream, + ): + plan = self.plan + assert mX.element_type == plan.dtype + assert mO.element_type == plan.dtype + # tiled_copy lays out `num_threads = threads_per_transform * rows_per_block` + # threads as `(rows_per_block, threads_per_transform)` rows. The kernel + # launch keeps the transform thread index in x and the packed row in y, + # so the copy path reconstructs the equivalent linear thread id as + # `row_in_cta * threads_per_transform + tidx`. + tiled_copy = copy_utils.tiled_copy_2d( + plan.dtype, plan.threads_per_transform, plan.num_threads, plan.copy_vecsize + ) + if const_expr(plan.use_tail_direct_store): + tail = plan.tail_store + tail_store_layout = cute.make_layout(tail.gmem_shape, stride=tail.gmem_stride) + tail_store_thr_layout = cute.make_layout(tail.thr_shape, stride=tail.thr_stride) + tail_store_val_layout = cute.make_layout(tail.val_shape, stride=tail.val_stride) + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + plan.dtype, + num_bits_per_copy=plan.copy_vecsize * plan.dtype.width, + ) + tail_store_copy = cute.make_tiled_copy_tv( + copy_atom, tail_store_thr_layout, tail_store_val_layout + ) + else: + tail_store_layout = cute.make_layout(1) + tail_store_copy = tiled_copy + tiler_mn = (plan.rows_per_block, plan.N_padded) + # Each CTA processes `rows_per_block` rows; ceil-div handles a trailing partial CTA. + num_blocks_m = cute.ceil_div(mX.shape[0], plan.rows_per_block) + # Persistent mode assumes this kernel is used only for shapes with occupancy 1 CTA/SM. + # Therefore one persistent CTA per SM is enough; do not query full occupancy here. + grid_x = cutlass.min(num_blocks_m, num_sms) if const_expr(self.persistent) else num_blocks_m + self.kernel(mX, mO, scale, tiler_mn, tiled_copy, tail_store_copy, tail_store_layout).launch( + grid=[grid_x, 1, 1], + block=[plan.threads_per_transform, plan.rows_per_block, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO: cute.Tensor, + scale: Float32, + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + tail_store_copy: cute.TiledCopy, + tail_store_layout: cute.Layout, + ): + plan = self.plan + tidx, _, _ = cute.arch.thread_idx() + row_in_cta = 0 if const_expr(plan.rows_per_block == 1) else cute.arch.thread_idx()[1] + block_row, _, _ = cute.arch.block_idx() + + smem = cutlass.utils.SmemAllocator() + + shape = mX.shape + # Each CTA, per iteration, processes a (rows_per_block, N_padded) tile starting at + # row block_row * rows_per_block. + gX, gO = [cute.local_tile(mT, tiler_mn, (None, 0)) for mT in (mX, mO)] + thr_copy = tiled_copy.get_slice(row_in_cta * plan.threads_per_transform + tidx) + tXgX = thr_copy.partition_S(gX) + tXgO = thr_copy.partition_D(gO) + cX = cute.make_identity_tensor(tiler_mn) + tXcX_full = thr_copy.partition_S(cX) + tXrX = cute.make_rmem_tensor_like(tXgX[..., 0]) + tXrO = cute.make_rmem_tensor_like(tXgO[..., 0]) + tXsX = None + if const_expr(self.async_load): + sX = smem.allocate_tensor( + plan.dtype, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 + ) + tXsX = thr_copy.partition_D(sX) + + num_rows = cute.size(mX.shape[0]) + is_even_N = const_expr(shape[1] == tiler_mn[1]) + tXpX = None if is_even_N else copy_utils.predicate_k(tXcX_full, limit=shape[1]) + copy = partial(copy_utils.copy, pred=tXpX) + + s_exchange_layout = cute.make_ordered_layout( + (plan.threads_per_transform * plan.ept, plan.rows_per_block), order=(0, 1) + ) + s_exchange = None + if const_expr(not self.use_shuffle): + s_exchange = smem.allocate_tensor(Float32, s_exchange_layout, byte_alignment=16) + s_exchange = s_exchange[None, row_in_cta] + sync_fn = ( + cute.arch.sync_warp if const_expr(plan.exchange_uses_syncwarp) else cute.arch.barrier + ) + + nblocks_m = cute.ceil_div(num_rows, plan.rows_per_block) + num_cta = cute.arch.grid_dim()[0] if const_expr(self.persistent) else nblocks_m + if const_expr(self.async_load): + if block_row < nblocks_m: + copy(tXgX[..., block_row], tXsX, is_async=True) + cute.arch.cp_async_commit_group() + num_iter = ( + 1 if const_expr(not self.persistent) else cute.ceil_div(nblocks_m - block_row, num_cta) + ) + for i in cutlass.range(num_iter, unroll=2 if const_expr(self.persistent) else 1): + row_block = block_row + i * num_cta + row = row_block * plan.rows_per_block + row_in_cta + row_is_valid = True if const_expr(tiler_mn[0] == 1) else row < num_rows + if const_expr(self.async_load): + cute.arch.cp_async_wait_group(0) + cute.autovec_copy(tXsX, tXrX) + next_row_block = row_block + num_cta + if next_row_block < nblocks_m: + copy(tXgX[..., next_row_block], tXsX, is_async=True) + cute.arch.cp_async_commit_group() + else: + if const_expr(not is_even_N): + tXrX.fill(tXrX.element_type.zero) + if row_is_valid: + copy(tXgX[..., row_block], tXrX) + + x_flat = cute.composition(tXrX, cute.make_layout((cute.size(tXrX), 1))) + x_vals = x_flat.to(Float32) + + if const_expr(self.use_shuffle): + _hadamard_thread_col(x_vals) + _hadamard_warp(x_vals, tidx, log_width=plan.log_threads_per_transform) + else: + for stage in cutlass.range_constexpr(plan.stages): + bit_shift = const_expr(plan.bit_shifts[stage]) + radix = const_expr(1 << bit_shift) + x_vals = cute.composition(x_vals, cute.make_layout((radix, plan.ept // radix))) + _hadamard_thread_col(x_vals) + if const_expr(stage < plan.stages - 1 or not plan.use_tail_direct_store): + if const_expr(stage > 0 or self.persistent): + # Before reusing the exchange buffer, wait for the previous + # exchange loads to finish. Warp-local exchange plans use a + # warp fence; cross-warp plans use a CTA barrier. + sync_fn() + x_vals = exchange(x_vals, s_exchange, tidx, sync_fn, bit_shift=bit_shift) + + if const_expr(not self.use_shuffle and plan.use_tail_direct_store): + tail_size = const_expr(1 << plan.tail_bit_shift) + rest_size = const_expr(plan.ept // (tail_size * plan.copy_vecsize)) + x_store = cute.composition( + x_vals, cute.make_layout((tail_size, plan.copy_vecsize, rest_size)) + ) + x_store = layout_utils.select(x_store, [1, 0, 2]).contiguous() + gO_store = gO[row_in_cta, None, row_block] + thr_store = tail_store_copy.get_slice(tidx) + tOgO = thr_store.partition_D(cute.composition(gO_store, tail_store_layout)) + tOrO = cute.make_rmem_tensor_like(tOgO, tXrO.element_type) + cute.coalesce(tOrO).store( + (cute.coalesce(x_store).load() * scale).to(tOrO.element_type) + ) + if row_is_valid: + if const_expr(is_even_N): + cute.copy(tail_store_copy, tOrO, tOgO) + else: + cO_store = cute.make_identity_tensor(tail_store_layout.shape) + tOcO = thr_store.partition_D(cO_store) + tOpO = _tail_store_pred(tOcO, tail_store_layout, shape[1]) + cute.copy(tail_store_copy, tOrO, tOgO, pred=tOpO) + else: + o_flat = cute.composition(tXrO, cute.make_layout((cute.size(tXrO), 1))) + o_flat.store((x_vals.load() * scale).to(tXrO.element_type)) + if row_is_valid: + copy(tXrO, tXgO[..., row_block]) + + @staticmethod + @jit_cache + def compile(dtype, N, persistent: bool = False): + batch_sym = cute.sym_int() + div = math.gcd(N, 128 // dtype.width) + x_cute = fake_tensor(dtype, (batch_sym, N), div) + out_cute = fake_tensor(dtype, (batch_sym, N), div) + return cute.compile( + HadamardTransform(dtype, N, persistent=persistent), + x_cute, + out_cute, + Float32(0.0), + 0, # num_sms, just for compilation + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +# ─── Autograd, public API ────────────────────────────────────────── + + +@cute_op( + add_op_namespace_prefix("_hadamard_transform_fwd"), + mutates_args={"out"}, + device_types="cuda", +) +def _hadamard_transform_fwd(x: Tensor, out: Tensor, scale: float) -> None: + """Custom-op binding: dispatch to the cached compiled kernel.""" + assert x.dim() == 2, "Input must be 2D" + assert out.shape == x.shape, "Output shape must match input" + assert x.dtype == out.dtype, "Output dtype must match input dtype" + assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype" + if x.numel() == 0: + return + N = x.size(1) + assert 2 <= N <= _MAX_N, f"Hadamard transform supports last dimension in [2, {_MAX_N}]" + dtype = torch2cute_dtype_map[x.dtype] + persistent = _should_use_persistent(x, N) + compiled = HadamardTransform.compile(dtype, N, persistent) + num_sms = _get_num_sms(x.device) if persistent else 0 + compiled(x, out, scale, num_sms) + + +def hadamard_transform_fwd(x: Tensor, scale: float = 1.0) -> Tensor: + """Forward-only entry point: normalize layout, dispatch to the custom op.""" + assert x.dim() >= 1, "Input must have at least one dimension" + x = _ensure_last_dim_contiguous(x) + N = x.size(-1) + assert 1 <= N <= _MAX_N, f"Hadamard transform supports last dimension in [1, {_MAX_N}]" + if x.numel() == 0: + return torch.empty_like(x) + if N == 1: + return x * float(scale) + x_2d = x.reshape(-1, N) + out_2d = torch.empty_like(x_2d) + _hadamard_transform_fwd(x_2d, out_2d, float(scale)) + return out_2d.reshape(x.shape) + + +def hadamard_transform_ref(x: Tensor, scale: float = 1.0) -> Tensor: + """PyTorch reference with the same zero-padding convention as fast-hadamard-transform.""" + assert x.dim() >= 1, "Input must have at least one dimension" + N = x.size(-1) + assert 1 <= N <= _MAX_N, f"Hadamard transform supports last dimension in [1, {_MAX_N}]" + N_padded = _next_power_of_2(N) + y = x.float().reshape(-1, N) + if N_padded != N: + y = torch.nn.functional.pad(y, (0, N_padded - N)) + h = 1 + while h < N_padded: + y = y.reshape(-1, N_padded // (2 * h), 2, h) + y0 = y[:, :, 0, :] + y1 = y[:, :, 1, :] + y = torch.stack((y0 + y1, y0 - y1), dim=2).reshape(-1, N_padded) + h *= 2 + return (y[:, :N] * scale).reshape(x.shape).to(x.dtype) + + +class HadamardTransformFunction(torch.autograd.Function): + """Autograd wrapper. The Hadamard transform is self-adjoint, so backward = forward.""" + + @staticmethod + def forward(ctx, x: Tensor, scale: float = 1.0): + ctx.scale = float(scale) + return hadamard_transform_fwd(x, ctx.scale) + + @staticmethod + def backward(ctx, dout: Tensor): + return hadamard_transform_fwd(dout, ctx.scale), None + + +def hadamard_transform(x: Tensor, scale: float = 1.0) -> Tensor: + """Apply a Sylvester Hadamard transform along the last dimension.""" + return HadamardTransformFunction.apply(x, scale) diff --git a/build/torch-cuda/quack/utils.py b/build/torch-cuda/quack/utils.py index 7039d8aeeae96ec0075f6549d9dd4b590702364b..7d2bb9c89b258194b860bd76cf387415b6303ed7 100644 --- a/build/torch-cuda/quack/utils.py +++ b/build/torch-cuda/quack/utils.py @@ -1,14 +1,13 @@ # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. import math -from typing import Optional, Tuple, Union +from typing import Optional, Tuple import cutlass import cutlass.cute as cute from cutlass import Float32, Int32, const_expr -from cutlass._mlir.dialects import arith as _arith -from cutlass._mlir.dialects import llvm, nvvm, vector +from cutlass._mlir.dialects import llvm, vector from cutlass.cutlass_dsl import T, dsl_user_op @@ -30,17 +29,8 @@ def set_block_rank( smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None ) -> Int32: """Map the given smem pointer to the address at another CTA rank in the cluster.""" - smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() - return Int32( - llvm.inline_asm( - T.i32(), - [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()], - "mapa.shared::cluster.u32 $0, $1, $2;", - "=r,r,r", - has_side_effects=False, - is_align_stack=False, - ) - ) + dsmem_ptr = cute.arch.map_dsmem_ptr(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip) + return Int32(dsmem_ptr.toint(loc=loc, ip=ip)) @dsl_user_op @@ -53,24 +43,18 @@ def store_shared_remote( loc=None, ip=None, ) -> None: - remote_smem_ptr_i32 = set_block_rank( - smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip - ).ir_value() - remote_mbar_ptr_i32 = set_block_rank( - mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip - ).ir_value() + remote_smem_ptr_i32 = set_block_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip) + remote_mbar_ptr_i32 = set_block_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip) if const_expr(isinstance(val, float)): val = Float32(val) assert isinstance(val, (Float32, Int32, cutlass.Int64)), "val must be Float32, Int32, or Int64" suffix = {Float32: "f32", Int32: "s32", cutlass.Int64: "s64"}[type(val)] - constraint = {Float32: "f", Int32: "r", cutlass.Int64: "l"}[type(val)] - llvm.inline_asm( - None, - [remote_smem_ptr_i32, val.ir_value(loc=loc, ip=ip), remote_mbar_ptr_i32], - f"st.async.shared::cluster.mbarrier::complete_tx::bytes.{suffix} [$0], $1, [$2];", - f"r,{constraint},r", - has_side_effects=True, - is_align_stack=False, + cute.arch.inline_ptx( + f"st.async.shared::cluster.mbarrier::complete_tx::bytes.{suffix} " + "[{$r0}], {$r1}, [{$r2}];", + read_only_args=[remote_smem_ptr_i32, val, remote_mbar_ptr_i32], + loc=loc, + ip=ip, ) @@ -87,87 +71,31 @@ def store_shared_remote_x4( loc=None, ip=None, ) -> None: - remote_smem_ptr_i32 = set_block_rank( - smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip - ).ir_value() - remote_mbar_ptr_i32 = set_block_rank( - mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip - ).ir_value() + remote_smem_ptr_i32 = set_block_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip) + remote_mbar_ptr_i32 = set_block_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip) assert isinstance(val0, (Float32, Int32)), "val must be Float32, or Int32" dtype = Float32 if isinstance(val0, Float32) else Int32 suffix = {Float32: "f32", Int32: "s32"}[dtype] - constraint = {Float32: "f", Int32: "r"}[dtype] - llvm.inline_asm( - None, - [ - remote_smem_ptr_i32, - remote_mbar_ptr_i32, - dtype(val0).ir_value(loc=loc, ip=ip), - dtype(val1).ir_value(loc=loc, ip=ip), - dtype(val2).ir_value(loc=loc, ip=ip), - dtype(val3).ir_value(loc=loc, ip=ip), - ], + cute.arch.inline_ptx( "{\n\t" f".reg .v4 .{suffix} abcd;\n\t" - f"mov.{suffix} abcd.x, $2;\n\t" - f"mov.{suffix} abcd.y, $3;\n\t" - f"mov.{suffix} abcd.z, $4;\n\t" - f"mov.{suffix} abcd.w, $5;\n\t" - f"st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.{suffix} [$0], abcd, [$1];\n\t" + f"mov.{suffix} abcd.x, {{$r2}};\n\t" + f"mov.{suffix} abcd.y, {{$r3}};\n\t" + f"mov.{suffix} abcd.z, {{$r4}};\n\t" + f"mov.{suffix} abcd.w, {{$r5}};\n\t" + f"st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.{suffix} " + "[{$r0}], abcd, [{$r1}];\n\t" "}\n", - f"r,r,{constraint},{constraint},{constraint},{constraint}", - has_side_effects=True, - is_align_stack=False, - ) - - -@dsl_user_op -def fmin(a: Union[float, Float32], b: Union[float, Float32], *, loc=None, ip=None) -> Float32: - if cutlass.const_expr(cutlass.CUDA_VERSION.major) == 12: - return Float32( - nvvm.fmin( - T.f32(), - Float32(a).ir_value(loc=loc, ip=ip), - Float32(b).ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - ) - return Float32( - nvvm.fmin( - Float32(a).ir_value(loc=loc, ip=ip), - Float32(b).ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def sqrt(a: float | Float32, *, loc=None, ip=None) -> Float32: - return Float32( - llvm.inline_asm( - T.f32(), - [Float32(a).ir_value(loc=loc, ip=ip)], - "sqrt.approx.f32 $0, $1;", - "=f,f", - has_side_effects=False, - is_align_stack=False, - ) - ) - - -@dsl_user_op -def ceil(a: float | Float32, *, loc=None, ip=None) -> Int32: - return Int32( - llvm.inline_asm( - T.i32(), - [Float32(a).ir_value(loc=loc, ip=ip)], - "cvt.rpi.ftz.s32.f32 $0, $1;", - "=r,f", - has_side_effects=False, - is_align_stack=False, - ) + read_only_args=[ + remote_smem_ptr_i32, + remote_mbar_ptr_i32, + dtype(val0), + dtype(val1), + dtype(val2), + dtype(val3), + ], + loc=loc, + ip=ip, ) @@ -209,10 +137,11 @@ def make_vector(elem_type, *values, loc=None, ip=None): vec_ty = ir.VectorType.get([n], mlir_ty) vec = llvm.mlir_undef(vec_ty, loc=loc, ip=ip) for i, v in enumerate(values): - vec = vector.insertelement( + vec = vector.insert( elem_type(v).ir_value(loc=loc, ip=ip), vec, - position=_arith.constant(T.i32(), i, loc=loc, ip=ip), + dynamic_position=[], + static_position=[i], loc=loc, ip=ip, ) @@ -258,62 +187,15 @@ def warp_prefix_sum(val: Int32, lane: Optional[Int32] = None) -> Int32: @dsl_user_op -def atomic_inc_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> Int32: - from cutlass import CUDA_VERSION - - # * NVVM call based on nvvm version - if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9: - # Old API: requires explicit result type as first positional argument - return nvvm.atomicrmw( - res=T.i32(), op=nvvm.AtomicOpKind.INC, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value() - ) - else: - # New API: infers result type automatically - return nvvm.atomicrmw( - op=nvvm.AtomicOpKind.INC, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value() - ) - - -@dsl_user_op -def atomic_add_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> Int32: - from cutlass import CUDA_VERSION - - # * NVVM call based on nvvm version - if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9: - # Old API: requires explicit result type as first positional argument - return nvvm.atomicrmw( - res=T.i32(), op=nvvm.AtomicOpKind.ADD, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value() - ) - else: - # New API: infers result type automatically - return nvvm.atomicrmw( - op=nvvm.AtomicOpKind.ADD, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value() - ) - - -@dsl_user_op -def issue_clc_query_nomulticast( - mbar_ptr: cute.Pointer, - clc_response_ptr: cute.Pointer, - loc=None, - ip=None, -) -> None: - """ - The clusterlaunchcontrol.try_cancel instruction requests atomically cancelling the launch - of a cluster that has not started running yet. It asynchronously writes an opaque response - to shared memory indicating whether the operation succeeded or failed. On success, the - opaque response contains the ctaid of the first CTA of the canceled cluster. - - :param mbar_ptr: A pointer to the mbarrier address in SMEM - :type mbar_ptr: Pointer - :param clc_response_ptr: A pointer to the cluster launch control response address in SMEM - :type clc_response_ptr: Pointer - """ - mbar_llvm_ptr = mbar_ptr.llvm_ptr - clc_response_llvm_ptr = clc_response_ptr.llvm_ptr - nvvm.clusterlaunchcontrol_try_cancel( - clc_response_llvm_ptr, - mbar_llvm_ptr, - loc=loc, - ip=ip, +def domain_offset_aligned( + coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None +) -> cute.Tensor: + assert isinstance(tensor.iterator, cute.Pointer) + # We assume that applying the offset does not change the pointer alignment + new_ptr = cute.make_ptr( + tensor.element_type, + elem_pointer(tensor, coord).toint(), + tensor.memspace, + assumed_align=tensor.iterator.alignment, ) + return cute.make_tensor(new_ptr, tensor.layout) diff --git a/build/torch-cuda/quack/varlen_utils.py b/build/torch-cuda/quack/varlen_utils.py index e8a45aa47434992668b42639ede52f3ba4681f69..032e4f22663e99a8024679139d0f98e34e56bdd3 100644 --- a/build/torch-cuda/quack/varlen_utils.py +++ b/build/torch-cuda/quack/varlen_utils.py @@ -123,7 +123,7 @@ class VarlenManager: return mAIdx_mk def offset_batch_SFA(self, mSFA_mkl: cute.Tensor, batch_idx: Int32) -> cute.Tensor: - """Offset SFA by padded per-expert offset (dQaccum-style). + """Offset SFA to this batch's tile-aligned region of the padded SF buffer. The padded offset, in tile units (128 source-M or source-K per tile), is simply `cu_seqlens[b] // 128 + b`. (Algebraically identical to diff --git a/build/torch-cuda/sonic_moe/__init__.py b/build/torch-cuda/sonic_moe/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch-cuda/sonic_moe/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))