add

w112009931-bit · w112009931-bit · commit 8c48f5314fb0 · 2026-05-29T20:13:43.000+08:00
diff --git a/benchmark_graph/consts.py b/benchmark_graph/consts.py
@@ -187,3 +187,9 @@ def compile_options():
     ((8, 8, 32, 32), 1, "AVG"),
     ((8, 4, 16, 16), 1, "MUL"),
 )
+
+
+RMSNORM_RHT_AMAX_SHAPES = (
+    (64, 2048, 2),
+    (128, 4096, 4),
+)
diff --git a/benchmark_graph/test_rmsnorm_rht_amax_perf.py b/benchmark_graph/test_rmsnorm_rht_amax_perf.py
@@ -0,0 +1,97 @@
+import pytest
+import torch
+
+import flag_dnn
+from benchmark_graph import consts
+from benchmark_graph.base import CudnnCompareBenchmark, get_cudnn
+
+EPS = 1e-5
+
+
+class RmsNormRhtAmaxBenchmark(CudnnCompareBenchmark):
+    op_name = "rmsnorm_rht_amax_wrapper_sm100"
+    dtypes = (torch.bfloat16,)
+    shapes = consts.RMSNORM_RHT_AMAX_SHAPES
+    shape_ids_env = "FLAGDNN_CUDNN_RMSNORM_RHT_AMAX_PERF_SHAPE_IDS"
+
+    def make_inputs(self, case, dtype):
+        m, n, rows_per_cta = case
+        self.rows_per_cta = rows_per_cta
+        x = torch.randn(
+            (m, n), device=flag_dnn.device, dtype=dtype
+        ).contiguous()
+        w = torch.randn((n,), device=flag_dnn.device, dtype=dtype).contiguous()
+        return x, w
+
+    def build_cudnn_runner(self, inputs):
+        cudnn = get_cudnn()
+        try:
+            wrapper = getattr(cudnn, "rmsnorm_rht_amax_wrapper_sm100")
+        except (AttributeError, ImportError, RuntimeError) as exc:
+            pytest.skip(
+                f"cuDNN RMSNorm RHT amax wrapper is unavailable: {exc}"
+            )
+        x, w = inputs
+
+        def run():
+            return wrapper(x, w, eps=EPS, rows_per_cta=self.rows_per_cta)
+
+        try:
+            run()
+            torch.cuda.synchronize()
+        except (AssertionError, ImportError, RuntimeError, ValueError) as exc:
+            pytest.skip(
+                "cuDNN RMSNorm RHT amax wrapper is unsupported here: " f"{exc}"
+            )
+        return run
+
+    def build_flag_dnn_runner(self, inputs):
+        x, w = inputs
+
+        @flag_dnn.graph
+        def flag_dnn_rmsnorm_rht_amax_graph(x, w):
+            return flag_dnn.rmsnorm_rht_amax_wrapper_sm100(
+                x, w, eps=EPS, rows_per_cta=self.rows_per_cta
+            )
+
+        compiled = flag_dnn.compile(
+            flag_dnn_rmsnorm_rht_amax_graph,
+            inputs=[
+                flag_dnn.TensorSpec.from_tensor(x, "x"),
+                flag_dnn.TensorSpec.from_tensor(w, "w"),
+            ],
+            options=consts.compile_options(),
+        )
+        assert [node.op_type for node in compiled.graph.nodes] == [
+            "rmsnorm_rht_amax_wrapper_sm100"
+        ]
+
+        def run():
+            return compiled.run(x, w)
+
+        return run
+
+    def transfer_bytes(self, inputs):
+        x, w = inputs
+        num_ctas = x.shape[0] // self.rows_per_cta
+        return (
+            x.numel() * x.element_size()
+            + w.numel() * w.element_size()
+            + x.numel() * x.element_size()
+            + num_ctas * torch.empty((), dtype=torch.float32).element_size()
+        )
+
+    def shape_detail(self, inputs):
+        x, w = inputs
+        return [x.size(), w.size(), self.rows_per_cta]
+
+
+@pytest.mark.cudnn_frontend
+@pytest.mark.rmsnorm_rht_amax
+@pytest.mark.graph
+@pytest.mark.perf
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+@pytest.mark.parametrize("dtype", RmsNormRhtAmaxBenchmark.dtypes)
+def test_perf_graph_rmsnorm_rht_amax_vs_cudnn_frontend(cudnn_handle, dtype):
+    torch.manual_seed(0)
+    RmsNormRhtAmaxBenchmark(cudnn_handle).run(dtype)
diff --git a/src/flag_dnn/graph/registry.py b/src/flag_dnn/graph/registry.py
@@ -931,6 +931,27 @@ def _layernorm_shape(
     ]
 
 
+_RMSNORM_RHT_AMAX_RPC_CANDIDATES = (2, 4, 8)
+_RMSNORM_RHT_AMAX_TARGET_MIN_CTAS = 148
+
+
+def _rmsnorm_rht_amax_pick_rows_per_cta(m: int) -> int:
+    for rows_per_cta in reversed(_RMSNORM_RHT_AMAX_RPC_CANDIDATES):
+        if m % rows_per_cta != 0:
+            continue
+        if m // rows_per_cta >= _RMSNORM_RHT_AMAX_TARGET_MIN_CTAS:
+            return rows_per_cta
+    return _RMSNORM_RHT_AMAX_RPC_CANDIDATES[0]
+
+
+def _squeeze_trailing_unit_spec_shape(
+    shape: tuple[Any, ...], expected_rank: int
+) -> tuple[Any, ...]:
+    if len(shape) == expected_rank + 1 and shape[-1] == 1:
+        return shape[:-1]
+    return shape
+
+
 def _rmsnorm_shape(
     input_specs: list[TensorSpec], attrs: dict[str, Any]
 ) -> list[TensorSpec]:
@@ -950,6 +971,68 @@ def _rmsnorm_shape(
     ]
 
 
+def _rmsnorm_rht_amax_shape(
+    input_specs: list[TensorSpec], attrs: dict[str, Any]
+) -> list[TensorSpec]:
+    inp = input_specs[0]
+    weight = input_specs[1]
+    x_shape = _squeeze_trailing_unit_spec_shape(tuple(inp.shape), 2)
+    w_shape = _squeeze_trailing_unit_spec_shape(tuple(weight.shape), 1)
+    if len(x_shape) != 2:
+        raise RuntimeError(
+            "rmsnorm_rht_amax_wrapper_sm100 x_tensor must be 2D"
+        )
+    if len(w_shape) != 1:
+        raise RuntimeError(
+            "rmsnorm_rht_amax_wrapper_sm100 w_tensor must be 1D"
+        )
+    m, n = x_shape
+    if isinstance(n, int) and isinstance(w_shape[0], int) and w_shape[0] != n:
+        raise RuntimeError(
+            "rmsnorm_rht_amax_wrapper_sm100 w_tensor length must match "
+            "x hidden dimension"
+        )
+    if isinstance(n, int) and n % 16 != 0:
+        raise RuntimeError(
+            "rmsnorm_rht_amax_wrapper_sm100 N must be divisible by 16"
+        )
+
+    rows_per_cta = attrs.get("rows_per_cta")
+    if rows_per_cta is None:
+        if not isinstance(m, int):
+            raise RuntimeError(
+                "rmsnorm_rht_amax_wrapper_sm100 requires concrete M when "
+                "rows_per_cta is omitted"
+            )
+        rows_per_cta = _rmsnorm_rht_amax_pick_rows_per_cta(m)
+    rows_per_cta = int(rows_per_cta)
+    if rows_per_cta <= 0:
+        raise RuntimeError(
+            "rmsnorm_rht_amax_wrapper_sm100 rows_per_cta must be positive"
+        )
+    if isinstance(m, int):
+        if m % rows_per_cta != 0:
+            raise RuntimeError(
+                "rmsnorm_rht_amax_wrapper_sm100 M must be divisible by "
+                "rows_per_cta"
+            )
+        amax_shape = (m // rows_per_cta,)
+    else:
+        amax_shape = (m,)
+
+    return [
+        TensorSpec(
+            name="o_tensor",
+            shape=x_shape,
+            dtype=inp.dtype,
+            layout="contiguous",
+            device=inp.device,
+            contiguous=True,
+        ),
+        _float32_spec(amax_shape, inp.device).with_name("amax_tensor"),
+    ]
+
+
 def _batchnorm_shape(
     input_specs: list[TensorSpec], attrs: dict[str, Any]
 ) -> list[TensorSpec]:
@@ -1068,6 +1151,51 @@ def _normalize_rmsnorm(
     return input_ids, attrs
 
 
+def _normalize_rmsnorm_rht_amax(
+    ctx: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+) -> tuple[list[int], dict]:
+    names = ("x_tensor", "w_tensor")
+    if len(args) > len(names):
+        raise TypeError(
+            "rmsnorm_rht_amax_wrapper_sm100 got too many positional args"
+        )
+    params = dict(kwargs)
+    values = {}
+    for name, value in zip(names, args):
+        values[name] = value
+    for name in names[len(args) :]:
+        if name in params:
+            values[name] = params.pop(name)
+    missing = [name for name in names if name not in values]
+    if missing:
+        raise TypeError(f"rmsnorm_rht_amax_wrapper_sm100 missing {missing[0]}")
+    current_stream = params.pop("current_stream", None)
+    if current_stream is not None:
+        raise TypeError(
+            "rmsnorm_rht_amax_wrapper_sm100 current_stream is not supported "
+            "in graph capture"
+        )
+    attrs = {
+        "eps": float(params.pop("eps", 1e-5)),
+        "num_threads": params.pop("num_threads", None),
+        "rows_per_cta": params.pop("rows_per_cta", None),
+        "name": params.pop("name", ""),
+    }
+    if attrs["num_threads"] is not None:
+        attrs["num_threads"] = int(attrs["num_threads"])
+    if attrs["rows_per_cta"] is not None:
+        attrs["rows_per_cta"] = int(attrs["rows_per_cta"])
+    if params:
+        raise TypeError(
+            "rmsnorm_rht_amax_wrapper_sm100 got unsupported graph attrs: "
+            f"{sorted(params)}"
+        )
+    return [
+        ctx.as_value(values["x_tensor"], "x_tensor"),
+        ctx.as_value(values["w_tensor"], "w_tensor"),
+    ], attrs
+
+
 def _normalize_batchnorm(
     ctx: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
 ) -> tuple[list[int], dict]:
@@ -2416,6 +2544,21 @@ def run(inputs: list[Any], attrs: dict[str, Any]) -> Any:
     return run
 
 
+def _run_rmsnorm_rht_amax(flag_ops: Any) -> RunFn:
+    def run(inputs: list[Any], attrs: dict[str, Any]) -> Any:
+        _require_runtime_backend(inputs[:2], "rmsnorm_rht_amax_wrapper_sm100")
+        result = flag_ops.rmsnorm_rht_amax_wrapper_sm100(
+            inputs[0],
+            inputs[1],
+            eps=attrs.get("eps", 1e-5),
+            num_threads=attrs.get("num_threads"),
+            rows_per_cta=attrs.get("rows_per_cta"),
+        )
+        return result["o_tensor"], result["amax_tensor"]
+
+    return run
+
+
 def _run_causal_conv1d(flag_ops: Any) -> RunFn:
     def run(inputs: list[Any], attrs: dict[str, Any]) -> torch.Tensor:
         _require_runtime_backend(inputs, "causal_conv1d")
@@ -2804,6 +2947,17 @@ def register_default_ops() -> None:
         )
     )
 
+    register_op(
+        OpSchema(
+            name="rmsnorm_rht_amax_wrapper_sm100",
+            normalize_fn=_normalize_rmsnorm_rht_amax,
+            shape_fn=_rmsnorm_rht_amax_shape,
+            run_fn=_run_rmsnorm_rht_amax(flag_ops),
+            num_outputs=2,
+            fusible=True,
+        )
+    )
+
     register_op(
         OpSchema(
             name="matmul",
diff --git a/src/flag_dnn/graph/wrappers.py b/src/flag_dnn/graph/wrappers.py
@@ -55,6 +55,7 @@
     "batchnorm_inference",
     "layernorm",
     "rmsnorm",
+    "rmsnorm_rht_amax_wrapper_sm100",
     "reduction",
     "sqrt",
     "square",
@@ -73,6 +74,11 @@
 )
 
 
+_DICT_OUTPUT_OPS = {
+    "rmsnorm_rht_amax_wrapper_sm100": ("o_tensor", "amax_tensor"),
+}
+
+
 def install_graph_wrappers(namespace: dict[str, Any]) -> None:
     for op_type in GRAPH_AWARE_OPS:
         eager_fn = namespace.get(op_type)
@@ -98,7 +104,18 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
                 raise RuntimeError(
                     f"FlagDNN graph op {op_type} used outside graph capture"
                 )
-            return ctx.add_op_call(op_type, args, kwargs)
+            outputs = ctx.add_op_call(op_type, args, kwargs)
+            output_keys = _DICT_OUTPUT_OPS.get(op_type)
+            if output_keys is not None:
+                if not isinstance(outputs, tuple):
+                    outputs = (outputs,)
+                if len(outputs) != len(output_keys):
+                    raise RuntimeError(
+                        f"FlagDNN graph op {op_type} returned "
+                        f"{len(outputs)} outputs, expected {len(output_keys)}"
+                    )
+                return dict(zip(output_keys, outputs))
+            return outputs
         return eager_fn(*args, **kwargs)
 
     setattr(wrapper, "__flagdnn_graph_wrapped__", True)
diff --git a/tests_graph/test_rmsnorm_rht_amax.py b/tests_graph/test_rmsnorm_rht_amax.py
