# CK Grouped GEMM 开发记录

## 最终架构

```
Python API (aiter/ops/grouped_gemm.py)
├── ck_grouped_gemm(a_tensors, b_tensors)          # alloc 输出
├── ck_grouped_gemm_out(a_tensors, b_tensors, c)   # prealloc 输出
├── ck_grouped_gemm_moe(...)                       # Stage 5: Python 侧 M padding（fallback）
└── ck_grouped_gemm_moe_out(...)

C++ JIT 模块 (csrc/py_itfs_ck/)
├── grouped_gemm_kernels.cu     # torch 封装，shape 校验，dtype 分发
├── ck_grouped_gemm_abi.h       # 纯 C 头文件，CK C ABI 声明
└── (link) CK instance 文件     # 各 dtype 的 instance，由 JIT 编译

CK 内核 (3rdparty/composable_kernel/example_hcu/ck_tile/19_grouped_gemm/)
├── grouped_gemm.hpp            # Config 定义 + C ABI 声明
├── grouped_gemm.cpp            # C ABI 入口，mpad 自动分发
├── instances/grouped_gemm_impl.hpp  # 模板实现 (Persistent tileloop)
├── instances/grouped_gemm_fp16.cpp  # V4 + V4Mpad
├── instances/grouped_gemm_bf16.cpp  # V4 + V4Mpad
├── instances/grouped_gemm_fp8.cpp   # V5 + V5Mpad
├── instances/grouped_gemm_bf8.cpp   # Bf8K64 + Bf8K64Mpad
├── instances/grouped_gemm_int8.cpp
└── instances/grouped_gemm_int4.cpp

测试 (op_tests/test_grouped_gemm.py)
├── 统一 shape（对齐 M，性能测试）   python ... --dtype fp16
├── 变长 shape（VARIBLE_SHAPES）     python ... --variable
├── 异构 shape（HETERO_SHAPES）      python ... --heterogeneous
├── MOE 任意 M（kPadM）              python ... --moe
└── 错误输入测试                     python ... --bad-input
```

## MOE 动态 M 实现：kPadM 自动分发

| dtype | 对齐配置 | Mpad 配置 | M tile | 分发条件 |
|-------|---------|----------|--------|---------|
| fp16  | V4 (kPadM=false)    | V4Mpad (kPadM=true)       | MPerBlock=64   | M % 64 != 0  |
| bf16  | V4 (kPadM=false)    | V4Mpad (kPadM=true)       | MPerBlock=64   | M % 64 != 0  |
| fp8   | V5 (kPadM=false)    | V5Mpad (kPadM=true)       | MPerBlock=128  | M % 128 != 0 |
| bf8   | Bf8K64 (kPadM=false)| Bf8K64Mpad (kPadM=true)   | MPerBlock=128  | M % 128 != 0 |

流程：`ck_tile_dcu_grouped_gemm_run` → `grouped_gemm_need_mpad()` 检测不对齐 → 自动选 V4/V5/Bf8K64 或对应 Mpad instance。
对齐时走原 instance（零开销），不对齐时 kernel 内 `pad_tensor_view` 将 M pad 到 tile 边界，epilogue 只写有效行。

核心修改（`universal_gemm_kernel.hpp` `MakeGemmPadViews`）：
- RowMajor A：`sequence<false, kPadK>` → `sequence<kPadM, kPadK>`
- RowMajor C：`sequence<false, kPadN>` → `sequence<kPadM, kPadN>`
- 当 kPadM=false（所有原 instance）行为不变；kPadM=true（Mpad instance）时 M 维也参与 padding

## 性能参考

| 场景 | fp16 | bf16 | fp8 |
|------|------|------|-----|
| 大 GEMM (1024³×3) | ~40 TFLOPS | ~40 TFLOPS | ~23 TFLOPS |
| MOE 小 M (1~100, N=K=128) | ~0.08 ms, ~0.12 TFLOPS | ~0.08 ms | ~0.08 ms |
| MOE kPadM vs Python padding | **3.8x 快** (0.08 vs 0.31 ms) | 同左 | — |

## Shape 约束

| dtype | N | K | M（不带 mpad）| M（带 mpad）|
|-------|---|---|--------------|-----------|
| fp16  | N % 128 == 0 | K % 64 == 0, K >= 128 | M % 64 == 0 | 任意 |
| bf16  | N % 128 == 0 | K % 64 == 0, K >= 128 | M % 64 == 0 | 任意 |
| fp8   | N % 128 == 0 | K % 128 == 0        | M % 128 == 0 | 任意 |
| bf8   | N % 128 == 0 | K % 128 == 0        | M % 128 == 0 | 任意（仅 CK C ABI，无 Python API）|
| int8  | N % 32 == 0  | K % 128 == 0        | M % 32 == 0  | 未实现 |

## 开发经验

### 性能
- CK V4 persistent tileloop 大 GEMM 天花板 ~40 TFLOPS（fp16/bf16），torch 约 90 TFLOPS。差距在 kernel 本身，不在 wrapper 开销。
- `ck_grouped_gemm_out`（prealloc C）相比 `ck_grouped_gemm`（alloc C）几乎无收益（~0.3%），瓶颈不在内存分配。
- V5/V3_2 等其他 config 对 fp16/bf16 大 GEMM 无提升；direct grid（非 persistent）更慢。
- 小 M MOE 场景 kPadM 收益巨大（3.8x），因为避免了 Python 侧 zero-pad + slice 的 memcpy 开销。

### CK 已知问题
- **fp16/bf16 K=64 不可用**：V4 K_Tile=64，恰好一个 tile 时会出错。wrapper 已强制 `K >= 128`。
- **fp16/bf16 组间 mixed K 会损坏结果**：即使每个 group 单独正确，组在一起 K 不同时会出错。异构测试固定 K，只变 M/N。
- **IsSupportedArgument 在 persistent tileloop 路径不被调用**，grid 直接取 MaxOccupancyGridSize。

### VGPR Spill 分析

全部 instance 均编译 `has_hot_loop=true`（K 足够大时触发 ping-pong double buffer 路径）和 `has_hot_loop=false` 两个分支。hot loop 路径寄存器压力远高于非 hot loop。

**汇总：各 dtype has_hot_loop=true + PassThrough epilogue + Row/Col layout 的 VGPR Spill**

| dtype | Config | Tile (M×N×K) | K_Warp_Tile | Acc/C | VGPR | VGPR Spill | Scratch (B/lane) | Instructions |
|-------|--------|-------------|-------------|-------|------|-----------|-----------------|-------------|
| bf8 | GemmConfigComputeBf8 (N_Tile=64) | 128×64×128 | 64 | float/float | 256 | **1292** | 1564 | 11331 |
| bf8 | GemmConfigComputeBf8V2 (N_Tile=128) | 128×128×128 | 32 | float/float | 256 | **1058** | 1368 | 9553 |
| int4 | GemmConfigComputeInt4 | 128×128×128 | 64 | int32/int32 | 256 | 311~322 | 608~636 | 6028~6130 |
| int4 | GemmConfigComputeInt4 + bias(1D) | 128×128×128 | 64 | int32/int32 | 256 | 139~143 | 468~472 | 3753~4064 |
| int4 | GemmConfigComputeInt4 + multiply(2D) | 128×128×128 | 64 | int32/int32 | 256 | 147~148 | 488 | 3255~3544 |

**has_hot_loop=false 时：**
| dtype | Config | VGPR | VGPR Spill |
|-------|--------|------|-----------|
| int4 + bias(1D) | GemmConfigComputeInt4 | 171 | **0** |
| int4 + multiply(2D) | GemmConfigComputeInt4 | 171 | **0** |

**结论：**

1. **bf8 has_hot_loop=true 是所有 dtype 里最严重的**，溢出 1058~1292，因为 Acc=CDataType=float（32-bit），累加器寄存器宽度是 int4 的 2 倍。
2. **int4 has_hot_loop=true + PassThrough 溢出 311~322**，加 bias/multiply epilogue 反而降到 140~150（epilogue 改变了寄存器分配策略）。**has_hot_loop=false 时零溢出**（VGPR=171）。
3. **Layout 间差异 < 20 spill**，不是瓶颈。
4. 所有变体 **Occupancy=1**（LDS 已吃满 49152~65536 bytes）。
5. **优化方向**：bf8 和 int4 的 has_hot_loop=true 路径是主要矛盾——考虑减小 M_Tile/N_Tile 降低寄存器压力，或使用 `__launch_bounds__` 显式限制 VGPR。如果实际 K 不大，has_hot_loop=false 路径零溢出，可接受。

### 8-bit 数值
- **fp8**：硬件 fp8×fp8 MMAC → float32 累加。参考 `a.float() @ b.float().T` 是先扩 float32 再乘。差异属于正常硬件舍入，rtol/atol=0.2。
- **int8**：硬件 int8×int8 MMAC → int32 累加。必须精确匹配（rtol=0, atol=0）。

### 远程工作
- 容器内有两份 aiter 代码（`/wksp/ai/aiter` 和 `/wksp/dcu-das/aiter`），必须设 `PYTHONPATH=/wksp/ai/aiter`。
- Windows → 远程的 shell 脚本会有 CRLF 问题，优先用 Python 脚本或 `sed -i 's/\r$//'`。
- JIT 首次编译 30~60 分钟，后续增量编译快。`AITER_REBUILD=1` 或删 `aiter/jit/build/module_grouped_gemm/` 强制重编。
- 单 GPU 命令设 `HIP_VISIBLE_DEVICES=<idle_device>`，多 GPU 不设。

### 常用命令

```bash
# 对齐 M（性能测试）
python test_grouped_gemm.py --dtype fp16                              # 1024³×3
python test_grouped_gemm.py --dtype fp16 --groups 1 --m 128 --n 128 --k 128
python test_grouped_gemm.py --dtype all --variable
python test_grouped_gemm.py --heterogeneous --dtype fp16

# 任意 M（MOE 测试）
python test_grouped_gemm.py --moe --dtype fp16
python test_grouped_gemm.py --moe --dtype fp16 --groups 3            # 只测前 3 个 M
python test_grouped_gemm.py --moe --dtype fp8

# CI 冒烟
python test_grouped_gemm.py --smoke --dtype fp16
python test_grouped_gemm.py --smoke --dtype all --variable

# 其他
python test_grouped_gemm.py --layout                                  # layout 变体
python test_grouped_gemm.py --bad-input                               # 错误输入
```
