diff --git a/bench/gemm/gemm_ops.py b/bench/gemm/gemm_ops.py index bace7c61..b731c77a 100644 --- a/bench/gemm/gemm_ops.py +++ b/bench/gemm/gemm_ops.py @@ -10,8 +10,15 @@ import torch from mslk.bench.common.utils import BenchOptions, do_bench -from mslk.flydsl.common import is_flydsl_version_at_least +from mslk.flydsl.common import is_flydsl_available from mslk.gemm.triton.fp8_gemm import matmul_fp8_block, matmul_fp8_row, to_mxfp8 + +if is_flydsl_available(): + from mslk.gemm.flydsl.preshuffle_gemm import ( + flydsl_preshuffle, + flydsl_preshuffle_batched_gemm, + flydsl_preshuffle_gemm, + ) from mslk.gemm.triton.grouped_gemm import grouped_gemm, grouped_gemm_fp8_rowwise from mslk.quantize.shuffle import ( ck_preshuffle, @@ -694,10 +701,6 @@ def compute_dtype(self) -> ComputeDtype: return ComputeDtype.FP8 -if is_flydsl_version_at_least(): - from mslk.gemm.flydsl import flydsl_preshuffle, flydsl_preshuffle_gemm - - @register_gemm_op class FP8RowwisePreshuffleFlyDSL(FP8Rowwise): """ @@ -722,7 +725,7 @@ def supported_accelerators(self) -> set[Accelerator]: def supported(self) -> bool: if not super().supported: return False - return is_flydsl_version_at_least() + return is_flydsl_available() @property def supported_gemm_types(self) -> set[GemmType]: @@ -1266,6 +1269,35 @@ def compute_dtype(self) -> ComputeDtype: return ComputeDtype.FP8 +@register_gemm_op +class FP8RowwiseBatchedPreshuffleFlyDSL(FP8RowwiseBatched): + """ + FP8 batched matmul with rowwise scaling and FlyDSL preshuffle kernel (gfx950). + """ + + def quantize(self, x, w): + xq, wq, x_scale, w_scale = super().quantize(x, w) + wq_shuf = torch.stack([flydsl_preshuffle(wq[i]) for i in range(wq.shape[0])]) + return xq, wq_shuf, x_scale, w_scale + + def compute(self, xq, wq, x_scale, w_scale): + return flydsl_preshuffle_batched_gemm(xq, wq, x_scale, w_scale) + + @property + def supported_accelerators(self) -> set[Accelerator]: + return {Accelerator.AMD_GFX950} + + @property + def supported(self) -> bool: + if get_current_accelerator() not in self.supported_accelerators: + return False + return is_flydsl_available() + + @property + def supported_gemm_types(self) -> set[GemmType]: + return {GemmType.GROUPED} + + # This kernel is broken and causes GPU to lock up, needs some investigation # @register_gemm_op class TritonFP8Rowwise(GemmOpBase): diff --git a/mslk/gemm/flydsl/__init__.py b/mslk/gemm/flydsl/__init__.py index 52f729af..581f84e4 100644 --- a/mslk/gemm/flydsl/__init__.py +++ b/mslk/gemm/flydsl/__init__.py @@ -5,8 +5,3 @@ # LICENSE file in the root directory of this source tree. # pyre-strict - -from mslk.gemm.flydsl.preshuffle_gemm import ( # noqa: F401 - flydsl_preshuffle, - flydsl_preshuffle_gemm, -) diff --git a/mslk/gemm/flydsl/_kernels/preshuffle_gemm.py b/mslk/gemm/flydsl/_kernels/preshuffle_gemm.py index b5d647a0..e2b92f07 100644 --- a/mslk/gemm/flydsl/_kernels/preshuffle_gemm.py +++ b/mslk/gemm/flydsl/_kernels/preshuffle_gemm.py @@ -140,6 +140,7 @@ def compile_preshuffle_gemm_a8( dvmem_preload: int = -1, epilogue: str = "none", # "none", "bias", "bias_relu", "bias_silu", "bias_gelu" xcd_swizzle: int = 0, + batched: bool = False, ): """Compile the preshuffle GEMM kernel using the @flyc.kernel API. @@ -199,6 +200,8 @@ def compile_preshuffle_gemm_a8( KERNEL_NAME += f"_ep_{epilogue}" if xcd_swizzle > 0: KERNEL_NAME += f"_xcd{xcd_swizzle}" + if batched: + KERNEL_NAME += "_batched" tile_k_bytes = int(tile_k) * int(elem_bytes) @@ -395,6 +398,8 @@ def kernel_gemm( tx = gpu.thread_id("x") bx = gpu.block_id("x") by = gpu.block_id("y") + if const_expr(batched): + bz = gpu.block_id("z") bx, by = xcd_remap_bx_by( bx, @@ -460,17 +465,33 @@ def kernel_gemm( _a_nrec = fx.Int64(c_m * (K * elem_bytes // a_elem_vec_pack)) _c_nrec = fx.Int64(c_m * c_n * 2) - def _ptr_buffer_resource(ptr, num_records_bytes=None): + # Grid-Z batch offset: each pointer advances by bz * batch_stride_bytes + _off_a = None + _off_b = None + _off_c = None + _off_sa = None + _off_sb = None + if const_expr(batched): + _bz_i64 = fx.Int64(bz) + _off_a = _bz_i64 * _a_nrec + _off_b = _bz_i64 * fx.Int64(fx.Index(N * K * elem_bytes // b_elem_vec_pack)) + _off_c = _bz_i64 * _c_nrec + _off_sa = _bz_i64 * fx.Int64(c_m * fx.Index(4)) + _off_sb = _bz_i64 * fx.Int64(fx.Index(N * 4)) + + def _ptr_buffer_resource(ptr, num_records_bytes=None, byte_offset=None): addr = fx.ptrtoint(ptr) addr_i64 = fx.arith.index_cast(T.i64, addr) + if byte_offset is not None: + addr_i64 = addr_i64 + byte_offset if num_records_bytes is None: return buffer_ops.create_buffer_resource_from_addr(addr_i64) return buffer_ops.create_buffer_resource_from_addr( addr_i64, num_records_bytes=num_records_bytes ) - a_rsrc = _ptr_buffer_resource(arg_a, _a_nrec) - c_rsrc = _ptr_buffer_resource(arg_c, _c_nrec) + a_rsrc = _ptr_buffer_resource(arg_a, _a_nrec, byte_offset=_off_a) + c_rsrc = _ptr_buffer_resource(arg_c, _c_nrec, byte_offset=_off_c) _needs_per_token_scale = not is_f16_or_bf16 and not is_fp4 scale_a_rsrc = None if const_expr(not is_f16_or_bf16): @@ -482,7 +503,9 @@ def _ptr_buffer_resource(ptr, num_records_bytes=None): ) else: _scale_a_nrec = fx.Int64(c_m * fx.Index(4)) - scale_a_rsrc = _ptr_buffer_resource(arg_scale_a, _scale_a_nrec) + scale_a_rsrc = _ptr_buffer_resource( + arg_scale_a, _scale_a_nrec, byte_offset=_off_sa + ) # ---- Bias buffer resource (for fused epilogue) ---- # Use max_size=True so the buffer descriptor's size is taken from the @@ -491,8 +514,12 @@ def _ptr_buffer_resource(ptr, num_records_bytes=None): bias_rsrc = None if const_expr(_has_bias): bias_rsrc = _ptr_buffer_resource(arg_bias) - b_rsrc = _ptr_buffer_resource(arg_b) - scale_b_rsrc = None if (is_f16_or_bf16) else _ptr_buffer_resource(arg_scale_b) + b_rsrc = _ptr_buffer_resource(arg_b, byte_offset=_off_b) + scale_b_rsrc = ( + None + if (is_f16_or_bf16) + else _ptr_buffer_resource(arg_scale_b, byte_offset=_off_sb) + ) bx_m = bx * tile_m by_n = by * tile_n @@ -2147,49 +2174,99 @@ def prefetch_a0_pack( store_output(final_accs, scales) # ── Host launcher ────────────────────────────────────────────────────── - @flyc.jit - def launch_gemm( - arg_c: fx.Pointer, - arg_a: fx.Pointer, - arg_b: fx.Pointer, - arg_scale_a: fx.Pointer, - arg_scale_b: fx.Pointer, - arg_bias: fx.Pointer, - i32_m: fx.Int32, - i32_n: fx.Int32, - stream: fx.Stream, - ): - allocator_pong.finalized = False - allocator_ping.finalized = False - ctx = CompilationContext.get_current() - from flydsl._mlir import ir + if batched: + + @flyc.jit + def launch_gemm( + arg_c: fx.Pointer, + arg_a: fx.Pointer, + arg_b: fx.Pointer, + arg_scale_a: fx.Pointer, + arg_scale_b: fx.Pointer, + arg_bias: fx.Pointer, + i32_m: fx.Int32, + i32_n: fx.Int32, + i32_b: fx.Int32, + stream: fx.Stream, + ): + allocator_pong.finalized = False + allocator_ping.finalized = False + ctx = CompilationContext.get_current() + from flydsl._mlir import ir - with ir.InsertionPoint(ctx.gpu_module_body): - allocator_pong.finalize() - allocator_ping.finalize() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator_pong.finalize() + allocator_ping.finalize() - gx = (i32_m + (tile_m - 1)) // tile_m - gy = i32_n // tile_n + gx = (i32_m + (tile_m - 1)) // tile_m + gy = i32_n // tile_n - kernel_gemm._func.__name__ = KERNEL_NAME - launcher = kernel_gemm( - arg_c, arg_a, arg_b, arg_scale_a, arg_scale_b, arg_bias, i32_m, i32_n - ) - if const_expr(waves_per_eu is not None): - _wpe = int(waves_per_eu) - if const_expr(_wpe >= 1): - for op in ctx.gpu_module_body.operations: - if const_expr( - hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func" - ): - op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( - fx.Int32.ir_type, _wpe - ) - launcher.launch( - grid=(gx, gy, 1), - block=(256, 1, 1), - stream=stream, - ) + kernel_gemm._func.__name__ = KERNEL_NAME + launcher = kernel_gemm( + arg_c, arg_a, arg_b, arg_scale_a, arg_scale_b, arg_bias, i32_m, i32_n + ) + if const_expr(waves_per_eu is not None): + _wpe = int(waves_per_eu) + if const_expr(_wpe >= 1): + for op in ctx.gpu_module_body.operations: + if const_expr( + hasattr(op, "attributes") + and op.OPERATION_NAME == "gpu.func" + ): + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + fx.Int32.ir_type, _wpe + ) + launcher.launch( + grid=(gx, gy, i32_b), + block=(256, 1, 1), + stream=stream, + ) + else: + + @flyc.jit + def launch_gemm( + arg_c: fx.Pointer, + arg_a: fx.Pointer, + arg_b: fx.Pointer, + arg_scale_a: fx.Pointer, + arg_scale_b: fx.Pointer, + arg_bias: fx.Pointer, + i32_m: fx.Int32, + i32_n: fx.Int32, + stream: fx.Stream, + ): + allocator_pong.finalized = False + allocator_ping.finalized = False + ctx = CompilationContext.get_current() + from flydsl._mlir import ir + + with ir.InsertionPoint(ctx.gpu_module_body): + allocator_pong.finalize() + allocator_ping.finalize() + + gx = (i32_m + (tile_m - 1)) // tile_m + gy = i32_n // tile_n + + kernel_gemm._func.__name__ = KERNEL_NAME + launcher = kernel_gemm( + arg_c, arg_a, arg_b, arg_scale_a, arg_scale_b, arg_bias, i32_m, i32_n + ) + if const_expr(waves_per_eu is not None): + _wpe = int(waves_per_eu) + if const_expr(_wpe >= 1): + for op in ctx.gpu_module_body.operations: + if const_expr( + hasattr(op, "attributes") + and op.OPERATION_NAME == "gpu.func" + ): + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + fx.Int32.ir_type, _wpe + ) + launcher.launch( + grid=(gx, gy, 1), + block=(256, 1, 1), + stream=stream, + ) return launch_gemm diff --git a/mslk/gemm/flydsl/preshuffle_gemm.py b/mslk/gemm/flydsl/preshuffle_gemm.py index 72248751..e9ae94d0 100644 --- a/mslk/gemm/flydsl/preshuffle_gemm.py +++ b/mslk/gemm/flydsl/preshuffle_gemm.py @@ -8,65 +8,88 @@ """FlyDSL preshuffle GEMM for FP8 rowwise scaling (gfx950). -Public API: - flydsl_preshuffle(src) -- shuffle weights for the FlyDSL layout - flydsl_preshuffle_gemm(XQ, WQ, x_scale, w_scale, ...) -- run the GEMM +Provides single and batched FP8 preshuffle GEMM via FlyDSL. Also registers +as the ROCm implementation of the ``mslk`` rowwise FP8 ops on gfx950. """ from dataclasses import dataclass from typing import Optional import torch -from mslk.flydsl.common import is_flydsl_available, require_flydsl from torch import Tensor # --------------------------------------------------------------------------- -# Default tile configurations for gfx950 preshuffle GEMM +# Tile configuration for the preshuffle GEMM kernel # --------------------------------------------------------------------------- @dataclass class KernelConfig: - """Tile configuration for the FlyDSL preshuffle GEMM kernel.""" + """Tile and launch parameters for one FlyDSL preshuffle GEMM variant.""" - tile_m: int - tile_n: int - tile_k: int - lds_stage: int - use_cshuffle_epilog: int = 0 - use_async_copy: int = 0 - waves_per_eu: int = 0 - xcd_swizzle: int = 0 + tile_m: int # M-dimension tile size (rows of activation per workgroup) + tile_n: int # N-dimension tile size (columns of weight per workgroup) + tile_k: int # K-dimension tile size (reduction loop step) + lds_stage: int # number of LDS pipeline stages (software pipelining depth) + use_cshuffle_epilog: int = 0 # 1 to use cross-lane shuffle in the epilogue + use_async_copy: int = 0 # 1 to use async global→LDS copy instructions + waves_per_eu: int = 0 # occupancy hint; 0 = let the compiler decide + xcd_swizzle: int = 0 # XCD (cross-chiplet die) swizzle pattern index -# Default configs for gfx950 heuristic selection. DEFAULT_CONFIGS_GFX950: dict[int, KernelConfig] = { - -1: KernelConfig(128, 256, 256, 2, 0, 0, 2, 0), - -2: KernelConfig(16, 64, 512, 2, 0, 0, 2, 0), - -3: KernelConfig(32, 64, 512, 2, 0, 0, 2, 0), - -4: KernelConfig(128, 128, 128, 2, 0, 0, 2, 0), + -1: KernelConfig(64, 128, 256, 2, xcd_swizzle=1, waves_per_eu=2), + -2: KernelConfig(16, 64, 512, 2, waves_per_eu=1), + -3: KernelConfig(32, 64, 512, 2, xcd_swizzle=1), + -4: KernelConfig(64, 64, 128, 2, xcd_swizzle=1), + -5: KernelConfig(128, 256, 128, 2, xcd_swizzle=1, waves_per_eu=2), + -6: KernelConfig(64, 128, 128, 2, xcd_swizzle=1), } - -def select_default_config(m: int, n: int, k: int) -> KernelConfig: +# Profile-guided overrides: (m_upper_bound, n, k) -> KernelConfig +# Entries are checked in order; first match where m <= m_upper_bound wins. +_SHAPE_OVERRIDES_GFX950: list[tuple[int, int, int, KernelConfig]] = [ + # N=1280, K=8192: sweep-tuned per M range + (1, 1280, 8192, KernelConfig(64, 128, 256, 2, xcd_swizzle=1, waves_per_eu=2)), + (64, 1280, 8192, KernelConfig(32, 128, 256, 2, xcd_swizzle=4, waves_per_eu=2)), + (256, 1280, 8192, KernelConfig(64, 64, 128, 2, xcd_swizzle=1)), + (512, 1280, 8192, KernelConfig(64, 256, 128, 2, xcd_swizzle=1)), + (2048, 1280, 8192, KernelConfig(128, 256, 128, 2, xcd_swizzle=1, waves_per_eu=2)), + (8192, 1280, 8192, KernelConfig(128, 256, 128, 2, waves_per_eu=2)), + # N=8192, K=1024: sweep-tuned per M range + (1, 8192, 1024, KernelConfig(16, 64, 512, 2, waves_per_eu=1)), + (256, 8192, 1024, KernelConfig(128, 256, 128, 2, xcd_swizzle=1, waves_per_eu=2)), + (8192, 8192, 1024, KernelConfig(128, 256, 128, 2, xcd_swizzle=1, waves_per_eu=2)), +] + + +def select_default_config(m: int, n: int, k: int, batch: int = 1) -> KernelConfig: """Select a default FlyDSL tile config based on shape heuristics.""" - fits = [ - c - for c in DEFAULT_CONFIGS_GFX950.values() - if n % c.tile_n == 0 and k % c.tile_k == 0 - ] + for m_upper, n_val, k_val, cfg in _SHAPE_OVERRIDES_GFX950: + if n == n_val and k == k_val and m <= m_upper: + return cfg + + configs = DEFAULT_CONFIGS_GFX950 + fits = [c for c in configs.values() if n % c.tile_n == 0 and k % c.tile_k == 0] if not fits: raise RuntimeError( f"No FlyDSL preshuffle config fits shape ({m}, {n}, {k}). " f"N must be divisible by tile_n and K by tile_k." ) want_tm = min(256, max(16, 1 << (m - 1).bit_length())) if m > 0 else 16 - return min(fits, key=lambda c: (abs(c.tile_m - want_tm), -c.tile_n, -c.tile_k)) + + def _sort_key(c: KernelConfig) -> tuple: + m_tiles = max(1, -(-m // c.tile_m)) + n_tiles = n // c.tile_n + low_occupancy = m_tiles * n_tiles * batch < 64 + return (low_occupancy, abs(c.tile_m - want_tm), -c.tile_n, -c.tile_k) + + return min(fits, key=_sort_key) # --------------------------------------------------------------------------- -# Weight shuffle +# Weight preshuffle # --------------------------------------------------------------------------- @@ -110,6 +133,8 @@ def _get_compile_fn(): # type: ignore[return] if _import_done: return _compile_fn _import_done = True + from mslk.flydsl.common import is_flydsl_available + if not is_flydsl_available(): return None try: @@ -121,8 +146,12 @@ def _get_compile_fn(): # type: ignore[return] return _compile_fn +def _as_i8(t: Tensor) -> Tensor: + return t.view(torch.int8) if "float8" in str(t.dtype) else t + + # --------------------------------------------------------------------------- -# GEMM entry point +# Single GEMM # --------------------------------------------------------------------------- @@ -156,6 +185,8 @@ def flydsl_preshuffle_gemm( Returns: Output tensor (M, N) in ``dtype``. """ + from mslk.flydsl.common import require_flydsl + require_flydsl() compile_fn = _get_compile_fn() @@ -210,14 +241,12 @@ def flydsl_preshuffle_gemm( ) import flydsl.expr as fx # pyre-ignore[21] - from mslk.gemm.flydsl._kernels.tensor_shim import _run_compiled, ptr_arg - - def _as_i8(t: Tensor) -> Tensor: - return t.view(torch.int8) if "float8" in str(t.dtype) else t + from mslk.flydsl.jit import run_compiled + from mslk.gemm.flydsl._kernels.tensor_shim import ptr_arg out_contig = out.contiguous() - _dummy_bias = torch.empty(0, dtype=out.dtype, device=out.device) - _run_compiled( + _dummy_bias = torch.empty(1, dtype=out.dtype, device=out.device) + run_compiled( exe, ptr_arg(out_contig.view(-1)), ptr_arg(_as_i8(XQ.contiguous()).view(-1)), @@ -236,15 +265,131 @@ def _as_i8(t: Tensor) -> Tensor: # --------------------------------------------------------------------------- -# Op registration (replaces CK on gfx950) +# Batched GEMM # --------------------------------------------------------------------------- + +def flydsl_preshuffle_batched_gemm( + XQ: Tensor, + WQ: Tensor, + x_scale: Tensor, + w_scale: Tensor, + out: Optional[Tensor] = None, + tile_m: Optional[int] = None, + tile_n: Optional[int] = None, + tile_k: Optional[int] = None, + lds_stage: int = 2, + use_cshuffle_epilog: int = 0, + use_async_copy: int = 0, + waves_per_eu: int = 0, + xcd_swizzle: int = 0, + dtype: torch.dtype = torch.bfloat16, +) -> Tensor: + """Batched FP8 preshuffle GEMM via Grid-Z batching. + + Args: + XQ: FP8 activation tensor (B, M, K). + WQ: FP8 weight tensor (B, N, K), pre-shuffled via ``flydsl_preshuffle``. + x_scale: Per-token activation scale (B, M) or (B, M, 1), float32. + w_scale: Per-channel weight scale (B, N) or (B, 1, N), float32. + out: Optional pre-allocated output tensor (B, M, N). + dtype: Output dtype (bfloat16 or float16). + """ + from mslk.flydsl.common import require_flydsl + + require_flydsl() + + assert XQ.dim() == 3, f"Expected 3D XQ, got {XQ.dim()}D" + assert WQ.dim() == 3, f"Expected 3D WQ, got {WQ.dim()}D" + B, M, K = XQ.shape + N = WQ.shape[1] + + if out is None: + out = torch.empty(B, M, N, dtype=dtype, device=XQ.device) + + compile_fn = _get_compile_fn() + if compile_fn is None: + raise RuntimeError("FlyDSL preshuffle kernel compiler not available") + + if tile_m is None or tile_n is None or tile_k is None: + cfg = select_default_config(M, N, K, batch=B) + tile_m = tile_m or cfg.tile_m + tile_n = tile_n or cfg.tile_n + tile_k = tile_k or cfg.tile_k + lds_stage = cfg.lds_stage + use_cshuffle_epilog = cfg.use_cshuffle_epilog + use_async_copy = cfg.use_async_copy + waves_per_eu = cfg.waves_per_eu + xcd_swizzle = cfg.xcd_swizzle + + if "float8" in str(XQ.dtype): + in_dtype = "fp8" + elif XQ.dtype == torch.int8: + in_dtype = "int8" + else: + raise ValueError(f"Unsupported input dtype {XQ.dtype}") + + out_dtype = "bf16" if dtype == torch.bfloat16 else "fp16" + wpe = None if waves_per_eu <= 0 else waves_per_eu + + exe = compile_fn( + N=N, + K=K, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + in_dtype=in_dtype, + out_dtype=out_dtype, + lds_stage=lds_stage, + use_cshuffle_epilog=bool(use_cshuffle_epilog), + use_async_copy=bool(use_async_copy), + waves_per_eu=wpe, + xcd_swizzle=int(xcd_swizzle), + batched=True, + ) + + import flydsl.expr as fx # pyre-ignore[21] + from mslk.flydsl.jit import run_compiled + from mslk.gemm.flydsl._kernels.tensor_shim import ptr_arg + + out_contig = out.contiguous() + XQ_i8 = _as_i8(XQ.contiguous()) + WQ_i8 = _as_i8(WQ.contiguous()) + xs_contig = x_scale.contiguous() + ws_contig = w_scale.contiguous() + + dummy_bias = torch.empty(1, dtype=dtype, device=XQ.device) + stream = fx.Stream(torch.cuda.current_stream()) + + run_compiled( + exe, + ptr_arg(out_contig.view(-1)), + ptr_arg(XQ_i8.view(-1)), + ptr_arg(WQ_i8.view(-1)), + ptr_arg(xs_contig.view(-1)), + ptr_arg(ws_contig.view(-1)), + ptr_arg(dummy_bias), + M, + N, + B, + stream, + ) + if out_contig is not out: + out.copy_(out_contig) + + return out + + +# --------------------------------------------------------------------------- +# Register FlyDSL as the ROCm implementation of mslk rowwise FP8 ops on gfx950 +# --------------------------------------------------------------------------- if torch.version.hip is not None and hasattr(torch.ops, "mslk"): + from mslk.flydsl.common import is_flydsl_available + if is_flydsl_available(): _preshuffle_cache: dict = {} def _get_preshuffled(WQ: Tensor) -> Tensor: - """Cache preshuffled weights by data pointer.""" key = WQ.data_ptr() cached = _preshuffle_cache.get(key) if cached is not None and cached.shape == WQ.shape: @@ -265,7 +410,12 @@ def _flydsl_rowwise_impl( ) -> Tensor: WQ_shuf = _get_preshuffled(WQ) return flydsl_preshuffle_gemm( - XQ, WQ_shuf, x_scale, w_scale, out=output, dtype=dtype + XQ, + WQ_shuf, + x_scale, + w_scale, + out=output, + dtype=dtype, ) if hasattr(torch.ops.mslk, "f8f8bf16_rowwise"): @@ -332,3 +482,31 @@ def _f8f8f16_rowwise_flydsl( use_fast_accum, dtype=torch.float16, ) + + # --- batched op --- + if hasattr(torch.ops.mslk, "f8f8bf16_rowwise_batched"): + from mslk.flydsl.common import is_flydsl_available as _is_flydsl_batched + + if _is_flydsl_batched(): + + @torch.library.impl("mslk::f8f8bf16_rowwise_batched", "CUDA") + def _f8f8bf16_rowwise_batched_flydsl( + XQ: Tensor, + WQ: Tensor, + x_scale: Tensor, + w_scale: Tensor, + bias: Optional[Tensor] = None, + use_fast_accum: bool = True, + output: Optional[Tensor] = None, + ) -> Tensor: + if bias is not None: + raise NotImplementedError( + "FlyDSL batched preshuffle GEMM does not support bias" + ) + return flydsl_preshuffle_batched_gemm( + XQ, + WQ, + x_scale, + w_scale, + out=output, + ) diff --git a/test/gemm/gemm_test.py b/test/gemm/gemm_test.py index 85d26056..b6dfb47f 100644 --- a/test/gemm/gemm_test.py +++ b/test/gemm/gemm_test.py @@ -39,6 +39,11 @@ from mslk.utils.device import compute_capability_in, supports_float8_fnuz if torch.cuda.is_available(): + from mslk.gemm.flydsl.preshuffle_gemm import ( + flydsl_preshuffle, + flydsl_preshuffle_batched_gemm, + flydsl_preshuffle_gemm, + ) from mslk.gemm.triton.fp8_gemm import matmul_fp8_block, matmul_fp8_row, to_mxfp8 from mslk.quantize.mx_mixed_dtype_utils import ( pack_fp6_e2m3, @@ -2915,7 +2920,7 @@ def test_gemm(self, M: int, N: int, K: int) -> None: B = torch.randn((N, K), dtype=torch.bfloat16, device=self.device) * 0.01 # Quantize A to MX8 with blocked scale layout - (a_scale_raw, aq) = to_mxfp8(A) + a_scale_raw, aq = to_mxfp8(A) a_scale = _to_blocked(a_scale_raw.view(torch.int8).reshape(M, -1)).view( torch.uint8 ) @@ -3071,7 +3076,7 @@ def test_gemm(self, M: int, N: int, K: int) -> None: B = torch.randn((N, K), dtype=torch.bfloat16, device=self.device) * 0.01 # MX8 activation with blocked scale layout. - (a_scale_raw, aq) = to_mxfp8(A) + a_scale_raw, aq = to_mxfp8(A) a_scale = _to_blocked(a_scale_raw.view(torch.int8).reshape(M, -1)).view( torch.uint8 ) @@ -3112,7 +3117,7 @@ def test_unpacked_weight_raises(self) -> None: A = torch.randn((M, K), dtype=torch.bfloat16, device=self.device) * 0.1 B = torch.randn((N, K), dtype=torch.bfloat16, device=self.device) * 0.01 - (a_scale_raw, aq) = to_mxfp8(A) + a_scale_raw, aq = to_mxfp8(A) a_scale = _to_blocked(a_scale_raw.view(torch.int8).reshape(M, -1)).view( torch.uint8 ) @@ -3365,10 +3370,6 @@ class FlyDSLPreshuffleGemmTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.device = torch.accelerator.current_accelerator() - from mslk.gemm.flydsl import flydsl_preshuffle, flydsl_preshuffle_gemm - - cls.flydsl_preshuffle = staticmethod(flydsl_preshuffle) - cls.flydsl_preshuffle_gemm = staticmethod(flydsl_preshuffle_gemm) @parameterized.expand( [ @@ -3385,13 +3386,54 @@ def test_gemm(self, M: int, N: int, K: int) -> None: xq, x_scale = quantize_fp8_row(x) wq, w_scale = quantize_fp8_row(w) - wq_shuffled = self.flydsl_preshuffle(wq) + wq_shuffled = flydsl_preshuffle(wq) - out = self.flydsl_preshuffle_gemm(xq, wq_shuffled, x_scale, w_scale) + out = flydsl_preshuffle_gemm(xq, wq_shuffled, x_scale, w_scale) ref = (x @ w.T).to(torch.bfloat16) torch.testing.assert_close(out, ref, atol=1.0, rtol=0.1) +@skipUnlessRocm() +@skipUnlessGfxArch("gfx950") +@unittest.skipUnless( + is_flydsl_version_at_least(), + f"requires FlyDSL >= {MIN_FLYDSL_VERSION}, found {flydsl_version()}", +) +class FlyDSLPreshuffleBatchedGemmTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.device = torch.accelerator.current_accelerator() + + @parameterized.expand( + [ + # small M (decode) + (16, 1, 1280, 8192), + (4, 1, 8192, 1024), + # medium M + (4, 32, 1280, 8192), + (8, 64, 8192, 1024), + (2, 128, 7424, 8192), + # large M (prefill) + (2, 1024, 1280, 8192), + (2, 4096, 8192, 1024), + # B=1 (degenerate batch) + (1, 64, 1280, 8192), + ] + ) + def test_gemm(self, B: int, M: int, N: int, K: int) -> None: + x = torch.randn(B, M, K, dtype=torch.bfloat16, device=self.device) * 0.1 + w = torch.randn(B, N, K, dtype=torch.bfloat16, device=self.device) * 0.01 + + xq, x_scale = quantize_fp8_row(x) + wq, w_scale = quantize_fp8_row(w) + wq_shuffled = torch.stack([flydsl_preshuffle(wq[i]) for i in range(B)]) + + out = flydsl_preshuffle_batched_gemm(xq, wq_shuffled, x_scale, w_scale) + + ref = torch.bmm(x, w.transpose(1, 2)).to(torch.bfloat16) + torch.testing.assert_close(out, ref, atol=1.0, rtol=0.1) + + if __name__ == "__main__": unittest.main()