fix

Author: w112009931-bit

cudnn_frontend_op_gap.md

@@ -156,17 +156,17 @@ FlagDNN 对标 cuDNN Frontend 时，不建议只补 eager 单算子。每个新
 
 | cuDNN op | 参数 | FlagDNN 状态 | 建议实现位置 |
 |---|---|---|---|
-| `matmul` | `A, B, compute_data_type=NOT_SET, padding=0.0, name=''` | 部分；已有 2D `mm`、3D batched Graph/Triton 和 general broadcast fallback，仍缺非零 `padding` 语义 | `ops/matmul.py` 继续补 padding |
+| `matmul` | `A, B, compute_data_type=NOT_SET, padding=0.0, name=''` | 已有；2D 走 `mm`，3D 同 batch 走 Triton 快路径，其它 rank/batch broadcast 也走 Triton batched kernel；`padding` 按 cuDNN backend padding-value 兼容参数接收，标准 K 匹配时不改变数学结果 | 已有；后续可补动态 K override/低精度专用 GEMM |
 | `moe_grouped_matmul` | `token, weight, first_token_offset, token_index=None, token_ks=None, mode=NONE, compute_data_type=FLOAT, top_k=0, name=''` | 缺失 | `ops/moe_grouped_matmul.py`，graph kernel |
 | `moe_grouped_matmul_bwd` | `doutput, token, first_token_offset, compute_data_type=FLOAT, name=''` | 缺失 | `ops/moe_grouped_matmul_bwd.py` |
 
-普通 `matmul` 已覆盖 2D、同 batch 3D 优化路径和其它 rank/broadcast fallback；cuDNN `padding` 是 backend plan padding value，仍需单独定义 FlagDNN 侧等价语义。
+普通 `matmul` 已覆盖 2D、同 batch 3D 优化路径和其它 rank/broadcast；cuDNN `padding` 是 backend matmul descriptor 的 padding value，FlagDNN 侧作为兼容参数接收，shape/K 仍按标准 matmul 校验。
 
 ### 4.7 Reduction / stats
 
 | cuDNN op | 参数 | FlagDNN 状态 | 建议实现位置 |
 |---|---|---|---|
-| `reduction` | `input, mode, compute_data_type=NOT_SET, name=''` | 已有 eager + Graph；支持 `ADD/AVG/MUL/MIN/MAX/AMAX/NORM1/NORM2/MUL_NO_ZEROS`，FlagDNN graph 需显式 `dim/keepdim`；当前 cuDNN standalone reduction 不给 engine | 已有；后续可补 cuDNN 输出 shape 推导式 API |
+| `reduction` | `input, mode, compute_data_type=NOT_SET, name=''` | 已有 eager + Graph；支持 `ADD/AVG/MUL/MIN/MAX/AMAX/NORM1/NORM2/MUL_NO_ZEROS`，其中 `MIN/MAX/AMAX/MUL_NO_ZEROS` 已改为 Triton 原生实现；FlagDNN graph 需显式 `dim/keepdim`；当前 cuDNN standalone reduction 不给 engine | 已有；后续可补 cuDNN 输出 shape 推导式 API |
 | `genstats` | `input, compute_data_type=NOT_SET, name=''` -> `[mean, inv_variance]` | 缺失 | `ops/genstats.py` |
 
 cuDNN `reduction_mode` 已映射到 FlagDNN graph 的显式 `dim/keepdim` reduction；本环境 cuDNN standalone reduction 无有效 engine，因此功能测试保留 torch-reference 覆盖。
@@ -175,7 +175,7 @@ cuDNN `reduction_mode` 已映射到 FlagDNN graph 的显式 `dim/keepdim` reduct
 
 | cuDNN op | 参数 | FlagDNN 状态 | 建议实现位置 |
 |---|---|---|---|
-| `batchnorm` | `input, scale, bias, in_running_mean, in_running_var, epsilon, momentum, peer_stats=[], compute_data_type=NOT_SET, name=''` -> list | 部分；已有 cudnn-style eager + Graph 多输出 forward，返回 `Y/mean/inv_var/next_running_mean/next_running_var`；`peer_stats` 未支持 | `ops/batchnorm.py`；后续补 peer_stats/backward |
+| `batchnorm` | `input, scale, bias, in_running_mean, in_running_var, epsilon, momentum, peer_stats=[], compute_data_type=NOT_SET, name=''` -> list | 部分；已有 cudnn-style eager + Graph 多输出 forward，返回 `Y/mean/inv_var/next_running_mean/next_running_var`；单设备 `peer_stats=[tensor]` 已可透传并按本地 BN 语义执行，多设备同步 peer stats 仍未支持 | `ops/batchnorm.py`；后续补多 GPU peer_stats/backward |
 | `batchnorm_inference` | `input, mean, inv_variance, scale, bias, compute_data_type=NOT_SET, name=''` | 已有 eager + Graph；功能/性能已改用 legacy cuDNN standalone `cudnnBatchNormalizationForwardInference` 对标；cuDNN Frontend graph standalone 在本环境无 plan | 已有 |
 | `batchnorm_backward` | `grad, input, scale, mean, inv_variance, peer_stats=[], compute_data_type=NOT_SET, name=''` | 缺失 | `ops/batch_norm_backward.py` |
 | `layernorm` | `norm_forward_phase, input, scale, bias, epsilon, compute_data_type=NOT_SET, name=''` -> list | 已有 cudnn-style eager + Graph 多输出 forward，返回 `Y/mean/inv_var`；保留原 `layer_norm` API | 已有；backward 另补 |
@@ -187,7 +187,7 @@ cuDNN `reduction_mode` 已映射到 FlagDNN graph 的显式 `dim/keepdim` reduct
 | `adalayernorm` | `norm_forward_phase, input, scale, bias=None, epsilon, compute_data_type=NOT_SET, name=''` | 缺失 | `ops/adalayer_norm.py` |
 | `adalayernorm_backward` | `grad, input, scale, mean, inv_variance, compute_data_type=NOT_SET, name=''` | 缺失 | `ops/adalayer_norm_backward.py` |
 
-Norm 的 cuDNN API 返回 list，通常包含 output 和保存给 backward 的 stats。FlagDNN graph 已启用 multi-output，并新增 cudnn-style `batchnorm/layernorm/rmsnorm` forward tuple 返回；backward 和少数分布式/peer stats 仍需单独补。
+Norm 的 cuDNN API 返回 list，通常包含 output 和保存给 backward 的 stats。FlagDNN graph 已启用 multi-output，并新增 cudnn-style `batchnorm/layernorm/rmsnorm` forward tuple 返回；backward 和真正跨设备同步 peer stats 仍需单独补。
 
 ### 4.9 Quant / FP8
 
@@ -239,7 +239,7 @@ Norm 的 cuDNN API 返回 list，通常包含 output 和保存给 backward 的 s
 | `gemm_swiglu_wrapper_sm100` | `a_tensor, b_tensor, alpha=1.0, c_major='n', ab12_dtype=float32, c_dtype=float16, acc_dtype=float32, mma_tiler_mn=(128,128), cluster_shape_mn=None, sfa_tensor=None, sfb_tensor=None, norm_const_tensor=None, sf_vec_size=16, vector_f32=False, ab12_stages=4, stream=None` | 缺失 | `ops/gemm_swiglu.py` |
 | `grouped_gemm_*_wrapper_sm100` | grouped/MoE variants with `padded_offsets`, `alpha_tensor`, optional pointer tensors, scale tensors, `mma_tiler_mn`, `cluster_shape_mn`, `sf_vec_size`, scheduler flags | 缺失 | `ops/grouped_gemm.py` plus specialized graph kernels |
 | `discrete_grouped_gemm_*_wrapper_sm100` | pointer-based discrete grouped GEMM variants | 缺失 | `ops/discrete_grouped_gemm.py` |
-| `rmsnorm_rht_amax_wrapper_sm100` | `x_tensor, w_tensor, eps=1e-5, num_threads=None, rows_per_cta=None, current_stream=None` | 部分；有 rms_norm，但缺 RHT 和 amax | `ops/rmsnorm_rht_amax.py` |
+| `rmsnorm_rht_amax_wrapper_sm100` | `x_tensor, w_tensor, eps=1e-5, num_threads=None, rows_per_cta=None, current_stream=None` | 已有；Triton 实现 `RMSNorm + 16 点 normalized RHT + per-CTA amax`，支持官方 wrapper 的 2D/尾部 1 维 squeeze、bf16 row-major 输入和 `rows_per_cta/num_threads` 校验；不是 SM100 CUTE 专用优化 kernel | 已有；后续可补 SM100/CUTE 专用性能路径 |
 
 ## 5. 建议新增文件清单
 

src/flag_dnn/__init__.py

@@ -48,6 +48,7 @@
     relu,
     rms_norm,
     rmsnorm,
+    rmsnorm_rht_amax_wrapper_sm100,
     silu,
     swish,
     softmax,
@@ -512,6 +513,7 @@ def all_registered_keys():
     "matmul",
     "batchnorm_inference",
     "reduction",
+    "rmsnorm_rht_amax_wrapper_sm100",
     "compile",
     "current_capture",
     "get_default_plan_cache",

src/flag_dnn/graph/registry.py

@@ -1093,19 +1093,25 @@ def _normalize_batchnorm(
     if missing:
         raise TypeError(f"batchnorm missing {missing[0]}")
     peer_stats = params.pop("peer_stats", [])
-    if peer_stats:
-        raise NotImplementedError(
-            "FlagDNN graph batchnorm does not support peer_stats"
-        )
+    if peer_stats is None:
+        peer_stats = []
+    if not isinstance(peer_stats, (list, tuple)):
+        raise TypeError("batchnorm peer_stats must be a list or tuple")
     attrs = {
         "compute_data_type": params.pop("compute_data_type", None),
         "name": params.pop("name", ""),
+        "peer_stats_count": len(peer_stats),
     }
     if params:
         raise TypeError(
             f"batchnorm got unsupported graph attrs: {sorted(params)}"
         )
-    return [ctx.as_value(values[name], name) for name in names], attrs
+    input_ids = [ctx.as_value(values[name], name) for name in names]
+    input_ids.extend(
+        ctx.as_value(peer_stat, f"peer_stats_{index}")
+        for index, peer_stat in enumerate(peer_stats)
+    )
+    return input_ids, attrs
 
 
 def _normalize_causal_conv1d(
@@ -2356,6 +2362,7 @@ def run(inputs: list[Any], attrs: dict[str, Any]) -> torch.Tensor:
 def _run_batchnorm(flag_ops: Any) -> RunFn:
     def run(inputs: list[Any], attrs: dict[str, Any]) -> Any:
         _require_runtime_backend(inputs[:5], "batchnorm")
+        peer_count = int(attrs.get("peer_stats_count", 0))
         return flag_ops.batchnorm(
             inputs[0],
             inputs[1],
@@ -2364,6 +2371,7 @@ def run(inputs: list[Any], attrs: dict[str, Any]) -> Any:
             inputs[4],
             inputs[5],
             inputs[6],
+            peer_stats=inputs[7 : 7 + peer_count],
             compute_data_type=attrs.get("compute_data_type"),
             name=attrs.get("name", ""),
         )

src/flag_dnn/ops/__init__.py

@@ -24,6 +24,7 @@
 from flag_dnn.ops.layernorm import layernorm
 from flag_dnn.ops.rms_norm import rms_norm
 from flag_dnn.ops.rmsnorm import rmsnorm
+from flag_dnn.ops.rmsnorm_rht_amax import rmsnorm_rht_amax_wrapper_sm100
 from flag_dnn.ops.group_norm import group_norm
 from flag_dnn.ops.max_pool2d import max_pool2d
 from flag_dnn.ops.avg_pool2d import avg_pool2d
@@ -216,6 +217,7 @@ def gelu_approx_tanh(input, *, compute_data_type=None, name=""):
     "layer_norm",
     "rms_norm",
     "rmsnorm",
+    "rmsnorm_rht_amax_wrapper_sm100",
     "group_norm",
     "max_pool2d",
     "avg_pool2d",

src/flag_dnn/ops/batch_norm.py

@@ -8,7 +8,6 @@
 from flag_dnn.runtime import torch_device_fn
 from flag_dnn.utils import triton_lang_extension as tle
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -243,10 +242,22 @@ def batchnorm_forward(
     momentum,
     peer_stats=None,
 ):
-    if peer_stats:
+    if peer_stats is None:
+        peer_stats = []
+    peer_stats = list(peer_stats)
+    if len(peer_stats) > 1:
         raise NotImplementedError(
-            "flag_dnn batchnorm does not support peer_stats"
+            "flag_dnn batchnorm supports only single-device peer_stats; "
+            "multi-device synchronized peer stats are not implemented"
         )
+    if peer_stats:
+        peer_stat = peer_stats[0]
+        if peer_stat.device != input.device:
+            raise RuntimeError(
+                "batchnorm peer_stats tensor must be on the input device"
+            )
+        if peer_stat.dtype != torch.float32:
+            raise RuntimeError("batchnorm peer_stats tensor must be float32")
     if input.dim() < 2:
         raise RuntimeError("batchnorm expects input rank >= 2")
     if not input.is_contiguous():

src/flag_dnn/ops/matmul.py

@@ -54,7 +54,7 @@ def _batched_matmul_kernel(
             mask=(k[:, None] < K) & (offs_n[None, :] < N),
             other=0.0,
         )
-        acc += tl.dot(a, b)
+        acc += tl.dot(a, b, input_precision="ieee")
 
     c_base = c_ptr + bid * stride_cb
     tl.store(
@@ -109,6 +109,90 @@ def _batched_matmul_3d(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
     return C
 
 
+_BATCHED_TRITON_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+
+
+def _prod(shape: tuple[int, ...]) -> int:
+    total = 1
+    for dim in shape:
+        total *= int(dim)
+    return total
+
+
+def _check_matmul_inputs(A: torch.Tensor, B: torch.Tensor) -> None:
+    if A.dim() < 2 or B.dim() < 2:
+        raise NotImplementedError(
+            "flag_dnn matmul requires tensors with rank >= 2"
+        )
+    if A.shape[-1] != B.shape[-2]:
+        raise RuntimeError(
+            "mat1 and mat2 shapes cannot be multiplied "
+            f"({A.shape[-2]}x{A.shape[-1]} and "
+            f"{B.shape[-2]}x{B.shape[-1]})"
+        )
+    if A.device != B.device:
+        raise RuntimeError(
+            "matmul: input tensors must be on the same device, "
+            f"but got {A.device} and {B.device}"
+        )
+    if A.dtype != B.dtype:
+        raise RuntimeError(
+            "expected mat1 and mat2 to have the same dtype, but got: "
+            f"{A.dtype} != {B.dtype}"
+        )
+    if A.layout != torch.strided or B.layout != torch.strided:
+        raise NotImplementedError(
+            "flag_dnn matmul supports dense strided tensors only"
+        )
+
+
+def _broadcast_batch_shape(
+    A: torch.Tensor, B: torch.Tensor
+) -> tuple[int, ...]:
+    try:
+        return tuple(
+            torch.broadcast_shapes(tuple(A.shape[:-2]), tuple(B.shape[:-2]))
+        )
+    except RuntimeError as exc:
+        raise RuntimeError(
+            "matmul batch dimensions are not broadcastable: "
+            f"{tuple(A.shape)} x {tuple(B.shape)}"
+        ) from exc
+
+
+def _as_batched_contiguous(
+    tensor: torch.Tensor, batch_shape: tuple[int, ...]
+) -> torch.Tensor:
+    matrix_shape = (int(tensor.shape[-2]), int(tensor.shape[-1]))
+    batch_count = _prod(batch_shape)
+    target_shape = batch_shape + matrix_shape
+
+    if tuple(tensor.shape[:-2]) == batch_shape:
+        batched = tensor.reshape(batch_count, *matrix_shape)
+    else:
+        batched = tensor.expand(target_shape).reshape(
+            batch_count, *matrix_shape
+        )
+    return batched if batched.is_contiguous() else batched.contiguous()
+
+
+def _batched_matmul_broadcast(
+    A: torch.Tensor, B: torch.Tensor
+) -> torch.Tensor:
+    if A.dtype not in _BATCHED_TRITON_DTYPES:
+        raise NotImplementedError(
+            "flag_dnn batched matmul supports fp16, bf16, and fp32 inputs"
+        )
+
+    batch_shape = _broadcast_batch_shape(A, B)
+    m = int(A.shape[-2])
+    n = int(B.shape[-1])
+    A3 = _as_batched_contiguous(A, batch_shape)
+    B3 = _as_batched_contiguous(B, batch_shape)
+    C3 = _batched_matmul_3d(A3, B3)
+    return C3.reshape(batch_shape + (m, n))
+
+
 def matmul(
     A: torch.Tensor,
     B: torch.Tensor,
@@ -117,19 +201,19 @@ def matmul(
     padding: float = 0.0,
     name: str = "",
 ) -> torch.Tensor:
-    del compute_data_type, name
-    if float(padding) != 0.0:
-        raise NotImplementedError("flag_dnn matmul does not support padding")
+    del compute_data_type, padding, name
+    _check_matmul_inputs(A, B)
+
     if A.dim() == 2 and B.dim() == 2:
         return mm(A, B)
-    if A.dim() == 3 and B.dim() == 3 and A.shape[0] == B.shape[0]:
+
+    if (
+        A.dim() == 3
+        and B.dim() == 3
+        and A.shape[0] == B.shape[0]
+        and A.is_contiguous()
+        and B.is_contiguous()
+    ):
         return _batched_matmul_3d(A, B)
-    if A.dim() < 2 or B.dim() < 2:
-        raise NotImplementedError(
-            "flag_dnn matmul requires tensors with rank >= 2"
-        )
-    # TODO: implement the remaining torch.matmul broadcasting cases in Triton.
-    raise NotImplementedError(
-        "flag_dnn matmul currently supports 2D x 2D and matching-batch "
-        "3D x 3D inputs only"
-    )
+
+    return _batched_matmul_broadcast(A, B)