diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md b/docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md new file mode 100644 index 000000000..697583f87 --- /dev/null +++ b/docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md @@ -0,0 +1,152 @@ +# GEMM + RoPE + MXFP8 Projection (SM100) + +**This is an experimental API and subject to change.** + +## Overview + +**Fused projection GEMM + per-head YARN RoPE + dual-direction MXFP8 quantize**: a persistent dense GEMM on NVIDIA Blackwell GPUs (SM100+) that projects bf16 activations, applies the Megatron MLA-YARN rotary embedding to each attention head's trailing rotary features, and MXFP8 (E4M3, block=32) quantizes the result in **both** the rowwise (D-direction) and columnwise (S-direction) layouts. Implemented with CUTLASS/CUTE. + +Emitting both scale directions makes the output directly consumable by block-scaled matmuls that need either operand orientation. For example, in DeepSeek-V3 the rowwise output feeds the forward `QK^T` and the columnwise output feeds the backward `dK = dS^T · Q` on the cuDNN `is_input_fp8` attention path. + +- **Inputs**: bf16 activations `x`, bf16 projection weight `w`, and bf16 rotary tables `cos`/`sin`. +- **Outputs**: rowwise and columnwise MXFP8 data (`out_fp8_row`, `out_fp8_col`) and their E8M0 scale factors (`out_scales_row`, `out_scales_col`). + +The kernel is tuned for the DeepSeek-V3 Q up-projection shapes: `NUM_HEADS=128`, `HEAD_DIM=192` (`QK_NOPE=128` + `QK_ROPE=64`), MXFP8 `BLOCK=32`, tile `TILE_M=128` (one head per CTA). + +### Shapes + +- **Inputs** + - `x`: `(tokens, Q_LORA)` — `tokens % TILE_M == 0` + - `w`: `(Q_LORA, NUM_HEADS·HEAD_DIM)` when `w_out_in=False`, or the transformer-engine-native transposed `(NUM_HEADS·HEAD_DIM, Q_LORA)` when `w_out_in=True` + - `cos`, `sin`: `(tokens, QK_ROPE)` + +- **Outputs** + - `out_fp8_row`, `out_fp8_col`: `(tokens, NUM_HEADS, HEAD_DIM)` + - `out_scales_row`: `(tokens, NUM_HEADS, HEAD_DIM // BLOCK)` + - `out_scales_col`: `(tokens // BLOCK, NUM_HEADS, HEAD_DIM)` + +### Equations + +Project and reshape per head, then apply the interleaved-in / halves-out YARN RoPE to the trailing `QK_ROPE` features of each head: + +$$ +Y[t, h, :] = (x \, W)\;\text{reshaped to}\;[\text{tokens}, \text{NUM\_HEADS}, \text{HEAD\_DIM}] +$$ + +$$ +Y_{\text{pe}} = \operatorname{RoPE}(Y[\ldots, \text{QK\_NOPE}:],\; \cos, \sin) +$$ + +MXFP8 quantize with block size `BLOCK=32`, independently for each direction (E8M0 per-block scale, E4M3 data): + +$$ +(\text{out\_fp8\_row}, \text{out\_scales\_row}) = \operatorname{MXFP8}_{\text{D}}(Y)\quad\text{(blocks along HEAD\_DIM)} +$$ + +$$ +(\text{out\_fp8\_col}, \text{out\_scales\_col}) = \operatorname{MXFP8}_{\text{S}}(Y)\quad\text{(blocks along tokens)} +$$ + +### Diagram + +```text +x (tokens x Q_LORA), w (Q_LORA x NUM_HEADS*HEAD_DIM) + | GEMM + v + Y (tokens x NUM_HEADS x HEAD_DIM) ---- per-head YARN RoPE on trailing QK_ROPE + | + +--> MXFP8 rowwise (D) -> out_fp8_row, out_scales_row + +--> MXFP8 columnwise(S)-> out_fp8_col, out_scales_col +``` + +## API Usage + +### High-level wrapper +```python +result = gemm_proj_rope_mxfp8_wrapper_sm100( + x, + w, + cos, + sin, + w_out_in=False, + stream=None, +) +out_fp8_row, out_scales_row, out_fp8_col, out_scales_col = result +# Key access: result["out_fp8_row"], result["out_scales_col"], ... +``` + +### Class API +```python +from cuda.bindings import driver as cuda + +op = GemmProjRopeMxfp8Sm100( + sample_x=x, + sample_w=w, + sample_cos=cos, + sample_sin=sin, + sample_out_fp8_row=out_fp8_row, + sample_out_scales_row=out_scales_row, + sample_out_fp8_col=out_fp8_col, + sample_out_scales_col=out_scales_col, + w_out_in=False, +) +assert op.check_support() +op.compile() +op.execute(x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col, current_stream=None) +``` + +--- + +## Parameters + +### Input/Output tensors +- Input **x**: `x` (wrapper) or `sample_x`/`x` (class) + - Shape: `(tokens, Q_LORA)`; Dtype: `bfloat16` +- Input **w**: `w` (wrapper) or `sample_w`/`w` (class) + - Shape: `(Q_LORA, NUM_HEADS·HEAD_DIM)` (`w_out_in=False`) or `(NUM_HEADS·HEAD_DIM, Q_LORA)` (`w_out_in=True`); Dtype: `bfloat16` +- Input **cos**, **sin**: Shape `(tokens, QK_ROPE)`; Dtype: `bfloat16` +- Output **out_fp8_row** / **out_fp8_col**: Shape `(tokens, NUM_HEADS, HEAD_DIM)`; Dtype: `float8_e4m3fn` +- Output **out_scales_row**: Shape `(tokens, NUM_HEADS, HEAD_DIM // BLOCK)`; Dtype: `uint8` (E8M0) +- Output **out_scales_col**: Shape `(tokens // BLOCK, NUM_HEADS, HEAD_DIM)`; Dtype: `uint8` (E8M0) + +### Common parameters +- `w_out_in: bool` + - Whether `w` is stored `[out, in]` (`True`) or `[in, out]` (`False`). The kernel consumes both as the logical `[out, in]` B operand via the cutlass major mode (like cuBLAS `transb`); no transposed copy is materialized. Default: `False` +- CUDA stream (`current_stream` in class API, `stream` in wrapper). Defaults to the current torch stream (required for CUDA-graph capture). + +### Wrapper return values + +Returns a `TupleDict` with keys `out_fp8_row`, `out_scales_row`, `out_fp8_col`, `out_scales_col`. Tuple unpacking order is `(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)`. + +--- + +## Support surface and constraints + +### Dtypes +- `x`, `w`, `cos`, `sin` must be `bfloat16`. +- `out_fp8_row`, `out_fp8_col` must be `float8_e4m3fn`; `out_scales_row`, `out_scales_col` must be `uint8` (E8M0). + +### Shapes and divisibility +- `tokens % TILE_M == 0` (`TILE_M = 128`); no tail handling. +- The projected weight dimension must equal `NUM_HEADS·HEAD_DIM`. + +### Environment +- Requires CUDA with SM100+ compute capability. + +--- + +## Source provenance + +Integrated from the DeepSeek-V3 MLA fused Q up-projection kernel developed for Megatron-LM MXFP8 training (Blackwell / customte CUTLASS 4.4.1). The public entry point `run(...)` and the pure-PyTorch oracle `gemm_proj_rope_mxfp8_reference(...)` live in `python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py`. + +## Installation + +Requires the optional CuTeDSL dependencies: + +```bash +pip install nvidia-cudnn-frontend[cutedsl] +``` + +## Usage examples + +For usage examples, see test cases in `test/python/fe_api/test_gemm_proj_rope_mxfp8.py`. diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 4a22e22d9..2bc3ccfb2 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -4,6 +4,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For details on currently implemented operations, see: - [GEMM + Amax](gemm_fusions/gemm_amax.md) +- [GEMM + RoPE + MXFP8 Projection](gemm_fusions/gemm_proj_rope_mxfp8.md) - [GEMM + SwiGLU](gemm_fusions/gemm_swiglu.md) - [GEMM + sReLU](gemm_fusions/gemm_srelu.md) - [GEMM + dsReLU](gemm_fusions/gemm_dsrelu.md) diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 89da943b1..6e8b99829 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -291,6 +291,8 @@ def _dlopen_cudnn(): "gemm_dsrelu_wrapper_sm100": (".gemm_dsrelu", "gemm_dsrelu_wrapper_sm100"), "GemmAmaxSm100": (".gemm_amax", "GemmAmaxSm100"), "gemm_amax_wrapper_sm100": (".gemm_amax", "gemm_amax_wrapper_sm100"), + "GemmProjRopeMxfp8Sm100": (".gemm_proj_rope_mxfp8", "GemmProjRopeMxfp8Sm100"), + "gemm_proj_rope_mxfp8_wrapper_sm100": (".gemm_proj_rope_mxfp8", "gemm_proj_rope_mxfp8_wrapper_sm100"), "RmsNormRhtAmaxSm100": (".rmsnorm_rht_amax", "RmsNormRhtAmaxSm100"), "rmsnorm_rht_amax_wrapper_sm100": (".rmsnorm_rht_amax", "rmsnorm_rht_amax_wrapper_sm100"), "grouped_gemm": (".grouped_gemm", None), diff --git a/python/cudnn/gemm_proj_rope_mxfp8/__init__.py b/python/cudnn/gemm_proj_rope_mxfp8/__init__.py new file mode 100644 index 000000000..c651a98ac --- /dev/null +++ b/python/cudnn/gemm_proj_rope_mxfp8/__init__.py @@ -0,0 +1,15 @@ +from .api import ( + GemmProjRopeMxfp8Sm100, + gemm_proj_rope_mxfp8_wrapper_sm100, +) +from .gemm_proj_rope_mxfp8 import ( + run, + gemm_proj_rope_mxfp8_reference, +) + +__all__ = [ + "GemmProjRopeMxfp8Sm100", + "gemm_proj_rope_mxfp8_reference", + "gemm_proj_rope_mxfp8_wrapper_sm100", + "run", +] diff --git a/python/cudnn/gemm_proj_rope_mxfp8/api.py b/python/cudnn/gemm_proj_rope_mxfp8/api.py new file mode 100644 index 000000000..fafa4d89c --- /dev/null +++ b/python/cudnn/gemm_proj_rope_mxfp8/api.py @@ -0,0 +1,315 @@ +from .gemm_proj_rope_mxfp8 import ( + gemm_proj_rope_mxfp8_host, + HEAD_DIM, + QK_ROPE, + BLOCK, + TILE_M, +) + +from cuda.bindings import driver as cuda +import logging +import torch +from typing import Optional + +import cutlass +import cutlass.utils +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + +from cudnn.api_base import APIBase, TupleDict + + +class GemmProjRopeMxfp8Sm100(APIBase): + """Fused projection GEMM + per-head YARN RoPE + dual-direction MXFP8 quantize (SM100).""" + + def __init__( + self, + sample_x: torch.Tensor, + sample_w: torch.Tensor, + sample_cos: torch.Tensor, + sample_sin: torch.Tensor, + sample_out_fp8_row: torch.Tensor, + sample_out_scales_row: torch.Tensor, + sample_out_fp8_col: torch.Tensor, + sample_out_scales_col: torch.Tensor, + w_out_in: bool = False, + ): + super().__init__() + + self._warn_experimental_api() + self._logger.debug("Entering __init__") + + self.x_desc = self._make_tensor_desc(sample_x, name="sample_x") + self.w_desc = self._make_tensor_desc(sample_w, name="sample_w") + self.cos_desc = self._make_tensor_desc(sample_cos, name="sample_cos") + self.sin_desc = self._make_tensor_desc(sample_sin, name="sample_sin") + self.out_fp8_row_desc = self._make_tensor_desc(sample_out_fp8_row, name="sample_out_fp8_row") + self.out_scales_row_desc = self._make_tensor_desc(sample_out_scales_row, name="sample_out_scales_row") + self.out_fp8_col_desc = self._make_tensor_desc(sample_out_fp8_col, name="sample_out_fp8_col") + self.out_scales_col_desc = self._make_tensor_desc(sample_out_scales_col, name="sample_out_scales_col") + + self.w_out_in = bool(w_out_in) + self.tokens = int(sample_x.shape[0]) + # Heads derived from the weight's projected dim (compile-time Constexpr for the kernel); + # check_support() validates that this divides evenly. + proj_dim = int(sample_w.shape[0] if self.w_out_in else sample_w.shape[1]) + self.num_heads = proj_dim // HEAD_DIM + + # The cute program is traced from real sample tensors at compile() time; kept only + # until then, then released (mirrors the sample_* teardown in the sibling APIs). + self._samples = ( + sample_x, + sample_w, + sample_cos, + sample_sin, + sample_out_fp8_row, + sample_out_scales_row, + sample_out_fp8_col, + sample_out_scales_col, + ) + self._logger.debug(f"__init__ completed: x {self.x_desc.shape}, w {self.w_desc.shape}, w_out_in {self.w_out_in}") + + def check_support(self) -> bool: + self._logger.debug("Entering check_support") + + self._check_dtype(self.x_desc, dtype=torch.bfloat16, name="x") + self._check_dtype(self.w_desc, dtype=torch.bfloat16, name="w") + self._check_dtype(self.cos_desc, dtype=torch.bfloat16, name="cos") + self._check_dtype(self.sin_desc, dtype=torch.bfloat16, name="sin") + self._check_dtype(self.out_fp8_row_desc, dtype=torch.float8_e4m3fn, name="out_fp8_row") + self._check_dtype(self.out_fp8_col_desc, dtype=torch.float8_e4m3fn, name="out_fp8_col") + self._check_dtype(self.out_scales_row_desc, dtype=torch.uint8, name="out_scales_row") + self._check_dtype(self.out_scales_col_desc, dtype=torch.uint8, name="out_scales_col") + + self._value_error_if( + self.tokens % TILE_M != 0, + f"tokens ({self.tokens}) must be a multiple of TILE_M ({TILE_M})", + ) + + # Shape / output-contract validation. NUM_HEADS is derived from the weight's projected + # dimension (matches the kernel's Constexpr); HEAD_DIM is the fixed per-head width the + # epilogue is specialized for, so the projected dim must be an integer multiple of it. + self._value_error_if( + len(self.x_desc.shape) != 2 or self.x_desc.shape[0] != self.tokens, + f"x must be [tokens, Q_LORA]; got {tuple(self.x_desc.shape)}", + ) + proj_dim = self.w_desc.shape[0] if self.w_out_in else self.w_desc.shape[1] + self._value_error_if( + len(self.w_desc.shape) != 2 or proj_dim % HEAD_DIM != 0, + f"w projected dim must be an integer multiple of HEAD_DIM ({HEAD_DIM}); got weight " + f"shape {tuple(self.w_desc.shape)} with w_out_in={self.w_out_in}", + ) + # GEMM contraction dim (Q_LORA): x's inner dim must match w's, or the fused kernel + # reads past the operands. w is [proj, K] when w_out_in else [K, proj]. + k_dim = self.w_desc.shape[1] if self.w_out_in else self.w_desc.shape[0] + self._value_error_if( + self.x_desc.shape[1] != k_dim, + f"x contraction dim ({self.x_desc.shape[1]}) must match w's ({k_dim}); " + f"x {tuple(self.x_desc.shape)}, w {tuple(self.w_desc.shape)}, w_out_in={self.w_out_in}", + ) + num_heads = self.num_heads + for name, desc in (("cos", self.cos_desc), ("sin", self.sin_desc)): + self._value_error_if( + tuple(desc.shape) != (self.tokens, QK_ROPE), + f"{name} must be [tokens, QK_ROPE] = [{self.tokens}, {QK_ROPE}]; got {tuple(desc.shape)}", + ) + expected = { + "out_fp8_row": (self.out_fp8_row_desc, (self.tokens, num_heads, HEAD_DIM)), + "out_scales_row": (self.out_scales_row_desc, (self.tokens, num_heads, HEAD_DIM // BLOCK)), + "out_fp8_col": (self.out_fp8_col_desc, (self.tokens, num_heads, HEAD_DIM)), + "out_scales_col": (self.out_scales_col_desc, (self.tokens // BLOCK, num_heads, HEAD_DIM)), + } + for name, (desc, shape) in expected.items(): + self._value_error_if( + tuple(desc.shape) != shape, + f"{name} must have shape {shape}; got {tuple(desc.shape)}", + ) + + # Device placement: every tensor must live on the same CUDA device. + all_descs = ( + self.x_desc, + self.w_desc, + self.cos_desc, + self.sin_desc, + self.out_fp8_row_desc, + self.out_scales_row_desc, + self.out_fp8_col_desc, + self.out_scales_col_desc, + ) + devices = {d.device for d in all_descs} + self._value_error_if( + len(devices) != 1 or next(iter(devices)).type != "cuda", + f"all tensors must be on a single CUDA device; got devices {sorted(str(d) for d in devices)}", + ) + + self._logger.debug("Checking environment") + self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") + device = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(device) + compute_capability = major * 10 + minor + self._runtime_error_if( + compute_capability < 100, + f"GemmProjRopeMxfp8 requires SM100+ compute capability, but found SM{compute_capability} on device {device}", + ) + + self._is_supported = True + self._logger.debug("check_support completed successfully") + return True + + def _to_cute_tensors( + self, + x, + w, + cos, + sin, + out_fp8_row, + out_scales_row, + out_fp8_col, + out_scales_col, + ): + """Wrap the torch tensors as (layout-dynamic) cute tensors for the host program. + + ``w`` may be stored [in, out] (default) or the TE-native [out, in] (``w_out_in``); + both are presented to the kernel as the logical [out, in] B operand. + """ + mA = from_dlpack(x.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + if self.w_out_in: + mB = from_dlpack(w.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + else: + mB = from_dlpack(w.detach().transpose(0, 1), assumed_align=16).mark_layout_dynamic(leading_dim=0) + mCos = from_dlpack(cos.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + mSin = from_dlpack(sin.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + mQrow = from_dlpack(out_fp8_row, assumed_align=16).mark_layout_dynamic(leading_dim=2) + mSrow = from_dlpack(out_scales_row, assumed_align=16).mark_layout_dynamic(leading_dim=2) + mQcol = from_dlpack(out_fp8_col, assumed_align=16).mark_layout_dynamic(leading_dim=2) + mScol = from_dlpack(out_scales_col, assumed_align=16).mark_layout_dynamic(leading_dim=2) + return mA, mB, mCos, mSin, mQrow, mSrow, mQcol, mScol + + def compile(self) -> None: + self._logger.debug("Entering compile") + self._ensure_support_checked() + if self._compiled_kernel is not None: + self._logger.debug("Kernel already compiled; skipping recompilation") + return + + cute_tensors = self._to_cute_tensors(*self._samples) + grid_m = self.tokens // TILE_M + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters(1) + swizzle_size = 8 + # Trace/compile on the current stream (a runtime argument; execute() supplies its own). + compile_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + self._logger.debug("Compiling gemm_proj_rope_mxfp8_host") + self._compiled_kernel = cute.compile( + gemm_proj_rope_mxfp8_host, + *cute_tensors, + grid_m, + self.num_heads, + max_active_clusters, + swizzle_size, + compile_stream, + ) + + self._samples = None + self._logger.debug("Kernel compiled successfully") + + def execute( + self, + x: torch.Tensor, + w: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + out_fp8_row: torch.Tensor, + out_scales_row: torch.Tensor, + out_fp8_col: torch.Tensor, + out_scales_col: torch.Tensor, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + self._logger.debug("Entering execute") + current_stream = self._get_default_stream(current_stream) + + self._runtime_error_if( + self._compiled_kernel is None, + "GemmProjRopeMxfp8Sm100 kernel not compiled; call compile() first", + ) + + cute_tensors = self._to_cute_tensors(x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col) + self._compiled_kernel(*cute_tensors, current_stream) + self._logger.debug("Executed with compiled kernel successfully") + + +_logger = logging.getLogger(__name__) +_cache_of_GemmProjRopeMxfp8Sm100Objects = {} + + +def gemm_proj_rope_mxfp8_wrapper_sm100( + x: torch.Tensor, + w: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + w_out_in: bool = False, + stream: Optional[cuda.CUstream] = None, +) -> TupleDict: + """Allocate outputs, (compile-and-)cache the kernel, run it, and return the MXFP8 Q tensors. + + Args: + x: ``[tokens, Q_LORA]`` bf16 activations (``tokens % TILE_M == 0``). + w: projection weight -- ``[Q_LORA, NUM_HEADS*HEAD_DIM]`` (``w_out_in=False``) or the + TE-native transposed ``[NUM_HEADS*HEAD_DIM, Q_LORA]`` (``w_out_in=True``). + cos, sin: ``[tokens, QK_ROPE]`` bf16 rotary tables. + w_out_in: whether ``w`` is stored ``[out, in]``. + stream: optional CUDA stream; defaults to the current torch stream. + + Returns: + ``TupleDict(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)`` -- rowwise / columnwise + MXFP8 (E4M3) data and E8M0 scales. + """ + tokens = x.shape[0] + device = x.device + # Heads derived from the weight's projected dimension (matches the kernel's Constexpr). + num_heads = (w.shape[0] if w_out_in else w.shape[1]) // HEAD_DIM + + out_fp8_row = torch.empty(tokens, num_heads, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device) + out_scales_row = torch.empty(tokens, num_heads, HEAD_DIM // BLOCK, dtype=torch.uint8, device=device) + out_fp8_col = torch.empty(tokens, num_heads, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device) + out_scales_col = torch.empty(tokens // BLOCK, num_heads, HEAD_DIM, dtype=torch.uint8, device=device) + + cache_key = ( + tuple(x.shape), + tuple(w.shape), + x.dtype, + w.dtype, + cos.dtype, + sin.dtype, + bool(w_out_in), + x.device, + ) + if cache_key in _cache_of_GemmProjRopeMxfp8Sm100Objects: + _logger.debug("gemm_proj_rope_mxfp8_wrapper_sm100: using cached GemmProjRopeMxfp8Sm100 object") + obj = _cache_of_GemmProjRopeMxfp8Sm100Objects[cache_key] + else: + _logger.debug("gemm_proj_rope_mxfp8_wrapper_sm100: creating new GemmProjRopeMxfp8Sm100 object") + obj = GemmProjRopeMxfp8Sm100( + sample_x=x, + sample_w=w, + sample_cos=cos, + sample_sin=sin, + sample_out_fp8_row=out_fp8_row, + sample_out_scales_row=out_scales_row, + sample_out_fp8_col=out_fp8_col, + sample_out_scales_col=out_scales_col, + w_out_in=w_out_in, + ) + assert obj.check_support() + obj.compile() + _cache_of_GemmProjRopeMxfp8Sm100Objects[cache_key] = obj + + obj.execute(x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col, current_stream=stream) + + return TupleDict( + out_fp8_row=out_fp8_row, + out_scales_row=out_scales_row, + out_fp8_col=out_fp8_col, + out_scales_col=out_scales_col, + ) diff --git a/python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py b/python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py new file mode 100644 index 000000000..189d57786 --- /dev/null +++ b/python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py @@ -0,0 +1,647 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Fused projection GEMM + per-head YARN RoPE + dual-direction MXFP8 quantize (Blackwell / SM100).""" + +import cuda.bindings.driver as cuda +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.runtime import from_dlpack +from cutlass._mlir import ir as _mlir_ir +from cutlass._mlir.dialects import llvm as _llvm_dialect +from cutlass.cutlass_dsl import Float32 as _Float32Sym, Int32 as _Int32Sym + +# ---- DSv3 constants ---- +# NUM_HEADS inferred +QK_NOPE = 128 +QK_ROPE = 64 +HALF = 32 +HEAD_DIM = 192 # 128 + 64 +BLOCK = 32 +FP8_MAX = 448.0 + +# ---- tile config ---- +io_dtype = cutlass.BFloat16 +acc_dtype = cutlass.Float32 +TILE_M = 128 +TILE_N = HEAD_DIM # 192, one head per CTA +K_TILE = 64 +COLBLK = TILE_M // BLOCK # col blocks per CTA along tokens +stage_dtype = cutlass.BFloat16 # SMEM staging dtype for post-rope tile + +mma_inst_shape_mnk = (TILE_M, TILE_N, 16) +mma_tiler_mnk = (TILE_M, TILE_N, K_TILE) + +ab_stages = 4 +acc_stages = 2 +NUM_EPI_WARPS = 12 # epilogue warps (rope + quant); 4 of them do T2R staging +T2R_WARPS = 4 # warps that drain TMEM->SMEM (fixed by 128-thread T2R) +threads_in_epilogue = NUM_EPI_WARPS * 32 +SACC_STRIDE = 196 # marginally better than 200 in testing +FEATCELL = 64 # VEC2: feature-cell width; lane owns 2 contiguous feats +N_FEATCELL = HEAD_DIM // FEATCELL # 3 +HALFW = 16 # lanes per 32-feature row-block within a featcell + +# ---- Structural specialization (compile-time) ---- +# Like the SDPA kernels' fixed head_dim, this kernel is specialized for the DeepSeek-V3 Q up-proj +# head geometry: HEAD_DIM=192 (QK_NOPE 128 + QK_ROPE 64), MXFP8 BLOCK=32, TILE_M=128. The epilogue's +# VEC2 feature-cell layout, warp specialization, and single trailing rope cell depend on these exact +# values -- changing them requires reworking the epilogue, not just editing a constant. These asserts +# fail loudly at import if the constants are set to an unsupported combination. +assert FEATCELL == 64, "FEATCELL is warp-size (32) x VEC2 (2); must be 64" +assert QK_NOPE + QK_ROPE == HEAD_DIM, "QK_NOPE + QK_ROPE must equal HEAD_DIM" +assert HEAD_DIM % FEATCELL == 0, "HEAD_DIM must be a whole number of 64-wide feature cells" +assert QK_ROPE == FEATCELL, "the rope occupies exactly the trailing feature cell; QK_ROPE must equal FEATCELL" +assert HALF == QK_ROPE // 2, "HALF must be QK_ROPE // 2" +assert TILE_M % BLOCK == 0, "TILE_M must be a whole number of MXFP8 blocks" +assert NUM_EPI_WARPS == COLBLK * N_FEATCELL, "epilogue warp count must equal COLBLK x N_FEATCELL" + + +@cute.struct +class SharedStorage: + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stages * 2] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buffer: cutlass.Int32 + + +@cute.jit +def _e8m0(amax): + scaled = amax * cutlass.Float32(1.0 / FP8_MAX) + packed_i16 = _llvm_dialect.inline_asm( + _mlir_ir.IntegerType.get_signless(16), + [_Float32Sym(scaled).ir_value()], + "cvt.rp.satfinite.ue8m0x2.f32 $0, 0f00000000, $1;", + "=h,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm_dialect.AsmDialect.AD_ATT, + ) + sbyte32 = cutlass.Int32(packed_i16) & cutlass.Int32(0xFF) + inv_bits = _llvm_dialect.inline_asm( + _mlir_ir.F32Type.get(), + [_Int32Sym(sbyte32).ir_value()], + "{ .reg .s32 t; sub.s32 t, 254, $1; shl.b32 t, t, 23; mov.b32 $0, t; }", + "=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm_dialect.AsmDialect.AD_ATT, + ) + return cutlass.Float32(inv_bits), sbyte32 + + +@cute.jit +def _e8m0_inv(sbyte32): + inv_bits = _llvm_dialect.inline_asm( + _mlir_ir.F32Type.get(), + [_Int32Sym(sbyte32).ir_value()], + "{ .reg .s32 t; sub.s32 t, 254, $1; shl.b32 t, t, 23; mov.b32 $0, t; }", + "=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm_dialect.AsmDialect.AD_ATT, + ) + return cutlass.Float32(inv_bits) + + +@cute.jit +def _e8m0_pair(amax0, amax1): + s0 = amax0 * cutlass.Float32(1.0 / FP8_MAX) + s1 = amax1 * cutlass.Float32(1.0 / FP8_MAX) + packed_i16 = _llvm_dialect.inline_asm( + _mlir_ir.IntegerType.get_signless(16), + [_Float32Sym(s0).ir_value(), _Float32Sym(s1).ir_value()], + "cvt.rp.satfinite.ue8m0x2.f32 $0, $1, $2;", + "=h,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm_dialect.AsmDialect.AD_ATT, + ) + p = cutlass.Int32(packed_i16) + sbyte0 = (p >> 8) & cutlass.Int32(0xFF) + sbyte1 = p & cutlass.Int32(0xFF) + return _e8m0_inv(sbyte0), sbyte0, _e8m0_inv(sbyte1), sbyte1 + + +@cute.jit +def _pack_e4m3x2(vlo, vhi): + packed_i16 = _llvm_dialect.inline_asm( + _mlir_ir.IntegerType.get_signless(16), + [_Float32Sym(vhi).ir_value(), _Float32Sym(vlo).ir_value()], + "cvt.rn.satfinite.e4m3x2.f32 $0, $1, $2;", + "=h,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm_dialect.AsmDialect.AD_ATT, + ) + return packed_i16 + + +@cute.kernel +def gemm_proj_rope_mxfp8_kernel( + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + a_smem_layout: cute.ComposedLayout, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + b_smem_layout: cute.ComposedLayout, + mCos: cute.Tensor, + mSin: cute.Tensor, + mQrow: cute.Tensor, + mSrow: cute.Tensor, + mQcol: cute.Tensor, + mScol: cute.Tensor, + epi_tile: cute.Tile, + cta_layout_vmnk: cute.Layout, + tile_sched_params: utils.PersistentTileSchedulerParams, + num_tmem_cols: cutlass.Constexpr, + num_heads: cutlass.Constexpr, +): + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + tidx, _, _ = cute.arch.thread_idx() + + epilogue_warp_ids = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11) + mma_warp_id = 12 + tma_warp_id = 13 + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sA = smem.allocate_tensor(element_type=io_dtype, layout=a_smem_layout.outer, byte_alignment=128, swizzle=a_smem_layout.inner) + sB = smem.allocate_tensor(element_type=io_dtype, layout=b_smem_layout.outer, byte_alignment=128, swizzle=b_smem_layout.inner) + sACC = smem.allocate_tensor(element_type=stage_dtype, layout=cute.make_layout((TILE_M, TILE_N), stride=(SACC_STRIDE, 1)), byte_alignment=128) + + if warp_idx == tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + + cta_rank_in_cluster = cute.arch.block_idx_in_cluster() + cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + + tma_mcast_mask_a = cpasync.create_tma_multicast_mask(cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2) + tma_mcast_mask_b = cpasync.create_tma_multicast_mask(cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1) + + gA = cute.local_tile(mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None)) + gB = cute.local_tile(mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None)) + + thr_mma = tiled_mma.get_slice(0) + tCgA = thr_mma.partition_A(gA) + tCgB = thr_mma.partition_B(gB) + + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + + acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, acc_stages)) + + epilogue_sync_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=threads_in_epilogue) + tmem_alloc_barrier = pipeline.NamedBarrier(barrier_id=2, num_threads=32 * len((mma_warp_id, *epilogue_warp_ids))) + tmem = utils.TmemAllocator(storage.tmem_holding_buffer, barrier_for_retrieve=tmem_alloc_barrier, allocator_warp_id=epilogue_warp_ids[0], is_two_cta=False) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + cta_in_cluster_coord_vmnk[2], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + cta_in_cluster_coord_vmnk[1], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + num_tma_copy_bytes = cute.size_in_bytes(io_dtype, cute.select(a_smem_layout, mode=[0, 1, 2])) + cute.size_in_bytes( + io_dtype, cute.select(b_smem_layout, mode=[0, 1, 2]) + ) + + mainloop_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + mainloop_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, size=1) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=ab_stages, + producer_group=mainloop_producer_group, + consumer_group=mainloop_consumer_group, + tx_count=num_tma_copy_bytes, + cta_layout_vmnk=cta_layout_vmnk, + ).make_participants() + + acc_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + acc_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, size=T2R_WARPS) + acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=acc_stages, + producer_group=acc_producer_group, + consumer_group=acc_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + ).make_participants() + + num_k_tiles = cute.size(tCgA, mode=[4]) + + tile_sched = utils.StaticPersistentTileScheduler.create(tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + + # ================= TMA load warp ================= + if warp_idx == tma_warp_id: + while work_tile.is_valid_tile: + coord = work_tile.tile_idx + m_idx = coord[0] + n_idx = coord[1] + tAgA_slice = tAgA[(None, m_idx, None)] + tBgB_slice = tBgB[(None, n_idx, None)] + for k_tile_idx in cutlass.range(num_k_tiles): + handle = ab_producer.acquire_and_advance() + cute.copy(tma_atom_a, tAgA_slice[(None, k_tile_idx)], tAsA[(None, handle.index)], tma_bar_ptr=handle.barrier, mcast_mask=tma_mcast_mask_a) + cute.copy(tma_atom_b, tBgB_slice[(None, k_tile_idx)], tBsB[(None, handle.index)], tma_bar_ptr=handle.barrier, mcast_mask=tma_mcast_mask_b) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + ab_producer.tail() + + # ================= MMA warp ================= + elif warp_idx == mma_warp_id: + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + while work_tile.is_valid_tile: + acc_empty = acc_producer.acquire_and_advance() + tCtAcc = tCtAcc_base[(None, None, None, acc_empty.index)] + for k_tile_idx in cutlass.range(num_k_tiles): + handle = ab_consumer.wait_and_advance() + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile_idx != 0) + tile_crd = (None, None, None, handle.index) + cute.gemm(tiled_mma, tCtAcc, tCrA[tile_crd], tCrB[tile_crd], tCtAcc) + handle.release() + acc_empty.commit() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + acc_producer.tail() + + # ================= Epilogue warps ================= + elif warp_idx < mma_warp_id: + tmem.allocate(num_tmem_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + copy_atom_t2r = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + + sACC_epi = cute.flat_divide(sACC, epi_tile) + buf0 = cute.make_rmem_tensor((BLOCK,), cutlass.Float32) + buf1 = cute.make_rmem_tensor((BLOCK,), cutlass.Float32) + rPr = cute.make_rmem_tensor((1,), cutlass.Uint16) + rPc = cute.make_rmem_tensor((1,), cutlass.Uint16) + st16 = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Uint16, num_bits_per_copy=16) + wid = warp_idx # epilogue warp id 0..NUM_EPI_WARPS-1 + lane = tidx % 32 + + while work_tile.is_valid_tile: + coord = work_tile.tile_idx + m_idx = coord[0] + head = coord[1] + token_base = m_idx * TILE_M + + if wid < T2R_WARPS: + acc_full = acc_consumer.wait_and_advance() + tCtAcc = tCtAcc_base[(None, None, None, acc_full.index)] + tCtAcc_epi = cute.flat_divide(tCtAcc[((None, None), 0, 0)], epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc_epi[(None, None, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc_epi) + tTR_sACC = thr_copy_t2r.partition_D(sACC_epi) + tTR_rAcc = cute.make_rmem_tensor(tTR_sACC[(None, None, None, 0, 0)].shape, cutlass.Float32) + tTR_rStg = cute.make_rmem_tensor(tTR_sACC[(None, None, None, 0, 0)].shape, stage_dtype) + tTR_tAcc_g = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_sACC_g = cute.group_modes(tTR_sACC, 3, cute.rank(tTR_sACC)) + subtile_cnt = cute.size(tTR_tAcc_g.shape, mode=[3]) + + for subtile_idx in cutlass.range_constexpr(subtile_cnt): + cute.copy(tiled_copy_t2r, tTR_tAcc_g[(None, None, None, subtile_idx)], tTR_rAcc) + tTR_rStg.store(tTR_rAcc.load().to(stage_dtype)) + cute.autovec_copy(tTR_rStg, tTR_sACC_g[(None, None, None, subtile_idx)]) + + with cute.arch.elect_one(): + acc_full.release() + + epilogue_sync_barrier.arrive_and_wait() + + # ---- VEC2: lane owns 2 contiguous features; 12 cells = 12 warps ---- + cell = wid + cb = cell // N_FEATCELL # token-block 0..3 + fc = cell % N_FEATCELL # feature-cell 0..2 + tok0 = cb * BLOCK + f0 = fc * FEATCELL + 2 * lane # global feature (even) + f1 = f0 + 1 + b = fc * 2 + (lane // HALFW) # 32-feature row-block 0..5 + col_amax0 = cutlass.Float32(0.0) + col_amax1 = cutlass.Float32(0.0) + is_rope = fc == (N_FEATCELL - 1) + if is_rope: + lib = lane % HALFW + pcol0 = QK_NOPE + 4 * lib + pcol1 = pcol0 + 2 + roff = HALF * (lane // HALFW) + rf = cutlass.Float32(lane // HALFW) + lf = cutlass.Float32(1.0) - rf + cidx0 = 2 * lib + roff + + cos_row_bytes = QK_ROPE * 2 + cos_base = mCos[token_base + tok0, None].iterator.toint() + cidx0 * 2 + sin_base = mSin[token_base + tok0, None].iterator.toint() + cidx0 * 2 + for r in cutlass.range_constexpr(BLOCK): + token = token_base + tok0 + r + p0 = sACC[tok0 + r, pcol0].to(cutlass.Float32) + q0 = sACC[tok0 + r, pcol0 + 1].to(cutlass.Float32) + p1 = sACC[tok0 + r, pcol1].to(cutlass.Float32) + q1 = sACC[tok0 + r, pcol1 + 1].to(cutlass.Float32) + tcos = cute.make_tensor(cute.make_ptr(cutlass.BFloat16, cos_base + r * cos_row_bytes, cute.AddressSpace.gmem, assumed_align=4), (2,)) + tsin = cute.make_tensor(cute.make_ptr(cutlass.BFloat16, sin_base + r * cos_row_bytes, cute.AddressSpace.gmem, assumed_align=4), (2,)) + c0 = tcos[0].to(cutlass.Float32) + s0 = tsin[0].to(cutlass.Float32) + c1 = tcos[1].to(cutlass.Float32) + s1 = tsin[1].to(cutlass.Float32) + # packed pair math: both features share lf/rf blend weights + pc0, pc1 = cute.arch.mul_packed_f32x2((p0, p1), (c0, c1)) + qs0, qs1 = cute.arch.mul_packed_f32x2((q0, q1), (s0, s1)) + lft0, lft1 = cute.arch.fma_packed_f32x2((qs0, qs1), (cutlass.Float32(-1.0), cutlass.Float32(-1.0)), (pc0, pc1)) + ps0, ps1 = cute.arch.mul_packed_f32x2((p0, p1), (s0, s1)) + qc0, qc1 = cute.arch.mul_packed_f32x2((q0, q1), (c0, c1)) + rgt0, rgt1 = cute.arch.add_packed_f32x2((ps0, ps1), (qc0, qc1)) + ll0, ll1 = cute.arch.mul_packed_f32x2((lft0, lft1), (lf, lf)) + v0, v1 = cute.arch.fma_packed_f32x2((rgt0, rgt1), (rf, rf), (ll0, ll1)) + buf0[r] = v0 + buf1[r] = v1 + col_amax0 = cute.arch.fmax(col_amax0, cute.arch.fmax(v0, -v0)) + col_amax1 = cute.arch.fmax(col_amax1, cute.arch.fmax(v1, -v1)) + else: + for r in cutlass.range_constexpr(BLOCK): + v0 = sACC[tok0 + r, f0].to(cutlass.Float32) + v1 = sACC[tok0 + r, f1].to(cutlass.Float32) + buf0[r] = v0 + buf1[r] = v1 + col_amax0 = cute.arch.fmax(col_amax0, cute.arch.fmax(v0, -v0)) + col_amax1 = cute.arch.fmax(col_amax1, cute.arch.fmax(v1, -v1)) + invc0, sbc0, invc1, sbc1 = _e8m0_pair(col_amax0, col_amax1) + scol_row = m_idx * COLBLK + cb + mScol[scol_row, head, f0] = cutlass.Uint8(sbc0) + mScol[scol_row, head, f1] = cutlass.Uint8(sbc1) + is_leader = (lane % HALFW) == 0 + vchunk = f0 // 2 # VEC=2 chunk index within HEAD_DIM + + row_bytes = num_heads * HEAD_DIM + pr_base = mQrow[token_base + tok0, head, None].iterator.toint() + pc_base = mQcol[token_base + tok0, head, None].iterator.toint() + for r in cutlass.range_constexpr(BLOCK): + token = token_base + tok0 + r + v0 = buf0[r] + v1 = buf1[r] + m = cute.arch.fmax(cute.arch.fmax(v0, -v0), cute.arch.fmax(v1, -v1)) + o8 = cute.arch.shuffle_sync_bfly(m, 8) + m = cute.arch.fmax(m, o8) + o4 = cute.arch.shuffle_sync_bfly(m, 4) + m = cute.arch.fmax(m, o4) + o2 = cute.arch.shuffle_sync_bfly(m, 2) + m = cute.arch.fmax(m, o2) + o1 = cute.arch.shuffle_sync_bfly(m, 1) + row_amax = cute.arch.fmax(m, o1) + inv_r, sbyte_r = _e8m0(row_amax) + if is_leader: + mSrow[token, head, b] = cutlass.Uint8(sbyte_r) + vr0, vr1 = cute.arch.mul_packed_f32x2((v0, v1), (inv_r, inv_r)) + vc0, vc1 = cute.arch.mul_packed_f32x2((v0, v1), (invc0, invc1)) + rPr[0] = cutlass.Uint16(_pack_e4m3x2(vr0, vr1)) + rPc[0] = cutlass.Uint16(_pack_e4m3x2(vc0, vc1)) + pr = cute.make_ptr(cutlass.Uint16, pr_base + r * row_bytes, cute.AddressSpace.gmem, assumed_align=16) + pc = cute.make_ptr(cutlass.Uint16, pc_base + r * row_bytes, cute.AddressSpace.gmem, assumed_align=16) + gr = cute.tiled_divide(cute.make_tensor(pr, (HEAD_DIM // 2,)), (1,)) + gc = cute.tiled_divide(cute.make_tensor(pc, (HEAD_DIM // 2,)), (1,)) + cute.copy(st16, rPr, gr[None, vchunk]) + cute.copy(st16, rPc, gc[None, vchunk]) + + epilogue_sync_barrier.arrive_and_wait() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + +@cute.jit +def gemm_proj_rope_mxfp8_host( + mA: cute.Tensor, + mB: cute.Tensor, + mCos: cute.Tensor, + mSin: cute.Tensor, + mQrow: cute.Tensor, + mSrow: cute.Tensor, + mQcol: cute.Tensor, + mScol: cute.Tensor, + grid_m: cutlass.Constexpr, + num_heads: cutlass.Constexpr, + max_active_clusters: cutlass.Constexpr, + swizzle_size: cutlass.Constexpr, + stream, +): + a_major = utils.LayoutEnum.from_tensor(mA).mma_major_mode() + b_major = utils.LayoutEnum.from_tensor(mB).mma_major_mode() + + op = tcgen05.MmaF16BF16Op(io_dtype, acc_dtype, mma_inst_shape_mnk, tcgen05.CtaGroup.ONE, tcgen05.OperandSource.SMEM, a_major, b_major) + tiled_mma = cute.make_tiled_mma(op) + + a_smem_layout = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, mA.element_type, ab_stages) + b_smem_layout = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, mB.element_type, ab_stages) + + cluster_shape_mnk = (1, 1, 1) + cta_layout_mnk = cute.make_layout(cluster_shape_mnk) + cta_layout_vmnk = cute.tiled_divide(cta_layout_mnk, (tiled_mma.thr_id,)) + + tma_op = cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.ONE) + + a_smem_layout_1 = cute.slice_(a_smem_layout, (None, None, None, 0)) + b_smem_layout_1 = cute.slice_(b_smem_layout, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(tma_op, mA, a_smem_layout_1, mma_tiler_mnk, tiled_mma, cta_layout_vmnk.shape) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(tma_op, mB, b_smem_layout_1, mma_tiler_mnk, tiled_mma, cta_layout_vmnk.shape) + + cta_tile_shape_mnk = (mma_tiler_mnk[0], mma_tiler_mnk[1], mma_tiler_mnk[2]) + c_layout_kind = utils.LayoutEnum.ROW_MAJOR + epi_tile = utils.compute_epilogue_tile_shape(cta_tile_shape_mnk, False, c_layout_kind, cutlass.Float32) + + acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, acc_stages)) + num_tmem_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100") + + num_ctas_mnl = (grid_m, num_heads, 1) + tile_sched_params = utils.PersistentTileSchedulerParams(num_ctas_mnl, cluster_shape_mnk, swizzle_size, True) + grid = utils.StaticPersistentTileScheduler.get_grid_shape(tile_sched_params, max_active_clusters) + + gemm_proj_rope_mxfp8_kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + a_smem_layout, + tma_atom_b, + tma_tensor_b, + b_smem_layout, + mCos, + mSin, + mQrow, + mSrow, + mQcol, + mScol, + epi_tile, + cta_layout_vmnk, + tile_sched_params, + num_tmem_cols, + num_heads, + ).launch( + grid=grid, + block=[(NUM_EPI_WARPS + 2) * 32, 1, 1], + cluster=cluster_shape_mnk, + stream=stream, + ) + + +_cache = {} + + +def run(x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col, w_out_in=False): + tokens = x.shape[0] + proj_dim = w.shape[0] if w_out_in else w.shape[1] + k_dim = w.shape[1] if w_out_in else w.shape[0] # GEMM contraction dim (Q_LORA) + # Lightweight structural validation (run() is a public entry that skips check_support); + # catch the common misuses with a clear error instead of an opaque CUTLASS/CUDA crash. + if tokens % TILE_M != 0: + raise ValueError(f"tokens ({tokens}) must be a multiple of TILE_M ({TILE_M})") + if proj_dim % HEAD_DIM != 0: + raise ValueError(f"weight projected dim ({proj_dim}) must be a multiple of HEAD_DIM ({HEAD_DIM})") + if x.shape[1] != k_dim: + raise ValueError(f"x contraction dim ({x.shape[1]}) must match w's ({k_dim}); " f"x {tuple(x.shape)}, w {tuple(w.shape)}, w_out_in={w_out_in}") + if tuple(cos.shape) != (tokens, QK_ROPE) or tuple(sin.shape) != (tokens, QK_ROPE): + raise ValueError(f"cos/sin must both be ({tokens}, {QK_ROPE}); got {tuple(cos.shape)}, {tuple(sin.shape)}") + # Every tensor must live on x's device: the stream and HardwareInfo below are read from + # the active device, and the kernel launch dereferences all operands on it. + operands = (x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col) + if any(t.device != x.device for t in operands): + raise ValueError(f"all tensors must be on the same device as x ({x.device}); got {[str(t.device) for t in operands]}") + # Number of heads is derived from the weight's projected dimension (compile-time Constexpr); + # HEAD_DIM is the fixed per-head width the epilogue is specialized for. + num_heads = proj_dim // HEAD_DIM + + # Bind the active device to x's device so the stream, HardwareInfo, and launch below all + # target the tensors' device (not whatever device happened to be current on entry). + with torch.cuda.device(x.device): + mA = from_dlpack(x.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + # Weight B operand. Two accepted layouts, both giving logical B=[N,K]=[out,in]: + # w_out_in=False (default): w is [K,N]=[in,out] contiguous -> transpose(0,1) view (N-contig). + # w_out_in=True: w is the NATIVE TE weight [N,K]=[out,in] (K-contig) -> consume directly, no + # transposed-contiguous copy needed (cutlass handles the major mode like cuBLAS's transb). + if w_out_in: + mB = from_dlpack(w.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + else: + mB = from_dlpack(w.detach().transpose(0, 1), assumed_align=16).mark_layout_dynamic(leading_dim=0) + mCos = from_dlpack(cos.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + mSin = from_dlpack(sin.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + mQrow = from_dlpack(out_fp8_row, assumed_align=16).mark_layout_dynamic(leading_dim=2) + mSrow = from_dlpack(out_scales_row, assumed_align=16).mark_layout_dynamic(leading_dim=2) + mQcol = from_dlpack(out_fp8_col, assumed_align=16).mark_layout_dynamic(leading_dim=2) + mScol = from_dlpack(out_scales_col, assumed_align=16).mark_layout_dynamic(leading_dim=2) + + # Launch on the CURRENT torch stream (required for CUDA-graph capture: null-stream + # work is not captured -> empty graph). Mirrors the v3 winner's stream handling. + current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + grid_m = tokens // TILE_M + # Key on every dimension the compiled fn specializes on: grid (tokens), B-operand layout + # (w_out_in), the num_heads Constexpr, and the device (HardwareInfo/max_active_clusters are + # queried from the current device at compile time). Dtypes are fixed for this entry point. + key = (tokens, bool(w_out_in), num_heads, x.device.index) + fn = _cache.get(key) + if fn is None: + hw = cutlass.utils.HardwareInfo() + max_active_clusters = hw.get_max_active_clusters(1) + swizzle_size = 8 + fn = cute.compile( + gemm_proj_rope_mxfp8_host, mA, mB, mCos, mSin, mQrow, mSrow, mQcol, mScol, grid_m, num_heads, max_active_clusters, swizzle_size, current_stream + ) + _cache[key] = fn + fn(mA, mB, mCos, mSin, mQrow, mSrow, mQcol, mScol, current_stream) + + +# --------------------------------------------------------------------------- +# PyTorch reference (oracle) for the fused kernel above. +# --------------------------------------------------------------------------- +def gemm_proj_rope_mxfp8_reference(x, w, cos, sin, w_out_in=False): + E8M0_BIAS = 127 + tokens = x.shape[0] + # Heads derived from the weight's projected dimension (matches the kernel's Constexpr). + num_heads = (w.shape[0] if w_out_in else w.shape[1]) // HEAD_DIM + + # Projection GEMM (fp32 accumulate), reshaped to per-head. + w_eff = w.float().t() if w_out_in else w.float() # -> [Q_LORA, num_heads*HEAD_DIM] + q = torch.matmul(x.float(), w_eff).view(tokens, num_heads, HEAD_DIM) + + # Per-head YARN RoPE on the trailing QK_ROPE (interleaved-in, halves-out). + q_nope, q_pe = q[..., :QK_NOPE], q[..., QK_NOPE:] + x1, x2 = q_pe[..., 0::2], q_pe[..., 1::2] + cl = cos[..., :HALF].unsqueeze(1).float() + sl = sin[..., :HALF].unsqueeze(1).float() + cr = cos[..., HALF:].unsqueeze(1).float() + sr = sin[..., HALF:].unsqueeze(1).float() + q_pe = torch.cat([x1 * cl - x2 * sl, x2 * cr + x1 * sr], dim=-1) + qf = torch.cat([q_nope, q_pe], dim=-1).contiguous() # [tokens, num_heads, HEAD_DIM] fp32 + + def _e8m0_quant(blocks, amax_dim): + amax = blocks.abs().amax(dim=amax_dim, keepdim=True).clamp(min=1e-30) + exp = torch.ceil(torch.log2(amax / FP8_MAX)).clamp(-127.0, 127.0) + data = (blocks * torch.pow(2.0, -exp)).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + scale = (exp + E8M0_BIAS).to(torch.uint8) + return data, scale + + # Rowwise (D-direction): 32-blocks along HEAD_DIM. + rb = qf.reshape(tokens, num_heads, HEAD_DIM // BLOCK, BLOCK) + rdata, rscale = _e8m0_quant(rb, amax_dim=-1) + out_fp8_row = rdata.reshape(tokens, num_heads, HEAD_DIM) + out_scales_row = rscale.squeeze(-1) # [tokens, num_heads, HEAD_DIM // BLOCK] + + # Columnwise (S-direction): 32-blocks along tokens. + cb = qf.reshape(tokens // BLOCK, BLOCK, num_heads, HEAD_DIM) + cdata, cscale = _e8m0_quant(cb, amax_dim=1) + out_fp8_col = cdata.reshape(tokens, num_heads, HEAD_DIM) + out_scales_col = cscale.squeeze(1) # [tokens // BLOCK, num_heads, HEAD_DIM] + + return out_fp8_row, out_scales_row, out_fp8_col, out_scales_col diff --git a/test/python/fe_api/test_gemm_proj_rope_mxfp8.py b/test/python/fe_api/test_gemm_proj_rope_mxfp8.py new file mode 100644 index 000000000..6522e4a1c --- /dev/null +++ b/test/python/fe_api/test_gemm_proj_rope_mxfp8.py @@ -0,0 +1,276 @@ +import torch + +import pytest + +from test_utils import torch_fork_set_rng +from fe_api.test_gemm_proj_rope_mxfp8_utils import with_gemm_proj_rope_mxfp8_params + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@with_gemm_proj_rope_mxfp8_params +def test_gemm_proj_rope_mxfp8_compile_execute(tokens, w_out_in, request): + _test_gemm_proj_rope_mxfp8_compile_execute(tokens=tokens, w_out_in=w_out_in, request=request) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@with_gemm_proj_rope_mxfp8_params +def test_gemm_proj_rope_mxfp8_wrapper(tokens, w_out_in, request): + _test_gemm_proj_rope_mxfp8_wrapper(tokens=tokens, w_out_in=w_out_in, request=request) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_gemm_proj_rope_mxfp8_check_support_rejects_unaligned_tokens(request): + """check_support() must reject a token count that is not a multiple of TILE_M.""" + try: + from cudnn import GemmProjRopeMxfp8Sm100 + from fe_api.test_gemm_proj_rope_mxfp8_utils import ( + allocate_input_tensors, + allocate_output_tensors, + gemm_proj_rope_mxfp8_init, + ) + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + gemm_proj_rope_mxfp8_init(request, tokens=2048, w_out_in=False) # arch skip if needed + tokens = 2048 + 32 # not a multiple of TILE_M (128) + x, w, cos, sin = allocate_input_tensors(tokens, w_out_in=False) + out_fp8_row, out_scales_row, out_fp8_col, out_scales_col = allocate_output_tensors(tokens - 32) # any outputs + obj = GemmProjRopeMxfp8Sm100( + sample_x=x, + sample_w=w, + sample_cos=cos, + sample_sin=sin, + sample_out_fp8_row=out_fp8_row, + sample_out_scales_row=out_scales_row, + sample_out_fp8_col=out_fp8_col, + sample_out_scales_col=out_scales_col, + w_out_in=False, + ) + with pytest.raises(ValueError): + obj.check_support() + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_gemm_proj_rope_mxfp8_check_support_rejects_kdim_mismatch(request): + """check_support() must reject x whose contraction dim does not match the weight's.""" + try: + import torch + + from cudnn import GemmProjRopeMxfp8Sm100 + from fe_api.test_gemm_proj_rope_mxfp8_utils import ( + Q_LORA, + allocate_input_tensors, + allocate_output_tensors, + gemm_proj_rope_mxfp8_init, + ) + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + gemm_proj_rope_mxfp8_init(request, tokens=2048, w_out_in=False) # arch skip if needed + tokens = 2048 + x, w, cos, sin = allocate_input_tensors(tokens, w_out_in=False) + x = torch.randn(tokens, Q_LORA + 8, dtype=x.dtype, device=x.device) # inner dim != w's K + out_fp8_row, out_scales_row, out_fp8_col, out_scales_col = allocate_output_tensors(tokens) + obj = GemmProjRopeMxfp8Sm100( + sample_x=x, + sample_w=w, + sample_cos=cos, + sample_sin=sin, + sample_out_fp8_row=out_fp8_row, + sample_out_scales_row=out_scales_row, + sample_out_fp8_col=out_fp8_col, + sample_out_scales_col=out_scales_col, + w_out_in=False, + ) + with pytest.raises(ValueError): + obj.check_support() + + +@pytest.mark.L0 +def test_gemm_proj_rope_mxfp8_run_rejects_kdim_mismatch(): + """run() is a public entry that skips check_support(); it must still reject an x + whose GEMM contraction dim does not match the weight's, before reaching the kernel. + + The structural checks at the top of run() fire before any CUDA work, so plain CPU + tensors are enough to exercise the rejection (no SM100 device required).""" + try: + from cudnn.gemm_proj_rope_mxfp8 import run + from fe_api.test_gemm_proj_rope_mxfp8_utils import ( + BLOCK, + HEAD_DIM, + NUM_HEADS, + QK_ROPE, + Q_LORA, + Q_OUT, + ) + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + tokens = 128 # a multiple of TILE_M so validation reaches the k-dim check + x = torch.empty(tokens, Q_LORA + 8, dtype=torch.bfloat16) # inner dim != w's K + w = torch.empty(Q_LORA, Q_OUT, dtype=torch.bfloat16) # [in, out], K = Q_LORA + cos = torch.empty(tokens, QK_ROPE, dtype=torch.bfloat16) + sin = torch.empty(tokens, QK_ROPE, dtype=torch.bfloat16) + out_fp8_row = torch.empty(tokens, NUM_HEADS, HEAD_DIM, dtype=torch.float8_e4m3fn) + out_scales_row = torch.empty(tokens, NUM_HEADS, HEAD_DIM // BLOCK, dtype=torch.uint8) + out_fp8_col = torch.empty(tokens, NUM_HEADS, HEAD_DIM, dtype=torch.float8_e4m3fn) + out_scales_col = torch.empty(tokens // BLOCK, NUM_HEADS, HEAD_DIM, dtype=torch.uint8) + + with pytest.raises(ValueError): + run(x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col, w_out_in=False) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize("tokens", [128, 256]) +@pytest.mark.parametrize("w_out_in", [False, True]) +def test_gemm_proj_rope_mxfp8_reference_contract(tokens, w_out_in): + """Exercise the PyTorch reference oracle directly (no SM100 kernel required). + + Gives ``gemm_proj_rope_mxfp8_reference`` contract coverage for both weight + layouts and valid token multiples: the four outputs must have the documented + shapes/dtypes and carry only finite (dequantized) values.""" + try: + from cudnn.gemm_proj_rope_mxfp8 import gemm_proj_rope_mxfp8_reference + from fe_api.test_gemm_proj_rope_mxfp8_utils import ( + BLOCK, + HEAD_DIM, + NUM_HEADS, + QK_ROPE, + Q_LORA, + Q_OUT, + ) + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + dev = "cuda" if torch.cuda.is_available() else "cpu" + x = torch.randn(tokens, Q_LORA, dtype=torch.bfloat16, device=dev) * 0.5 + # w is stored [out, in] for w_out_in else [in, out]; both project Q_LORA -> Q_OUT. + w_shape = (Q_OUT, Q_LORA) if w_out_in else (Q_LORA, Q_OUT) + w = torch.randn(*w_shape, dtype=torch.bfloat16, device=dev) * 0.02 + cos = torch.randn(tokens, QK_ROPE, dtype=torch.bfloat16, device=dev) + sin = torch.randn(tokens, QK_ROPE, dtype=torch.bfloat16, device=dev) + + qr, sr, qc, sc = gemm_proj_rope_mxfp8_reference(x, w, cos, sin, w_out_in=w_out_in) + + assert qr.shape == (tokens, NUM_HEADS, HEAD_DIM) + assert sr.shape == (tokens, NUM_HEADS, HEAD_DIM // BLOCK) + assert qc.shape == (tokens, NUM_HEADS, HEAD_DIM) + assert sc.shape == (tokens // BLOCK, NUM_HEADS, HEAD_DIM) + assert qr.dtype == torch.float8_e4m3fn and qc.dtype == torch.float8_e4m3fn + assert sr.dtype == torch.uint8 and sc.dtype == torch.uint8 + # No NaN/Inf should leak through the e8m0/e4m3 quantization path. + assert torch.isfinite(qr.float()).all() and torch.isfinite(qc.float()).all() + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize("tokens", [128, 256]) +def test_gemm_proj_rope_mxfp8_reference_w_out_in_equivalence(tokens): + """The oracle must be layout-invariant: feeding the transposed weight with + w_out_in=True reproduces the w_out_in=False result (both map to the same + logical B=[out, in] operand).""" + try: + from cudnn.gemm_proj_rope_mxfp8 import gemm_proj_rope_mxfp8_reference + from fe_api.test_gemm_proj_rope_mxfp8_utils import QK_ROPE, Q_LORA, Q_OUT + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + dev = "cuda" if torch.cuda.is_available() else "cpu" + x = torch.randn(tokens, Q_LORA, dtype=torch.bfloat16, device=dev) * 0.5 + w_in_out = torch.randn(Q_LORA, Q_OUT, dtype=torch.bfloat16, device=dev) * 0.02 + cos = torch.randn(tokens, QK_ROPE, dtype=torch.bfloat16, device=dev) + sin = torch.randn(tokens, QK_ROPE, dtype=torch.bfloat16, device=dev) + + ref_in_out = gemm_proj_rope_mxfp8_reference(x, w_in_out, cos, sin, w_out_in=False) + ref_out_in = gemm_proj_rope_mxfp8_reference(x, w_in_out.t().contiguous(), cos, sin, w_out_in=True) + + names = ("out_fp8_row", "out_scales_row", "out_fp8_col", "out_scales_col") + for name, a, b in zip(names, ref_in_out, ref_out_in): + # Identical logical operands; allow a tiny fraction of e4m3 boundary flips from + # differing fp32 matmul reduction orders, but the layouts must otherwise agree. + frac_equal = (a.float() == b.float()).float().mean().item() + assert frac_equal >= 0.999, f"{name} differs between w_out_in layouts: matched={frac_equal:.4f}" + + +""" +GemmProjRopeMxfp8 API with explicit check_support, compile, and execute paths. +Use this method when running one static configuration per object. +""" + + +def _test_gemm_proj_rope_mxfp8_compile_execute(tokens, w_out_in, request): + try: + from cudnn import GemmProjRopeMxfp8Sm100 + from cuda.bindings import driver as cuda + from fe_api.test_gemm_proj_rope_mxfp8_utils import ( + allocate_input_tensors, + allocate_output_tensors, + check_ref_gemm_proj_rope_mxfp8, + gemm_proj_rope_mxfp8_init, + ) + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + cfg = gemm_proj_rope_mxfp8_init(request, tokens, w_out_in) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + x, w, cos, sin = allocate_input_tensors(cfg["tokens"], cfg["w_out_in"]) + outputs = allocate_output_tensors(cfg["tokens"]) + + gemm = GemmProjRopeMxfp8Sm100( + sample_x=x, + sample_w=w, + sample_cos=cos, + sample_sin=sin, + sample_out_fp8_row=outputs[0], + sample_out_scales_row=outputs[1], + sample_out_fp8_col=outputs[2], + sample_out_scales_col=outputs[3], + w_out_in=cfg["w_out_in"], + ) + try: + assert gemm.check_support(), "Unsupported testcase" + except (ValueError, NotImplementedError) as e: + pytest.skip(f"Unsupported testcase: {e}") + gemm.compile() + gemm.execute(x, w, cos, sin, *outputs, current_stream=stream) + + check_ref_gemm_proj_rope_mxfp8(x, w, cos, sin, outputs, cfg["w_out_in"], skip_ref=cfg["skip_ref"]) + + +""" +GemmProjRopeMxfp8 API via the high-level wrapper (no explicit setup/compile). +""" + + +def _test_gemm_proj_rope_mxfp8_wrapper(tokens, w_out_in, request): + try: + from cudnn import gemm_proj_rope_mxfp8_wrapper_sm100 + from cuda.bindings import driver as cuda + from fe_api.test_gemm_proj_rope_mxfp8_utils import ( + allocate_input_tensors, + check_ref_gemm_proj_rope_mxfp8, + gemm_proj_rope_mxfp8_init, + ) + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + cfg = gemm_proj_rope_mxfp8_init(request, tokens, w_out_in) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + x, w, cos, sin = allocate_input_tensors(cfg["tokens"], cfg["w_out_in"]) + + try: + for _ in range(2): # run twice to exercise the caching path + out = gemm_proj_rope_mxfp8_wrapper_sm100(x, w, cos, sin, w_out_in=cfg["w_out_in"], stream=stream) + except (ValueError, NotImplementedError) as e: + pytest.skip(f"Unsupported testcase: {e}") + + # TupleDict supports both key access and tuple unpacking. + outputs = (out["out_fp8_row"], out["out_scales_row"], out["out_fp8_col"], out["out_scales_col"]) + check_ref_gemm_proj_rope_mxfp8(x, w, cos, sin, outputs, cfg["w_out_in"], skip_ref=cfg["skip_ref"]) diff --git a/test/python/fe_api/test_gemm_proj_rope_mxfp8_utils.py b/test/python/fe_api/test_gemm_proj_rope_mxfp8_utils.py new file mode 100644 index 000000000..739da58e8 --- /dev/null +++ b/test/python/fe_api/test_gemm_proj_rope_mxfp8_utils.py @@ -0,0 +1,103 @@ +""" +Utilities and parameterization for the fused GEMM + per-head RoPE + MXFP8 projection tests. +Contains test configuration, tensor creation, and reference comparison helpers. +""" + +import torch +import pytest + +# DSv3 Q up-proj shapes the kernel is specialized for. +NUM_HEADS = 128 +QK_ROPE = 64 +HEAD_DIM = 192 +Q_LORA = 1536 +Q_OUT = NUM_HEADS * HEAD_DIM +BLOCK = 32 +TILE_M = 128 +E8M0_BIAS = 127 + + +GEMM_PROJ_ROPE_MXFP8_PARAM_MARKS = [ + pytest.mark.parametrize("tokens", [2048, 4096]), + pytest.mark.parametrize("w_out_in", [False, True]), +] + + +def with_gemm_proj_rope_mxfp8_params(func): + """Apply all parameterization marks to a test function.""" + for mark in reversed(GEMM_PROJ_ROPE_MXFP8_PARAM_MARKS): + func = mark(func) + return func + + +def gemm_proj_rope_mxfp8_init(request, tokens, w_out_in): + """Build test config; skip on unsupported architecture.""" + major, minor = torch.cuda.get_device_capability() + compute_capability = major * 10 + minor + if compute_capability < 100: + pytest.skip(f"Environment not supported: requires compute capability >= 100, found {compute_capability}") + + skip_ref = request.config.getoption("--skip-ref", default=False) + + return { + "tokens": tokens, + "w_out_in": bool(w_out_in), + "skip_ref": skip_ref, + } + + +def allocate_input_tensors(tokens, w_out_in): + """Allocate bf16 activations, projection weight (in the requested layout), and rope tables.""" + dev = "cuda" + x = torch.randn(tokens, Q_LORA, dtype=torch.bfloat16, device=dev) * 0.5 + if w_out_in: + w = torch.randn(Q_OUT, Q_LORA, dtype=torch.bfloat16, device=dev) * 0.02 # [out, in] + else: + w = torch.randn(Q_LORA, Q_OUT, dtype=torch.bfloat16, device=dev) * 0.02 # [in, out] + cos = torch.randn(tokens, QK_ROPE, dtype=torch.bfloat16, device=dev) + sin = torch.randn(tokens, QK_ROPE, dtype=torch.bfloat16, device=dev) + return x, w, cos, sin + + +def allocate_output_tensors(tokens): + """Allocate the four MXFP8 output tensors (rowwise + columnwise data and E8M0 scales).""" + dev = "cuda" + out_fp8_row = torch.empty(tokens, NUM_HEADS, HEAD_DIM, dtype=torch.float8_e4m3fn, device=dev) + out_scales_row = torch.empty(tokens, NUM_HEADS, HEAD_DIM // BLOCK, dtype=torch.uint8, device=dev) + out_fp8_col = torch.empty(tokens, NUM_HEADS, HEAD_DIM, dtype=torch.float8_e4m3fn, device=dev) + out_scales_col = torch.empty(tokens // BLOCK, NUM_HEADS, HEAD_DIM, dtype=torch.uint8, device=dev) + return out_fp8_row, out_scales_row, out_fp8_col, out_scales_col + + +def _deq_row(data, scale): + t = data.shape[0] + inv = torch.pow(2.0, scale.float() - E8M0_BIAS).unsqueeze(-1) + return (data.float().reshape(t, NUM_HEADS, HEAD_DIM // BLOCK, BLOCK) * inv).reshape(t, NUM_HEADS, HEAD_DIM) + + +def _deq_col(data, scale): + t = data.shape[0] + inv = torch.pow(2.0, scale.float() - E8M0_BIAS).reshape(t // BLOCK, 1, NUM_HEADS, HEAD_DIM) + return (data.float().reshape(t // BLOCK, BLOCK, NUM_HEADS, HEAD_DIM) * inv).reshape(t, NUM_HEADS, HEAD_DIM) + + +def _matched(got, ref, atol=0.1, rtol=0.1): + diff = (got.float() - ref.float()).abs() + return (diff <= atol + rtol * ref.float().abs()).float().mean().item() + + +def check_ref_gemm_proj_rope_mxfp8(x, w, cos, sin, outputs, w_out_in, skip_ref=False, need=0.95): + """Compare each of the four public outputs against the PyTorch reference oracle.""" + if skip_ref: + print("Skipping reference check") + return + + from cudnn.gemm_proj_rope_mxfp8 import gemm_proj_rope_mxfp8_reference + + out_fp8_row, out_scales_row, out_fp8_col, out_scales_col = outputs + ref_qr, ref_sr, ref_qc, ref_sc = gemm_proj_rope_mxfp8_reference(x, w, cos, sin, w_out_in=w_out_in) + + row = _matched(_deq_row(out_fp8_row, out_scales_row), _deq_row(ref_qr, ref_sr)) + col = _matched(_deq_col(out_fp8_col, out_scales_col), _deq_col(ref_qc, ref_sc)) + assert row >= need, f"rowwise MXFP8 mismatch vs reference: matched={row:.4f}" + assert col >= need, f"columnwise MXFP8 mismatch vs reference: matched={col:.4f}"