Merge branch 'flagos-ai:master' into w

Author: YY4994

benchmark/bench_fused_add_rms_norm.py

@@ -0,0 +1,118 @@
+"""
+Benchmark: fused_add_rms_norm
+Compares: FlagGems vs torch.compile vs vLLM (if available)
+"""
+
+import time
+
+import torch
+
+import flag_gems
+
+
+# ── reference: naive torch ──────────────────────────────────────────────
+def torch_fused_add_rms_norm(x, residual, weight, eps=1e-5):
+    x = x + residual
+    variance = x.pow(2).mean(-1, keepdim=True)
+    return x * torch.rsqrt(variance + eps) * weight, x
+
+
+# ── reference: torch.compile ────────────────────────────────────────────
+@torch.compile
+def compiled_fused_add_rms_norm(x, residual, weight, eps=1e-5):
+    x = x + residual
+    variance = x.pow(2).mean(-1, keepdim=True)
+    return x * torch.rsqrt(variance + eps) * weight, x
+
+
+# ── reference: vLLM ─────────────────────────────────────────────────────
+try:
+    import os
+
+    os.environ["VLLM_CONFIGURE_LOGGING"] = "0"
+    from vllm._custom_ops import fused_add_rms_norm as vllm_fused_add_rms_norm
+
+    HAS_VLLM = True
+except (ImportError, AttributeError):
+    HAS_VLLM = False
+    print("vLLM not available, skipping vLLM baseline\n")
+
+
+# ── benchmark helper ────────────────────────────────────────────────────
+def bench_fn(fn, warmup=20, rep=100):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(rep):
+        fn()
+    torch.cuda.synchronize()
+    t1 = time.perf_counter()
+    return (t1 - t0) / rep * 1000  # ms
+
+
+# ── main ────────────────────────────────────────────────────────────────
+shapes = [
+    (1, 4096),
+    (32, 4096),
+    (128, 4096),
+    (512, 4096),
+    (1024, 4096),
+    (4096, 4096),
+    (128, 8192),
+    (128, 11008),
+]
+dtypes = [torch.float16, torch.bfloat16]
+
+print(f"{'shape':>18s} {'dtype':>10s} | {'naive':>8s} {'compile':>8s}", end="")
+if HAS_VLLM:
+    print(f" {'vllm':>8s}", end="")
+print(f" {'flaggems':>8s}  (ms)")
+print("-" * 80)
+
+for shape in shapes:
+    for dtype in dtypes:
+        M, N = shape
+        device = "cuda"
+        eps = 1e-5
+
+        x_ref = torch.randn(M, N, dtype=dtype, device=device)
+        r_ref = torch.randn(M, N, dtype=dtype, device=device)
+        w = torch.randn(N, dtype=dtype, device=device)
+
+        # ── naive torch ──
+        t_naive = bench_fn(
+            lambda: torch_fused_add_rms_norm(x_ref.clone(), r_ref.clone(), w, eps)
+        )
+
+        # ── torch.compile ──
+        # warmup compile
+        _ = compiled_fused_add_rms_norm(x_ref.clone(), r_ref.clone(), w, eps)
+        t_compile = bench_fn(
+            lambda: compiled_fused_add_rms_norm(x_ref.clone(), r_ref.clone(), w, eps)
+        )
+
+        # ── vLLM ──
+        if HAS_VLLM:
+            # vLLM's fused_add_rms_norm is in-place: (x, residual) modified
+            def run_vllm():
+                xc = x_ref.clone()
+                rc = r_ref.clone()
+                vllm_fused_add_rms_norm(xc, rc, w, eps)
+
+            t_vllm = bench_fn(run_vllm)
+
+        # ── FlagGems ──
+        def run_gems():
+            xc = x_ref.clone()
+            rc = r_ref.clone()
+            flag_gems.fused_add_rms_norm(xc, rc, (N,), w, eps)
+
+        t_gems = bench_fn(run_gems)
+
+        # ── print ──
+        tag = f"({M}, {N})"
+        print(f"{tag:>18s} {str(dtype):>10s} | {t_naive:8.3f} {t_compile:8.3f}", end="")
+        if HAS_VLLM:
+            print(f" {t_vllm:8.3f}", end="")
+        print(f" {t_gems:8.3f}")

benchmark/conftest.py

@@ -4,6 +4,7 @@
 
 import pytest
 import torch
+import yaml
 
 import flag_gems
 from benchmark.attri_util import (
@@ -58,6 +59,7 @@ def __init__(self):
 
 
 Config = BenchConfig()
+REGISTERED_MARKS = []
 
 
 def pytest_addoption(parser):
@@ -156,9 +158,21 @@ def pytest_addoption(parser):
         ),
     )
 
+    parser.addoption(
+        "--collect-marks",
+        action="store_true",
+        help="Collect the tests with marker information without executing them",
+    )
+
 
 def pytest_configure(config):
     global Config  # noqa: F824
+    global REGISTERED_MARKS
+
+    REGISTERED_MARKS = {
+        marker.split(":")[0].strip() for marker in config.getini("markers")
+    }
+
     mode_value = config.getoption(
         "--mode" if vendor_name != "kunlunxin" else "--fg_mode"
     )
@@ -278,3 +292,33 @@ def extract_and_log_op_attributes(request):
     yield
     if Config.record_log and op_attributes:
         emit_record_logger(json.dumps(op_attributes, indent=2))
+
+
+def pytest_collection_modifyitems(session, config, items):
+    if config.getoption("--collect-marks"):
+        report = []
+        for item in items:
+            data = {}
+
+            # Collect some general information
+            if item.cls:
+                data["class"] = item.cls.__name__
+            data["test_case"] = item.name
+            if item.originalname:
+                data["function"] = item.originalname
+            data["file"] = item.location[0]
+
+            all_marks = list(item.iter_markers())
+            op_marks = [
+                mark.name
+                for mark in all_marks
+                if mark.name not in BUILTIN_MARKS and mark.name not in REGISTERED_MARKS
+            ]
+
+            data["marks"] = op_marks
+            report.append(data)
+
+        print(yaml.dump(report, indent=2))
+
+        # Skip all tests
+        items.clear()

benchmark/core_shapes.yaml

@@ -54,6 +54,13 @@ BlasBenchmark:
     - [16, 4096, 4096, 4096]
   shape_desc: "B, M, N, K"   # shapes are defined as (B, M, N, K)
 
+groupmm:
+  shapes:
+    - [16, 512, 2048]
+    - [16, 2560, 2048]
+    - [64, 2048, 128]
+  shape_desc: "Groups, N, K"   # shapes are defined as (Groups, N, K)
+
 MvAndOuterBenchmark:
   shapes:
     - [384, 384]