Add an mxfp8-input version of the gemm_proj_rope_mxfp8 kernel.

docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md

@@ -4,21 +4,27 @@
 
 ## 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.
+**Fused projection GEMM + per-head YARN RoPE + dual-direction MXFP8 quantize**: a persistent dense GEMM on NVIDIA Blackwell GPUs (SM100+) that projects 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.
+Two sibling kernels implement the same operation, differing only in the **GEMM input precision**, and are selected by the dtype of `x`/`w`:
 
-- **Inputs**: bf16 activations `x`, bf16 projection weight `w`, and bf16 rotary tables `cos`/`sin`.
+- **BF16 input** (`GemmProjRopeMxfp8Bf16InSm100`): `x` and `w` are `bfloat16`; the GEMM runs in bf16.
+- **MXFP8 input** (`GemmProjRopeMxfp8Mxfp8InSm100`): `x` and `w` are pre-quantized MXFP8 (E4M3 codes + E8M0 rowwise block scales); the GEMM runs in MXFP8.
+
+The **output is dual-direction MXFP8 in both cases**. 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**: activations `x`, projection weight `w`, and bf16 rotary tables `cos`/`sin`; for the MXFP8-input path, the E8M0 block scales `x_scale`/`w_scale` as well.
 - **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)`
+  - `x`: `(tokens, Q_LORA)` — `tokens % TILE_M == 0`. Dtype `bfloat16` (bf16 path) or `float8_e4m3fn` (MXFP8 path).
+  - `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`. Same dtype as `x`.
+  - `x_scale`, `w_scale` (MXFP8 path only): E8M0 rowwise block scales, `(tokens, Q_LORA // BLOCK)` and `(NUM_HEADS·HEAD_DIM, Q_LORA // BLOCK)`, `uint8`.
+  - `cos`, `sin`: `(tokens, QK_ROPE)`, `bfloat16`.
 
 - **Outputs**
   - `out_fp8_row`, `out_fp8_col`: `(tokens, NUM_HEADS, HEAD_DIM)`
@@ -50,8 +56,8 @@ $$
 ### Diagram
 
 ```text
-x (tokens x Q_LORA), w (Q_LORA x NUM_HEADS*HEAD_DIM)
-     |  GEMM
+x (tokens x Q_LORA), w (Q_LORA x NUM_HEADS*HEAD_DIM)     [bf16, or MXFP8 codes+scales]
+     |  GEMM (bf16 or MXFP8, per input dtype)
      v
    Y (tokens x NUM_HEADS x HEAD_DIM)  ---- per-head YARN RoPE on trailing QK_ROPE
      |
@@ -61,57 +67,68 @@ x (tokens x Q_LORA), w (Q_LORA x NUM_HEADS*HEAD_DIM)
 
 ## API Usage
 
-### High-level wrapper
+### High-level wrapper (dtype-dispatch)
+
+Selects the bf16- or mxfp8-input kernel by `x`/`w` dtype (which must match); pass `x_scale`/`w_scale` only for the MXFP8 path.
+
 ```python
+# BF16 inputs
+result = gemm_proj_rope_mxfp8_wrapper_sm100(x, w, cos, sin, w_out_in=True, stream=None)
+
+# MXFP8 inputs (E4M3 codes + E8M0 scales)
 result = gemm_proj_rope_mxfp8_wrapper_sm100(
-    x,
-    w,
-    cos,
-    sin,
-    w_out_in=False,
-    stream=None,
+    x_code, w_code, cos, sin, x_scale=x_scale, w_scale=w_scale, w_out_in=True, 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
+### Class API — BF16 input
 ```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,
+op = GemmProjRopeMxfp8Bf16InSm100(
+    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)
 ```
 
+### Class API — MXFP8 input
+```python
+op = GemmProjRopeMxfp8Mxfp8InSm100(
+    sample_x_code=x_code, sample_x_scale=x_scale, sample_w_code=w_code, sample_w_scale=w_scale,
+    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,
+)   # weight is TE-native [out, in]; the wrapper transposes the [in, out] layout for you
+assert op.check_support()
+op.compile()
+op.execute(x_code, x_scale, w_code, w_scale, 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)
+- Input **x**: `(tokens, Q_LORA)`; Dtype `bfloat16` (bf16 path) or `float8_e4m3fn` (MXFP8 path).
+- Input **w**: `(Q_LORA, NUM_HEADS·HEAD_DIM)` (`w_out_in=False`) or `(NUM_HEADS·HEAD_DIM, Q_LORA)` (`w_out_in=True`); same dtype as **x**.
+- Input **x_scale**, **w_scale** (MXFP8 path only): `(tokens, Q_LORA // BLOCK)` and `(NUM_HEADS·HEAD_DIM, Q_LORA // BLOCK)`; Dtype `uint8` (E8M0 rowwise block scales).
+- Input **cos**, **sin**: `(tokens, QK_ROPE)`; Dtype `bfloat16`.
+- Output **out_fp8_row** / **out_fp8_col**: `(tokens, NUM_HEADS, HEAD_DIM)`; Dtype `float8_e4m3fn`.
+- Output **out_scales_row**: `(tokens, NUM_HEADS, HEAD_DIM // BLOCK)`; Dtype `uint8` (E8M0).
+- Output **out_scales_col**: `(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`
+- `w_out_in: bool` — whether `w` is stored `[out, in]` (`True`) or `[in, out]` (`False`). Wrapper default: `True`. On the bf16 path the kernel consumes both via the cutlass major mode (no transposed copy); on the MXFP8 path the wrapper transposes the code + scale for `[in, out]`.
+- `x_scale`, `w_scale: Optional[Tensor]` — required for `float8_e4m3fn` inputs; must be `None` for `bfloat16`. `x.dtype == w.dtype` is asserted.
 - CUDA stream (`current_stream` in class API, `stream` in wrapper). Defaults to the current torch stream (required for CUDA-graph capture).
 
 ### Wrapper return values
@@ -123,12 +140,13 @@ Returns a `TupleDict` with keys `out_fp8_row`, `out_scales_row`, `out_fp8_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).
+- BF16 path: `x`, `w`, `cos`, `sin` are `bfloat16`.
+- MXFP8 path: `x`, `w` are `float8_e4m3fn` codes; `x_scale`, `w_scale` are `uint8` (E8M0); `cos`, `sin` are `bfloat16`.
+- Outputs (both paths): `out_fp8_row`, `out_fp8_col` are `float8_e4m3fn`; `out_scales_row`, `out_scales_col` are `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`.
+- The projected weight dimension must equal `NUM_HEADS·HEAD_DIM`. On the MXFP8 path, `Q_LORA % BLOCK == 0`.
 
 ### Environment
 - Requires CUDA with SM100+ compute capability.
@@ -137,7 +155,12 @@ Returns a `TupleDict` with keys `out_fp8_row`, `out_scales_row`, `out_fp8_col`,
 
 ## 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`.
+Integrated from the DeepSeek-V3 MLA fused Q up-projection kernel developed for Megatron-LM MXFP8 training (Blackwell / customte CUTLASS 4.4.1); originally added in the GEMM+RoPE+MXFP8 fusion commit. The BF16-GEMM and MXFP8-GEMM variants were consolidated into two input-precision kernel modules selected by input dtype:
+
+- `python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8_bf16in.py` — BF16-input kernel; also hosts the pure-PyTorch oracle `gemm_proj_rope_mxfp8_reference(...)`.
+- `python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.py` — MXFP8-input kernel.
+
+The compiled-kernel lifecycle (`check_support`/`compile`/`execute`) lives in the APIBase classes in `api.py`; the wrapper caches the compiled objects (matching the sibling GEMM-fusion packages).
 
 ## Installation
 

python/cudnn/__init__.py

@@ -316,7 +316,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"),
+    "GemmProjRopeMxfp8Bf16InSm100": (".gemm_proj_rope_mxfp8", "GemmProjRopeMxfp8Bf16InSm100"),
+    "GemmProjRopeMxfp8Mxfp8InSm100": (".gemm_proj_rope_mxfp8", "GemmProjRopeMxfp8Mxfp8InSm100"),
     "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"),

python/cudnn/gemm_proj_rope_mxfp8/__init__.py

@@ -1,15 +1,13 @@
 from .api import (
-    GemmProjRopeMxfp8Sm100,
+    GemmProjRopeMxfp8Bf16InSm100,
+    GemmProjRopeMxfp8Mxfp8InSm100,
     gemm_proj_rope_mxfp8_wrapper_sm100,
 )
-from .gemm_proj_rope_mxfp8 import (
-    run,
-    gemm_proj_rope_mxfp8_reference,
-)
+from .gemm_proj_rope_mxfp8_bf16in import gemm_proj_rope_mxfp8_reference
 
 __all__ = [
-    "GemmProjRopeMxfp8Sm100",
-    "gemm_proj_rope_mxfp8_reference",
+    "GemmProjRopeMxfp8Bf16InSm100",
+    "GemmProjRopeMxfp8Mxfp8InSm100",
     "gemm_proj_rope_mxfp8_wrapper_sm100",
-    "run",
+    "gemm_proj_rope_mxfp8_reference",
 ]