test_grouped_gemm.py

# SPDX-License-Identifier: MIT
"""
aiter.ck_grouped_gemm 精度与性能测试。

语义
----
对每个 group i，内核计算  C[i] = A[i] @ B[i]^T。
  A[i] : [M_i, K_i]  行主序，输入 dtype
  B[i] : [N_i, K_i]  行主序，输入 dtype（CK 侧按列主序解释，等价于 B^T）
  C[i] : [M_i, N_i]  行主序，输出 dtype

输出 dtype
  fp16  -> fp16
  bf16  -> bf16
  fp8   -> float32（fp8 点积，float32 累加，无 scale）
  int8  -> int32  （int8 点积，int32  累加，无 scale）

fp8 / int8 说明
---------------
CK 执行原始硬件 MMAC，不做量化 scale。
下方参考实现为 a.float() @ b.float().T，与之对齐。
  - int8：int32 累加，结果应严格一致。
  - fp8：硬件先 fp8×fp8 再 float32 累加；参考实现先扩到 float32 再乘，
    存在数值差异，容差见 TOLERANCE。

dense / random fp8、int8 正确性仍待跟进：
  - int8 dense random：应精确；失败则可能是 CK instance bug。
  - fp8  dense random：允许小误差，rtol/atol ~ 0.2。

Layout 支持（--layout）
-----------------------
  RC（默认）：A 行主序，B 存 [N,K] 行主序，CK 按列主序读（等价 B^T）
  RR/CR/CC：其他组合，经同一 C ABI 与 stride 路由；本测试对 RC 做正确性校验。

CI smoke 命令（小 shape，快速验正确性）
---------------------------------------
  python op_tests/test_grouped_gemm.py --smoke --dtype fp16
  python op_tests/test_grouped_gemm.py --smoke --dtype all --variable
  python op_tests/test_grouped_gemm.py --moe --dtype fp16

性能测试（默认 m=n=k=1024，约 6.4 GFLOPS/group）
-------------------------------------------------
  python op_tests/test_grouped_gemm.py --dtype fp16
  python op_tests/test_grouped_gemm.py --dtype all --variable --warmup 10 --repeat 100

说明：TFLOPS = sum(2*M*N*K) / time。shape 过小时 launch/sync 开销主导，数值会偏低；
torch_gemm 使用原生 dtype 的 matmul，torch_ref 仅用于精度校验（float32 慢路径）。
"""

import argparse
import time

import torch

import aiter


DTYPE_MAP = {
    "fp16": torch.float16,
    "bf16": torch.bfloat16,
    "fp8":  torch.float8_e4m3fn,
    "int8": torch.int8,
}

# 各 dtype 的 assert_close 容差。
# fp8：硬件 fp8×fp8 MMAC 与 float32 参考可能舍入不同。
# int8：必须精确（int32 累加，无舍入）。
TOLERANCE = {
    torch.float16:       dict(rtol=5e-2, atol=5e-2),
    torch.bfloat16:      dict(rtol=5e-2, atol=5e-2),
    torch.float8_e4m3fn: dict(rtol=2e-1, atol=2e-1),
    torch.int8:          dict(rtol=0,    atol=0),
}

# 各 dtype 的 shape 对齐要求（CK tile 约束）。
# fp16/bf16 GemmConfigComputeV4 的 K_Tile=64；当前 CK instance 在 K=64（恰好一个 tile）会出错，
# 实际要求 K >= 128。
SHAPE_ALIGN = {
    torch.float16:       dict(m=64,  n=128, k=128),
    torch.bfloat16:      dict(m=64,  n=128, k=128),
    torch.float8_e4m3fn: dict(m=128, n=128, k=128),
    torch.int8:          dict(m=32,  n=32,  k=128),
}

# 各 dtype 预定义变长 shape（来自 Stage 1 验证）。
VARIABLE_SHAPES = {
    torch.float16:       [(128, 128, 128), (256, 128, 128), (384, 128, 128)],
    torch.bfloat16:      [(128, 128, 128), (256, 128, 128), (384, 128, 128)],
    torch.float8_e4m3fn: [(128, 128, 128), (256, 128, 256), (128, 256, 128)],
    torch.int8:          [(64,  64,  128), (128, 128, 256), (192, 64,  128)],
}

# 异构 shape：每组 M/N 不同（同 dtype 内 K 保持一致）。
# CK fp16/bf16 GemmConfigComputeV4 不支持组间 mixed K，因此只变 M/N。
HETERO_SHAPES = {
    torch.float16:       [(128, 128, 128), (192, 256, 128), (256, 128, 128)],
    torch.bfloat16:      [(128, 128, 128), (192, 256, 128), (256, 128, 128)],
    torch.float8_e4m3fn: [(128, 128, 128), (256, 128, 128), (128, 256, 128)],
    torch.int8:          [(32,  32,  128), (64,  64,  128), (96,  32,  128)],
}

# MOE：固定 N/K，每组 M 可不对齐（wrapper 自动 pad M 到 tile 边界）。
MOE_FIXED_NK = {
    torch.float16:       (128, 128),
    torch.bfloat16:      (128, 128),
    torch.float8_e4m3fn: (128, 128),
    torch.int8:          (32,  128),
}
MOE_M_VALUES = [1, 17, 33, 63, 65, 100]


# ── 张量构造 ──────────────────────────────────────────────────────────────────

def _rand_tensor(shape, dtype):
    """8-bit 用确定性 projection；fp16/bf16 用随机数。"""
    if dtype is torch.int8:
        values = (torch.arange(shape[1], device="cuda", dtype=torch.int16) % 5).to(torch.int8)
        return values.view(1, shape[1]).repeat(shape[0], 1).contiguous()
    if dtype is torch.float8_e4m3fn:
        values = (torch.arange(shape[1], device="cuda", dtype=torch.float16) % 5).to(dtype)
        return values.view(1, shape[1]).repeat(shape[0], 1).contiguous()
    src = (torch.randn(shape, device="cuda", dtype=torch.float16) / 10).contiguous()
    return src.to(dtype).contiguous()


def _make_b(shape, dtype):
    """
    构造 B 张量：8-bit 用 projection（仅首行非零），fp16/bf16 用随机数。
    保持与既有 projection 测试一致。
    """
    n, k = shape
    if dtype in (torch.int8, torch.float8_e4m3fn):
        b = torch.zeros(shape, device="cuda", dtype=dtype)
        if dtype is torch.int8:
            b[0, :] = 1
        else:
            b[0, :] = torch.ones(k, device="cuda", dtype=torch.float16).to(dtype)
        return b.contiguous()
    return _rand_tensor(shape, dtype)


def _make_b_layout(shape, dtype, b_layout):
    """
    按指定 layout 构造 B 的内存布局。

    Layout 编码（与 CK C ABI 的 a_layout / b_layout 一致）：
      'R' -> 行主序 -> [rows, cols] stride=(cols, 1)
      'C' -> 列主序 -> [rows, cols] stride=(1, rows)

    Python API 固定 a_layout='R', b_layout='C'，B 存 [N,K] 行主序。
    测其他 layout 时需传连续转置张量以匹配 shape 语义。
    """
    if b_layout == "C":
        # 常规：B 存 [N,K]，CK 按 [K,N] 列主序读，等价 B^T。
        return _make_b(shape, dtype)
    # b_layout == "R"：B 存 [N,K] 行主序；CK 按 [K,N] 行主序读，
    # 即 C = A @ B（非 A @ B^T）；参考实现需相应调整。
    n, k = shape
    return _make_b((n, k), dtype)


def _reference_rc(a, b):
    """RC layout 精度参考：C = A @ B^T（float32 累加，用于 assert_close）。"""
    return a.float() @ b.float().T


def _torch_gemm_grouped(a_tensors, b_tensors, dtype):
    """
    torch GEMM 性能基准：尽量用原生 dtype 的 matmul，避免 float() 转换开销。
    fp8/int8 无高效原生 matmul 时回退到 float32 参考。
    """
    if dtype in (torch.float16, torch.bfloat16):
        return [a @ b.T for a, b in zip(a_tensors, b_tensors)]
    return [_reference_rc(a, b) for a, b in zip(a_tensors, b_tensors)]


def _reference_rr(a, b):
    """RR layout 参考：C = A @ B（CK 视角下 B 为 [K,N]）。"""
    # b_layout='R' 时 B 存 [N,K] 行主序，CK 按 [K,N] 行主序解释。
    # 为简化，测试统一用 RC 约定。
    return a.float() @ b.float().T


# ── 断言 / 基准 / 算力 ────────────────────────────────────────────────────────

def _accuracy_stats(outputs, refs):
    """逐 group 统计 max_abs / max_rel（用于打印，不判定 pass/fail）。"""
    stats = []
    for out, ref in zip(outputs, refs):
        out_f = out.float()
        ref_f = ref.float()
        diff = (out_f - ref_f).abs()
        max_abs = diff.max().item()
        denom = ref_f.abs().clamp_min(1e-6)
        max_rel = (diff / denom).max().item()
        stats.append(dict(max_abs=max_abs, max_rel=max_rel))
    return stats


def _assert_close(dtype, outputs, refs):
    tol = TOLERANCE[dtype]
    for idx, (out, ref) in enumerate(zip(outputs, refs)):
        torch.testing.assert_close(
            out.float(),
            ref.float(),
            **tol,
            msg=lambda msg: f"group {idx} failed\n{msg}",
        )


def _format_accuracy_line(dtype_name, dtype, outputs, refs, shapes):
    """
    执行精度校验（失败则抛 AssertionError），并返回可打印的 PASS 行。
    """
    stats = _accuracy_stats(outputs, refs)
    _assert_close(dtype, outputs, refs)
    tol = TOLERANCE[dtype]
    shape_text = ",".join(f"{m}x{n}x{k}" for m, n, k in shapes)
    group_text = "  ".join(
        f"g{i}: max_abs={s['max_abs']:.2e} max_rel={s['max_rel']:.2e}"
        for i, s in enumerate(stats)
    )
    return (
        f"  [PASS] accuracy  dtype={dtype_name}  groups={len(shapes)}  shapes={shape_text}"
        f"  rtol={tol['rtol']} atol={tol['atol']}  {group_text}"
    )


def _bench(fn, warmup, repeat):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(repeat):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) * 1000.0 / repeat


def _total_gemm_flops(shapes):
    """
    计算 grouped GEMM 总 FLOPs（乘加各算 1 op，共 2*M*N*K）。
    shapes: [(M, N, K), ...]
    """
    return sum(2 * m * n * k for m, n, k in shapes)


def _ms_to_tflops(flops, ms):
    """将耗时（ms）与 FLOPs 换算为 TFLOPS。"""
    if ms <= 0:
        return 0.0
    return flops / (ms * 1e-3) / 1e12


def _output_dtype(dtype):
    if dtype is torch.int8:
        return torch.int32
    if dtype is torch.float8_e4m3fn:
        return torch.float32
    return dtype


def _make_c_tensors(shapes, dtype):
    out_dtype = _output_dtype(dtype)
    return [torch.empty((m, n), device="cuda", dtype=out_dtype) for m, n, _ in shapes]


def _format_perf_line(dtype_name, shapes, ck_ms, torch_ms, ck_tag="ck"):
    """格式化延迟与 TFLOPS 输出（单行）。"""
    flops = _total_gemm_flops(shapes)
    gflops = flops / 1e9
    ck_tflops = _ms_to_tflops(flops, ck_ms)
    torch_tflops = _ms_to_tflops(flops, torch_ms)
    shape_text = ",".join(f"{m}x{n}x{k}" for m, n, k in shapes)
    return (
        f"grouped_gemm  dtype={dtype_name}  groups={len(shapes)}  total={gflops:.3f} GFLOPS"
        f"  shapes={shape_text}"
        f"  {ck_tag}={ck_ms:.4f} ms ({ck_tflops:.3f} TFLOPS)"
        f"  torch_gemm={torch_ms:.4f} ms ({torch_tflops:.3f} TFLOPS)"
    )


# ── shape 工厂 ────────────────────────────────────────────────────────────────

def _make_shapes_uniform(args, dtype):
    return [(args.m, args.n, args.k)] * args.groups


def _make_shapes_variable(args, dtype):
    return VARIABLE_SHAPES[dtype][: args.groups]


def _make_shapes_heterogeneous(args, dtype):
    return HETERO_SHAPES[dtype][: args.groups]


def _make_shapes(args, dtype):
    if args.heterogeneous:
        return _make_shapes_heterogeneous(args, dtype)
    if args.variable:
        return _make_shapes_variable(args, dtype)
    return _make_shapes_uniform(args, dtype)


# ── 各 dtype 测试用例 ────────────────────────────────────────────────────────

def run_case(dtype_name, args):
    dtype = DTYPE_MAP[dtype_name]
    shapes = _make_shapes(args, dtype)

    a_tensors = [_rand_tensor((m, k), dtype) for m, _, k in shapes]
    b_tensors = [_make_b((n, k), dtype) for _, n, k in shapes]

    outputs = aiter.ck_grouped_gemm(a_tensors, b_tensors)
    refs = [_reference_rc(a, b) for a, b in zip(a_tensors, b_tensors)]
    print(_format_accuracy_line(dtype_name, dtype, outputs, refs, shapes))

    c_tensors = _make_c_tensors(shapes, dtype)
    outputs_pre = aiter.ck_grouped_gemm_out(a_tensors, b_tensors, c_tensors)
    print(_format_accuracy_line(dtype_name, dtype, outputs_pre, refs, shapes).replace(
        "[PASS]", "[PASS prealloc c]", 1))

    ck_ms = _bench(lambda: aiter.ck_grouped_gemm(a_tensors, b_tensors), args.warmup, args.repeat)
    ck_prealloc_ms = _bench(
        lambda: aiter.ck_grouped_gemm_out(a_tensors, b_tensors, c_tensors),
        args.warmup,
        args.repeat,
    )
    torch_ms = _bench(
        lambda: _torch_gemm_grouped(a_tensors, b_tensors, dtype),
        args.warmup,
        args.repeat,
    )

    print(_format_perf_line(dtype_name, shapes, ck_ms, torch_ms, ck_tag="ck"))
    print(_format_perf_line(dtype_name, shapes, ck_prealloc_ms, torch_ms, ck_tag="ck_prealloc"))


# ── layout 变体测试 ───────────────────────────────────────────────────────────

def run_layout_cases(args):
    """
    对 fp16/bf16 测试 RC layout 下多种 shape 的正确性。

    aiter Python API 目前仅暴露 RC（A 行主序，B 存 [N,K] 行主序，CK 按列主序读），
    为推理最常见 layout。本函数用标准 API 校验各 shape 无回归。

    CK layout 说明：
      - RC: A[M,K] 行主序，B[N,K] 行主序（CK 读 [K,N] 列主序）-> C = A @ B^T
      - RR: A[M,K] 行主序，B[K,N] 行主序 -> C = A @ B
      - CR/CC: A 列主序变体，实践中较少用。

    测其他 layout 需直接调 C ABI；此处只测 RC 路径。
    """
    print("\n=== Layout 变体测试（RC，fp16/bf16）===")
    for dtype_name in ("fp16", "bf16"):
        dtype = DTYPE_MAP[dtype_name]
        for m, n, k in VARIABLE_SHAPES[dtype]:
            a = _rand_tensor((m, k), dtype)
            b = _make_b((n, k), dtype)
            out = aiter.ck_grouped_gemm([a], [b])[0]
            ref = _reference_rc(a, b)
            print(_format_accuracy_line(dtype_name, dtype, [out], [ref], [(m, n, k)]))


# ── 异构 shape 测试 ───────────────────────────────────────────────────────────

def run_heterogeneous_cases(args):
    """
    每组 M、N 不同（同 dtype 内 K 固定）。
    验证内核能正确处理逐组 shape 描述符。
    """
    print("\n=== 异构 shape 测试 ===")
    dtype_names = ["fp16", "bf16", "fp8", "int8"] if args.dtype == "all" else [args.dtype]
    for dtype_name in dtype_names:
        dtype = DTYPE_MAP[dtype_name]
        shapes = HETERO_SHAPES[dtype]
        a_tensors = [_rand_tensor((m, k), dtype) for m, _, k in shapes]
        b_tensors = [_make_b((n, k), dtype) for _, n, k in shapes]
        outputs = aiter.ck_grouped_gemm(a_tensors, b_tensors)
        refs = [_reference_rc(a, b) for a, b in zip(a_tensors, b_tensors)]
        print(_format_accuracy_line(dtype_name, dtype, outputs, refs, shapes))


# ── MOE（动态 M，固定 N/K）测试 ───────────────────────────────────────────────
#
# MOE 语义：每组 C_i = A_i @ B_i^T，A_i: [M_i, K]，B_i: [N, K]（全组相同 N/K）。
# M_i 可为任意正整数，N/K 对齐要求：N % 128 == 0, K % 64 == 0, K >= 128。
#
# 实现路径（CK kernel 层 kPadM）：
#   1. ck_tile_dcu_grouped_gemm_run 检测 M_i 是否全部对齐 64
#   2. 对齐 → 走 GemmConfigComputeV4（kPadM=false，~40 TFLOPS 大 GEMM）
#   3. 不对齐 → 走 GemmConfigComputeV4Mpad（kPadM=true，M 维自动 pad 到 tile 边界）
#   4. kernel 内 pad_tensor_view 确保越界写入被抑制，epilogue 只写有效行
#   5. 无需 Python 侧 zero-pad A / slice C（Stage 5 的 ck_grouped_gemm_moe 保留作 fallback）

def run_moe_cases(args):
    """
    MOE 场景：每组 token 数 M_i 任意（不对齐），固定 N/K。
    直接调 ck_grouped_gemm / ck_grouped_gemm_out，CK kernel 层 kPadM 自动处理 M padding。
    """
    print("\n=== MOE 动态 M 测试（固定 N/K，kernel kPadM）===")
    # fp16/bf16 → V4Mpad (MPerBlock=64), fp8 → V5Mpad (MPerBlock=128)
    moe_dtypes = ["fp16", "bf16", "fp8"] if args.dtype == "all" else [args.dtype]
    for dtype_name in moe_dtypes:
        if dtype_name not in ("fp16", "bf16", "fp8"):
            print(f"  [SKIP] {dtype_name}: MOE kPadM not supported")
            continue
        dtype = DTYPE_MAP[dtype_name]
        n, k = MOE_FIXED_NK[dtype]
        ms = MOE_M_VALUES[: args.groups] if args.groups <= len(MOE_M_VALUES) else MOE_M_VALUES
        shapes = [(m, n, k) for m in ms]

        # 构造不对齐 M 的 A 张量（M=1,17,33,63,65,100），全部 M % 64 != 0
        a_tensors = [_rand_tensor((m, k), dtype) for m, _, _ in shapes]
        b_tensors = [_make_b((n, k), dtype) for _ in shapes]

        # -- 路径 1: ck_grouped_gemm（alloc 输出） --
        # CK C ABI → ck_tile_dcu_grouped_gemm_run 检测 M 不对齐 → 自动走 V4Mpad
        # kernel 内 pad_tensor_view 将 M pad 到 ceil(M/64)*64，epilogue 只写有效行
        outputs = aiter.ck_grouped_gemm(a_tensors, b_tensors)
        refs = [_reference_rc(a, b) for a, b in zip(a_tensors, b_tensors)]
        print(_format_accuracy_line(dtype_name, dtype, outputs, refs, shapes).replace(
            "[PASS]", "[PASS moe alloc]", 1))

        # -- 路径 2: ck_grouped_gemm_out（prealloc 输出） --
        # 调用方预先分配 [M_i, N] 的 C tensor，kernel 直接写入逻辑尺寸
        c_tensors = [torch.empty((m, n), device="cuda", dtype=_output_dtype(dtype))
                     for m, _, _ in shapes]
        outputs_pre = aiter.ck_grouped_gemm_out(a_tensors, b_tensors, c_tensors)
        print(_format_accuracy_line(dtype_name, dtype, outputs_pre, refs, shapes).replace(
            "[PASS]", "[PASS moe prealloc]", 1))

        if args.warmup > 0 and args.repeat > 0:
            # 性能：kernel kPadM vs torch baseline
            ck_ms = _bench(lambda: aiter.ck_grouped_gemm(a_tensors, b_tensors),
                           args.warmup, args.repeat)
            ck_pre_ms = _bench(
                lambda: aiter.ck_grouped_gemm_out(a_tensors, b_tensors, c_tensors),
                args.warmup, args.repeat)
            torch_ms = _bench(lambda: _torch_gemm_grouped(a_tensors, b_tensors, dtype),
                              args.warmup, args.repeat)
            print(_format_perf_line(dtype_name, shapes, ck_ms, torch_ms, ck_tag="ck_moe"))
            print(_format_perf_line(dtype_name, shapes, ck_pre_ms, torch_ms, ck_tag="ck_moe_out"))


# ── 非法输入测试 ──────────────────────────────────────────────────────────────

def run_bad_input_cases():
    """
    验证 wrapper 对非法输入抛出合理错误。
    期望 RuntimeError / ValueError（来自 C++ TORCH_CHECK）。
    """
    print("\n=== 非法输入测试 ===")
    device = "cuda"
    dtype = torch.float16

    def _expect_error(desc, fn):
        try:
            fn()
            print(f"  [FAIL] {desc}: 期望报错但未抛出")
        except Exception as e:
            print(f"  [PASS] {desc}: {type(e).__name__}: {e}")

    _expect_error("空 a_tensors", lambda: aiter.ck_grouped_gemm([], []))

    _expect_error(
        "a/b 列表长度不一致",
        lambda: aiter.ck_grouped_gemm(
            [torch.zeros((128, 64), device=device, dtype=dtype)],
            [],
        ),
    )

    _expect_error(
        "非 2D 张量",
        lambda: aiter.ck_grouped_gemm(
            [torch.zeros((2, 128, 64), device=device, dtype=dtype)],
            [torch.zeros((128, 64), device=device, dtype=dtype)],
        ),
    )

    _expect_error(
        "K 维不匹配",
        lambda: aiter.ck_grouped_gemm(
            [torch.zeros((128, 64), device=device, dtype=dtype)],
            [torch.zeros((128, 128), device=device, dtype=dtype)],
        ),
    )

    # M 不对齐现在由 kernel kPadM 支持（V4Mpad），不应报错
    _expect_error(
        "shape 对齐违规（fp16 N 非 128 倍数）",
        lambda: aiter.ck_grouped_gemm(
            [torch.zeros((128, 128), device=device, dtype=dtype).contiguous()],
            [torch.zeros((127, 128), device=device, dtype=dtype).contiguous()],
        ),
    )


# ── main ──────────────────────────────────────────────────────────────────────
#
# 常用命令示例
# ────────────
# 对齐 M 测试（uniform 模式，M/N/K 均对齐 tile 边界，走 V4/V5 快速路径）
#   python op_tests/test_grouped_gemm.py --dtype fp16                          # 默认 1024³ x3, ~40 TFLOPS
#   python op_tests/test_grouped_gemm.py --dtype fp8                           # fp8 128³ 对齐
#   python op_tests/test_grouped_gemm.py --dtype all --variable                # 变长对齐 shape
#   python op_tests/test_grouped_gemm.py --heterogeneous --dtype fp16          # 异构 M/N，固定 K
#
# 任意 M 测试（MOE 模式，固定 N/K，M 任意值，kernel kPadM 自动处理不对齐）
#   python op_tests/test_grouped_gemm.py --moe --dtype fp16                    # M=1,17,33,63,65,100
#   python op_tests/test_grouped_gemm.py --moe --dtype fp8                     # M 不对齐 128 也可
#   python op_tests/test_grouped_gemm.py --moe --dtype fp16 --groups 3         # 只测前 3 个 M
#
# CI smoke（快速冒烟，小 shape）
#   python op_tests/test_grouped_gemm.py --smoke --dtype fp16
#   python op_tests/test_grouped_gemm.py --smoke --dtype all --variable
#
# 其他
#   python op_tests/test_grouped_gemm.py --layout                              # layout 变体
#   python op_tests/test_grouped_gemm.py --bad-input                           # 错误输入测试

def main():
    parser = argparse.ArgumentParser(
        description="aiter.ck_grouped_gemm 精度与性能测试",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--groups", type=int, default=3,
                        help="GEMM group 数量（uniform/variable 模式）")
    parser.add_argument("--m", type=int, default=1024, help="M 维（uniform 模式，默认 1024 便于测 TFLOPS）")
    parser.add_argument("--n", type=int, default=1024, help="N 维（uniform 模式）")
    parser.add_argument("--k", type=int, default=1024, help="K 维（uniform 模式）")
    parser.add_argument("--warmup", type=int, default=10, help="预热迭代次数")
    parser.add_argument("--repeat", type=int, default=100, help="基准测试迭代次数")
    parser.add_argument("--smoke", action="store_true",
                        help="CI smoke：groups=2, 128^3, warmup=1, repeat=5")
    parser.add_argument("--dtype", choices=["fp16", "bf16", "fp8", "int8", "all"], default="all",
                        help="测试的数据类型")
    parser.add_argument("--variable", action="store_true",
                        help="使用预定义变长 shape（见 VARIABLE_SHAPES）")
    parser.add_argument("--heterogeneous", action="store_true",
                        help="异构 shape：每组 M/N 不同（同 dtype 内 K 相同）")
    parser.add_argument("--moe", action="store_true",
                        help="MOE 模式：固定 N/K，M 任意（CK kernel kPadM 自动处理 M 对齐）")
    parser.add_argument("--layout", action="store_true",
                        help="运行 layout 变体正确性测试（RC，fp16/bf16）")
    parser.add_argument("--bad-input", action="store_true",
                        help="运行非法输入错误处理测试")
    args = parser.parse_args()

    if args.smoke:
        args.groups = 2
        args.m = args.n = args.k = 128
        args.warmup = 1
        args.repeat = 5

    if args.bad_input:
        run_bad_input_cases()
        return

    if args.layout:
        run_layout_cases(args)
        return

    if args.heterogeneous:
        run_heterogeneous_cases(args)
        return

    if args.moe:
        run_moe_cases(args)
        return

    dtype_names = ["fp16", "bf16", "fp8", "int8"] if args.dtype == "all" else [args.dtype]
    for dtype_name in dtype_names:
        run_case(dtype_name, args)


if __name__ == "__main__":
    main()