diff --git a/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md b/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md new file mode 100644 index 00000000..f32dc992 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md @@ -0,0 +1,266 @@ +# DCU Kernel Auto-Opt 工作流速查（给接手 AI 的快速上手） + +> 维护说明：本文档由人工/AI 维护，内容基于 2026-08 的实际代码与现场任务。 +> 代码演进后请同步更新；涉及版本、容器名、具体数值的地方会标注“以实际为准”。 +> 遇到不确定的信息，先查代码/任务现场，不要凭记忆下结论。 + +## 0. 这个功能是什么 + +MetaInfer 的一个 task 插件：在一个节点（worker29，4×K500SM_AI/gfx928）上用多 Agent + +4 GPU 对 DCU 算子做“生成/优化 → 可信正确性+性能基准 → 串行回归验证”的闭环。 +当前主算子是 **INT8 W8A8 GEMM**（DeepSeek-V4 TP4/TP8 的六个逻辑 GEMM），内核语言 HIP C++（DUMMA Tensor Core）。 + +一句话流程：`新建任务表单 → 解析配置并冻结 API 契约 → 固定 Triton Graph 基线 → 每个 GPU 一个 Agent worker 迭代优化 → 合并技能 → 串行验证全部 API shape（含回归） → 出报告`。 + +## 1. 目录地图（先认路） + +``` +metainfer/tasks/dcu_kernel_auto_opt/ +├── form.yaml # 新建任务表单定义（所有可配置项） +├── api/ # 接入的算子 API 文件（plugin 本地权威契约） +│ └── int8w8a8gemm/int8_w8a8_gemm_api.py +├── variant/ # 参考变体 HIP 代码 +│ └── w8a8_gemm_variants.hip +├── assets/ +│ ├── smoke_harness.cpp # infra smoke 模式的向量 kernel +│ ├── w8a8_baseline/ # W8A8 基线扩展模板（bindings/w8a8_gemm_hip/profile_pmc.sh…） +│ └── w8a8_bench.py # 可信 harness（任务仓库里会拷贝一份） +├── orchestrator/ +│ ├── config.py # load_config：解析表单 + shape/GPU 分配校验 +│ ├── api_contracts.py # 解析权威 API、冻结快照、default_optimization_shapes +│ ├── w8a8_baselines.py # 固定 Triton Graph 基线表（缺条目直接报错） +│ ├── w8a8_pipeline.py # 真实 worker 生命周期：W8A8Runner/benchmark/PMC/验收 +│ ├── gen_and_opt_pipeline.py # Generate 模式 + 最终 synthesis + 串行验证 +│ ├── real_pipeline.py # 真实（非 mock）流程骨架 +│ ├── worker.py / phases.py # mock worker / 状态机阶段 +│ ├── guidance.py / skill_store.py / result_store.py / gpu_binding.py / pmc_profile.py +│ ├── cli.py # 命令行入口：dcu-kernel-auto-opt run requirements.json ... +│ └── adapters/ # mock kernel adapter +├── server/ # Web 插件路由（summary/iterations/guidance…） +├── static/ # 前端（dkao-shape-input.js 里硬编码了 shape 常量！） +├── bridge/ # agent bridge（控制面↔agent） +└── tests/ # 单元测试（改行为后必须跑） +``` + +## 2. 端到端工作流 + +阶段状态机见 `orchestrator/phases.py`： +`prepare → generate_kernel_repo → baseline → parallel_explore → skill_synthesis → serial_validate → report → finished`。 + +### 2.1 新建任务（表单字段，见 form.yaml） + +关键字段：`operator`（Quantized GEMM）、`kernel_language`（HIP C++）、`target_hardware`（K500SM_AI/gfx928）、 +`dtype`（INT8 W8A8）、`agent_framework`（ccb / dsh）、`agent_model`（ccb: Opus/Sonnet；dsh: deepseek-v4-flash）、`execution_mode`（Mock / Real INT8 W8A8 GEMM / +Generate & optimize / Infra smoke）、`target_repo_path`、`shape_assignment_mode`（AI automatic / Manual by GPU）、 +`shape_scope`（All API shapes / Selected shapes only）、`shape_config`、`max_iterations`、`minimum_improvement_percent`、`extra_notes`。 + +`shape_config` 是 YAML（config.py 解析，最多 4 个 worker、每 GPU 一个、shape 必须恰好分配一次）： + +```yaml +shapes: + - {id: tp4_wqkv_a_m4096, tp_size: 4, operator: wqkv_a, M: 4096, N: 1536, K: 4096} + # ... 其余 shape +assignments: + worker_0: {gpu: 0, shapes: [tp4_wqkv_a_m4096]} +``` + +提交后 `load_config` 会调用 `api_contracts.validate_contract_shapes`，用**冻结契约**校验每个 shape 的 +M/N/K 是否合法（M 范围、K%32==0、N%16==0、(K,N) 是否在 TP4/TP8 表内）。 + +### 2.2 固定接口（先读这三个文件） + +1. **权威 API 契约**：`metainfer/tasks/dcu_kernel_auto_opt/api/int8w8a8gemm/int8_w8a8_gemm_api.py` + （orchestrator 从这里 resolve；`METAINFER_OPERATOR_API_ROOT` 可覆盖，测试用临时目录走覆盖路径）。 + 参考变体 HIP 代码放在 `metainfer/tasks/dcu_kernel_auto_opt/variant/w8a8_gemm_variants.hip`。 +2. **任务内冻结快照**：任务仓库 `kernel-repos//int8_w8a8_gemm_api.py`，其 sha256 记录在 + `scaffold_manifest.json` 的 `control_plane_files` 中，`gen_and_opt_pipeline._task_local_api_contract` + 每次运行都校验 digest。**改权威 API 只影响新任务**（控制面会重新 staging 新 digest）；**不要手改旧任务仓库里的快照**。 +3. **固定调用面**：`w8a8_gemm_out(x_q[M,K] int8, packed_weight, x_scale[M,1] fp32, + packed_weight_scale[N,1], out[M,N] bf16, workspace) -> out`，底层是 `torch.ops.zth_w8a8.gemm_out`； + 可选 `pack_weight`。语义： + + `out[m,n] = bf16( int32_dot(x_q[m,:], weight[:,n]) * x_scale[m] * weight_scale[n] )` + + 计时区只包含 `w8a8_gemm_out`；`prepare_weight`/`allocate_workspace` 在 Graph capture 之前、不计时。 + +### 2.3 Baseline（固定表 + 可自测） + +`w8a8_baselines.py::fixed_triton_graph_baseline(shape_id, shape)` 按 `(tp, M, N, K)` 查表， +**查不到就抛 ValueError**（baseline 阶段直接失败），所以新 shape 必须先补表。 +表值是 Triton Graph 基线（µs），TP4 M=4096 条目是 2026-08-06 在 worker29 用 +`baseline/int8_utils.py`（lmslim）实测的（graph replay median：wqkv_a 13247、wq_b 20590、 +wo_b 19882、gate_up 8790、down 5546；wq_b 与 indexer.wq_b 共用一条）。 + +要自己测 Triton baseline：用 `matmul_int8`（即 SGLang/lmslim 实际调用路径），M>1024 默认 config +是 `BM256/BN256/BK64/GROUP8/SPLIT_K1/warps8`，GPU event、预分配 out（排除分配）、 +热缓存协议建议 `warmups=10, samples=20, launches_per_sample=5`；可参考 +`zth_infer/baseline/bench_triton_tp4_m4096.py`。 + +### 2.4 Parallel explore（worker 生命周期，w8a8_pipeline.py） + +- 每个 assignment 一个 worker（`worker_N ↔ GPU N`，最多 4 个）；每个 worker 负责若干 shape。 +- **Agent 每轮只能改 `csrc/w8a8_gemm_hip.hip` 和 `proposal.json`**（控制面拥有其余文件）。 +- 阶段：bootstrap（iteration 0，正确性优先；大 M 要求直接上 DUMMA tile kernel，标量只做 + unmatched/small-M fallback）→ 多轮迭代。 +- 每轮：控制面用 `W8A8Runner.benchmark` 跑 `w8a8_bench.py` + （CPU int64 exact reference + CUDA Graph replay 计时，median/P90）→ + `evaluate_candidate_acceptance`（median 提升 ≥1%（`ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT=1.0`） + 且无 P90 回归）→ 可选 `profile_pmc.sh`（hipprof PMC）。 +- 接受的 artifact 落在 `workers//accepted//`（`kernel.hip`、`kernel.cuda.o`、 + `manifest.json`、`gfx928.co`、`isa.txt`），`result.json` 里记录 `source_sha256/object_sha256`（后续会校验）。 + +### 2.5 Synthesis + Serial validate（gen_and_opt_pipeline._synthesize_final_candidate） + +1. 用 worker 的 accepted objects 重建 `final/source`：`git restore` → 拷贝 object/HIP → + `llvm-objcopy` 符号命名（`mi__` 前缀）→ 渲染 `csrc/w8a8_dispatch.cpp` → + 写 `artifact_manifest.json`。 +2. **dispatch 只按 (N,K) 路由、忽略 M**（见 w8a8_dispatch.cpp），所以回归 shape 也会走自定义 object。 +3. 验证集合 = 6 个优化 shape（默认 protocol）+ 全部 API fallback shapes（`warmups=2, samples=3`）。 + 优化 shape 还有 `final ≤ 1.05× worker best` 的性能门。 +4. 任一 shape 正确性失败 → `Final prebuilt-object W8A8 validation failed: correctness failed for : `。 +5. 全部通过 → 提交 final worktree、cherry-pick 回 seed、写 `final_report.json`，任务 finished/success。 + +### 2.6 主 agent 与子 agent 的职责边界（严格区分，不是靠自觉） + +**是严格区分的**，由“提示词 + agent 运行时 + 控制面代码”三层强制，任何越界都会被控制面判失败并重试： + +| 角色 | 允许写 | 禁止 | +|---|---|---| +| 主协调 agent（Generate 阶段，`kernel_coordinator`） | 只写 `proposal.json` | 任何 HIP/C++/CUDA/Python backend、setup/build/test/benchmark/API 文件；不得编译或跑 harness | +| 子 worker agent（bootstrap + 每轮迭代） | 只创建/修改 `csrc/w8a8_gemm_hip.hip`（+ `proposal.json`） | 其它一切；API/backend/bindings/setup/harness 都是控制面所有 | +| Final synthesis agent（主侧合并） | 只改 `csrc/w8a8_gemm_hip.hip` + `proposal.json` | 其它一切；不得自己编译/跑 benchmark | + +三层强制： + +1. **提示词层**（`orchestrator/prompts.py`）： + - `main_coordinator_prompt`：`Hard role boundary: You are not a kernel implementation agent... Write only proposal.json`； + - `bootstrap_worker_prompt`：`You own and must create csrc/w8a8_gemm_hip.hip`； + - synthesis prompt：`edit only csrc/w8a8_gemm_hip.hip and proposal.json`。 +2. **运行时层**：所有 agent 都以 + `_SOURCE_ONLY_AGENT_ARGS = ["--disallowedTools", "Bash,Skill,WebFetch,WebSearch"]` + 启动 → 不能执行 shell/编译/benchmark/联网，只能用文件工具编辑源码。 +3. **控制面层（最终裁决，`gen_and_opt_pipeline.py` / `w8a8_pipeline.py`）**： + - 主协调返回后：API 契约与参考文件 digest 校验 + `git status --porcelain --untracked-files=all`， + 除 `proposal.json` 外任何改动 → `generate_role_violation`，该 attempt 失败； + - worker bootstrap：控制面文件 digest 锁定 + `git status`，检出除 `csrc/w8a8_gemm_hip.hip` + 外改动 → `bootstrap Agent changed files outside its HIP ownership`； + - worker 迭代：契约 digest + `git diff --name-only`，只允许 `csrc/w8a8_gemm_hip.hip` + → `agent changed control-plane-owned extension infrastructure`。 + +控制面（非 agent）负责：任务 staging、API 快照、编译构建、`w8a8_bench` 正确性/性能、 +hipprof PMC、验收门槛、final synthesis（符号命名 + dispatch 渲染）、串行验证、报告。 +子 agent 的结果互不直接合并——合并由控制面/synthesis agent 完成。 + +## 3. 关键常量（2026-08 现状，改前先读 API 文件） + +TP4 (K,N)： + +| operator | K | N | +|---|---:|---:| +| wqkv_a | 4096 | 1536 | +| wq_b / indexer.wq_b | 1024 | 8192 | +| wo_b | 2048 | 4096 | +| shared_gate_up_proj | 4096 | 1024 | +| shared_down_proj | 512 | 4096 | + +TP8 另有 wq_b/wo_b (1024,4096)、gate_up (4096,512)、down (256,4096)、indexer (1024,8192) 等。 + +- `DEFAULT_OPTIMIZATION_M_VALUES = (2, 16, 3072)`；`TP4_EXTRA_OPTIMIZATION_M_VALUES = (4096,)` + （TP4 专属，2026-08-06 加）→ 默认 42 个 shape（TP4 24 + TP8 18）。 +- **模型目录同步**：前端 `static/dkao-shape-input.js` 的 `MODEL_WORKLOADS`、基线表 + `orchestrator/w8a8_baselines.py`、权威契约 `api/int8w8a8gemm/int8_w8a8_gemm_api.py` + 三处必须一致。2026-08-12 前端和基线表加了 Hy3/MiniMax M3/GLM5.2 的 TP1/TP4/TP8 目录， + 但契约只补了 Hy3 TP4；MiniMax M3 TP4 任务在 prepare 直接 + `infra_fail`（`qkv_proj (6144,2304)` 不在 `TP4_OPERATOR_ALLOWED_KN` 内）。 + 已把目录 TP4/TP8 形状并入契约的 ALLOWED 集合（`HY3_TP8/MINIMAX_TP4/MINIMAX_TP8/GLM52_TP4/GLM52_TP8_OPERATOR_KN`）， + 但不进 `TP_OPERATOR_KN`，默认 42 shape 不变。**新模型/新 shape 上线时三处一起改**； + 注意 TP1 目录前端可选、契约仍只收 tp_size 4/8，选 TP1 会同样在 prepare 被拒。 +- **模型目录 TP8 大 prefill 边界（2026-08-27）**：Hy3/MiniMax M3/GLM5.2 的 TP8 形状支持 + M=4096 优化（"Selected shapes only" 手动选择，如 minimax 任务）。契约加 + `MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES=(4096,)`（**不进** `DEFAULT_OPTIMIZATION_SHAPES`， + 默认 42 不变，串行验证 fallback 不受影响）；前端三个模型的 TP8 topology 加 + `mValues: [2,16,3072,4096]`；基线表 15 条 `(8,4096,…)` 已实测 + （`baseline/bench_triton_tp8_m4096.py`，int8_utils.matmul_kernel + CUDA-graph replay， + 结果存 `baseline/tp8_m4096_graph.json`）。DeepSeek TP8 不加 M=4096。 +- `MIN_M=1, MAX_M=4096`；`WORKSPACE_BUDGET_BYTES=16MB`。 +- **M=4096 时大部分 (N,K) 的 split-K workspace 容量为 0** → 大 M kernel 必须走 2D M-tile 路径， + 不能依赖 split-K workspace。 + +## 4. 环境（重要） + +- **本机就是 worker29**（hostname=worker29，10.18.17.80，4×K500SM_AI/gfx928，DTK 26.04）。 +- **MetaInfer 相关容器一律 `zth_meta`**：挂载 `zth_infer → /workspace`，PID1=`serve.py --port 8765`。 + GPU python 需要： + `source /opt/dtk/env.sh` + `HIP_VISIBLE_DEVICES=` + `PYTHONPATH=/workspace/MetaInfer`； + 任务工作区在容器里是 `/workspace/MetaInfer/nodes/worker29/workspaces//`。 +- serving/benchmark 容器 `zth1-sglang-deepseek-v4-flash-tracing`（torch2.9/triton3.3/lmslim） + 只用于 Triton 基线测量等独立用途，**不要改它代码、不要在里面跑 MetaInfer 任务流程**。 + （容器名历史上变过，以 `docker ps` 为准。） +- GPU 绑定：`gpu_binding.py` 只设 `HIP_VISIBLE_DEVICES`，不要同时设 `ROCR_VISIBLE_DEVICES`。 + +## 5. 常见坑（都是真实踩过的） + +1. **布局一致性（最容易翻车）**：object 的 `pack_weight` 如果重排了 B（例：down_proj (512,4096) + 打成 N-blocked 64 列），它内部**所有按 M 分发的路径必须读同一个布局**。dispatch 按 (N,K) 忽略 M， + 回归 M（2/16/3072）会走进通用路径；通用 DUMMA 若按 row-major 读 N-blocked B 就全错。 + 修复示例：`dumma_eligible && (n,k)==(4096,512)` 时改用 N-blocked-aware 的 64x64 tile kernel， + 精确 M=4096 的 128x64 路径保持不变。 +2. **worker 只验证分配到的 M** → 回归 bug 拖到 serial validate 才暴露。 + 建议：worker 验收时对每个 (K,N) 额外跑 M=2/16/3072 正确性（M-sweep）。 +3. **不要手改任务仓库的 API 快照**（digest mismatch）；改 + `metainfer/tasks/dcu_kernel_auto_opt/api/` 下的权威契约后新建任务。 +4. **CPU int64 reference 在 M≥3072 很慢**（分钟级）→ 用 reference 缓存目录 + （`final/cache/references/`，key 是 m-n-k），w8a8_bench 用 `--reference-cache-dir` 命中。 + **2026-08-23 起控制面自动预置**：`W8A8Runner.benchmark`（check_correctness 时）发现 + `workers//cache/references/exact-int64-v1-m-n-k.pt` 缺失，会先以 + `--prepare-reference` 模式（独立 3600s 超时 `_REFERENCE_PREPARE_TIMEOUT_S`）生成缓存， + 再跑 900s 超时的正式 bench（命中缓存）。**坑**：torch.mm 的 int64 GEMM 只有 ~0.2 GFLOPS， + M=4096/K=6144 单个参考就要 ~10 分钟 solo、4 worker 并发 CPU 争抢更久——没预置缓存时 + bootstrap 的 900s 子进程必然超时（minimaxm3-dsh-tp4-m4096-1-0c2f84a9 的 worker_0/1/3 就因此 + 3 次全超时 failed）。老任务仓库里已 staging 的 w8a8_bench.py 是旧版（无 `--prepare-reference`）， + 恢复老任务需把新 harness 提交进任务仓库（同步 `scaffold_manifest.json` 的 + `control_plane_files.w8a8_bench.py` digest）再重启 worker；新建任务自动用新版。 +5. **serial validate 失败后的恢复流程**（有先例 `zth_infer/recover_c9_serial_validate.py`、 + `recover_cc86e2b2_serial_validate.py`）： + patch `workers//accepted//kernel.hip` → 用 hipcc 重编 `kernel.cuda.o` + （`-O3 --offload-arch=gfx928 -std=c++17 -fPIC -fno-gpu-rdc`）→ 更新 + `result.json`/`manifest.json` 里的 `source_sha256/object_sha256` → 写恢复驱动 + （`GenAndOptPipeline._phase(VALIDATE,...) → _synthesize_final_candidate → _phase(REPORT)`）。 + WebUI 的 `POST /workers/{id}/restart` 注意两点（2026-08-23 修过）： + - **不再强制传 `--claude-bin`**（原来传 `METAINFER_CLAUDE_BIN`，对 dsh 框架会拉起 + Claude 二进制而不是 `bridge/dsh/dsh_agent.py`）；`restart_worker.py` / + `integrate_restarted_worker.py` 自己按 `agent_framework` 解析。 + - `restart_worker.py` 归档 `failure.json` 用 `shutil.move` 而非 `Path.replace`： + overlayfs 下老文件在 lower layer、新 `restarts/` 目录在 upper layer， + `os.rename` 抛 EXDEV（[Errno 18] Invalid cross-device link）。 + 注意：容器内后台跑长任务用 `docker exec -d`，否则进程会随 exec 会话结束被杀。 +6. **前端 shape 常量** `static/dkao-shape-input.js` 与 API 默认 shape 要同步 + （TP4 专属 M=4096 是分别维护的）。 +7. **角色越界会自动判失败**：`generate_role_violation` / `changed files outside its HIP ownership` + / `agent changed control-plane-owned extension infrastructure` 都是控制面在 agent 返回后 + 用 git diff/digest 检出的。修复方式是恢复被改文件、重试该 attempt，**不要绕过检查**。 + +## 6. 新 AI 快速开始 checklist + +1. 先读：`form.yaml` → `api_contracts.py` → 权威 `int8_w8a8_gemm_api.py` → + `w8a8_pipeline.py`（W8A8Runner/worker loop）→ `gen_and_opt_pipeline.py` + （`_synthesize_final_candidate`/serial validate）。 +2. 跑测试：`python3 -m pytest metainfer/tasks/dcu_kernel_auto_opt/tests/`。 +3. 看一个真实任务现场（已成功的例子：`dcu-kernel-auto-opt-cc86e2b2`）： + `plan.json`（配置）、`shared_baseline/results.json`（基线）、`workers/*/accepted/*`（kernel）、 + `final/source/csrc/w8a8_dispatch.cpp`（路由）、`final_report.json`（结果）。 +4. 内核级工作参考 skill：W8A8 调优已按阶段拆成 skill 家族（2026-08-23）—— + `int8-w8a8-gemm-decode`（M≤32）+ `int8-w8a8-gemm-prefill`（M>32）+ 共享地基 + `int8-w8a8-gemm-foundations`；旧名 `int8-w8a8-quantized-gemm-optimization` 保留为路由器。 + 其余：`dcu-kernel-tuning`、`hygon-dcu-kernel`、`hygon-gfx928-memory-isa`、 + `sglang-custom-kernel-integration`；环境/SSH/容器细节参考 `remote-dcu-env`。 + 改动 skill 库（`~/.dsh/skills/`）后记得跑 `sync_skill_libraries()` 镜像到 + `~/.claude/skills/`（skill_store 测试里有覆盖）。 +5. 改动任何行为后：更新本文件相关段落 + 跑 tests + 用真实任务验证（优先在 zth_meta 里）。 + +## 7. 不确定性标注 + +- 版本号、容器名、基线数值均为 2026-08 观察值，动手前以实际代码/`docker ps`/现场数据为准。 +- “恢复流程”是手工驱动（复用 `_synthesize_final_candidate`），不是 UI 一键重试；UI 是否提供重试以 + `server/routes.py` 实际实现为准。 +- 本文档不替代 skill 里的性能调优细节（tile 选择、LDS、DUMMA API、hipprof 用法），那些看对应 skill。 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/__init__.py b/metainfer/tasks/dcu_kernel_auto_opt/__init__.py new file mode 100644 index 00000000..9e186b92 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/__init__.py @@ -0,0 +1,4 @@ +"""DCU multi-agent kernel auto-optimization task package.""" + +from .orchestrator import plugin as _task_plugin # noqa: F401 +from .server import plugin as _web_plugin # noqa: F401 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/api/int8w8a8gemm/int8_w8a8_gemm_api.py b/metainfer/tasks/dcu_kernel_auto_opt/api/int8w8a8gemm/int8_w8a8_gemm_api.py new file mode 100644 index 00000000..2951f48c --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/api/int8w8a8gemm/int8_w8a8_gemm_api.py @@ -0,0 +1,556 @@ +"""Stable Python contract for TP=4/8 INT8 W8A8 GEMM. + +This file is the boundary between benchmarks/framework code and an optimizable +backend. Kernel authors may change packing, dispatch, HIP/DUMMA kernels, and +workspace use, but should not change the public function signatures or tensor +semantics in this file. + +Logical operation: + out[m, n] = bf16( + int32_dot(x_q[m, :], raw_weight[:, n]) + * x_scale[m, 0] + * weight_scale[n, 0] + ) + +Only ``w8a8_gemm_out`` belongs inside the timed/CUDA-Graph region. +``prepare_weight`` and ``allocate_workspace`` must run before graph capture and +are excluded from GEMM timing. + +Required backend operator schema: + zth_w8a8::gemm_out( + Tensor x_q, + Tensor packed_weight, + Tensor x_scale, + Tensor packed_weight_scale, + Tensor(a!) out, + Tensor(b!) workspace + ) -> Tensor(a!) + +Optional backend packing schema: + zth_w8a8::pack_weight( + Tensor raw_weight, + Tensor weight_scale + ) -> (Tensor, Tensor) +""" + +from __future__ import annotations + +from typing import Final, Mapping + +try: + import torch +except ModuleNotFoundError: # Metadata validation runs in CPU-only CI. + torch = None # type: ignore[assignment] + + +def _torch_no_grad(): + """Return PyTorch's decorator, or an import-only CI fallback.""" + if torch is None: + return lambda function: function + return torch.no_grad() + + +def _require_torch() -> None: + if torch is None: + raise RuntimeError( + "the INT8 W8A8 runtime API requires PyTorch in the DCU environment" + ) + + +# These are logical (K, N) dimensions after TP partitioning, not checkpoint +# storage shapes. Operator identity is retained even when two operators share +# the same numerical shape, so TP=4 and TP=8 task metadata cannot be confused. +TP4_OPERATOR_KN: Final[Mapping[str, tuple[int, int]]] = { + "wqkv_a": (4096, 1536), + "wq_b": (1024, 8192), + "indexer.wq_b": (1024, 8192), + "wo_b": (2048, 4096), + "shared_gate_up_proj": (4096, 1024), + "shared_down_proj": (512, 4096), +} +TP8_OPERATOR_KN: Final[Mapping[str, tuple[int, int]]] = { + "wqkv_a": (4096, 1536), + "wq_b": (1024, 4096), + "indexer.wq_b": (1024, 8192), + "wo_b": (1024, 4096), + "shared_gate_up_proj": (4096, 512), + "shared_down_proj": (256, 4096), +} +HY3_TP4_OPERATOR_KN: Final[Mapping[str, tuple[int, int]]] = { + "qkv_proj": (4096, 2560), + "o_proj": (2048, 4096), + "shared_gate_up_proj": (4096, 768), + "shared_down_proj": (384, 4096), +} +# Additional TP4/TP8 model-catalog operators (Hy3, MiniMax M3, GLM5.2). +# These mirror the frontend workload catalog (static/dkao-shape-input.js) +# and the measured Triton baseline table (orchestrator/w8a8_baselines.py) so +# tasks created from the model catalog pass the fixed contract validation. +# They are deliberately excluded from TP_OPERATOR_KN below, so +# DEFAULT_OPTIMIZATION_SHAPES (DeepSeek TP4/TP8 only) is unchanged. +HY3_TP8_OPERATOR_KN: Final[Mapping[str, tuple[int, int]]] = { + "qkv_proj": (4096, 1280), + "o_proj": (1024, 4096), + "shared_gate_up_proj": (4096, 384), + "shared_down_proj": (192, 4096), +} +MINIMAX_TP4_OPERATOR_KN: Final[Mapping[str, tuple[int, int]]] = { + "qkv_proj": (6144, 2304), + "qkv_proj_and_indexer_qk": (6144, 2560), + "o_proj": (2048, 6144), + "shared_gate_up_proj": (6144, 1536), + "shared_down_proj": (768, 6144), +} +MINIMAX_TP8_OPERATOR_KN: Final[Mapping[str, tuple[int, int]]] = { + "qkv_proj": (6144, 1280), + "qkv_proj_and_indexer_qk": (6144, 1536), + "o_proj": (1024, 6144), + "shared_gate_up_proj": (6144, 768), + "shared_down_proj": (384, 6144), +} +GLM52_TP4_OPERATOR_KN: Final[Mapping[str, tuple[int, int]]] = { + "fused_qkv_a_proj": (6144, 2624), + "q_b_proj": (2048, 4096), + "kv_b_proj": (512, 7168), + "o_proj": (4096, 6144), + "shared_gate_up_proj": (6144, 1024), + "shared_down_proj": (512, 6144), +} +GLM52_TP8_OPERATOR_KN: Final[Mapping[str, tuple[int, int]]] = { + "fused_qkv_a_proj": (6144, 2624), + "q_b_proj": (2048, 2048), + "kv_b_proj": (512, 3584), + "o_proj": (2048, 6144), + "shared_gate_up_proj": (6144, 512), + "shared_down_proj": (256, 6144), +} + + +def _merge_kn_tables( + *tables: Mapping[str, tuple[int, int]], +) -> Mapping[str, tuple[tuple[int, int], ...]]: + """Merge (K, N) tables per operator, keeping every distinct pair. + + A shared operator name (e.g. ``o_proj``) may carry different (K, N) + pairs per model; each pair stays individually allowed. + """ + merged: dict[str, list[tuple[int, int]]] = {} + for table in tables: + for operator, kn in table.items(): + if kn not in merged.setdefault(operator, []): + merged[operator].append(kn) + return {operator: tuple(kns) for operator, kns in merged.items()} + + +_TP4_CATALOG_OPERATOR_KN: Final[Mapping[str, tuple[tuple[int, int], ...]]] = ( + _merge_kn_tables( + TP4_OPERATOR_KN, + HY3_TP4_OPERATOR_KN, + MINIMAX_TP4_OPERATOR_KN, + GLM52_TP4_OPERATOR_KN, + ) +) +_TP8_CATALOG_OPERATOR_KN: Final[Mapping[str, tuple[tuple[int, int], ...]]] = ( + _merge_kn_tables( + TP8_OPERATOR_KN, + HY3_TP8_OPERATOR_KN, + MINIMAX_TP8_OPERATOR_KN, + GLM52_TP8_OPERATOR_KN, + ) +) +TP_OPERATOR_KN: Final[Mapping[int, Mapping[str, tuple[int, int]]]] = { + 4: TP4_OPERATOR_KN, + 8: TP8_OPERATOR_KN, +} +TP4_OPERATOR_ALLOWED_KN: Final[Mapping[str, frozenset[tuple[int, int]]]] = { + operator: frozenset(kns) + for operator, kns in _TP4_CATALOG_OPERATOR_KN.items() +} +TP8_OPERATOR_ALLOWED_KN: Final[Mapping[str, frozenset[tuple[int, int]]]] = { + operator: frozenset(kns) + for operator, kns in _TP8_CATALOG_OPERATOR_KN.items() +} +TP_OPERATOR_ALLOWED_KN: Final[ + Mapping[int, Mapping[str, frozenset[tuple[int, int]]]] +] = { + 4: TP4_OPERATOR_ALLOWED_KN, + 8: TP8_OPERATOR_ALLOWED_KN, +} +TP4_W8A8_KN: Final[frozenset[tuple[int, int]]] = frozenset( + kn for allowed in TP4_OPERATOR_ALLOWED_KN.values() for kn in allowed +) +TP8_W8A8_KN: Final[frozenset[tuple[int, int]]] = frozenset( + kn for allowed in TP8_OPERATOR_ALLOWED_KN.values() for kn in allowed +) +SUPPORTED_W8A8_KN: Final[frozenset[tuple[int, int]]] = ( + TP4_W8A8_KN | TP8_W8A8_KN +) +# Compatibility alias for repositories generated with the older API. +TP4_DECODE_KN: Final[frozenset[tuple[int, int]]] = TP4_W8A8_KN + +DEFAULT_OPTIMIZATION_M_VALUES: Final[tuple[int, ...]] = (2, 16, 3072) +# TP4-only large-prefill boundary added on 2026-08-06. The TP8 workload keeps +# the original three M values, so each TP4 operator gets exactly one M=4096 +# shape while TP8 defaults are unchanged. +TP4_EXTRA_OPTIMIZATION_M_VALUES: Final[tuple[int, ...]] = (4096,) +# Model-catalog TP8 large-prefill boundary added on 2026-08-27. The DeepSeek +# TP8 default workload keeps the original three M values; Hy3 / MiniMax M3 / +# GLM5.2 TP8 operators are additionally optimizable at M=4096 via +# "Selected shapes only" (their (K,N) pairs are already in the allowed sets +# and their Triton CUDA-graph baselines are measured). This constant is not +# applied by _default_optimization_shapes(), so DEFAULT_OPTIMIZATION_SHAPES +# (DeepSeek-only) and its serial-validation fallback scope are unchanged. +MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES: Final[tuple[int, ...]] = (4096,) + + +def _default_optimization_shapes() -> tuple[dict[str, int | str], ...]: + """Return logical TP/operator shapes with collision-free task IDs. + + Identical numerical GEMMs may appear more than once because the task ID and + metadata deliberately preserve their model call site. Each shape still + produces an independent accepted artifact, so no runtime dispatch conflict + is introduced. + """ + shapes: list[dict[str, int | str]] = [] + for tp_size, operators in TP_OPERATOR_KN.items(): + m_values = ( + DEFAULT_OPTIMIZATION_M_VALUES + + TP4_EXTRA_OPTIMIZATION_M_VALUES + if tp_size == 4 + else DEFAULT_OPTIMIZATION_M_VALUES + ) + for operator, (k, n) in operators.items(): + operator_id = operator.replace(".", "_") + for m in m_values: + shapes.append( + { + "id": f"tp{tp_size}_{operator_id}_m{m}", + "tp_size": tp_size, + "operator": operator, + "M": m, + "N": n, + "K": k, + } + ) + return tuple(shapes) + + +# Default workload used when MetaInfer New Task leaves shapes empty. It covers +# decode, short-token DUMMA, and the requested large-prefill boundary. +DEFAULT_OPTIMIZATION_SHAPES: Final[tuple[dict[str, int | str], ...]] = ( + _default_optimization_shapes() +) + +MIN_M: Final[int] = 1 +MAX_M: Final[int] = 4096 +WORKSPACE_BUDGET_BYTES: Final[int] = 16 * 1024 * 1024 +WORKSPACE_MAX_SPLIT_K: Final[int] = 16 +# Backward-compatible name used by existing generated repositories. Capacity +# is now calculated from the byte budget for each shape rather than fixed at 8. +WORKSPACE_SPLIT_K_CAP: Final[int] = WORKSPACE_MAX_SPLIT_K +WORKSPACE_ALIGNMENT: Final[int] = 256 +INT32_BYTES: Final[int] = 4 + + +def _check_cuda_tensor( + name: str, + tensor: torch.Tensor, + dtype: torch.dtype, + *, + contiguous: bool = True, +) -> None: + _require_torch() + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if not tensor.is_cuda: + raise ValueError(f"{name} must be on a CUDA/HIP device") + if tensor.dtype != dtype: + raise TypeError(f"{name}.dtype must be {dtype}, got {tensor.dtype}") + if contiguous and not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + +def _check_same_device(reference: torch.Tensor, **tensors: torch.Tensor) -> None: + for name, tensor in tensors.items(): + if tensor.device != reference.device: + raise ValueError( + f"{name}.device must be {reference.device}, got {tensor.device}" + ) + + +def _check_target_shape(m: int, n: int, k: int) -> None: + if not MIN_M <= m <= MAX_M: + raise ValueError(f"M must be in [{MIN_M}, {MAX_M}], got {m}") + if (k, n) not in SUPPORTED_W8A8_KN: + raise ValueError( + f"unsupported TP4/TP8 logical shape (K, N)=({k}, {n}); " + f"supported={sorted(SUPPORTED_W8A8_KN)}" + ) + if k % 32 != 0: + raise ValueError(f"K must be divisible by the INT8 DUMMA K tile 32, got {k}") + if n % 16 != 0: + raise ValueError(f"N must be divisible by the DUMMA N tile 16, got {n}") + + +def validate_optimization_shape(shape: Mapping[str, object]) -> None: + """Validate one topology-qualified MetaInfer optimization shape.""" + try: + tp_size = int(shape["tp_size"]) + operator = str(shape["operator"]) + m = int(shape["M"]) + n = int(shape["N"]) + k = int(shape["K"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError( + "optimization shape requires tp_size, operator, M, N and K" + ) from exc + operators = TP_OPERATOR_ALLOWED_KN.get(tp_size) + if operators is None: + raise ValueError(f"tp_size must be 4 or 8, got {tp_size}") + allowed = operators.get(operator) + if allowed is None: + raise ValueError( + f"unsupported TP={tp_size} operator {operator!r}; " + f"supported={sorted(operators)}" + ) + if (k, n) not in allowed: + raise ValueError( + f"TP={tp_size} {operator} requires (K, N) in {sorted(allowed)}, " + f"got ({k}, {n})" + ) + _check_target_shape(m, n, k) + + +def _optional_op(namespace: str, name: str): + _require_torch() + try: + return getattr(getattr(torch.ops, namespace), name) + except AttributeError: + return None + + +def _required_op(namespace: str, name: str): + op = _optional_op(namespace, name) + if op is None: + raise RuntimeError( + f"required custom operator torch.ops.{namespace}.{name} is not " + "registered; load/build the W8A8 HIP extension before calling it" + ) + return op + + +@_torch_no_grad() +def prepare_weight( + raw_weight: torch.Tensor, + weight_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare a logical [K, N] INT8 weight outside the timed region. + + The optional backend op may return any contiguous opaque packed layout. + Until a packing op is registered, the identity contiguous layout is used. + The returned tensors must remain alive and at stable addresses throughout + CUDA/HIP Graph capture and replay. + """ + _check_cuda_tensor("raw_weight", raw_weight, torch.int8) + _check_cuda_tensor("weight_scale", weight_scale, torch.float32) + if raw_weight.ndim != 2: + raise ValueError( + f"raw_weight must have logical shape [K, N], got {raw_weight.shape}" + ) + k, n = raw_weight.shape + _check_target_shape(MAX_M, n, k) + if weight_scale.shape != (n, 1): + raise ValueError( + f"weight_scale must have shape ({n}, 1), got {weight_scale.shape}" + ) + _check_same_device(raw_weight, weight_scale=weight_scale) + + pack_op = _optional_op("zth_w8a8", "pack_weight") + if pack_op is None: + return raw_weight.contiguous(), weight_scale.contiguous() + + packed_weight, packed_scale = pack_op(raw_weight, weight_scale) + _check_cuda_tensor("packed_weight", packed_weight, torch.int8) + _check_cuda_tensor("packed_weight_scale", packed_scale, torch.float32) + _check_same_device( + raw_weight, + packed_weight=packed_weight, + packed_weight_scale=packed_scale, + ) + if packed_weight.numel() < k * n: + raise ValueError( + "packed_weight must contain at least K*N int8 elements; " + f"got {packed_weight.numel()} for K*N={k*n}" + ) + if packed_scale.numel() < n: + raise ValueError( + "packed_weight_scale must contain at least N fp32 elements; " + f"got {packed_scale.numel()} for N={n}" + ) + return packed_weight, packed_scale + + +def allocate_workspace( + m: int, + n: int, + k: int, + device: torch.device | str | int, +) -> torch.Tensor: + """Allocate an opaque, reusable workspace before graph capture. + + Capacity is shape-aware and capped by a fixed byte budget. This permits + non-power-of-two split-K choices while keeping allocation bounded. + """ + _check_target_shape(m, n, k) + required_bytes = max( + WORKSPACE_ALIGNMENT, + workspace_split_k_capacity(m, n, k) * m * n * INT32_BYTES, + ) + aligned_bytes = ( + (required_bytes + WORKSPACE_ALIGNMENT - 1) + // WORKSPACE_ALIGNMENT + * WORKSPACE_ALIGNMENT + ) + return torch.empty(aligned_bytes, dtype=torch.uint8, device=device) + + +def workspace_split_k_capacity(m: int, n: int, k: int) -> int: + """Return legal split-K partial planes, or zero when none fit the budget.""" + _check_target_shape(m, n, k) + bytes_per_partial = m * n * INT32_BYTES + budget_capacity = WORKSPACE_BUDGET_BYTES // bytes_per_partial + stage_capacity = k // 32 + return min(WORKSPACE_MAX_SPLIT_K, budget_capacity, stage_capacity) + + +def validate_gemm_out_inputs( + x_q: torch.Tensor, + packed_weight: torch.Tensor, + x_scale: torch.Tensor, + packed_weight_scale: torch.Tensor, + out: torch.Tensor, + workspace: torch.Tensor, +) -> tuple[int, int, int]: + """Validate the frozen logical contract and return (M, N, K).""" + _check_cuda_tensor("x_q", x_q, torch.int8) + _check_cuda_tensor("packed_weight", packed_weight, torch.int8) + _check_cuda_tensor("x_scale", x_scale, torch.float32) + _check_cuda_tensor("packed_weight_scale", packed_weight_scale, torch.float32) + _check_cuda_tensor("out", out, torch.bfloat16) + _check_cuda_tensor("workspace", workspace, torch.uint8) + + if x_q.ndim != 2: + raise ValueError(f"x_q must have shape [M, K], got {x_q.shape}") + m, k = x_q.shape + if out.ndim != 2: + raise ValueError(f"out must have shape [M, N], got {out.shape}") + if out.shape[0] != m: + raise ValueError(f"out.shape[0] must equal M={m}, got {out.shape[0]}") + n = out.shape[1] + _check_target_shape(m, n, k) + + if x_scale.shape != (m, 1): + raise ValueError(f"x_scale must have shape ({m}, 1), got {x_scale.shape}") + if packed_weight.numel() < k * n: + raise ValueError( + "packed_weight must contain at least K*N int8 elements; " + f"got {packed_weight.numel()} for K*N={k*n}" + ) + if packed_weight_scale.numel() < n: + raise ValueError( + "packed_weight_scale must contain at least N fp32 elements; " + f"got {packed_weight_scale.numel()} for N={n}" + ) + minimum_workspace_bytes = ( + workspace_split_k_capacity(m, n, k) * m * n * INT32_BYTES + ) + if workspace.numel() < minimum_workspace_bytes: + raise ValueError( + f"workspace requires at least {minimum_workspace_bytes} uint8 " + f"elements, got {workspace.numel()}" + ) + _check_same_device( + x_q, + packed_weight=packed_weight, + x_scale=x_scale, + packed_weight_scale=packed_weight_scale, + out=out, + workspace=workspace, + ) + return m, n, k + + +@_torch_no_grad() +def w8a8_gemm_out( + x_q: torch.Tensor, + packed_weight: torch.Tensor, + x_scale: torch.Tensor, + packed_weight_scale: torch.Tensor, + out: torch.Tensor, + workspace: torch.Tensor, +) -> torch.Tensor: + """Run the graph-capturable W8A8 GEMM into preallocated ``out``. + + Backend requirements: + * launch on PyTorch's current stream; + * perform no allocation, compilation, autotuning, packing, or host sync; + * include every required split-K/combine/epilogue kernel in this call; + * return the exact same tensor/storage as ``out``. + """ + validate_gemm_out_inputs( + x_q, + packed_weight, + x_scale, + packed_weight_scale, + out, + workspace, + ) + gemm_op = _required_op("zth_w8a8", "gemm_out") + result = gemm_op( + x_q, + packed_weight, + x_scale, + packed_weight_scale, + out, + workspace, + ) + if not isinstance(result, torch.Tensor): + raise TypeError("torch.ops.zth_w8a8.gemm_out must return a Tensor") + if result.data_ptr() != out.data_ptr(): + raise RuntimeError("gemm_out must return the same storage as out") + return result + + +__all__ = [ + "DEFAULT_OPTIMIZATION_M_VALUES", + "DEFAULT_OPTIMIZATION_SHAPES", + "GLM52_TP4_OPERATOR_KN", + "GLM52_TP8_OPERATOR_KN", + "HY3_TP4_OPERATOR_KN", + "HY3_TP8_OPERATOR_KN", + "MAX_M", + "MIN_M", + "MINIMAX_TP4_OPERATOR_KN", + "MINIMAX_TP8_OPERATOR_KN", + "MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES", + "TP4_EXTRA_OPTIMIZATION_M_VALUES", + "TP4_DECODE_KN", + "TP4_OPERATOR_ALLOWED_KN", + "TP4_OPERATOR_KN", + "TP4_W8A8_KN", + "TP8_OPERATOR_ALLOWED_KN", + "TP8_OPERATOR_KN", + "TP8_W8A8_KN", + "TP_OPERATOR_ALLOWED_KN", + "TP_OPERATOR_KN", + "SUPPORTED_W8A8_KN", + "WORKSPACE_BUDGET_BYTES", + "WORKSPACE_MAX_SPLIT_K", + "WORKSPACE_SPLIT_K_CAP", + "allocate_workspace", + "prepare_weight", + "validate_optimization_shape", + "validate_gemm_out_inputs", + "w8a8_gemm_out", + "workspace_split_k_capacity", +] diff --git a/metainfer/tasks/dcu_kernel_auto_opt/assets/smoke_harness.cpp b/metainfer/tasks/dcu_kernel_auto_opt/assets/smoke_harness.cpp new file mode 100644 index 00000000..0c7ca7ff --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/assets/smoke_harness.cpp @@ -0,0 +1,134 @@ +#include + +#include +#include +#include +#include +#include +#include + +#define HIP_CHECK(call) do { \ + hipError_t error = (call); \ + if (error != hipSuccess) { \ + std::fprintf(stderr, "%s\n", hipGetErrorString(error)); \ + return 2; \ + } \ +} while (0) + +__global__ void scalar_kernel(const float* input, float* output, size_t count) { + size_t index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < count) output[index] = input[index] * 2.0f + 1.0f; +} + +__global__ void vector4_kernel( + const float4* input, float4* output, size_t count4) { + size_t index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < count4) { + float4 value = input[index]; + output[index] = make_float4( + value.x * 2.0f + 1.0f, value.y * 2.0f + 1.0f, + value.z * 2.0f + 1.0f, value.w * 2.0f + 1.0f); + } +} + +static void launch( + const std::string& variant, const float* input, float* output, + size_t count, hipStream_t stream) { + constexpr int threads = 256; + if (variant == "vector4") { + size_t count4 = count / 4; + hipLaunchKernelGGL( + vector4_kernel, dim3((count4 + threads - 1) / threads), + dim3(threads), 0, stream, reinterpret_cast(input), + reinterpret_cast(output), count4); + } else { + hipLaunchKernelGGL( + scalar_kernel, dim3((count + threads - 1) / threads), + dim3(threads), 0, stream, input, output, count); + } +} + +int main(int argc, char** argv) { + int count_devices = 0; + HIP_CHECK(hipGetDeviceCount(&count_devices)); + hipDeviceProp_t properties{}; + if (count_devices > 0) HIP_CHECK(hipGetDeviceProperties(&properties, 0)); + if (argc == 2 && std::string(argv[1]) == "--probe") { + std::printf( + "{\"visible_devices\":%d,\"logical_device\":0," + "\"device_name\":\"%s\",\"warp_size\":%d}\n", + count_devices, count_devices ? properties.name : "", + count_devices ? properties.warpSize : 0); + return count_devices == 1 ? 0 : 3; + } + if (argc != 3) { + std::fprintf(stderr, "usage: smoke_harness ELEMENTS scalar|vector4\n"); + return 2; + } + size_t count = std::strtoull(argv[1], nullptr, 10); + std::string variant = argv[2]; + if (count < 4 || count % 4 != 0 || + (variant != "scalar" && variant != "vector4")) { + std::fprintf(stderr, "invalid arguments\n"); + return 2; + } + + std::vector host_input(count); + std::vector host_output(count); + for (size_t i = 0; i < count; ++i) { + host_input[i] = static_cast(static_cast(i % 257) - 128) / 17; + } + float *input = nullptr, *output = nullptr; + HIP_CHECK(hipMalloc(&input, count * sizeof(float))); + HIP_CHECK(hipMalloc(&output, count * sizeof(float))); + HIP_CHECK(hipMemcpy( + input, host_input.data(), count * sizeof(float), hipMemcpyHostToDevice)); + + for (int i = 0; i < 10; ++i) launch(variant, input, output, count, nullptr); + HIP_CHECK(hipDeviceSynchronize()); + std::vector samples; + samples.reserve(50); + hipEvent_t begin, end; + HIP_CHECK(hipEventCreate(&begin)); + HIP_CHECK(hipEventCreate(&end)); + for (int i = 0; i < 50; ++i) { + HIP_CHECK(hipEventRecord(begin)); + launch(variant, input, output, count, nullptr); + HIP_CHECK(hipGetLastError()); + HIP_CHECK(hipEventRecord(end)); + HIP_CHECK(hipEventSynchronize(end)); + float milliseconds = 0; + HIP_CHECK(hipEventElapsedTime(&milliseconds, begin, end)); + samples.push_back(milliseconds * 1000.0f); + } + HIP_CHECK(hipMemcpy( + host_output.data(), output, count * sizeof(float), hipMemcpyDeviceToHost)); + bool passed = true; + for (size_t i = 0; i < count; ++i) { + float expected = host_input[i] * 2.0f + 1.0f; + if (std::fabs(host_output[i] - expected) > 1e-5f) { + passed = false; + break; + } + } + std::sort(samples.begin(), samples.end()); + float median = samples[samples.size() / 2]; + float p90 = samples[static_cast(samples.size() * 0.9)]; + double seconds = median * 1e-6; + double tflops = (2.0 * static_cast(count)) / seconds / 1e12; + double bandwidth = (8.0 * static_cast(count)) / seconds / 1e9; + std::printf( + "{\"passed\":%s,\"variant\":\"%s\",\"elements\":%zu," + "\"visible_devices\":%d,\"device_name\":\"%s\"," + "\"median_us\":%.6f,\"p90_us\":%.6f,\"min_us\":%.6f," + "\"max_us\":%.6f,\"tflops\":%.9f,\"bandwidth_gb_s\":%.6f," + "\"warmup\":10,\"samples\":50}\n", + passed ? "true" : "false", variant.c_str(), count, count_devices, + properties.name, median, p90, samples.front(), samples.back(), + tflops, bandwidth); + hipEventDestroy(begin); + hipEventDestroy(end); + hipFree(input); + hipFree(output); + return passed ? 0 : 4; +} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/csrc/bindings.cpp b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/csrc/bindings.cpp new file mode 100644 index 00000000..1765cf88 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/csrc/bindings.cpp @@ -0,0 +1,104 @@ +#include +#include + +#include +#include + +#include +#include + + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream); + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream); + + +namespace { + +std::tuple pack_weight_impl( + const at::Tensor& raw_weight, + const at::Tensor& weight_scale) { + auto packed_weight = at::empty_like(raw_weight); + auto packed_weight_scale = at::empty_like(weight_scale); + const int k = static_cast(raw_weight.size(0)); + const int n = static_cast(raw_weight.size(1)); + const auto device_index = raw_weight.device().index(); + const hipStream_t stream = + c10::cuda::getCurrentCUDAStream(device_index).stream(); + launch_pack_w8a8_weight( + raw_weight.data_ptr(), + weight_scale.data_ptr(), + packed_weight.data_ptr(), + packed_weight_scale.data_ptr(), + k, + n, + stream); + return std::make_tuple(packed_weight, packed_weight_scale); +} + +at::Tensor gemm_out_impl( + const at::Tensor& x_q, + const at::Tensor& packed_weight, + const at::Tensor& x_scale, + const at::Tensor& packed_weight_scale, + at::Tensor out, + const at::Tensor& workspace) { + const int m = static_cast(x_q.size(0)); + const int k = static_cast(x_q.size(1)); + const int n = static_cast(out.size(1)); + const auto device_index = x_q.device().index(); + const hipStream_t stream = + c10::cuda::getCurrentCUDAStream(device_index).stream(); + launch_w8a8_gemm( + x_q.data_ptr(), + packed_weight.data_ptr(), + x_scale.data_ptr(), + packed_weight_scale.data_ptr(), + out.data_ptr(), + workspace.data_ptr(), + static_cast(workspace.numel()), + m, + n, + k, + stream); + return out; +} + +} // namespace + + +TORCH_LIBRARY(zth_w8a8, m) { + m.def( + "pack_weight(Tensor raw_weight, Tensor weight_scale) " + "-> (Tensor, Tensor)"); + m.def( + "gemm_out(Tensor x_q, Tensor packed_weight, Tensor x_scale, " + "Tensor packed_weight_scale, Tensor(a!) out, Tensor(b!) workspace) " + "-> Tensor(a!)"); +} + +TORCH_LIBRARY_IMPL(zth_w8a8, CUDA, m) { + m.impl("pack_weight", pack_weight_impl); + m.impl("gemm_out", gemm_out_impl); +} + +// Keep setup.py builds importable as well as torch.ops.load_library compatible. +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/csrc/w8a8_gemm_hip.hip b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/csrc/w8a8_gemm_hip.hip new file mode 100644 index 00000000..d04ccc91 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/csrc/w8a8_gemm_hip.hip @@ -0,0 +1,100 @@ +#include +#include + +#include + + +namespace { + +constexpr int kBlockThreads = 256; + +__global__ void w8a8_gemm_scalar( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t output_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t output_elements = static_cast(m) * n; + if (output_index >= output_elements) { + return; + } + + const int row = static_cast(output_index / n); + const int col = static_cast(output_index - static_cast(row) * n); + const int8_t* a_row = a + static_cast(row) * k; + + int32_t accumulator = 0; + for (int inner = 0; inner < k; ++inner) { + accumulator += static_cast(a_row[inner]) * + static_cast( + b[static_cast(inner) * n + col]); + } + const float scaled = static_cast(accumulator) * x_scale[row] * + weight_scale[col]; + out[output_index] = hip_bfloat16(scaled); +} + +} // namespace + + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + const int64_t output_elements = static_cast(m) * n; + const dim3 block(kBlockThreads); + const dim3 grid( + static_cast( + (output_elements + kBlockThreads - 1) / kBlockThreads)); + hipLaunchKernelGGL( + w8a8_gemm_scalar, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast(out), + m, + n, + k); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + hipMemcpyAsync( + packed_weight, + raw_weight, + static_cast(k) * n * sizeof(int8_t), + hipMemcpyDeviceToDevice, + stream); + hipMemcpyAsync( + packed_weight_scale, + weight_scale, + static_cast(n) * sizeof(float), + hipMemcpyDeviceToDevice, + stream); +} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/profile_pmc.sh b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/profile_pmc.sh new file mode 100644 index 00000000..7ac66a50 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/profile_pmc.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [ "$#" -ne 6 ]; then + echo "usage: profile_pmc.sh HARNESS SOURCE M N K OUTPUT_DIR" >&2 + exit 2 +fi + +harness=$1 +source_dir=$2 +m=$3 +n=$4 +k=$5 +output_dir=$6 + +run_profile() { + mode=$1 + name=$2 + csv_path="$output_dir/$name.csv" + mkdir -p "$output_dir/data-$name" + /opt/dtk/bin/hipprof \ + "$mode" \ + --pmc-type 3 \ + --flush-interval 1000 \ + --exit-cleanup \ + -o "$output_dir/$name" \ + -d "$output_dir/data-$name" \ + python3 "$harness" \ + --source "$source_dir" \ + --m "$m" \ + --n "$n" \ + --k "$k" \ + --warmups 0 \ + --samples 1 \ + --replays-per-sample 1 \ + --profile-only + + # pmc-read/pmc-write use replay and their launcher may return before the + # replay child flushes the CSV. Serialize the counter groups and make a + # missing result a hard profiling failure. + deadline=$((SECONDS + 300)) + while [ ! -s "$csv_path" ]; do + if [ "$SECONDS" -ge "$deadline" ]; then + echo "hipprof did not produce $csv_path within 300 seconds" >&2 + return 1 + fi + sleep 1 + done +} + +# General resources/instructions and memory traffic are separate hipprof +# counter groups. Keep all three raw CSV files so the control plane can apply +# DTK's documented request-size formulas without trusting Agent arithmetic. +run_profile --pmc pmc +run_profile --pmc-read pmc-read +run_profile --pmc-write pmc-write diff --git a/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/setup.py b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/setup.py new file mode 100644 index 00000000..e77a8e99 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/setup.py @@ -0,0 +1,32 @@ +from pathlib import Path + +from setuptools import setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + +ROOT = Path(__file__).resolve().parent +PREBUILT = sorted((ROOT / "prebuilt").glob("*.o")) +SOURCES = ["csrc/bindings.cpp"] +if PREBUILT: + SOURCES.append("csrc/w8a8_dispatch.cpp") +else: + SOURCES.append("csrc/w8a8_gemm_hip.hip") + + +setup( + name="metainfer_w8a8_backend", + ext_modules=[ + CUDAExtension( + name="metainfer_w8a8_backend", + sources=SOURCES, + extra_objects=[str(path) for path in PREBUILT], + extra_compile_args={ + "cxx": ["-O3"], + "nvcc": ["-O3", "--offload-arch=gfx928"], + }, + ) + ], + cmdclass={ + "build_ext": BuildExtension.with_options(no_python_abi_suffix=True) + }, +) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/w8a8_backend.py b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/w8a8_backend.py new file mode 100644 index 00000000..7038fd03 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/w8a8_backend.py @@ -0,0 +1,80 @@ +"""Trusted loader for the generated W8A8 HIP extension. + +The control plane owns this file. Optimization agents edit the HIP kernel, +not the extension-loading contract. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import torch +from torch.utils.cpp_extension import load + + +_SOURCE_DIR = Path(__file__).resolve().parent +_LOADED = False + + +def _compile_source_dir() -> Path: + configured = os.environ.get("METAINFER_W8A8_COMPILE_SOURCE_DIR") + return Path(configured).resolve() if configured else _SOURCE_DIR + + +def _extension_inputs() -> tuple[list[str], list[str]]: + """Select either an exploration HIP source or final prebuilt objects.""" + compile_source = _compile_source_dir() + prebuilt_dir = compile_source / "prebuilt" + prebuilt = sorted(prebuilt_dir.glob("*.o")) + dispatch = compile_source / "csrc" / "w8a8_dispatch.cpp" + if prebuilt: + if not dispatch.is_file(): + raise RuntimeError( + "prebuilt W8A8 objects exist without csrc/w8a8_dispatch.cpp" + ) + return ( + [ + str(compile_source / "csrc" / "bindings.cpp"), + str(dispatch), + ], + [str(path) for path in prebuilt], + ) + return ( + [ + str(compile_source / "csrc" / "bindings.cpp"), + str(compile_source / "csrc" / "w8a8_gemm_hip.hip"), + ], + [], + ) + + +def load_extension() -> None: + """Build and load the TORCH_LIBRARY extension once per process.""" + global _LOADED + if _LOADED: + return + sources, prebuilt_objects = _extension_inputs() + build_key = os.environ.get("METAINFER_W8A8_BUILD_KEY", "default") + safe_key = "".join( + char for char in build_key.lower() if char in "0123456789abcdef" + )[:24] or "default" + load( + name=f"metainfer_w8a8_backend_{safe_key}", + sources=sources, + extra_cflags=["-O3"], + extra_cuda_cflags=["-O3", "--offload-arch=gfx928"], + extra_ldflags=prebuilt_objects, + is_python_module=False, + with_cuda=True, + verbose=False, + ) + if not hasattr(torch.ops.zth_w8a8, "gemm_out"): + raise RuntimeError( + "W8A8 extension loaded without registering " + "torch.ops.zth_w8a8.gemm_out" + ) + _LOADED = True + + +__all__ = ["load_extension"] diff --git a/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/w8a8_graph.py b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/w8a8_graph.py new file mode 100644 index 00000000..a9d716d4 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_baseline/w8a8_graph.py @@ -0,0 +1,67 @@ +"""Python entry point for capturing the fixed W8A8 API in a CUDA/HIP Graph.""" + +from __future__ import annotations + +from typing import Any + +import torch + + +class W8A8GraphRunner: + """Replay a captured W8A8 call and return its caller-owned output.""" + + def __init__( + self, + graph: torch.cuda.CUDAGraph, + output: torch.Tensor, + stream: torch.cuda.Stream, + ) -> None: + self.graph = graph + self.output = output + self.stream = stream + + def replay(self) -> torch.Tensor: + self.graph.replay() + torch.cuda.current_stream().wait_stream(self.stream) + return self.output + + +def capture_w8a8_graph( + api: Any, + a: torch.Tensor, + packed_weight: torch.Tensor, + a_scale: torch.Tensor, + packed_weight_scale: torch.Tensor, + out: torch.Tensor, + workspace: torch.Tensor, +) -> W8A8GraphRunner: + """Warm up, capture and return a Python-callable W8A8 Graph runner.""" + + def invoke() -> None: + returned = api.w8a8_gemm_out( + a, + packed_weight, + a_scale, + packed_weight_scale, + out, + workspace, + ) + if returned.data_ptr() != out.data_ptr(): + raise RuntimeError( + "w8a8_gemm_out must return the caller-provided out tensor" + ) + + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + invoke() + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + invoke() + runner = W8A8GraphRunner(graph, out, capture_stream) + runner.replay() + torch.cuda.current_stream().wait_stream(capture_stream) + torch.cuda.synchronize() + return runner diff --git a/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_bench.py b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_bench.py new file mode 100644 index 00000000..51a3b879 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_bench.py @@ -0,0 +1,523 @@ +#!/usr/bin/env python3 +"""Trusted correctness and performance harness for the gfx928 W8A8 adapter.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import math +import os +import statistics +import sys +from pathlib import Path + +try: + import torch +except ModuleNotFoundError: # Allow CPU-only CI to import pure helpers. + torch = None # type: ignore[assignment] + + +def load_module(module_path: Path, name: str): + spec = importlib.util.spec_from_file_location( + name, module_path + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot import {module_path}") + module = importlib.util.module_from_spec(spec) + sys.path.insert(0, str(module_path.parent)) + spec.loader.exec_module(module) + return module + + +def percentile(values: list[float], fraction: float) -> float: + ordered = sorted(values) + return ordered[max(0, math.ceil(fraction * len(ordered)) - 1)] + + +def validate_profile_protocol( + profile_only: bool, + warmups: int, + samples: int, + replays_per_sample: int, +) -> None: + """Keep one unambiguous operator replay after the PMC marker.""" + if profile_only and (warmups, samples, replays_per_sample) != (0, 1, 1): + raise ValueError( + "--profile-only requires --warmups 0 --samples 1 " + "--replays-per-sample 1" + ) + + +class CapturedGraphRunner: + """Small internal runner used by the trusted benchmark.""" + + def __init__( + self, + graph: torch.cuda.CUDAGraph, + output: torch.Tensor, + stream: torch.cuda.Stream, + ) -> None: + self.graph = graph + self.output = output + self.stream = stream + + def replay(self) -> torch.Tensor: + self.graph.replay() + return self.output + + +def capture_candidate_graph( + candidate, output: torch.Tensor +) -> CapturedGraphRunner: + """Capture a zero-argument candidate on a non-default HIP stream.""" + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + candidate() + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + candidate() + runner = CapturedGraphRunner(graph, output, capture_stream) + runner.replay() + torch.cuda.current_stream().wait_stream(capture_stream) + torch.cuda.synchronize() + return runner + + +def exact_w8a8_reference( + a: torch.Tensor, + b: torch.Tensor, + a_scale: torch.Tensor, + b_scale: torch.Tensor, +) -> torch.Tensor: + """Compute the contract's integer dot exactly before float scaling. + + CPU int64 avoids treating a float32 GEMM's accumulation order as the + W8A8 contract. Large-prefill callers should run this once per candidate, + then use ``--skip-correctness`` only for repeated timing/profiling of the + exact same source and deterministic inputs. + """ + device = a.device + dot = torch.mm( + a.to(device="cpu", dtype=torch.int64), + b.to(device="cpu", dtype=torch.int64), + ) + scaled = ( + dot.to(torch.float32) + * a_scale.to(device="cpu", dtype=torch.float32) + * b_scale.to(device="cpu", dtype=torch.float32).T + ) + return scaled.to(torch.bfloat16).to(device) + + +def w8a8_seed(m: int, n: int, k: int) -> int: + """Deterministic seed shared by the timed run and reference preparation.""" + return 20260724 + m + n + k + + +def generate_w8a8_inputs(m: int, n: int, k: int): + """Generate the exact deterministic inputs a timed run would use. + + Inputs are produced on CUDA/HIP so the RNG stream matches the timed + benchmark; a reference pre-seeded from these inputs is therefore + bit-identical to one computed inside the benchmark itself. + """ + torch.manual_seed(w8a8_seed(m, n, k)) + a = torch.randint( + -127, 128, (m, k), dtype=torch.int8, device="cuda" + ) + b = torch.randint( + -127, 128, (k, n), dtype=torch.int8, device="cuda" + ) + a_scale = torch.rand( + (m, 1), dtype=torch.float32, device="cuda" + ) * 0.01 + b_scale = torch.rand( + (n, 1), dtype=torch.float32, device="cuda" + ) * 0.01 + return a, b, a_scale, b_scale + + +def reference_cache_path( + m: int, n: int, k: int, cache_dir: Path +) -> Path: + """Path of the exact int64 reference cache for one (M, N, K) shape.""" + return cache_dir / f"exact-int64-v1-m{m}-n{n}-k{k}.pt" + + +def save_reference_cache( + reference: torch.Tensor, reference_path: Path +) -> None: + """Persist a computed reference atomically (tmp + replace). + + The parent directory is created on demand: the timed benchmark can reach + this path with a fresh ``--reference-cache-dir`` (e.g. serial validation + uses ``final/cache/references``) whose parent has never been created. + """ + reference_path.parent.mkdir(parents=True, exist_ok=True) + temporary = reference_path.with_name( + f"{reference_path.name}.tmp-{os.getpid()}" + ) + torch.save(reference.to("cpu"), temporary) + temporary.replace(reference_path) + + +def prepare_reference(m: int, n: int, k: int, cache_dir: Path) -> Path: + """Compute and cache the exact reference for one shape. + + This is the slow part (CPU int64 GEMM) of a correctness-checked run for + M>=3072; call it once per shape outside the timed benchmark budget. + """ + cache_dir.mkdir(parents=True, exist_ok=True) + reference_path = reference_cache_path(m, n, k, cache_dir) + if reference_path.is_file(): + return reference_path + a, b, a_scale, b_scale = generate_w8a8_inputs(m, n, k) + reference = exact_w8a8_reference(a, b, a_scale, b_scale) + save_reference_cache(reference, reference_path) + return reference_path + + +def main() -> int: + if torch is None: + raise RuntimeError( + "w8a8_bench.py requires PyTorch in the DCU benchmark environment" + ) + parser = argparse.ArgumentParser() + parser.add_argument("--source", type=Path) + parser.add_argument("--m", type=int) + parser.add_argument("--n", type=int) + parser.add_argument("--k", type=int) + parser.add_argument("--warmups", type=int, default=100) + parser.add_argument("--samples", type=int, default=30) + parser.add_argument("--replays-per-sample", type=int, default=100) + parser.add_argument("--reference-cache-dir", type=Path) + parser.add_argument("--probe", action="store_true") + parser.add_argument("--self-test", action="store_true") + parser.add_argument( + "--skip-correctness", + action="store_true", + help=( + "Skip the CPU int64 reference when the exact same source and " + "deterministic inputs already passed trusted correctness." + ), + ) + parser.add_argument( + "--profile-only", + action="store_true", + help="Alias for --skip-correctness used by trusted PMC profiling.", + ) + parser.add_argument( + "--prepare-reference", + action="store_true", + help=( + "Compute and cache the exact CPU int64 reference for (m, n, k) " + "using the same deterministic inputs as the timed run, then exit " + "without building or running any backend. For M>=3072 this is " + "the slow part of a correctness-checked benchmark, so the " + "control plane runs it once per shape outside the benchmark " + "subprocess timeout." + ), + ) + args = parser.parse_args() + + validate_profile_protocol( + args.profile_only, + args.warmups, + args.samples, + args.replays_per_sample, + ) + + if args.self_test: + a = torch.tensor( + [[1, -2, 3], [-4, 5, -6]], dtype=torch.int8 + ) + b = torch.tensor( + [[7, -8], [9, 10], [-11, 12]], dtype=torch.int8 + ) + a_scale = torch.tensor( + [[0.5], [0.25]], dtype=torch.float32 + ) + b_scale = torch.tensor( + [[2.0], [4.0]], dtype=torch.float32 + ) + actual = exact_w8a8_reference(a, b, a_scale, b_scale) + expected = torch.tensor( + [[-44.0, 16.0], [41.5, 10.0]], dtype=torch.bfloat16 + ) + passed = bool(torch.equal(actual, expected)) + print(json.dumps({ + "self_test": "exact_w8a8_reference", + "passed": passed, + "device": "cpu", + "torch_version": torch.__version__, + "actual": actual.float().tolist(), + "expected": expected.float().tolist(), + }, sort_keys=True)) + return 0 if passed else 4 + + missing = [ + name for name in ("source", "m", "n", "k") + if getattr(args, name) is None + ] + if missing: + parser.error( + "the following arguments are required unless --self-test is " + f"used: {', '.join('--' + name for name in missing)}" + ) + + count = torch.cuda.device_count() + props = torch.cuda.get_device_properties(0) if count else None + if args.probe: + print(json.dumps({ + "visible_devices": count, + "logical_device": 0, + "device_name": props.name if props else "", + "multi_processor_count": ( + props.multi_processor_count if props else 0 + ), + "cudagraph_available": hasattr(torch.cuda, "CUDAGraph"), + "python_graph_api": "torch.cuda.CUDAGraph", + })) + return ( + 0 + if count == 1 and hasattr(torch.cuda, "CUDAGraph") + else 3 + ) + + if min(args.m, args.n, args.k) <= 0: + raise ValueError("M, N and K must be positive") + + if args.prepare_reference: + if args.reference_cache_dir is None: + parser.error( + "--prepare-reference requires --reference-cache-dir" + ) + reference_path = reference_cache_path( + args.m, args.n, args.k, args.reference_cache_dir + ) + cache_hit = reference_path.is_file() + if not cache_hit: + prepare_reference( + args.m, args.n, args.k, args.reference_cache_dir + ) + print(json.dumps({ + "reference_prepared": True, + "reference_cache_hit": cache_hit, + "shape": {"M": args.m, "N": args.n, "K": args.k}, + "reference_cache_path": str(reference_path), + }, sort_keys=True)) + return 0 + + source = args.source.resolve() + fixed_contract = source / "int8_w8a8_gemm_api.py" + if fixed_contract.is_file(): + backend = load_module( + source / "w8a8_backend.py", "metainfer_w8a8_backend" + ) + backend.load_extension() + api = load_module(fixed_contract, "metainfer_w8a8_contract") + fixed_api = True + else: + # Backward compatibility for pre-contract extracted repositories. + api = load_module( + source / "w8a8_gemm.py", "metainfer_w8a8_candidate" + ) + api.load_extension() + fixed_api = False + + a, b, a_scale, b_scale = generate_w8a8_inputs( + args.m, args.n, args.k + ) + out = torch.empty( + (args.m, args.n), dtype=torch.bfloat16, device="cuda" + ) + + if fixed_api: + packed_weight, packed_weight_scale = api.prepare_weight(b, b_scale) + workspace = api.allocate_workspace( + args.m, args.n, args.k, a.device + ) + + def candidate() -> None: + api.w8a8_gemm_out( + a, + packed_weight, + a_scale, + packed_weight_scale, + out, + workspace, + ) + path = "w8a8_gemm_out" + elif args.m <= 16: + workspace = api.empty_optimized_workspace(a, b) + + def candidate() -> None: + api.gemm_out_optimized(a, b, a_scale, b_scale, out, workspace) + path = "gemm_out_optimized" + else: + def candidate() -> None: + api.gemm_out_prefill(a, b, a_scale, b_scale, out) + path = "gemm_out_prefill" + + try: + graph_runner = capture_candidate_graph(candidate, out) + except Exception as exc: + print(json.dumps({ + "passed": False, + "operator": "int8_w8a8_gemm", + "path": path, + "shape": {"M": args.m, "N": args.n, "K": args.k}, + "visible_devices": count, + "device_name": props.name if props else "", + "graph_capture_passed": False, + "timing_mode": "cuda_graph_replay", + "python_callable": True, + "graph_error": f"{type(exc).__name__}: {exc}", + "mismatch_count": None, + "first_mismatch": None, + }, sort_keys=True)) + return 0 + + correctness_checked = not ( + args.skip_correctness or args.profile_only + ) + reference_cache_hit = False + mismatch_count = None + passed = True + max_abs_error = None + first_mismatch = None + if correctness_checked: + reference_path = None + if args.reference_cache_dir is not None: + reference_path = reference_cache_path( + args.m, args.n, args.k, args.reference_cache_dir + ) + if reference_path is not None and reference_path.is_file(): + cached_reference = torch.load( + reference_path, map_location="cpu", weights_only=True + ) + if ( + cached_reference.shape != out.shape + or cached_reference.dtype != torch.bfloat16 + ): + raise RuntimeError( + f"invalid cached W8A8 reference: {reference_path}" + ) + reference = cached_reference.to(a.device) + reference_cache_hit = True + else: + reference = exact_w8a8_reference(a, b, a_scale, b_scale) + if reference_path is not None: + save_reference_cache(reference, reference_path) + mismatch_mask = out != reference + mismatch_count = int(mismatch_mask.sum().item()) + passed = mismatch_count == 0 + absolute_error = (out.float() - reference.float()).abs() + max_abs_error = float(absolute_error.max().item()) + if mismatch_count: + first_flat = int( + mismatch_mask.reshape(-1).nonzero()[0].item() + ) + first_m = first_flat // args.n + first_n = first_flat % args.n + first_mismatch = { + "flat_index": first_flat, + "m": first_m, + "n": first_n, + "actual": float(out[first_m, first_n].float().item()), + "expected": float(reference[first_m, first_n].float().item()), + "abs_error": float(absolute_error[first_m, first_n].item()), + } + + if args.replays_per_sample <= 0: + raise ValueError("replays-per-sample must be positive") + for _ in range(args.warmups): + graph_runner.replay() + graph_runner.stream.synchronize() + profile_marker_emitted = False + if args.profile_only: + # Graph creation performs an eager warmup and a validation replay. + # Emit a non-W8A8 dispatch after both so the PMC parser can identify + # the single timed operator replay that follows. The marker is outside + # the timing events and does not change candidate inputs or workspace. + with torch.cuda.stream(graph_runner.stream): + out.zero_() + graph_runner.stream.synchronize() + profile_marker_emitted = True + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + samples_us: list[float] = [] + for _ in range(args.samples): + with torch.cuda.stream(graph_runner.stream): + begin.record() + for _ in range(args.replays_per_sample): + graph_runner.replay() + end.record() + end.synchronize() + samples_us.append( + float(begin.elapsed_time(end)) + * 1000.0 + / args.replays_per_sample + ) + + median_us = statistics.median(samples_us) + seconds = median_us * 1.0e-6 + logical_ops = 2.0 * args.m * args.n * args.k + algorithmic_bytes = ( + args.m * args.k + + args.k * args.n + + 4 * (args.m + args.n) + + 2 * args.m * args.n + ) + print(json.dumps({ + "passed": passed, + "operator": "int8_w8a8_gemm", + "path": path, + "shape": {"M": args.m, "N": args.n, "K": args.k}, + "visible_devices": count, + "device_name": props.name if props else "", + "graph_capture_passed": True, + "timing_mode": "cuda_graph_replay", + "python_callable": True, + "python_graph_api": "torch.cuda.CUDAGraph", + "median_us": median_us, + "p90_us": percentile(samples_us, 0.9), + "min_us": min(samples_us), + "max_us": max(samples_us), + "latency_samples_us": samples_us, + "logical_ops": logical_ops, + "logical_tops": logical_ops / seconds / 1.0e12, + "algorithmic_bytes": algorithmic_bytes, + "algorithmic_bandwidth_gb_s": ( + algorithmic_bytes / seconds / 1.0e9 + ), + "metric_semantics": { + "logical_tops": ( + "INT8 GEMM logical operation rate; one multiply and one add " + "count as two operations." + ), + "algorithmic_bandwidth_gb_s": ( + "Algorithmic minimum bytes divided by unprofiled median " + "latency; this is not measured HBM traffic." + ), + }, + "max_abs_error": max_abs_error, + "mismatch_count": mismatch_count, + "first_mismatch": first_mismatch, + "correctness_checked": correctness_checked, + "profile_only": args.profile_only, + "profile_replay_marker_emitted": profile_marker_emitted, + "reference_cache_hit": reference_cache_hit, + "warmup": args.warmups, + "samples": args.samples, + "replays_per_sample": args.replays_per_sample, + }, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/agent_bridge_client.py b/metainfer/tasks/dcu_kernel_auto_opt/bridge/agent_bridge_client.py new file mode 100644 index 00000000..aa5af565 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/agent_bridge_client.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Container-side executable compatible with SubAgentManager's Claude CLI.""" + +from __future__ import annotations + +import json +import os +import socket +import struct +import sys + + +SOCKET_PATH = os.environ.get( + "METAINFER_AGENT_BRIDGE_SOCKET", + "/workspace/MetaInfer/.metainfer-agent-bridge.sock", +) +FORWARDED_ENV = ( + "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "TORCH_EXTENSIONS_DIR", + "TRITON_CACHE_DIR", "XDG_CACHE_HOME", "TMPDIR", + "DISABLE_INTERACTIVITY", +) + + +def _recv_exact(conn: socket.socket, size: int) -> bytes: + chunks = [] + while size: + chunk = conn.recv(size) + if not chunk: + raise EOFError("agent bridge closed unexpectedly") + chunks.append(chunk) + size -= len(chunk) + return b"".join(chunks) + + +def main() -> int: + prompt = sys.stdin.buffer.read() + request = json.dumps({ + "args": sys.argv[1:], + "cwd": os.getcwd(), + "task_id": os.environ.get("METAINFER_TASK_ID"), + "env": { + key: os.environ[key] for key in FORWARDED_ENV + if key in os.environ + }, + }).encode() + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as conn: + conn.connect(SOCKET_PATH) + conn.sendall(struct.pack("!I", len(request)) + request) + conn.sendall(struct.pack("!I", len(prompt)) + prompt) + while True: + kind = conn.recv(1) + if not kind: + return 1 + size = struct.unpack("!I", _recv_exact(conn, 4))[0] + payload = _recv_exact(conn, size) + if kind == b"O": + sys.stdout.buffer.write(payload) + sys.stdout.buffer.flush() + elif kind == b"E": + return struct.unpack("!i", payload)[0] + else: + sys.stderr.buffer.write(payload + b"\n") + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/agent_bridge_server.py b/metainfer/tasks/dcu_kernel_auto_opt/bridge/agent_bridge_server.py new file mode 100644 index 00000000..ca113423 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/agent_bridge_server.py @@ -0,0 +1,404 @@ +#!/usr/bin/env python3 +"""Host-side, credential-preserving Claude CLI bridge over a Unix socket.""" + +from __future__ import annotations + +import json +import os +import re +import signal +import socket +import struct +import subprocess +import threading +from pathlib import Path + + +_DEFAULT_META_ROOT = Path(__file__).resolve().parents[4] +HOST_ROOT = Path(os.environ.get( + "METAINFER_AGENT_BRIDGE_ROOT", + str(_DEFAULT_META_ROOT), +)).resolve() +HOST_WORKSPACE_ROOT = Path(os.environ.get( + "METAINFER_AGENT_BRIDGE_WORKSPACE_ROOT", + str(HOST_ROOT.parent), +)).resolve() +SOCKET_PATH = Path(os.environ.get( + "METAINFER_AGENT_BRIDGE_SOCKET", + str(HOST_ROOT / ".metainfer-agent-bridge.sock"), +)) +ALLOWED_ROOTS = ( + HOST_ROOT, + (HOST_WORKSPACE_ROOT / "kernel-repos").resolve(), + (HOST_WORKSPACE_ROOT / "API").resolve(), +) +CLAUDE_BIN = os.environ.get( + "METAINFER_HOST_CLAUDE_BIN", + "/usr/lib/node_modules/@anthropic-ai/claude-code/bin/claude.exe", +) +ALLOWED_ENV = { + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "TORCH_EXTENSIONS_DIR", + "TRITON_CACHE_DIR", + "XDG_CACHE_HOME", + "TMPDIR", + "DISABLE_INTERACTIVITY", +} +PATH_ENV = { + "TORCH_EXTENSIONS_DIR", "TRITON_CACHE_DIR", "XDG_CACHE_HOME", "TMPDIR", +} +VALUE_FLAGS = { + "--output-format", "--input-format", "--permission-mode", "--add-dir", + "--model", "--effort", "--resume", "--session-id", "--setting-sources", + "--tools", "--disallowedTools", "--disallowed-tools", +} +BARE_FLAGS = {"-p", "--verbose"} +SOURCE_ONLY_TOOLS = "Read,Glob,Grep,Write,Edit" +CLI_SLOT = threading.Semaphore(int(os.environ.get( + # worker29 has four physical GPUs and the control plane assigns at most + # one child per GPU. A single slot serializes independent workers and + # makes their absolute bootstrap timeout include queueing behind peers. + "METAINFER_AGENT_BRIDGE_CONCURRENCY", "4" +))) +TASK_LOCK = threading.Lock() +TASK_CONNECTIONS: dict[str, set[socket.socket]] = {} +TASK_PROCESSES: dict[str, set[subprocess.Popen[bytes]]] = {} +TASK_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") + + +def _recv_exact(conn: socket.socket, size: int) -> bytes: + chunks = [] + remaining = size + while remaining: + chunk = conn.recv(remaining) + if not chunk: + raise EOFError("bridge request ended early") + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) + + +def _frame(conn: socket.socket, kind: bytes, payload: bytes) -> None: + conn.sendall(kind + struct.pack("!I", len(payload)) + payload) + + +def _validated_task_id(value: object) -> str | None: + if value is None: + return None + task_id = str(value) + if not TASK_ID_RE.fullmatch(task_id): + raise ValueError(f"invalid task id: {task_id!r}") + return task_id + + +def _track_connection(task_id: str, conn: socket.socket) -> None: + with TASK_LOCK: + TASK_CONNECTIONS.setdefault(task_id, set()).add(conn) + + +def _track_process( + task_id: str, process: subprocess.Popen[bytes], +) -> None: + with TASK_LOCK: + TASK_PROCESSES.setdefault(task_id, set()).add(process) + + +def _untrack_task_resources( + task_id: str | None, + conn: socket.socket, + process: subprocess.Popen[bytes] | None, +) -> None: + if task_id is None: + return + with TASK_LOCK: + connections = TASK_CONNECTIONS.get(task_id) + if connections is not None: + connections.discard(conn) + if not connections: + TASK_CONNECTIONS.pop(task_id, None) + if process is not None: + processes = TASK_PROCESSES.get(task_id) + if processes is not None: + processes.discard(process) + if not processes: + TASK_PROCESSES.pop(task_id, None) + + +def _terminate_process( + process: subprocess.Popen[bytes], *, force: bool = False, +) -> bool: + if process.poll() is not None: + return False + try: + os.killpg( + os.getpgid(process.pid), + signal.SIGKILL if force else signal.SIGTERM, + ) + process.wait(timeout=1 if force else 5) + except ProcessLookupError: + return False + except subprocess.TimeoutExpired: + try: + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=5) + return True + + +def _task_process_groups( + task_id: str, + tracked_pids: set[int], +) -> set[int]: + """Resolve tracked descendants plus exact task-marker process groups.""" + rows: dict[int, tuple[int, int]] = {} + marker = task_id.encode() + marker_pids: set[int] = set() + for entry in os.scandir("/proc"): + if not entry.name.isdigit(): + continue + pid = int(entry.name) + try: + raw_stat = Path( + f"/proc/{pid}/stat" + ).read_text(encoding="utf-8", errors="replace") + rparen = raw_stat.rfind(")") + fields = raw_stat[rparen + 2:].split() + if rparen == -1 or len(fields) < 3: + continue + ppid = int(fields[1]) + pgid = os.getpgid(pid) + rows[pid] = (ppid, pgid) + cmdline = Path(f"/proc/{pid}/cmdline").read_bytes() + if marker in cmdline: + marker_pids.add(pid) + except (OSError, ValueError, ProcessLookupError): + continue + + owned = set(tracked_pids) | marker_pids + changed = True + while changed: + changed = False + for pid, (ppid, _) in rows.items(): + if ppid in owned and pid not in owned: + owned.add(pid) + changed = True + server_pgid = os.getpgrp() + return { + rows[pid][1] + for pid in owned + if pid in rows and rows[pid][1] != server_pgid + } + + +def _cancel_task(task_id: str, *, force: bool = False) -> dict[str, int]: + """Stop host agents and disconnect queued clients owned by one task.""" + with TASK_LOCK: + processes = list(TASK_PROCESSES.get(task_id, ())) + connections = list(TASK_CONNECTIONS.get(task_id, ())) + process_groups = _task_process_groups( + task_id, {process.pid for process in processes} + ) + group_signal = signal.SIGKILL if force else signal.SIGTERM + killed_groups = 0 + for pgid in process_groups: + try: + os.killpg(pgid, group_signal) + killed_groups += 1 + except ProcessLookupError: + pass + killed = sum( + _terminate_process(process, force=force) + for process in processes + ) + disconnected = 0 + for client in connections: + try: + client.shutdown(socket.SHUT_RDWR) + disconnected += 1 + except OSError: + pass + return { + "killed_processes": killed, + "killed_process_groups": killed_groups, + "disconnected_clients": disconnected, + } + + +def _translate_prompt_paths(prompt: bytes) -> bytes: + """Translate only the /workspace path segment, not /workspaces names.""" + return re.sub( + rb"/workspace(?=/|$)", + lambda _: str(HOST_WORKSPACE_ROOT).encode(), + prompt, + ) + + +def _host_path(value: str) -> Path: + if value == "/workspace" or value.startswith("/workspace/"): + value = str(HOST_WORKSPACE_ROOT) + value[len("/workspace"):] + path = Path(value).resolve() + if not any(path == root or root in path.parents for root in ALLOWED_ROOTS): + raise ValueError( + f"path outside allowed MetaInfer workspace roots: {path}" + ) + return path + + +def _host_agent_cwd(path: Path) -> Path: + """Use a neutral non-Git cwd for a container-created worktree.""" + git_marker = path / ".git" + if not git_marker.is_file(): + return path + # The common repo contains worktree metadata with container-only + # /workspace paths, while the worker path sits below the large MetaInfer + # checkout. Both make Claude Code repository discovery stall. The + # kernel-repos parent is allowed, writable, and is not itself a Git repo. + return (HOST_WORKSPACE_ROOT / "kernel-repos").resolve() + + +def _validated_args(args: list[str]) -> list[str]: + out = [] + index = 0 + while index < len(args): + flag = args[index] + if flag in BARE_FLAGS: + out.append(flag) + index += 1 + continue + if flag not in VALUE_FLAGS or index + 1 >= len(args): + raise ValueError(f"unsupported Claude argument: {flag}") + value = args[index + 1] + # Claude Code 2.1.161 on worker29 fails to consume stdin when its + # default text input mode is repeated explicitly. Omit only this + # redundant pair; stream-json input remains forwarded unchanged. + if flag == "--input-format" and value == "text": + index += 2 + continue + if flag == "--add-dir": + add_dir = _host_path(value) + # Container-created Git worktrees have a .git indirection file + # containing a /workspace/... gitdir that is invalid on the host. + # Grant the worker root instead; the source directory remains + # inside that exact allowed tree and is named in the prompt. + if (add_dir / ".git").is_file(): + add_dir = add_dir.parent + value = str(add_dir) + out.extend([flag, value]) + index += 2 + if "-p" not in out: + raise ValueError("print mode is required") + if not any( + flag in out + for flag in ("--tools", "--disallowedTools", "--disallowed-tools") + ): + out.extend(["--tools", SOURCE_ONLY_TOOLS]) + return out + + +def _handle(conn: socket.socket) -> None: + slot_acquired = False + process: subprocess.Popen[bytes] | None = None + task_id: str | None = None + try: + header_size = struct.unpack("!I", _recv_exact(conn, 4))[0] + if header_size > 128 * 1024: + raise ValueError("bridge header too large") + request = json.loads(_recv_exact(conn, header_size)) + prompt_size = struct.unpack("!I", _recv_exact(conn, 4))[0] + if prompt_size > 128 * 1024: + raise ValueError("bridge prompt too large") + prompt = _translate_prompt_paths(_recv_exact(conn, prompt_size)) + if request.get("action") == "cancel_task": + cancel_id = _validated_task_id(request.get("task_id")) + if cancel_id is None: + raise ValueError("cancel_task requires task_id") + result = json.dumps(_cancel_task( + cancel_id, + force=bool(request.get("force", False)), + )).encode() + b"\n" + _frame(conn, b"O", result) + _frame(conn, b"E", struct.pack("!i", 0)) + return + task_id = _validated_task_id(request.get("task_id")) + if task_id is not None: + _track_connection(task_id, conn) + cwd = _host_path(str(request["cwd"])) + # Container-created git worktrees store /workspace/... in their + # .git indirection file. That path does not exist on the host where + # Claude runs, and Claude Code stalls during repository discovery. + # Use the worker root as cwd while retaining source via --add-dir. + cwd = _host_agent_cwd(cwd) + args = _validated_args([str(v) for v in request.get("args", [])]) + prompt_text = prompt.decode("utf-8", errors="strict") + print_index = args.index("-p") + args.insert(print_index + 1, prompt_text) + overrides = { + str(key): str(value) + for key, value in (request.get("env") or {}).items() + if key in ALLOWED_ENV + } + for key in PATH_ENV & overrides.keys(): + overrides[key] = str(_host_path(overrides[key])) + env = dict(os.environ) + env.update(overrides) + env["PWD"] = str(cwd) + while not CLI_SLOT.acquire(timeout=5): + heartbeat = json.dumps({ + "type": "system", + "subtype": "bridge_queued", + "message": "Waiting for the host Claude CLI slot", + }).encode() + b"\n" + _frame(conn, b"O", heartbeat) + slot_acquired = True + process = subprocess.Popen( + [CLAUDE_BIN, *args], + cwd=cwd, + env=env, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + if task_id is not None: + _track_process(task_id, process) + assert process.stdout is not None + for chunk in iter(process.stdout.readline, b""): + _frame(conn, b"O", chunk) + returncode = process.wait() + _frame(conn, b"E", struct.pack("!i", returncode)) + except Exception as exc: + if process is not None and process.poll() is None: + try: + _terminate_process(process) + except (ProcessLookupError, subprocess.TimeoutExpired): + pass + try: + _frame(conn, b"X", str(exc).encode("utf-8", errors="replace")) + except OSError: + pass + finally: + _untrack_task_resources(task_id, conn, process) + if slot_acquired: + CLI_SLOT.release() + conn.close() + + +def main() -> int: + SOCKET_PATH.parent.mkdir(parents=True, exist_ok=True) + try: + SOCKET_PATH.unlink() + except FileNotFoundError: + pass + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as server: + server.bind(str(SOCKET_PATH)) + os.chmod(SOCKET_PATH, 0o660) + server.listen(16) + while True: + conn, _ = server.accept() + threading.Thread(target=_handle, args=(conn,), daemon=True).start() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/control.py b/metainfer/tasks/dcu_kernel_auto_opt/bridge/control.py new file mode 100644 index 00000000..4823e083 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/control.py @@ -0,0 +1,59 @@ +"""Task-scoped control client for the host Claude bridge.""" + +from __future__ import annotations + +import json +import os +import socket +import struct +from typing import Any, Dict + + +SOCKET_PATH = os.environ.get( + "METAINFER_AGENT_BRIDGE_SOCKET", + "/workspace/MetaInfer/.metainfer-agent-bridge.sock", +) + + +def _recv_exact(conn: socket.socket, size: int) -> bytes: + chunks = [] + while size: + chunk = conn.recv(size) + if not chunk: + raise EOFError("agent bridge closed unexpectedly") + chunks.append(chunk) + size -= len(chunk) + return b"".join(chunks) + + +def cancel_task( + task_id: str, + timeout_s: float = 10.0, + *, + force: bool = False, +) -> Dict[str, Any]: + """Terminate host agents and queued bridge clients for ``task_id``.""" + request = json.dumps({ + "action": "cancel_task", + "task_id": task_id, + "force": force, + }).encode() + payload: Dict[str, Any] = {} + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as conn: + conn.settimeout(timeout_s) + conn.connect(SOCKET_PATH) + conn.sendall(struct.pack("!I", len(request)) + request) + conn.sendall(struct.pack("!I", 0)) + while True: + kind = _recv_exact(conn, 1) + size = struct.unpack("!I", _recv_exact(conn, 4))[0] + body = _recv_exact(conn, size) + if kind == b"O": + payload = json.loads(body) + elif kind == b"E": + return { + "ok": struct.unpack("!i", body)[0] == 0, + **payload, + } + else: + raise RuntimeError(body.decode("utf-8", errors="replace")) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/README.md b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/README.md new file mode 100644 index 00000000..22aace2d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/README.md @@ -0,0 +1,102 @@ +# MetaInfer → DSH agent（dcu-kernel-auto-opt） + +把 dcu-kernel-auto-opt 任务的 agent 执行链路从 Claude Code 换成 **DeepSeek +Harness（DSH）**：主 agent（kernel coordinator）与子 agent（worker / repair / +synthesis）全部由 DSH Python SDK 驱动，MetaInfer 的编排与验证 harness 语义 +不变。 + +## 原理 + +MetaInfer 的 `SubAgentManager` 把每个 agent 当作一个 CLI 进程调用： + +``` + -p --output-format stream-json --input-format text --verbose \ + --permission-mode bypassPermissions --add-dir [--add-dir ...] \ + [--model ] [--effort ] [--resume | --session-id ] [extra...] +``` + +prompt 从 stdin 传入，stdout 输出逐行 stream-json 事件 +（`system` / `assistant` / `result`），退出码 0 + `result` 事件 = 成功。 + +本目录提供： + +| 文件 | 作用 | +|---|---| +| `dsh_agent.py` | ccb 兼容的 CLI 包装器：解析上述参数、stdin 读 prompt、经 DSH Python SDK 跑一个 DSH agent、把结果重新输出为 stream-json 事件流 | +| `cordis.yml` | DSH runtime 的自定义组合：agent spine + tool-fs + tool-bash + tool-subagent + tool-todo + skills + 会话持久化 | +| `run_dsh_task.sh` | 便捷启动脚本：把 `METAINFER_CLAUDE_BIN` 指向本包装器后直接跑 orchestrator CLI | +| `tests/smoke_sdk.py` | SDK 冒烟测试（runtime 启动 / 模型调用 / 会话 resume / 文件工具） | + +## 本机接入方式（worker29） + +WebUI new-task 表单中 dcu-kernel-auto-opt 增加 **Agent framework** 选择： + +- `ccb` → Claude Code（模型 `Sonnet` / `Opus`） +- `dsh` → DeepSeek Harness（模型 `deepseek-v4-flash`，即 + `deepseek/deepseek-v4-flash-0731`） + +orchestrator CLI 在 `agent_framework=dsh` 时自动把 `claude_bin` 指向本目录的 +`dsh_agent.py`（`--claude-bin` 显式传入时优先），无需改动 MetaInfer 共享代码。 + +命令行直接运行（headless）： + +```bash +python3 -m metainfer.tasks.dcu_kernel_auto_opt.orchestrator.cli run \ + --state-dir ... --workspace-dir ... +``` + +requirements.json 中 `answers.agent_framework = "dsh"` 即可。 + +## 模型端点（本机） + +本机 DSH profile（`~/.dsh/settings.yaml`）使用 TokenHub 网关，dev-checkout +runtime 只内置 `deepseek-official` adapter（`llm-deepseek`），它通过 +`DEEPSEEK_BASE_URL` 环境变量覆盖端点——所以用官方 adapter 指向同一网关： + +- provider：`deepseek-official`（runtime adapter；TokenHub 经 base_url 覆盖） +- baseURL：`https://tokenhub.tencentmaas.com/plan/v3` +- 模型：`deepseek/deepseek-v4-flash-0731`（默认） +- API key：`~/.dsh/.credentials.yaml` 的 `TENCENT_API_KEY` + +`dsh_agent.py` 的环境变量： + +| 变量 | 默认 | 说明 | +|---|---|---| +| `DSH_AGENT_PROVIDER` | `deepseek-official` | DSH runtime provider（dev runtime 仅内置此 adapter） | +| `DSH_AGENT_MODEL` | `deepseek/deepseek-v4-flash-0731` | 模型覆盖 | +| `DSH_AGENT_BASE_URL` | `https://tokenhub.tencentmaas.com/plan/v3` | 模型端点（`DEEPSEEK_BASE_URL` 优先） | +| `TENCENT_API_KEY` / `DEEPSEEK_API_KEY` | 凭据文件兜底 | API key | +| `DSH_AGENT_CORDIS` | 本目录 `cordis.yml` | 自定义组合路径 | +| `DSH_AGENT_SESSION_ROOT` | `{最后一个 --add-dir}/.dsh-sessions` | 会话 JSONL 持久化根（跨迭代 resume 需要稳定路径） | +| `DSH_AGENT_MAX_TOKENS` | 65536 | 每次请求输出上限 | +| `DSH_AGENT_DEBUG` | — | 保留 runtime 日志 | + +## 会话连续性 + +`SubAgentManager` 用 `--session-id`（首轮）与 `--resume `（后续轮）延续 +agent 会话；DSH runtime 按 session id 把会话 JSONL 持久化在 +`DSH_AGENT_SESSION_ROOT` 下，同一个 id 再次 `run()` 即恢复上下文。默认 +session root 取 orchestrator 传入的 workspace_dir（`--add-dir` 的最后一个）， +跨迭代稳定。 + +**已知限制（0.1.0rc6 runtime）**：小会话 resume 可用（SDK 冒烟测试验证），但 +大会话（如一次完整 kernel 编辑迭代，~400KB zstd）resume 时 runtime 报 +`corrupt session log` / `unsupported flat-file layout` 并返回空回复。 +`dsh_agent.py` 已做兜底：resume 失败（finish_reason != completed/max-tokens） +时自动改用全新 session 重跑一次（prompt 自带全部上下文，效果等价），并把最终 +session id 透出给编排器。代价是每个 resume 失败的迭代多一次快速失败尝试，不影响 +正确性。 + +## 已知注意点 + +- `dsh_agent.py` 需要有执行位（`chmod 755`），否则 `SubAgentManager` 的 + Popen 会报 `PermissionError`。 +- SDK 为 dev checkout（`/root/deepseek-harness/python/sdk`），provider + `tencent` + TokenHub 端点在 runtime 侧解析；改 `cordis.yml` 时以 runtime + 报错为准逐步对齐。 +- 迭代循环在 agent 失败时会把 `attempt_limit` 递增（`w8a8_pipeline.py` 的 + `replacement_iteration` 逻辑），失败多会显著拉长任务；resume 回退已消除 + 空回复型失败。 +- 最终验证（serial_validate）对 GPU 争抢敏感：本机其他容器跑重负载时会把 + benchmark 拉高数倍，导致 `performance regressed` 误报。clean 复测（GPU + 空闲窗口）同一 prebuilt object 可测到与 worker best 一致甚至更优的延迟。 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/cordis.yml b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/cordis.yml new file mode 100644 index 00000000..da88e0d7 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/cordis.yml @@ -0,0 +1,104 @@ +# MetaInfer -> DSH agent composition for the dsh-jsonrpc-agent runtime. +# +# Unattended coding-agent deployment: stdout is reserved for JSON-RPC (the +# SDK client owns it); no console logger or terminal UI. +# +# This is the SDK example composition (examples/jsonrpc-agent/cordis.yml) +# adapted for kernel-optimization workers: +# * skills enabled (MetaInfer injects per-task skill bundles the agents +# must be able to load) +# * tool-fs / tool-bash / tool-subagent / tool-todo available so agents can +# read/edit HIP sources, run hipcc/harness, and fan out work +# * DSH_CWD / DSH_SESSION_ROOT drive the agent workspace and persistence + +- id: sdk-jsonrpc-server + name: '@deepseek-ai/dsh-sdk-jsonrpc-server' + config: + maxTokensAsSuccess: !!js "process.env.DSH_MAX_TOKENS_AS_SUCCESS === undefined ? true : JSON.parse(process.env.DSH_MAX_TOKENS_AS_SUCCESS)" + +# The DeepSeek adapter. Shipped default: full thinking at max effort on every +# request; the model arrives per session over JSON-RPC, so it is not pinned +# here. +- id: llm-deepseek + name: '@deepseek-ai/dsh-llm-deepseek' + config: + thinking: enabled + reasoningEffort: max + +# Managed child-process groups for the bash executor (spawn/kill/output plumbing). +- id: subprocess + name: '@deepseek-ai/dsh-subprocess-local' + +- id: bash + name: '@deepseek-ai/dsh-bash-local' + config: + cwd: !!js process.env.DSH_CWD ?? process.cwd() + timeoutMs: 120000 + +- id: agent-spine + name: '@deepseek-ai/dsh-agent-spine-demo' + config: + persona: !!js process.env.DSH_SYSTEM_PROMPT ?? 'You are a coding agent.' + workspaceContext: false + skills: + enabled: true + toolBash: + enableRunInBackground: false + toolJobs: false + +- id: sessions + name: '@deepseek-ai/dsh-session-persistence-jsonl' + config: + root: !!js process.env.DSH_SESSION_ROOT ?? './.sessions' + compression: !!js "process.env.DSH_SNAPSHOT === undefined ? 'zstd' : 'none'" + +- id: session-checkpoints + name: '@deepseek-ai/dsh-session-checkpoint-policy' + +- id: subagent + name: '@deepseek-ai/dsh-subagent' + +- id: subagent-spawn-in-process + name: '@deepseek-ai/dsh-subagent-spawn-in-process' + config: + providerName: spawn + +- id: tool-subagent + name: '@deepseek-ai/dsh-tool-subagent' + config: + provider: spawn + toolName: subagent + enableRunInBackground: false + +- id: tool-todo + name: '@deepseek-ai/dsh-tool-todo' + config: + allowParallelInProgress: true + +- id: fs-local + name: '@deepseek-ai/dsh-fs-local' + config: + cwd: !!js process.env.DSH_CWD ?? process.cwd() + +- id: fs-observation-policy + name: '@deepseek-ai/dsh-fs-observation-policy' + +- id: tool-fs + name: '@deepseek-ai/dsh-tool-fs' + +- id: tool-str-replace-editor + name: '@deepseek-ai/dsh-tool-str-replace-editor' + +- id: skill-filesystem + name: '@deepseek-ai/dsh-skill-filesystem' + +- id: token-meter + name: '@deepseek-ai/dsh-token-meter' + +- id: compaction-basic + name: '@deepseek-ai/dsh-compaction-basic' + config: + thresholdRatio: 0.8 + retainRatio: 0.16 + maxTokens: 8192 + compactionRetries: 1 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/dsh_agent.py b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/dsh_agent.py new file mode 100755 index 00000000..ce292b9d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/dsh_agent.py @@ -0,0 +1,366 @@ +#!/usr/bin/env python3 +"""MetaInfer -> DSH agent driver (ccb-compatible CLI). + +SubAgentManager spawns sub-agents as a CLI process (``claude_bin``, default +``ccb``) with this contract: + + -p --output-format stream-json --input-format text --verbose \\ + --permission-mode --add-dir [--add-dir ...] \\ + [--model ] [--effort ] [--resume | --session-id ] \\ + [extra_args...] + +with the agent prompt piped on **stdin**. The process must emit a +line-delimited stream-json event stream on **stdout** and exit 0 on success: + + {"type":"system", "session_id": "", ...} (first) + {"type":"assistant", "message": {"content": [{"type":"text","text": "..."}]}} + {"type":"result", "session_id": "", "result": "", + "usage": {...}} (last) + +This wrapper speaks that exact protocol but runs a **DeepSeek Harness agent** +through the Python SDK (``deepseek_harness``) instead of Claude Code. The +orchestrator therefore needs zero code changes: point ``claude_bin`` at this +script (``METAINFER_CLAUDE_BIN`` or ``--claude-bin``) and both the coordinator +(main agent) and the kernel workers (sub-agents) run on DSH. + +Environment: + DSH_AGENT_PROVIDER provider name for the DSH runtime + (default: deepseek-official — the only adapter + shipped by the dev-checkout runtime) + DSH_AGENT_MODEL model override (default deepseek/deepseek-v4-flash-0731) + DSH_AGENT_BASE_URL model endpoint; falls back to DEEPSEEK_BASE_URL + (default: https://tokenhub.tencentmaas.com/plan/v3) + TENCENT_API_KEY preferred API key (matches ~/.dsh/settings.yaml) + DEEPSEEK_API_KEY fallback API key + DSH_AGENT_CORDIS custom cordis.yml for the SDK runtime + (default: dsh/cordis.yml next to this file) + DSH_AGENT_SESSION_ROOT session JSONL persistence root; stable across + resume chains (default: /.dsh-sessions) + DSH_AGENT_MAX_TOKENS per-request output cap (default 65536) + DSH_AGENT_DEBUG set to 1 to keep the SDK runtime log lines +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional + +# Dev-checkout carrier: this host has no compiled ``deepseek-harness-runtime-bin`` +# executable, so the SDK runtime is launched through the dev-only node carrier +# (see resolve_bundled_launch_args). Callers with the exe carrier installed can +# override with DSH_RUNTIME_MODE=exe. +os.environ.setdefault("DSH_RUNTIME_MODE", "node") + +# --------------------------------------------------------------------------- # +# Host wiring (worker29: TokenHub DSV4-Flash via the local dsh CLI profile) +# --------------------------------------------------------------------------- # + +# Provider name understood by the DSH runtime. The dev-checkout runtime only +# ships the llm-deepseek adapter (provider "deepseek-official"); the TokenHub +# endpoint is selected via DEEPSEEK_BASE_URL below, so the official adapter +# talks to the same gateway the local dsh CLI uses. +def default_provider() -> str: + return ( + os.environ.get("DSH_AGENT_PROVIDER", "deepseek-official").strip() + or "deepseek-official" + ) + + +def default_model() -> str: + return ( + os.environ.get("DSH_AGENT_MODEL", "deepseek/deepseek-v4-flash-0731") + .strip() or "deepseek/deepseek-v4-flash-0731" + ) + + +def default_base_url() -> str: + env = os.environ.get("DEEPSEEK_BASE_URL") + if env and env.strip(): + return env.strip() + return ( + os.environ.get( + "DSH_AGENT_BASE_URL", + "https://tokenhub.tencentmaas.com/plan/v3", + ).strip() + or "https://tokenhub.tencentmaas.com/plan/v3" + ) + + +def _credentials_api_key() -> str: + """Read the API key from the local dsh credentials file when the env + variables are not set (e.g. when this wrapper runs standalone).""" + try: + import yaml + + path = Path( + os.environ.get("DSH_CREDENTIALS", "~/.dsh/.credentials.yaml") + ).expanduser() + if not path.is_file(): + return "" + data = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + for key in ("TENCENT_API_KEY", "DEEPSEEK_API_KEY"): + value = data.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + except Exception: # noqa: BLE001 - credentials are best-effort + pass + return "" + + +def default_api_key() -> str: + """Resolve the API key for the configured gateway. + + The TokenHub (tencent) gateway is authenticated by the TENCENT key. The + DEEPSEEK key is also present in ~/.dsh/.credentials.yaml and is exported + into the metainfer server env by start_webui.sh, but TokenHub rejects it; + for that gateway the tencent credential therefore takes priority over the + DEEPSEEK env fallback. + """ + tencent_gateway = ( + "tencentmaas" in default_base_url().lower() + or default_provider().strip().lower() == "tencent" + ) + if tencent_gateway: + value = os.environ.get("TENCENT_API_KEY") + if value and value.strip(): + return value.strip() + cred = _credentials_api_key() + if cred: + return cred + value = os.environ.get("DEEPSEEK_API_KEY") + if value and value.strip(): + return value.strip() + else: + for env in ("DEEPSEEK_API_KEY", "TENCENT_API_KEY"): + value = os.environ.get(env) + if value and value.strip(): + return value.strip() + return _credentials_api_key() + return "" + + +# --------------------------------------------------------------------------- # +# Model mapping (legacy Claude labels -> this host's DSH model ids) +# --------------------------------------------------------------------------- # + +_MODEL_MAP = { + "opus": "deepseek-v4-pro", + "sonnet": "deepseek-v4-flash", + "haiku": "deepseek-v4-flash", +} + + +def map_model(requested: Optional[str]) -> str: + if not requested: + return default_model() + key = requested.strip().lower() + if key in _MODEL_MAP: + return _MODEL_MAP[key] + if key == "deepseek-v4-flash": + # Bare user-facing label -> the pinned host model id. + return default_model() + if key.startswith("deepseek"): + # Full model id (e.g. deepseek/deepseek-v4-flash-0731): pass through. + return key + # Unknown label: keep the caller's intent but stay on a known model id. + return default_model() + + +# --------------------------------------------------------------------------- # +# stream-json event emission (SubAgentManager wire protocol) +# --------------------------------------------------------------------------- # + +def emit(event: Dict[str, Any]) -> None: + sys.stdout.write(json.dumps(event, ensure_ascii=False) + "\n") + sys.stdout.flush() + + +def emit_system(session_id: str, model: str) -> None: + emit({ + "type": "system", + "session_id": session_id, + "model": model, + "subagent_id": session_id, + "cwd": os.getcwd(), + }) + + +def extract_text_blocks(content: Any) -> List[str]: + """Pull text blocks from an assistant message content array.""" + if not isinstance(content, list): + return [] + out: List[str] = [] + for blk in content: + if isinstance(blk, dict) and blk.get("type") == "text": + text = blk.get("text") + if isinstance(text, str): + out.append(text) + return out + + +# --------------------------------------------------------------------------- # +# CLI +# --------------------------------------------------------------------------- # + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(prog="dsh_agent", add_help=False) + p.add_argument("-p", action="store_true") + p.add_argument("--output-format") + p.add_argument("--input-format") + p.add_argument("--verbose", action="store_true") + p.add_argument("--permission-mode") + p.add_argument("--add-dir", action="append", default=[]) + p.add_argument("--model") + p.add_argument("--effort") + p.add_argument("--resume") + p.add_argument("--session-id") + p.add_argument("--max-turns") + # Anything else (claude-specific) is ignored. + return p + + +def main(argv: Optional[List[str]] = None) -> int: + args, _unknown = build_parser().parse_known_args(argv) + + prompt = sys.stdin.buffer.read().decode("utf-8", errors="replace").strip() + if not prompt: + sys.stderr.write("dsh_agent: empty prompt on stdin\n") + return 1 + + try: + from deepseek_harness import DeepSeekHarness, DeepSeekHarnessConfig + except ImportError as exc: # pragma: no cover - environment check + sys.stderr.write( + "dsh_agent: deepseek_harness SDK not installed " + f"({exc}); run: pip install deepseek-harness-sdk\n" + ) + return 1 + + model = map_model(args.model) + # Session continuity: --resume continues an existing DSH conversation; + # --session-id pins the id on the first turn so the orchestrator can + # resume by it later. Both map to the SDK's per-session id. + requested_session = args.resume or args.session_id + session_id = requested_session or f"session-{os.urandom(8).hex()}" + + add_dirs = [str(Path(d).resolve()) for d in args.add_dir if d] + # Stable session persistence: prefer env, else the last --add-dir (the + # orchestrator passes workspace_dir after the per-agent workdir), else cwd. + session_root = os.environ.get("DSH_AGENT_SESSION_ROOT") + if not session_root and add_dirs: + session_root = str(Path(add_dirs[-1]) / ".dsh-sessions") + if not session_root: + session_root = str(Path.cwd() / ".dsh-sessions") + + cordis = os.environ.get("DSH_AGENT_CORDIS") + if not cordis: + cordis = str(Path(__file__).resolve().parent / "cordis.yml") + + max_tokens = int(os.environ.get("DSH_AGENT_MAX_TOKENS", "65536")) + api_key = default_api_key() + if not api_key: + sys.stderr.write( + "dsh_agent: no API key found (set TENCENT_API_KEY or " + "DEEPSEEK_API_KEY)\n" + ) + return 1 + + def run_agent(sid: str): + config = DeepSeekHarnessConfig( + provider=default_provider(), + model=model, + max_tokens=max_tokens, + cwd=os.getcwd(), + session_root=session_root, + cordis=cordis, + # API key / base url are injected explicitly; the runtime also + # inherits them from the process environment by default. + base_url=default_base_url(), + api_key=api_key, + env={ + "DSH_SESSION_ROOT": session_root, + "DSH_CWD": os.getcwd(), + }, + request_timeout_seconds=3600.0, + shutdown_timeout_seconds=15.0, + ) + + def on_notification(notification: Any) -> None: + # Stream assistant text live so the SubAgentManager stuck-watchdog + # sees fresh stdout and the WebUI log stays readable. + try: + if getattr(notification, "method", None) != "session.event": + return + payload = notification.payload or {} + if payload.get("sessionId") != sid: + return + event = payload.get("event") + if not isinstance(event, dict): + return + if event.get("type") != "assistant/message": + return + data = event.get("data") or {} + message = data.get("message") + content = message.get("content") if isinstance(message, dict) else data.get("content") + blocks = extract_text_blocks(content) + if blocks: + emit({ + "type": "assistant", + "session_id": sid, + "message": {"content": [{"type": "text", "text": b} for b in blocks]}, + }) + except Exception: # pragma: no cover - observability must not kill the run + pass + + try: + result = DeepSeekHarness(config).run( + prompt, + session_id=sid, + on_notification=on_notification, + ) + return result, None + except Exception as exc: # pragma: no cover - surfaced to the orchestrator + return None, exc + + # Resume is best-effort: large prior sessions can fail to reload in the + # runtime (turn/end reason "error" with an empty response). Fall back to a + # fresh session so every iteration still produces real agent work; the + # prompt MetaInfer passes is self-contained and includes continuation + # context, so the fresh run remains effective. + result, exc = run_agent(session_id) + if result is not None and result.finish_reason not in (None, "completed", "max-tokens"): + sys.stderr.write( + f"dsh_agent: session {session_id} resume failed " + f"(finish_reason={result.finish_reason!r}); retrying as a fresh session\n" + ) + session_id = f"session-{os.urandom(8).hex()}" + result, exc = run_agent(session_id) + if exc is not None: + sys.stderr.write(f"dsh_agent: DSH run failed: {exc!r}\n") + return 1 + if result.finish_reason not in (None, "completed", "max-tokens"): + sys.stderr.write( + f"dsh_agent: DSH run finished with reason {result.finish_reason!r}\n" + ) + return 1 + + # Emit the system event only after the successful run so the stream's + # first session_id is the id the orchestrator should resume from later. + emit_system(result.session_id, model) + final_text = result.final_response or "" + emit({ + "type": "result", + "session_id": result.session_id, + "result": final_text, + "finish_reason": result.finish_reason, + "usage": {}, # optional; token-budget accounting skips when absent + }) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/run_dsh_task.sh b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/run_dsh_task.sh new file mode 100755 index 00000000..49196165 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/run_dsh_task.sh @@ -0,0 +1,45 @@ +#!/bin/bash +# Launch a dcu-kernel-auto-opt task with DSH agents instead of Claude Code. +# +# Usage: +# TENCENT_API_KEY= bash run_dsh_task.sh +# +# The requirements.json must contain task_type=dcu-kernel-auto-opt and the +# standard answers; agent_framework should be "dsh" (see the WebUI new-task +# form). This script only sets the claude_bin override — the orchestrator CLI +# resolves the DSH wrapper automatically when agent_framework=dsh, so this +# script is a convenience for headless runs. +# +# Optional env: +# DSH_AGENT_MAX_TOKENS per-request output cap (default 65536) +# DSH_AGENT_MODEL model override (default deepseek/deepseek-v4-flash-0731) +# METAINFER_GPU_IDS restrict GPUs, e.g. 4,5,6,7 + +set -euo pipefail + +REQ="${1:?requirements.json path required}" +TASK_ID="${2:?task id required}" + +# Resolve the MetaInfer root from this script's location: +# /metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/run_dsh_task.sh +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +META_ROOT="$(cd "$HERE/../../../../.." && pwd)" +DSH_DIR="$HERE" +NODE_ROOT="${METAINFER_NODE_ROOT:-$META_ROOT/nodes/worker29}" + +export METAINFER_CLAUDE_BIN="${METAINFER_CLAUDE_BIN:-$DSH_DIR/dsh_agent.py}" +export DSH_AGENT_MAX_TOKENS="${DSH_AGENT_MAX_TOKENS:-65536}" + +STATE_DIR="$NODE_ROOT/.metainfer/tasks/$TASK_ID" +WORKSPACE_DIR="$NODE_ROOT/workspaces/$TASK_ID" +mkdir -p "$STATE_DIR" "$WORKSPACE_DIR" + +echo "==> DSH agent: $METAINFER_CLAUDE_BIN" +echo "==> state: $STATE_DIR" +echo "==> workspace: $WORKSPACE_DIR" + +cd "$META_ROOT" +exec python3 -m metainfer.tasks.dcu_kernel_auto_opt.orchestrator.cli run \ + "$REQ" \ + --state-dir "$STATE_DIR" \ + --workspace-dir "$WORKSPACE_DIR" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/smoke_sdk.py b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/smoke_sdk.py new file mode 100644 index 00000000..94f29330 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/smoke_sdk.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +"""Smoke test: drive a real DSH agent via the Python SDK inside zth_meta. + +Verifies, in order: + 1. runtime boots with the MetaInfer cordis composition + 2. a trivial prompt returns a final response (model endpoint works) + 3. session id is stable and a second prompt resumes the same conversation + 4. the agent can create a file through its tools (fs/bash) in DSH_CWD + 5. the agent can read a file outside DSH_CWD (absolute path) — no sandbox + +Usage (inside the container): + DEEPSEEK_API_KEY=... python3 /workspace/MetaInfer/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/smoke_sdk.py +""" + +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path + +from deepseek_harness import DeepSeekHarness, DeepSeekHarnessConfig + +CORDIS = Path(__file__).resolve().parents[1] / "cordis.yml" +WORK = Path("/tmp/dsh-smoke") +SESSION_ROOT = WORK / ".sessions" + + +def main() -> int: + WORK.mkdir(parents=True, exist_ok=True) + cfg = DeepSeekHarnessConfig( + provider="deepseek-official", + model="deepseek-v4-flash", + max_tokens=2048, + cwd=str(WORK), + session_root=str(SESSION_ROOT), + cordis=str(CORDIS), + env={"DSH_CWD": str(WORK), "DSH_SESSION_ROOT": str(SESSION_ROOT)}, + request_timeout_seconds=600.0, + shutdown_timeout_seconds=15.0, + ) + sid = f"session-smoke-{os.urandom(4).hex()}" + print(f"[smoke] session={sid} cordis={CORDIS}", flush=True) + with DeepSeekHarness(cfg) as h: + r1 = h.run("Reply with exactly: DSH_OK", session_id=sid) + print("[smoke] run1 session_id:", r1.session_id, flush=True) + print("[smoke] run1 finish_reason:", r1.finish_reason, flush=True) + print("[smoke] run1 final_response:", repr(r1.final_response), flush=True) + if "DSH_OK" not in (r1.final_response or ""): + print("[smoke] FAIL: unexpected run1 response", flush=True) + return 1 + + r2 = h.run( + "Continuing the same conversation: what was the exact token you " + "were asked to reply with? Reply with that token only.", + session_id=sid, + ) + print("[smoke] run2 session_id:", r2.session_id, flush=True) + print("[smoke] run2 final_response:", repr(r2.final_response), flush=True) + if "DSH_OK" not in (r2.final_response or ""): + print("[smoke] WARN: resume did not preserve context", flush=True) + + r3 = h.run( + "Create a file named smoke.txt in the current working directory " + "containing exactly the text FILE_OK. Then print the file path.", + session_id=sid, + ) + print("[smoke] run3 final_response:", repr(r3.final_response), flush=True) + created = (WORK / "smoke.txt").exists() + content = (WORK / "smoke.txt").read_text() if created else "" + print("[smoke] file created:", created, "content:", repr(content), flush=True) + if not created or "FILE_OK" not in content: + print("[smoke] FAIL: agent could not create file in DSH_CWD", flush=True) + return 1 + + r4 = h.run( + "Read the file /workspace/MetaInfer/metainfer/tasks/dcu_kernel_auto_opt/" + "bridge/dsh/cordis.yml and report the id of the first entry.", + session_id=sid, + ) + print("[smoke] run4 final_response:", repr(r4.final_response), flush=True) + + print("[smoke] OK", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/summarize_iters.py b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/summarize_iters.py new file mode 100644 index 00000000..1f0ff14a --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/summarize_iters.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 +"""Summarize dsh-mini-smoke worker_0 iteration outcomes.""" +import json +import os +import sys + +base = sys.argv[1] if len(sys.argv) > 1 else ( + "/workspace/MetaInfer/nodes/worker29/workspaces/dsh-mini-smoke/" + "workers/worker_0/iterations/tp8_wqkv_a_m32" +) +print("=== iteration results ===") +if not os.path.isdir(base): + print("no iterations dir:", base) + sys.exit(0) +rows = [] +for d in sorted(os.listdir(base), key=lambda x: int(x[9:]) if x[9:].isdigit() else 0): + f = os.path.join(base, d, "iteration.json") + if not os.path.exists(f): + continue + j = json.load(open(f)) + acc = j.get("acceptance") or {} + rows.append(( + d, + j.get("build_success"), + j.get("correctness_passed"), + acc.get("accepted"), + acc.get("best_us"), + acc.get("candidate_us"), + j.get("baseline_us"), + str(j.get("failure_reason"))[:70], + )) +for r in rows: + print( + f"{r[0]}: build={r[1]} correct={r[2]} accepted={r[3]} " + f"best={r[4]} cand={r[5]} bl={r[6]} err={r[7]}" + ) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/test_dsh_agent.py b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/test_dsh_agent.py new file mode 100644 index 00000000..0580f2b3 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/test_dsh_agent.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python3 +"""Unit tests for dsh_agent.py (SubAgentManager stream-json protocol). + +Runs without a real model endpoint: the ``deepseek_harness`` SDK module is +replaced by a fake whose behavior is scripted per test. Verify: + + 1. arg parsing tolerates every SubAgentManager flag (incl. unknown ones) + 2. model mapping: sonnet/haiku -> deepseek-v4-flash, opus -> deepseek-v4-pro + 3. the emitted stream contains system -> assistant* -> result, exit 0 + 4. a resume failure (finish_reason == "error") falls back to a fresh + session and the stream's first session_id is the fallback one + 5. an SDK exception exits nonzero without a result event + +Usage (inside the container where deepseek-harness-sdk is installed, or with +PYTHONPATH pointing at the SDK source): + python3 tests/test_dsh_agent.py +""" + +from __future__ import annotations + +import importlib.util +import io +import json +import os +import sys +import types +import unittest +from contextlib import redirect_stderr, redirect_stdout +from pathlib import Path + +WRAPPER = Path(__file__).resolve().parents[1] / "dsh_agent.py" + + +class FakeRunResult: + def __init__(self, session_id, final_response, finish_reason): + self.session_id = session_id + self.final_response = final_response + self.finish_reason = finish_reason + + +class FakeDeepSeekHarness: + """Scripted harness: returns queued results in order, or raises.""" + + instances = [] + + def __init__(self, config=None, **kwargs): + self.config = config + self.runs = [] + FakeDeepSeekHarness.instances.append(self) + + def run(self, prompt, session_id=None, on_notification=None): + self.runs.append({"prompt": prompt, "session_id": session_id}) + script = FakeDeepSeekHarness.script + if isinstance(script, Exception): + raise script + item = script.pop(0) + if item == "raise": + raise RuntimeError("boom") + return FakeRunResult( + session_id=item[0], final_response=item[1], finish_reason=item[2] + ) + + def __enter__(self): + return self + + def __exit__(self, *args): + return None + + +class FakeConfig: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +def install_fake_sdk(script): + FakeDeepSeekHarness.script = list(script) + FakeDeepSeekHarness.instances = [] + fake = types.ModuleType("deepseek_harness") + fake.DeepSeekHarness = FakeDeepSeekHarness + fake.DeepSeekHarnessConfig = FakeConfig + sys.modules["deepseek_harness"] = fake + # Reload the wrapper so its import picks up the fake. + for name in list(sys.modules): + if name == "dsh_agent": + del sys.modules[name] + spec = importlib.util.spec_from_file_location("dsh_agent", WRAPPER) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class FakeStdin: + """Exposes .buffer (BytesIO) like a real process stdin.""" + + def __init__(self, text): + self.buffer = io.BytesIO(text.encode("utf-8")) + + +def run_wrapper(module, argv, prompt): + """Run the wrapper with a scripted prompt/stdin, plus a fake API key. + + ``dsh_agent.main()`` refuses to start without credentials (TENCENT_API_KEY + / DEEPSEEK_API_KEY env or ~/.dsh/.credentials.yaml). CI runners have none + of those, so inject a dummy key for the duration of the call to keep the + tests hermetic and environment-independent. + """ + out, err = io.StringIO(), io.StringIO() + old_stdin = sys.stdin + saved_key = os.environ.get("TENCENT_API_KEY") + os.environ["TENCENT_API_KEY"] = saved_key or "ci-test-key" + sys.stdin = FakeStdin(prompt) + try: + with redirect_stdout(out), redirect_stderr(err): + code = module.main(argv) + finally: + sys.stdin = old_stdin + if saved_key is None: + os.environ.pop("TENCENT_API_KEY", None) + else: + os.environ["TENCENT_API_KEY"] = saved_key + return code, out.getvalue(), err.getvalue() + + +BASE_ARGS = [ + "-p", "--output-format", "stream-json", "--input-format", "text", + "--verbose", "--permission-mode", "bypassPermissions", + "--add-dir", "/tmp/wtest", "--model", "sonnet", "--effort", "low", + "--max-turns", "50", "--disallowedTools", "Edit,Write", +] + + +class DshAgentTests(unittest.TestCase): + def test_success_stream_protocol(self): + mod = install_fake_sdk([ + ("session-aaa", "The answer is 42", "completed"), + ]) + code, out, err = run_wrapper(mod, BASE_ARGS, "Do the thing.") + self.assertEqual(code, 0) + events = [json.loads(l) for l in out.splitlines() if l.strip()] + types_seen = [e["type"] for e in events] + # Stream order: system (with final session id) -> optional assistant + # text -> result. The fake emits no assistant events; the real SDK + # streams them via on_notification. + self.assertEqual(types_seen[0], "system") + self.assertEqual(types_seen[-1], "result") + result = events[-1] + self.assertEqual(result["result"], "The answer is 42") + self.assertEqual(result["session_id"], "session-aaa") + self.assertEqual(result["finish_reason"], "completed") + # The system event carries the same session id. + self.assertEqual(events[0]["session_id"], "session-aaa") + + def test_model_mapping(self): + mod = install_fake_sdk([ + ("s1", "ok", "completed"), + ]) + # opus -> deepseek-v4-pro + run_wrapper(mod, [a if a != "sonnet" else "opus" for a in BASE_ARGS], "hi") + cfg = FakeDeepSeekHarness.instances[0].config + self.assertEqual(cfg.model, "deepseek-v4-pro") + + def test_resume_fallback(self): + mod = install_fake_sdk([ + ("session-old", "", "error"), # resume attempt fails + ("session-new", "recovered", "completed"), # fresh session works + ]) + args = BASE_ARGS + ["--resume", "session-old"] + code, out, err = run_wrapper(mod, args, "Continue the work.") + self.assertEqual(code, 0) + events = [json.loads(l) for l in out.splitlines() if l.strip()] + result = events[-1] + self.assertEqual(result["result"], "recovered") + self.assertEqual(result["session_id"], "session-new") + # The first session_id seen by the parser must be the fallback one. + first_sid = next(e["session_id"] for e in events if "session_id" in e) + self.assertEqual(first_sid, "session-new") + self.assertIn("retrying as a fresh session", err) + + def test_sdk_exception_exits_nonzero(self): + mod = install_fake_sdk(["raise"]) + code, out, err = run_wrapper(mod, BASE_ARGS, "hi") + self.assertNotEqual(code, 0) + events = [l for l in out.splitlines() if l.strip()] + # No result event on failure. + self.assertFalse(any("result" in json.loads(l) for l in events)) + + def test_empty_prompt_fails(self): + mod = install_fake_sdk([("s1", "x", "completed")]) + code, out, err = run_wrapper(mod, BASE_ARGS, " ") + self.assertNotEqual(code, 0) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/form.yaml b/metainfer/tasks/dcu_kernel_auto_opt/form.yaml new file mode 100644 index 00000000..96215dc2 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/form.yaml @@ -0,0 +1,202 @@ +- key: operator + question: "Which operator will this task optimize? The concrete adapter contract is supplied when the operator is connected." + header: "Operator" + required: true + multi: false + default: "Quantized GEMM" + options: + - label: "Quantized GEMM" + description: "INT8/BF16/FP16 matrix multiplication, including W8A8" + - label: "Attention" + description: "MHA, MQA, GQA or MLA kernels" + - label: "RMSNorm / LayerNorm" + description: "Normalization kernels" + - label: "RoPE" + description: "Rotary positional embedding" + - label: "Custom operator" + description: "An extracted DCU kernel with a custom adapter" + +- key: kernel_language + question: "Kernel implementation stack." + header: "Stack" + required: true + multi: false + default: "HIP C++" + options: + - label: "HIP C++" + description: "Native DTK/HIP kernel for gfx928" + - label: "PyTorch HIP extension" + description: "HIP source loaded through a PyTorch extension" + - label: "TileLang / Triton" + description: "Generated kernel stack when supported by the target project" + +- key: target_hardware + question: "Target DCU platform." + header: "Hardware" + required: true + multi: false + default: "K500SM_AI / gfx928" + options: + - label: "K500SM_AI / gfx928" + description: "worker29, four DCUs, wavefront 64" + - label: "Other DCU" + description: "Requires an explicit hardware adapter" + +- key: dtype + question: "Compute and storage dtype." + header: "Dtype" + required: true + multi: false + default: "INT8 W8A8" + options: + - label: "INT8 W8A8" + description: "INT8 activation and weight" + - label: "FP16 / BF16" + description: "Half-precision floating point" + - label: "Other" + description: "Defined by the concrete operator adapter" + +- key: agent_framework + question: "Agent framework used by coordinator, kernel workers, repairs, synthesis, and validation." + header: "Agent framework" + required: true + multi: false + default: "ccb" + options: + - label: "ccb" + description: "Claude Code (claude-sonnet-5 / claude-opus-5)" + - label: "dsh" + description: "DeepSeek Harness (deepseek-v4-flash)" + +- key: agent_model + question: "Model used by the selected agent framework (ccb: Sonnet/Opus; dsh: deepseek-v4-flash)." + header: "Agent model" + required: true + multi: false + default: "Opus" + override_component: "agent-model" + override_module: "app/dkao-agent-fields" + options: + - label: "Opus" + description: "Pinned to Claude Opus 5 (claude-opus-5) for difficult kernel architecture and correctness work" + - label: "Sonnet" + description: "Pinned to Claude Sonnet 5 (claude-sonnet-5) for faster, lower-cost optimization runs" + - label: "deepseek-v4-flash" + description: "DeepSeek Harness model (deepseek-v4-flash)" + +- key: execution_mode + question: "Task execution mode for this run." + header: "Mode" + required: true + multi: false + default: "Generate & optimize (auto-create kernel repo)" + options: + - label: "Generate & optimize (auto-create kernel repo)" + description: "Agent generates the complete kernel repo from scratch, then optimizes it. No pre-existing kernel code required." + +- key: target_repo_path + question: "Repository folder name under kernel-repos next to MetaInfer. Example: int8 test2 creates and directly uses kernel-repos/int8 test2." + header: "Kernel repo" + required: false + form: text + default: "" + +- key: model + question: "Which model's INT8 W8A8 GEMM workload will this task optimize? Shape selection is grouped by TP size (1/4/8)." + header: "Model" + required: true + multi: false + default: "DeepSeek V4 Flash" + options: + - label: "DeepSeek V4 Flash" + description: "SGLang DeepSeek-V4-Flash Channel INT8 W8A8 (TP1/4/8), current default workload" + - label: "Hy3 (Hunyuan 3)" + description: "Hy3 Channel INT8 W8A8 via vLLM v0.18.0" + - label: "MiniMax M3" + description: "MiniMax M3 W8A8; currently not deployable for testing" + - label: "GLM5.2" + description: "GLM-5.2 INT4/INT8 mix via SGLang v0.5.15" + +- key: shape_assignment_mode + question: "Choose whether the control plane or you assign shapes to physical GPUs." + header: "Shape assignment" + required: true + multi: false + default: "AI automatic" + options: + - label: "AI automatic" + description: "The control plane deterministically balances exact shapes across available GPUs" + - label: "Manual by GPU" + description: "Use four GPU cards below to explicitly choose where every shape is optimized" + +- key: shape_scope + question: "Choose whether this task optimizes the complete API workload or only a selected subset." + header: "Optimization scope" + required: true + multi: false + default: "All API shapes" + options: + - label: "All API shapes" + description: "Optimize every default shape exposed by the immutable operator API" + - label: "Selected shapes only" + description: "Optimize a task-sized subset; unselected shapes remain on the trusted fallback and are regression-tested" + +- key: shape_config + question: "Reads default optimization shapes from the fixed operator API. In manual mode, assign every shape to one GPU." + header: "Workload" + required: false + form: textarea + default: "" + override_component: "shape-input" + override_module: "app/dkao-shape-input" + +- key: correctness_ref + question: "Trusted correctness reference owned by the MetaInfer adapter." + header: "Reference" + required: true + multi: false + default: "PyTorch eager" + options: + - label: "PyTorch eager" + description: "Compare with a deterministic PyTorch implementation" + - label: "SGLang operator" + description: "Compare with the extracted SGLang reference path" + - label: "Adapter-provided reference" + description: "The concrete adapter supplies the immutable reference" + +- key: perf_target + question: "Optional performance target, for example >=1.2x baseline." + header: "Perf bar" + required: false + default: "Improve stable median without P90 regression" + +- key: max_iterations + question: "Maximum optimization rounds per worker." + header: "Max iters" + required: true + multi: false + default: "5" + options: + - label: "1" + description: "Single-round validation" + - label: "3" + description: "Quick MVP exercise" + - label: "5" + description: "Default" + - label: "10" + description: "Longer optimization loop" + - label: "20" + description: "Extended exploration" + +- key: minimum_improvement_percent + question: "Final target improvement over the fixed baseline. This is not the per-round candidate acceptance threshold." + header: "Final improve" + required: true + form: number + default: 1.0 + +- key: extra_notes + question: "Additional constraints, such as determinism or forbidden libraries." + header: "Notes" + required: false + form: textarea diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/__init__.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/__init__.py new file mode 100644 index 00000000..fb262b80 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/__init__.py @@ -0,0 +1,9 @@ +"""Orchestrator registration for dcu-kernel-auto-opt.""" + +from metainfer.orchestrator.tasks import register + +from .plugin import PLUGIN + +register(PLUGIN) + +__version__ = "0.1.0" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/adapters/__init__.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/adapters/__init__.py new file mode 100644 index 00000000..481c89f7 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/adapters/__init__.py @@ -0,0 +1,6 @@ +"""Kernel adapter interfaces and MVP implementations.""" + +from .base import AdapterResult, KernelAdapter +from .mock import MockKernelAdapter + +__all__ = ["AdapterResult", "KernelAdapter", "MockKernelAdapter"] diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/adapters/base.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/adapters/base.py new file mode 100644 index 00000000..3ee462be --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/adapters/base.py @@ -0,0 +1,50 @@ +"""Operator-agnostic seam between orchestration and a concrete kernel.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict + +from ..config import ShapeSpec + + +@dataclass(frozen=True) +class AdapterResult: + success: bool + metrics: Dict[str, float] = field(default_factory=dict) + evidence: Dict[str, Any] = field(default_factory=dict) + error: str | None = None + + +class KernelAdapter(ABC): + """A trusted adapter owned by MetaInfer, not by optimization agents.""" + + requires_gpu = True + + @abstractmethod + def describe_environment(self) -> Dict[str, Any]: + raise NotImplementedError + + @abstractmethod + def prepare(self, workspace: Path) -> AdapterResult: + raise NotImplementedError + + @abstractmethod + def build(self, workspace: Path) -> AdapterResult: + raise NotImplementedError + + @abstractmethod + def correctness(self, workspace: Path, shape: ShapeSpec) -> AdapterResult: + raise NotImplementedError + + @abstractmethod + def benchmark( + self, workspace: Path, shape: ShapeSpec, *, iteration: int = 0 + ) -> AdapterResult: + raise NotImplementedError + + @abstractmethod + def profile(self, workspace: Path, shape: ShapeSpec) -> AdapterResult: + raise NotImplementedError diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/adapters/mock.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/adapters/mock.py new file mode 100644 index 00000000..bcbafa5f --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/adapters/mock.py @@ -0,0 +1,84 @@ +"""Deterministic no-GPU adapter used to validate the orchestration MVP.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict + +from .base import AdapterResult, KernelAdapter +from ..config import ShapeSpec + + +class MockKernelAdapter(KernelAdapter): + requires_gpu = False + + def describe_environment(self) -> Dict[str, Any]: + return { + "adapter": "mock", + "requires_gpu": False, + "gpu_runtime_loaded": False, + } + + def prepare(self, workspace: Path) -> AdapterResult: + workspace.mkdir(parents=True, exist_ok=True) + return AdapterResult(True, evidence={"workspace": str(workspace)}) + + def build(self, workspace: Path) -> AdapterResult: + return AdapterResult(True, evidence={"build": "mock-success"}) + + def correctness(self, workspace: Path, shape: ShapeSpec) -> AdapterResult: + return AdapterResult( + True, + evidence={ + "reference": "mock-trusted-reference", + "seed": 20260724, + "shape": {"id": shape.id, **shape.params}, + }, + ) + + def benchmark( + self, workspace: Path, shape: ShapeSpec, *, iteration: int = 0 + ) -> AdapterResult: + seed = sum(ord(ch) for ch in shape.id) + baseline_us = 80.0 + float(seed % 70) + factor = max(0.70, 1.0 - 0.025 * max(0, iteration)) + median_us = round(baseline_us * factor, 4) + try: + m = float(shape.params.get("M", 0)) + n = float(shape.params.get("N", 0)) + k = float(shape.params.get("K", 0)) + except (TypeError, ValueError): + m = n = k = 0.0 + elapsed_s = median_us * 1e-6 + tflops = ( + (2.0 * m * n * k) / elapsed_s / 1e12 + if elapsed_s > 0 and m > 0 and n > 0 and k > 0 else 0.0 + ) + # Mock W8A8 traffic model: int8 A/B plus int32 output. + bytes_moved = m * k + k * n + 4.0 * m * n + bandwidth_gb_s = ( + bytes_moved / elapsed_s / 1e9 + if elapsed_s > 0 and bytes_moved > 0 else 0.0 + ) + return AdapterResult( + True, + metrics={ + "median_us": median_us, + "p90_us": round(median_us * 1.015, 4), + "min_us": round(median_us * 0.99, 4), + "max_us": round(median_us * 1.03, 4), + "tflops": round(tflops, 4), + "bandwidth_gb_s": round(bandwidth_gb_s, 4), + }, + evidence={"samples": 30, "warmup": 10, "mock": True}, + ) + + def profile(self, workspace: Path, shape: ShapeSpec) -> AdapterResult: + return AdapterResult( + True, + evidence={ + "tool": "mock-profiler", + "bottleneck": "synthetic-dispatch-overhead", + "profile_is_not_benchmark": True, + }, + ) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/api_contracts.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/api_contracts.py new file mode 100644 index 00000000..fe06262b --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/api_contracts.py @@ -0,0 +1,255 @@ +"""Resolve and stage user-owned, immutable operator API contracts.""" + +from __future__ import annotations + +import hashlib +import importlib.util +import os +import shutil +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import Any, Mapping, Sequence + + +W8A8_API_FILENAME = "int8_w8a8_gemm_api.py" +W8A8_BACKEND_FILENAME = "w8a8_backend.py" +W8A8_VARIANTS_FILENAME = "w8a8_gemm_variants.hip" +W8A8_VARIANTS_RELATIVE = Path("references") / W8A8_VARIANTS_FILENAME + + +@dataclass(frozen=True) +class OperatorAPIContract: + operator: str + dtype: str + source: Path + destination_name: str + reference_sources: tuple[Path, ...] = () + + +_CONTRACT_PATHS = { + ("Quantized GEMM", "INT8 W8A8"): ( + Path("int8w8a8gemm") / W8A8_API_FILENAME + ), +} + + +def _default_api_root() -> Path: + override = os.environ.get("METAINFER_OPERATOR_API_ROOT") + if override: + return Path(override).expanduser().resolve() + # Plugin-local integrated API files: + # .../metainfer/tasks/dcu_kernel_auto_opt/api + return Path(__file__).resolve().parents[1] / "api" + + +def _load_module(path: Path) -> ModuleType: + name = "metainfer_operator_api_" + hashlib.sha256( + str(path).encode("utf-8") + ).hexdigest()[:12] + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ValueError(f"cannot import operator API contract: {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def resolve_operator_api( + operator: str, + dtype: str, +) -> OperatorAPIContract: + relative = _CONTRACT_PATHS.get((operator, dtype)) + if relative is None: + supported = ", ".join( + f"{op} / {dt}" for op, dt in sorted(_CONTRACT_PATHS) + ) + raise ValueError( + f"no operator API contract registered for {operator} / {dtype}; " + f"supported: {supported}" + ) + source = (_default_api_root() / relative).resolve() + if not source.is_file(): + raise FileNotFoundError( + f"operator API contract not found: {source}; place the manually " + "maintained interface under " + "metainfer/tasks/dcu_kernel_auto_opt/api or set " + "METAINFER_OPERATOR_API_ROOT" + ) + module = _load_module(source) + required = ( + "prepare_weight", + "allocate_workspace", + "validate_gemm_out_inputs", + "w8a8_gemm_out", + ) + missing = [name for name in required if not callable(getattr(module, name, None))] + if missing: + raise ValueError( + f"operator API contract {source} is missing callables: {missing}" + ) + if not callable(getattr(module, "_check_target_shape", None)): + raise ValueError( + f"operator API contract {source} must expose a shape validator " + "named _check_target_shape" + ) + # Reference variant HIP code lives in the plugin's variant/ directory; + # fall back to the legacy location next to the API file. The fine-grained + # variant TREE (variant/////.hip) is + # staged alongside the legacy single file so agents can navigate both. + variant_root = Path(__file__).resolve().parents[1] / "variant" + legacy = next( + (path for path in (variant_root / W8A8_VARIANTS_FILENAME,) + if path.is_file()), + next((path for path in (source.parent / W8A8_VARIANTS_FILENAME,) + if path.is_file()), None), + ) + variant_sources: tuple = () + if legacy is not None: + variant_sources += (legacy,) + # every subdirectory of variant/ is a fine-grained variant family tree + # (e.g. int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip) + if variant_root.is_dir(): + variant_sources += tuple( + path for path in sorted(variant_root.iterdir()) if path.is_dir() + ) + return OperatorAPIContract( + operator=operator, + dtype=dtype, + source=source, + destination_name=W8A8_API_FILENAME, + reference_sources=variant_sources, + ) + + +def validate_contract_shapes( + contract: OperatorAPIContract, + shapes: Mapping[str, Any], +) -> None: + module = _load_module(contract.source) + validate_shape = getattr(module, "_check_target_shape") + validate_logical_shape = getattr( + module, "validate_optimization_shape", None + ) + for shape_id, shape in shapes.items(): + params = getattr(shape, "params", shape) + try: + m = int(params["M"]) + n = int(params["N"]) + k = int(params["K"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError( + f"{shape_id}: operator API requires integer M, N and K" + ) from exc + try: + if callable(validate_logical_shape) and ( + "tp_size" in params or "operator" in params + ): + validate_logical_shape(params) + else: + validate_shape(m, n, k) + except Exception as exc: + raise ValueError( + f"{shape_id}: shape (M={m}, N={n}, K={k}) is outside " + f"the fixed API contract {contract.source}: {exc}" + ) from exc + + +def default_optimization_shapes( + contract: OperatorAPIContract, +) -> list[dict[str, Any]]: + module = _load_module(contract.source) + raw = getattr(module, "DEFAULT_OPTIMIZATION_SHAPES", None) + if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)): + raise ValueError( + f"operator API contract {contract.source} must define " + "DEFAULT_OPTIMIZATION_SHAPES as a sequence" + ) + shapes: list[dict[str, Any]] = [] + seen_ids: set[str] = set() + for index, item in enumerate(raw): + if not isinstance(item, Mapping): + raise ValueError( + f"DEFAULT_OPTIMIZATION_SHAPES[{index}] must be a mapping" + ) + shape = {str(key): value for key, value in item.items()} + shape_id = str(shape.get("id") or "").strip() + if not shape_id: + raise ValueError( + f"DEFAULT_OPTIMIZATION_SHAPES[{index}] requires an id" + ) + if shape_id in seen_ids: + raise ValueError( + f"duplicate default optimization shape id: {shape_id}" + ) + seen_ids.add(shape_id) + shapes.append(shape) + if not shapes: + raise ValueError("DEFAULT_OPTIMIZATION_SHAPES must not be empty") + validate_contract_shapes( + contract, + { + str(shape["id"]): { + key: value for key, value in shape.items() if key != "id" + } + for shape in shapes + }, + ) + return shapes + + +def stage_operator_api( + contract: OperatorAPIContract, + destination_dir: Path, +) -> Path: + destination = destination_dir / contract.destination_name + destination_dir.mkdir(parents=True, exist_ok=True) + if destination.exists(): + destination.chmod(0o644) + # A New Task receives a new repository artifact. Preserve the contract + # bytes, but not the source asset's historical timestamp, so provenance + # cannot be mistaken for code copied from an earlier task repository. + shutil.copyfile(contract.source, destination) + destination.chmod(0o444) + return destination + + +def stage_operator_references( + contract: OperatorAPIContract, + destination_dir: Path, +) -> list[Path]: + """Stage optional read-only evidence without adding it to the build. + + Files are copied into ``references/`` as-is; directories (the fine-grained + variant tree) are copied recursively under ``references//``. + """ + staged: list[Path] = [] + references_dir = destination_dir / "references" + for source in contract.reference_sources: + references_dir.mkdir(parents=True, exist_ok=True) + if source.is_dir(): + destination = references_dir / source.name + if destination.exists(): + shutil.rmtree(destination) + shutil.copytree( + source, destination, + ignore=shutil.ignore_patterns("*.bak-*"), + ) + for path in destination.rglob("*"): + if path.is_file(): + path.chmod(0o444) + staged.append(destination) + continue + destination = references_dir / source.name + if destination.exists(): + destination.chmod(0o644) + shutil.copyfile(source, destination) + destination.chmod(0o444) + staged.append(destination) + return staged + + +def file_digest(path: Path) -> str | None: + if not path.is_file(): + return None + return hashlib.sha256(path.read_bytes()).hexdigest() diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/cli.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/cli.py new file mode 100644 index 00000000..30f086d4 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/cli.py @@ -0,0 +1,46 @@ +"""CLI entry point for dcu-kernel-auto-opt.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="dcu-kernel-auto-opt") + sub = parser.add_subparsers(dest="command", required=True) + run_p = sub.add_parser("run") + run_p.add_argument("requirements", type=Path) + run_p.add_argument("--state-dir", type=Path, required=True) + run_p.add_argument("--workspace-dir", type=Path, required=True) + run_p.add_argument("--dry-run", action="store_true") + run_p.add_argument( + "--claude-bin", + default=None, + help=( + "Agent binary override. Defaults are resolved from the task's " + "agent_framework answer: ccb -> METAINFER_CLAUDE_BIN (or 'ccb'), " + "dsh -> the bundled bridge/dsh/dsh_agent.py wrapper." + ), + ) + args = parser.parse_args(argv) + if args.command == "run": + from .config import resolve_agent_framework, resolve_claude_bin + from .orchestrator import run_with_requirements + + req = json.loads(args.requirements.read_text(encoding="utf-8")) + framework = resolve_agent_framework(req) + claude_bin = resolve_claude_bin(framework, explicit=args.claude_bin) + return run_with_requirements( + args.requirements, + state_dir=args.state_dir, + workspace_dir=args.workspace_dir, + dry_run=args.dry_run, + claude_bin=claude_bin, + ) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/config.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/config.py new file mode 100644 index 00000000..d908ba13 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/config.py @@ -0,0 +1,492 @@ +"""Validated, operator-agnostic configuration for the optimizer MVP.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Mapping + +import yaml + +MOCK_MODE = "Mock (no GPU)" +LEGACY_SMOKE_MODE = "Real agents + DCU (smoke harness)" +SMOKE_MODE = "Infrastructure smoke (not operator optimization)" +W8A8_MODE = "Real INT8 W8A8 GEMM" +GEN_AND_OPT_MODE = "Generate & optimize (auto-create kernel repo)" +CLAUDE_MODELS = { + "Sonnet": "claude-sonnet-5", + "Opus": "claude-opus-5", +} +# Agent frameworks selectable from the new-task form. Each framework maps a +# user-facing model label to the model id handed to SubAgentManager (and from +# there to the agent binary via --model). ccb keeps Claude Code semantics; +# dsh routes claude_bin to bridge/dsh/dsh_agent.py (see resolve_claude_bin). +CCB_FRAMEWORK = "ccb" +DSH_FRAMEWORK = "dsh" +DSH_DEFAULT_MODEL_ID = "deepseek/deepseek-v4-flash-0731" + + +def dsh_model_id() -> str: + """Model id used for the dsh framework (env DSH_AGENT_MODEL overrides).""" + override = os.environ.get("DSH_AGENT_MODEL") + return (override.strip() if override and override.strip() + else DSH_DEFAULT_MODEL_ID) + + +def agent_framework_models(framework: str) -> Dict[str, str]: + """Label -> model id for one agent framework.""" + if framework == DSH_FRAMEWORK: + return {"deepseek-v4-flash": dsh_model_id()} + return dict(CLAUDE_MODELS) + + +def agent_framework_default_model(framework: str) -> str: + return "deepseek-v4-flash" if framework == DSH_FRAMEWORK else "Opus" + + +def resolve_agent_framework(req: Mapping[str, Any]) -> str: + """Read + validate the agent_framework answer (default ccb).""" + answers = _answers(req) + framework = str( + answers.get("agent_framework") or CCB_FRAMEWORK + ).strip().lower() + if framework not in {CCB_FRAMEWORK, DSH_FRAMEWORK}: + raise ValueError( + "agent_framework must be one of " + f"{sorted([CCB_FRAMEWORK, DSH_FRAMEWORK])}" + ) + return framework + + +def resolve_model_id( + req: Mapping[str, Any], framework: str | None = None +) -> str: + """Resolve the model id for the chosen framework. + + Prefers the ``agent_model`` answer; falls back to the legacy + ``claude_model`` answer (pre-framework requirements files) and finally to + the framework default. + """ + framework = framework or resolve_agent_framework(req) + answers = _answers(req) + models = agent_framework_models(framework) + label = str(answers.get("agent_model") or "").strip() + if not label: + label = str(answers.get("claude_model") or "").strip() + if not label: + label = agent_framework_default_model(framework) + try: + return models[label] + except KeyError as exc: + raise ValueError( + f"agent_model must be one of {sorted(models)} for " + f"framework {framework!r}" + ) from exc + + +def resolve_claude_bin(framework: str, explicit: str | None = None) -> str: + """Pick the agent binary for a framework. + + An explicit ``--claude-bin`` always wins. Otherwise dsh uses the bundled + ccb-compatible DSH wrapper (bridge/dsh/dsh_agent.py) and ccb falls back to + ``METAINFER_CLAUDE_BIN`` (or ``ccb``). + """ + if explicit: + return explicit + if framework == DSH_FRAMEWORK: + return str( + Path(__file__).resolve().parents[1] + / "bridge" / "dsh" / "dsh_agent.py" + ) + return os.environ.get("METAINFER_CLAUDE_BIN", "ccb") +# A small, stable gain may be accumulated across rounds. The user-facing +# minimum_improvement_percent is the final target versus the fixed baseline. +ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT = 1.0 + + +def _kernel_repos_root() -> Path: + """Root directory for operator kernel repositories. + + ``kernel-repos/`` sits next to the MetaInfer root dir. When a user + types a relative name in the Kernel repo field (e.g. "int8 test3"), + it resolves to ``/int8 test3/``. + """ + override = os.environ.get("METAINFER_KERNEL_REPOS") + if override: + return Path(override).expanduser().resolve() + # Derive from METAINFER_ROOT or this installed MetaInfer package. Do not + # use cwd: orchestrators run with their task state directory as cwd. + meta_root = os.environ.get("METAINFER_ROOT") + base = ( + Path(meta_root).expanduser().resolve() + if meta_root + else Path(__file__).resolve().parents[4] + ) + return (base / ".." / "kernel-repos").resolve() + + +def _kernel_repo_from_name(name: str) -> Path: + """Resolve one repository folder name below the sibling kernel-repos.""" + repo_name = name.strip() + candidate = Path(repo_name) + if ( + not repo_name + or repo_name in {".", ".."} + or candidate.is_absolute() + or len(candidate.parts) != 1 + or "/" in repo_name + or "\\" in repo_name + ): + raise ValueError( + "Kernel repo must be a single folder name under kernel-repos " + "(for example: int8 test2)" + ) + kernel_root = _kernel_repos_root() + target = (kernel_root / repo_name).resolve() + if target.parent != kernel_root: + raise ValueError("Kernel repo resolves outside kernel-repos") + return target + + +@dataclass(frozen=True) +class ShapeSpec: + id: str + params: Dict[str, Any] + + +@dataclass(frozen=True) +class WorkerAssignment: + worker_id: str + gpu: int + shape_ids: List[str] + + +@dataclass(frozen=True) +class OptimizerConfig: + operator: str + dtype: str + hardware: str + kernel_language: str + claude_model: str + execution_mode: str + target_repo_path: Path | None + shapes: Dict[str, ShapeSpec] + assignments: List[WorkerAssignment] + assignment_mode: str + shape_scope: str + mock_iterations: int + minimum_improvement_percent: float + # Agent framework (ccb | dsh) that produced claude_model; the pipeline + # uses it to pick the agent binary via resolve_claude_bin. + agent_framework: str = CCB_FRAMEWORK + + +def _answers(req: Mapping[str, Any]) -> Mapping[str, Any]: + value = req.get("answers") + return value if isinstance(value, Mapping) else req + + +def _parse_shape_config(raw: Any) -> Dict[str, Any]: + if isinstance(raw, Mapping): + parsed = dict(raw) + elif isinstance(raw, str): + try: + parsed = yaml.safe_load(raw) or {} + except yaml.YAMLError as exc: + raise ValueError(f"shape_config is not valid YAML: {exc}") from exc + else: + raise ValueError("shape_config must be YAML text or an object") + if not isinstance(parsed, dict): + raise ValueError("shape_config must decode to an object") + # Backward compatibility for the first UI prototype, which nested full + # shape objects under each worker: + # + # workers: + # worker_0: {gpu: 0, shapes: [{id: m2, M: 2, ...}]} + # + # Normalize it to the current canonical shapes + assignments schema. + legacy_workers = parsed.get("workers") + if ( + "shapes" not in parsed + and "assignments" not in parsed + and isinstance(legacy_workers, Mapping) + ): + shapes: list[Dict[str, Any]] = [] + assignments: Dict[str, Any] = {} + for worker_id, worker in legacy_workers.items(): + if not isinstance(worker, Mapping): + raise ValueError(f"worker {worker_id!r} must be an object") + shape_ids: list[str] = [] + for shape in worker.get("shapes") or []: + if not isinstance(shape, Mapping): + raise ValueError( + f"{worker_id}: legacy shapes must be full objects" + ) + item = dict(shape) + shape_id = str(item.get("id") or "").strip() + if not shape_id: + raise ValueError( + f"{worker_id}: every legacy shape requires an id" + ) + shapes.append(item) + shape_ids.append(shape_id) + assignments[str(worker_id)] = { + "gpu": worker.get("gpu", -1), + "shapes": shape_ids, + } + parsed = {"shapes": shapes, "assignments": assignments} + return parsed + + +def load_config(req: Mapping[str, Any]) -> OptimizerConfig: + answers = _answers(req) + agent_framework = resolve_agent_framework(req) + claude_model = resolve_model_id(req, agent_framework) + mode = str(answers.get("execution_mode", MOCK_MODE)) + if mode not in { + MOCK_MODE, + LEGACY_SMOKE_MODE, + SMOKE_MODE, + W8A8_MODE, + GEN_AND_OPT_MODE, + }: + raise ValueError("unsupported execution_mode") + + target_raw = str(answers.get("target_repo_path") or "").strip() + target = Path(target_raw).expanduser() if target_raw else None + + # Smoke modes don't consume an external operator repository. + if mode in {LEGACY_SMOKE_MODE, SMOKE_MODE}: + target = None + elif mode == GEN_AND_OPT_MODE: + # Generate mode always owns a concrete repository under the + # kernel-repos directory next to MetaInfer. The user-facing field is + # a folder name, not an arbitrary filesystem path. + repo_name = target_raw or str(req.get("task_id") or "generated-kernel") + target = _kernel_repo_from_name(repo_name) + elif target is not None: + if target.is_absolute(): + # Absolute path: use as-is if it exists. + if not target.is_dir(): + target = None + else: + target = _kernel_repo_from_name(target_raw) + + raw_shape_config = answers.get("shape_config") + using_api_defaults = False + if ( + raw_shape_config is None + or ( + isinstance(raw_shape_config, str) + and not raw_shape_config.strip() + ) + ): + using_api_defaults = True + from .api_contracts import ( + default_optimization_shapes, + resolve_operator_api, + ) + + contract = resolve_operator_api( + str(answers.get("operator") or "Custom operator"), + str(answers.get("dtype") or "Other"), + ) + raw_shape_config = { + "shapes": default_optimization_shapes(contract), + } + parsed = _parse_shape_config(raw_shape_config) + shape_scope = str( + parsed.get("shape_scope") + or ("all" if using_api_defaults else "custom") + ).strip().lower() + if shape_scope not in {"all", "subset", "custom"}: + raise ValueError( + "shape_config.shape_scope must be all, subset or custom" + ) + raw_shapes = parsed.get("shapes") + if not isinstance(raw_shapes, list) or not raw_shapes: + raise ValueError("shape_config.shapes must be a non-empty list") + + shapes: Dict[str, ShapeSpec] = {} + for item in raw_shapes: + if not isinstance(item, Mapping): + raise ValueError("every shape must be an object") + shape_id = str(item.get("id") or "").strip() + if not shape_id: + raise ValueError("every shape requires a non-empty id") + if shape_id in shapes: + raise ValueError(f"duplicate shape id: {shape_id}") + shapes[shape_id] = ShapeSpec( + id=shape_id, + params={str(k): v for k, v in item.items() if k != "id"}, + ) + + raw_assignments = parsed.get("assignments") + assignment_mode = str( + parsed.get("assignment_mode") + or ("manual" if raw_assignments else "ai") + ).strip().lower() + if assignment_mode not in {"ai", "manual"}: + raise ValueError("shape_config.assignment_mode must be ai or manual") + assignments_omitted = ( + assignment_mode == "ai" + and (mode == GEN_AND_OPT_MODE or using_api_defaults) + and (not isinstance(raw_assignments, Mapping) or not raw_assignments) + ) + if not assignments_omitted: + if not isinstance(raw_assignments, Mapping) or not raw_assignments: + raise ValueError("shape_config.assignments must be a non-empty object") + if isinstance(raw_assignments, Mapping) and len(raw_assignments) > 4: + raise ValueError("worker29 MVP supports at most four workers") + + assignments: List[WorkerAssignment] = [] + seen_gpus: set[int] = set() + assigned_shapes: set[str] = set() + + if assignments_omitted: + # API-default workloads do not carry machine-specific GPU placement. + # Start with one deterministic worker. Generate mode replaces this + # placeholder with the generate agent's validated proposal. + assignments.append( + WorkerAssignment("worker_0", 0, list(shapes.keys())) + ) + assigned_shapes.update(shapes.keys()) + else: + assert isinstance(raw_assignments, Mapping) + for worker_id, item in raw_assignments.items(): + if not isinstance(item, Mapping): + raise ValueError(f"assignment {worker_id!r} must be an object") + gpu = int(item.get("gpu", -1)) + if gpu not in range(4): + raise ValueError(f"{worker_id}: gpu must be one of 0,1,2,3") + if gpu in seen_gpus: + raise ValueError(f"physical GPU {gpu} is assigned more than once") + seen_gpus.add(gpu) + shape_ids = [str(v) for v in (item.get("shapes") or [])] + if not shape_ids: + raise ValueError(f"{worker_id}: shapes must not be empty") + unknown = [sid for sid in shape_ids if sid not in shapes] + if unknown: + raise ValueError(f"{worker_id}: unknown shapes: {unknown}") + duplicate = [sid for sid in shape_ids if sid in assigned_shapes] + if duplicate: + raise ValueError(f"shapes assigned more than once: {duplicate}") + assigned_shapes.update(shape_ids) + assignments.append( + WorkerAssignment(str(worker_id), gpu, shape_ids) + ) + + missing = sorted(set(shapes) - assigned_shapes) + if missing: + raise ValueError(f"unassigned shapes: {missing}") + + try: + iterations = int( + answers.get("max_iterations", answers.get("mock_iterations", 3)) + ) + except (TypeError, ValueError) as exc: + raise ValueError("mock_iterations must be an integer") from exc + if iterations < 1 or iterations > 50: + raise ValueError("mock_iterations must be between 1 and 50") + + try: + threshold = float(answers.get("minimum_improvement_percent", 1.0)) + except (TypeError, ValueError) as exc: + raise ValueError("minimum_improvement_percent must be numeric") from exc + if threshold < 0: + raise ValueError("minimum_improvement_percent must be non-negative") + + return OptimizerConfig( + operator=str(answers.get("operator") or "Custom operator"), + dtype=str(answers.get("dtype") or "Other"), + hardware=str(answers.get("target_hardware") or "Other DCU"), + kernel_language=str(answers.get("kernel_language") or "HIP C++"), + claude_model=claude_model, + execution_mode=mode, + target_repo_path=target, + shapes=shapes, + assignments=assignments, + assignment_mode=assignment_mode, + shape_scope=shape_scope, + mock_iterations=iterations, + minimum_improvement_percent=threshold, + agent_framework=agent_framework, + ) + + +def replace_assignments( + config: OptimizerConfig, + assignments: List[WorkerAssignment], +) -> OptimizerConfig: + """Return a new OptimizerConfig with the given assignments.""" + return OptimizerConfig( + operator=config.operator, + dtype=config.dtype, + hardware=config.hardware, + kernel_language=config.kernel_language, + claude_model=config.claude_model, + execution_mode=config.execution_mode, + target_repo_path=config.target_repo_path, + shapes=config.shapes, + assignments=assignments, + assignment_mode=config.assignment_mode, + shape_scope=config.shape_scope, + mock_iterations=config.mock_iterations, + minimum_improvement_percent=config.minimum_improvement_percent, + agent_framework=config.agent_framework, + ) + + +def validate_gpu_assignment( + shapes: Dict[str, ShapeSpec], + raw: Dict[str, Any], +) -> List[WorkerAssignment]: + """Validate and parse a GPU assignment dict from the generate agent. + + Args: + shapes: The known shape specs. + raw: A dict mapping worker_id → {gpu: int, shapes: [str, ...]}. + + Returns: + A validated list of WorkerAssignment. + """ + if not isinstance(raw, dict) or not raw: + raise ValueError("gpu_assignment must be a non-empty object") + if len(raw) > 4: + raise ValueError("at most four workers supported") + + assignments: List[WorkerAssignment] = [] + seen_gpus: set[int] = set() + assigned_shapes: set[str] = set() + + for worker_id, item in raw.items(): + if not isinstance(item, dict): + raise ValueError(f"{worker_id}: must be an object") + gpu = int(item.get("gpu", -1)) + if gpu not in range(4): + raise ValueError(f"{worker_id}: gpu must be 0-3, got {gpu}") + if gpu in seen_gpus: + raise ValueError(f"GPU {gpu} assigned to more than one worker") + seen_gpus.add(gpu) + + shape_ids = [str(s) for s in (item.get("shapes") or [])] + if not shape_ids: + raise ValueError(f"{worker_id}: shapes must not be empty") + unknown = [sid for sid in shape_ids if sid not in shapes] + if unknown: + raise ValueError(f"{worker_id}: unknown shapes: {unknown}") + duplicate = [sid for sid in shape_ids if sid in assigned_shapes] + if duplicate: + raise ValueError(f"shapes assigned more than once: {duplicate}") + assigned_shapes.update(shape_ids) + + assignments.append( + WorkerAssignment(str(worker_id), gpu, shape_ids) + ) + + missing = sorted(set(shapes) - assigned_shapes) + if missing: + raise ValueError(f"unassigned shapes: {missing}") + + return assignments diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/experience_store.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/experience_store.py new file mode 100644 index 00000000..b5afc07f --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/experience_store.py @@ -0,0 +1,96 @@ +"""Cross-task experience derived only from trusted iteration records.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Dict + + +def _same_shape(record: Dict[str, Any], shape: Dict[str, Any]) -> bool: + recorded = record.get("shape") + if not isinstance(recorded, dict): + return False + try: + return all( + int(recorded[key]) == int(shape[key]) + for key in ("M", "N", "K") + ) + except (KeyError, TypeError, ValueError): + return False + + +def _classification(record: Dict[str, Any]) -> str | None: + if ( + record.get("accepted") is True + and record.get("correctness_passed") is True + ): + return "accepted_correct" + if ( + record.get("build_success") is True + and record.get("correctness_passed") is False + and isinstance(record.get("speedup"), (int, float)) + and float(record["speedup"]) > 1.0 + ): + return "faster_incorrect_repairable" + if ( + record.get("build_success") is True + and record.get("correctness_passed") is True + ): + return "measured_correct_rejection" + if record.get("build_success") is False: + return "compile_or_agent_failure" + return None + + +def load_verified_experience( + kernel_repos_root: Path | None, + shape: Dict[str, Any], + *, + exclude_repo: Path | None = None, + limit: int = 12, +) -> list[Dict[str, Any]]: + """Load exact-shape facts without trusting generated Skill prose.""" + if kernel_repos_root is None or not kernel_repos_root.is_dir(): + return [] + candidates: list[tuple[float, Path]] = [] + for path in kernel_repos_root.glob( + "*/candidates/**/iteration.json" + ): + try: + if exclude_repo is not None and path.is_relative_to(exclude_repo): + continue + candidates.append((path.stat().st_mtime, path)) + except OSError: + continue + evidence: list[Dict[str, Any]] = [] + for _, path in sorted(candidates, reverse=True): + try: + record = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + continue + if not isinstance(record, dict) or not _same_shape(record, shape): + continue + classification = _classification(record) + if classification is None: + continue + evidence.append({ + "classification": classification, + "shape_id": record.get("shape_id"), + "iteration": record.get("iteration"), + "proposed_change_not_verified_fact": record.get("hypothesis"), + "metrics": record.get("metrics") or {}, + "baseline_us": record.get("baseline_us"), + "speedup": record.get("speedup"), + "accepted": record.get("accepted"), + "build_success": record.get("build_success"), + "correctness_passed": record.get("correctness_passed"), + "record_path": str(path), + "trust_rule": ( + "Only classification and measured fields are trusted. " + "The proposed-change text may contain an Agent hypothesis." + ), + }) + if len(evidence) >= limit: + break + return evidence diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gen_and_opt_pipeline.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gen_and_opt_pipeline.py new file mode 100644 index 00000000..daa10c5a --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gen_and_opt_pipeline.py @@ -0,0 +1,2425 @@ +"""Generate + Optimize pipeline with strict coordinator/worker ownership. + +Every New Task creates a clean repository and trusted test scaffold during +GENERATE. The control plane owns shape/GPU placement and the main Agent +reviews it. Child Agents create and optimize their assigned HIP kernels +during parallel EXPLORE. +""" + +from __future__ import annotations + +import json +import os +import shutil +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import Any, Dict, List + +from metainfer.orchestrator.state import StateStore +from metainfer.orchestrator.subagent_manager import AgentSpec, SubAgentManager + +from . import phases +from .api_contracts import ( + OperatorAPIContract, + W8A8_API_FILENAME, + W8A8_BACKEND_FILENAME, + W8A8_VARIANTS_RELATIVE, + default_optimization_shapes, + file_digest, + resolve_operator_api, + stage_operator_api, + stage_operator_references, + validate_contract_shapes, +) +from .config import ( + GEN_AND_OPT_MODE, + ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT, + OptimizerConfig, + WorkerAssignment, + load_config, + replace_assignments, + validate_gpu_assignment, +) +from .prompts import ( + HARNESS_PATH, + bootstrap_worker_prompt, + generate_kernel_prompt, + shape_balanced_assignment, +) +from .real_pipeline import _run, _safe, _status +from .result_store import SCHEMA_VERSION, write_json +from .skill_store import generate_merged_skill, generate_worker_skill +from .w8a8_pipeline import ( + RealW8A8OptimizationPipeline, + W8A8Runner, + _SOURCE_ONLY_AGENT_ARGS, + _sha256_file, + evaluate_final_target, + snapshot_accepted_kernel_artifact, +) +from .w8a8_baselines import fixed_triton_graph_baseline + + +_MAX_GENERATE_RETRIES = 3 +_COORDINATOR_AGENT_ARGS = ["--tools", "Read,Glob,Grep,Write"] +_MAX_BOOTSTRAP_RETRIES = 3 +_MAX_SYNTHESIS_RETRIES = 3 +# µs-scale decode kernels are one-sided noisy on shared GPUs: transient +# co-tenant load / clock / thermal state only inflates the median, never +# deflates it. When the final performance gate (final <= 1.05 x worker best) +# trips with no source change, re-measure up to _PERF_GATE_MAX_RETRIES times, +# waiting _PERF_GATE_RETRY_INTERVAL_S between attempts, and accept the best +# (min) median; only fail when every attempt still exceeds the gate. +_PERF_GATE_MAX_RETRIES = 3 +_PERF_GATE_RETRY_INTERVAL_S = 300 +_PMC_PLAN = { + "script": "profile_pmc.sh", + "mode": "hipprof_pmc_csv", + "trigger": ( + "usable DUMMA bootstrap, newly accepted official best, or " + "late-round plateau/ISA decision" + ), + "reuse_when_source_digest_matches": True, + "skip_scalar_bootstrap": True, + "acceptance_timing": ( + "unprofiled_cuda_graph_replay_median_p90" + ), +} +_BASELINE_SOURCE_DIR = ( + Path(__file__).resolve().parent.parent / "assets" / "w8a8_baseline" +) +_TRUSTED_HARNESS_SOURCE = ( + Path(__file__).resolve().parent.parent / "assets" / "w8a8_bench.py" +) +_HARNESS_FILENAME = "w8a8_bench.py" + + +def _task_local_api_contract( + origin: OperatorAPIContract, + source_dir: Path, +) -> OperatorAPIContract: + """Resolve and verify the API snapshot committed into one task repo.""" + source = source_dir / origin.destination_name + if not source.is_file(): + raise RuntimeError(f"task-local API contract is missing: {source}") + manifest_path = source_dir / "scaffold_manifest.json" + if not manifest_path.is_file(): + raise RuntimeError( + f"task scaffold manifest is missing: {manifest_path}" + ) + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + expected = manifest["control_plane_files"][origin.destination_name] + except (KeyError, TypeError, ValueError, OSError) as exc: + raise RuntimeError( + f"task scaffold manifest has no API digest: {manifest_path}" + ) from exc + actual = file_digest(source) + if not isinstance(expected, str) or actual != expected: + raise RuntimeError( + "task-local API contract digest mismatch: " + f"expected={expected}, actual={actual}, source={source}" + ) + references = tuple( + candidate + for item in origin.reference_sources + for candidate in (source_dir / "references" / item.name,) + if candidate.is_file() + ) + return OperatorAPIContract( + operator=origin.operator, + dtype=origin.dtype, + source=source, + destination_name=origin.destination_name, + reference_sources=references, + ) +_SCAFFOLD_FILES = ( + W8A8_BACKEND_FILENAME, + "setup.py", + "csrc/bindings.cpp", + "profile_pmc.sh", + "w8a8_graph.py", + _HARNESS_FILENAME, +) +_GENERATED_KERNEL_FILE = "csrc/w8a8_gemm_hip.hip" +_GIT_ATTRIBUTES_FILE = ".gitattributes" +_GIT_ATTRIBUTES = """\ +* text=auto eol=lf +*.sh text eol=lf +*.py text eol=lf +*.cpp text eol=lf +*.hip text eol=lf +*.json text eol=lf +""" +_CONTROL_PLANE_GENERATED_FILES = frozenset({ + "csrc/bindings_hip.cpp", +}) + + +def _is_control_plane_artifact(path: str) -> bool: + """Return whether *path* is generated by Python/HIPify, not an Agent.""" + normalized = path.replace("\\", "/") + return ( + normalized in _CONTROL_PLANE_GENERATED_FILES + or normalized.endswith(".pyc") + or normalized.startswith("__pycache__/") + or "/__pycache__/" in normalized + ) + + +def _stage_trusted_baseline(destination: Path) -> list[str]: + """Install control-plane-owned loader/build scaffolding, never HIP code. + + A brand-new task must visibly generate its kernel. Only a future explicit + continuation mode may start from an existing HIP implementation. + """ + staged: list[str] = [] + for relative in _SCAFFOLD_FILES: + source = ( + _TRUSTED_HARNESS_SOURCE + if relative == _HARNESS_FILENAME + else _BASELINE_SOURCE_DIR / relative + ) + target = destination / relative + if target.exists(): + continue + target.parent.mkdir(parents=True, exist_ok=True) + # copyfile intentionally does not preserve the asset timestamp. The + # destination is a newly-created task repository, not a checkout or a + # continuation of the asset (or of any previous task). + shutil.copyfile(source, target) + staged.append(relative) + return staged + + +def _last_json_object(text: str) -> Dict[str, Any]: + for line in reversed(text.splitlines()): + try: + value = json.loads(line) + except ValueError: + continue + if isinstance(value, dict): + return value + raise ValueError(f"no JSON object in output: {text[-1000:]}") + + +def _validate_generate_scaffold( + source: Path, + *, + contract_sha256: str, + shapes: Dict[str, Dict[str, Any]], +) -> Dict[str, Any]: + """Run trusted, implementation-free Generate preflight checks.""" + kernel_path = source / _GENERATED_KERNEL_FILE + if kernel_path.exists(): + raise RuntimeError( + "Generate scaffold unexpectedly contains a HIP implementation" + ) + required = [ + source / _GIT_ATTRIBUTES_FILE, + source / W8A8_API_FILENAME, + *[source / relative for relative in _SCAFFOLD_FILES], + source / "scaffold_manifest.json", + ] + missing = [ + str(path.relative_to(source)) + for path in required + if not path.is_file() + ] + if missing: + raise RuntimeError( + f"Generate scaffold is missing trusted files: {missing}" + ) + actual_contract_sha256 = file_digest(source / W8A8_API_FILENAME) + if actual_contract_sha256 != contract_sha256: + raise RuntimeError("staged API contract digest is not trusted") + + harness = source / _HARNESS_FILENAME + self_test_result = _run( + ["python3", str(harness), "--self-test"], + cwd=source, + timeout=120, + ) + self_test = _last_json_object(self_test_result.stdout) + if self_test.get("passed") is not True: + raise RuntimeError( + f"trusted PyTorch reference self-test failed: {self_test}" + ) + + probe_env = dict(os.environ) + probe_env.update({ + "HIP_VISIBLE_DEVICES": "0", + "ROCR_VISIBLE_DEVICES": "0", + "PYTHONDONTWRITEBYTECODE": "1", + }) + probe_result = _run( + [ + "python3", str(harness), + "--source", str(source), + "--m", "1", "--n", "1", "--k", "1", + "--probe", + ], + cwd=source, + env=probe_env, + timeout=120, + ) + gpu_probe = _last_json_object(probe_result.stdout) + if gpu_probe.get("visible_devices") != 1: + raise RuntimeError( + f"Generate GPU probe failed: {gpu_probe}" + ) + if gpu_probe.get("cudagraph_available") is not True: + raise RuntimeError( + f"Generate CUDA/HIP Graph probe failed: {gpu_probe}" + ) + + profile_script = source / "profile_pmc.sh" + _run(["bash", "-n", str(profile_script)], cwd=source, timeout=30) + profile_text = profile_script.read_text(encoding="utf-8") + required_pmc_tokens = ( + "/opt/dtk/bin/hipprof", + "--pmc", + "--pmc-type 3", + '"$source_dir"', + '"$output_dir', + ) + missing_tokens = [ + token for token in required_pmc_tokens + if token not in profile_text + ] + hipprof = Path("/opt/dtk/bin/hipprof") + if missing_tokens or not hipprof.is_file() or not os.access( + hipprof, os.X_OK + ): + raise RuntimeError( + "trusted PMC entry point validation failed: " + f"missing_tokens={missing_tokens}, hipprof={hipprof}" + ) + + return { + "schema_version": SCHEMA_VERSION, + "status": "passed", + "implementation_present": False, + "api_contract_sha256": actual_contract_sha256, + "shape_ids": sorted(shapes), + "harness": { + "path": _HARNESS_FILENAME, + "sha256": file_digest(harness), + "reference_self_test": self_test, + }, + "gpu_probe": gpu_probe, + "graph": { + "required": True, + "python_api": "torch.cuda.CUDAGraph", + "wrapper": "w8a8_graph.py", + "capture_execution_deferred_until_child_kernel": True, + }, + "pmc": { + "script": "profile_pmc.sh", + "sha256": file_digest(profile_script), + "bash_syntax_passed": True, + "hipprof_path": str(hipprof), + "hipprof_executable": True, + "command_tokens_verified": list(required_pmc_tokens), + "profile_execution_deferred_until_child_kernel": True, + }, + "timestamp": time.time(), + } + + +def _match_tree_owner(path: Path, owner_source: Path) -> None: + """Make bind-mounted agent files writable by the host workspace owner.""" + if os.geteuid() != 0: + return + owner = owner_source.stat() + for root, dirs, files in os.walk(path): + os.chown(root, owner.st_uid, owner.st_gid) + for name in dirs: + os.chown(Path(root) / name, owner.st_uid, owner.st_gid) + for name in files: + os.chown(Path(root) / name, owner.st_uid, owner.st_gid) + + +def _require_valid_child_assignments( + assignments: List[WorkerAssignment], +) -> None: + """Enforce one-to-four workers whose IDs match their physical GPUs.""" + if not 1 <= len(assignments) <= 4: + raise ValueError( + "gpu_assignment must use between one and four workers" + ) + actual = { + item.worker_id: item.gpu for item in assignments + } + invalid = { + worker_id: gpu for worker_id, gpu in actual.items() + if worker_id != f"worker_{gpu}" + } + if invalid: + raise ValueError( + "gpu_assignment must map each worker_N to physical GPU N; " + f"invalid mappings: {invalid}" + ) + + +def _final_synthesis_prompt( + *, + worker_inputs: list[Dict[str, Any]], + shapes: list[Dict[str, Any]], + source: Path, + proposal_path: Path, + previous_failure: str | None, + fallback_shapes: list[Dict[str, Any]] | None = None, +) -> str: + failure_block = ( + "\nPrevious trusted synthesis failure:\n" + f"{previous_failure}\n" + if previous_failure else "" + ) + fallback_block = "" + if fallback_shapes: + fallback_block = f""" +The following API shapes are outside this task's optimization scope. Preserve +the generic fallback for them; the trusted control plane will run correctness +regression checks before publishing: +{json.dumps(fallback_shapes, indent=2)} +""" + return f"""You are the FINAL W8A8 HIP synthesis agent. + +The control plane has four independently measured worker branches: +{json.dumps(worker_inputs, indent=2)} + +Build one deployable implementation in `{source}`. Read each worker's +`csrc/w8a8_gemm_hip.hip`, preserve the trusted loader/bindings, and merge only +shape-specific kernel/launch dispatch that is supported by measured worker +results. The final `torch.ops.zth_w8a8.gemm_out` must support every shape below +through one compiled extension: +{json.dumps(shapes, indent=2)} +{fallback_block} + +Hard constraints: +- edit only `csrc/w8a8_gemm_hip.hip` and `proposal.json`; +- keep the generic scalar path as a correctness fallback; +- keep M=2/M=16 variants of each operator in one explicit dispatch family; +- wavefront=64, blockDim a multiple of 64, current stream only; +- DUMMA INT8 is m16n16k32 with int32 accumulation on gfx928; installed DTK + uses `` and namespace `du::dumma`; +- final per-shape median must stay within 5% of that worker's measured best; +- do not run Docker, the harness, compilation, benchmarks, or environment + probes; trusted control plane validation runs after you return. +{failure_block} +Write strict JSON to `{proposal_path}` with `hypothesis`, `workers_merged`, +`dispatch_families`, and `files_changed`. +""" + + +def _artifact_symbol_prefix(shape_id: str) -> str: + safe_symbol = "".join( + char if char.isalnum() or char == "_" else "_" + for char in str(shape_id) + ) + return f"mi_{safe_symbol}_" + + +def _namespace_prebuilt_object( + source_object: Path, + destination_object: Path, + shape_id: str, +) -> Dict[str, Any]: + """Give every defined host symbol a shape namespace without recompiling.""" + destination_object.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source_object, destination_object) + nm_output = _run( + ["nm", "-g", "--defined-only", str(destination_object)], + cwd=destination_object.parent, + ).stdout + symbols: list[str] = [] + for line in nm_output.splitlines(): + fields = line.split() + if len(fields) >= 2: + symbols.append(fields[-1]) + symbols = sorted(set(symbols)) + required = {"launch_w8a8_gemm", "launch_pack_w8a8_weight"} + missing = sorted(required - set(symbols)) + if missing: + raise RuntimeError( + f"accepted object for {shape_id} is missing ABI symbols: " + f"{missing}" + ) + prefix = _artifact_symbol_prefix(shape_id) + mapping_path = destination_object.with_suffix(".symbols") + mapping_path.write_text( + "".join(f"{symbol} {prefix}{symbol}\n" for symbol in symbols), + encoding="utf-8", + ) + objcopy = next( + ( + candidate for candidate in ( + shutil.which("llvm-objcopy"), + "/opt/dtk/aillvm/bin/llvm-objcopy", + "/opt/dtk/dcc/bin/llvm-objcopy", + ) + if candidate and Path(candidate).is_file() + ), + None, + ) + if objcopy is None: + raise RuntimeError("llvm-objcopy is required for artifact linking") + _run([ + str(objcopy), + f"--redefine-syms={mapping_path}", + str(destination_object), + ], cwd=destination_object.parent) + mapping_path.unlink(missing_ok=True) + return { + "symbol_prefix": prefix, + "launch_symbol": f"{prefix}launch_w8a8_gemm", + "pack_symbol": f"{prefix}launch_pack_w8a8_weight", + "defined_symbols_namespaced": len(symbols), + "linked_object_sha256": _sha256_file(destination_object), + } + + +def _render_prebuilt_dispatch( + artifacts: list[Dict[str, Any]], +) -> str: + """Render control-plane C++ dispatch; this file contains no HIP kernel.""" + declarations: list[str] = [] + gemm_routes: list[str] = [] + pack_routes: list[str] = [] + seen_pairs: set[tuple[int, int]] = set() + for artifact in artifacts: + launch = artifact["launch_symbol"] + pack = artifact["pack_symbol"] + shape = artifact["shape"] + n = int(shape["N"]) + k = int(shape["K"]) + declarations.extend([ + f'extern "C" void {launch}(', + " const int8_t*, const int8_t*, const float*, const float*,", + " void*, void*, int64_t, int, int, int, hipStream_t);", + f'extern "C" void {pack}(', + " const int8_t*, const float*, int8_t*, float*,", + " int, int, hipStream_t);", + "", + ]) + pair = (n, k) + if pair in seen_pairs: + continue + seen_pairs.add(pair) + gemm_routes.extend([ + f" if (n == {n} && k == {k}) {{", + f" {launch}(a, b, x_scale, weight_scale, out, workspace,", + " workspace_bytes, m, n, k, stream);", + " return;", + " }", + ]) + pack_routes.extend([ + f" if (n == {n} && k == {k}) {{", + f" {pack}(raw_weight, weight_scale, packed_weight,", + " packed_weight_scale, k, n, stream);", + " return;", + " }", + ]) + if not artifacts: + raise RuntimeError("cannot render dispatch without worker artifacts") + fallback_launch = artifacts[0]["launch_symbol"] + fallback_pack = artifacts[0]["pack_symbol"] + return "\n".join([ + "// Generated by the trusted control plane. Contains no HIP kernel.", + "#include ", + "#include ", + "", + *declarations, + 'extern "C" void launch_w8a8_gemm(', + " const int8_t* a, const int8_t* b,", + " const float* x_scale, const float* weight_scale,", + " void* out, void* workspace, int64_t workspace_bytes,", + " int m, int n, int k, hipStream_t stream) {", + *gemm_routes, + f" {fallback_launch}(a, b, x_scale, weight_scale, out, workspace,", + " workspace_bytes, m, n, k, stream);", + "}", + "", + 'extern "C" void launch_pack_w8a8_weight(', + " const int8_t* raw_weight, const float* weight_scale,", + " int8_t* packed_weight, float* packed_weight_scale,", + " int k, int n, hipStream_t stream) {", + *pack_routes, + f" {fallback_pack}(raw_weight, weight_scale, packed_weight,", + " packed_weight_scale, k, n, stream);", + "}", + "", + ]) + + +def _final_performance_gate( + *, + shape_id: str, + best_median: float, + metrics: Dict[str, Any], + benchmark, + max_retries: int, + retry_interval_s: float, + store: StateStore, +) -> Dict[str, Any]: + """Accept the final validation measurement for one optimized shape. + + µs-scale decode kernels are one-sided noisy on shared GPUs: transient + co-tenant load / clock / thermal state only inflates the median, never + deflates it (decode skill §5). When ``median > 1.05 x best`` with no + source change, re-measure with the same protocol up to ``max_retries`` + times, waiting ``retry_interval_s`` between attempts, and accept the best + (min) median. Raise only when every attempt still exceeds the gate. + + Returns the metrics dict of the accepted (best) measurement. + """ + final_median = float(metrics.get("median_us") or float("inf")) + accepted = metrics + retries_left = max_retries + while final_median > best_median * 1.05 and retries_left > 0: + retries_left -= 1 + attempt = max_retries - retries_left + store.append_timeline( + "final_perf_gate_retry", + { + "shape_id": shape_id, + "attempt": attempt, + "median_us": final_median, + "best_median_us": best_median, + "wait_s": retry_interval_s, + }, + ) + time.sleep(retry_interval_s) + retry = benchmark() + if not retry.get("passed"): + raise RuntimeError( + f"correctness failed for {shape_id}: {json.dumps(retry)}" + ) + accepted = retry + final_median = min( + final_median, float(retry.get("median_us") or float("inf")) + ) + if final_median > best_median * 1.05: + raise RuntimeError( + f"performance regressed for {shape_id}: final " + f"{final_median:.3f} us vs worker best {best_median:.3f} us " + f"after {max_retries} re-measures" + ) + return accepted + + +class GenAndOptPipeline(RealW8A8OptimizationPipeline): + """Coordinate a clean kernel repo, then let child Agents implement it. + + GENERATE owns repository scaffolding and GPU assignment only. Up to four + child lanes independently create, validate, and optimize HIP in EXPLORE. + """ + + # ------------------------------------------------------------------ # + # Overrides from RealW8A8OptimizationPipeline + # ------------------------------------------------------------------ # + + def _validate_contract(self, config: OptimizerConfig) -> None: + """Override: no pre-existing repo validation in generate mode.""" + if config.operator != "Quantized GEMM": + raise ValueError( + "Generate mode requires operator=Quantized GEMM" + ) + if config.dtype != "INT8 W8A8": + raise ValueError( + "Generate mode requires dtype=INT8 W8A8" + ) + for shape in config.shapes.values(): + for key in ("M", "N", "K"): + if key not in shape.params: + raise ValueError(f"{shape.id} is missing {key}") + contract = resolve_operator_api(config.operator, config.dtype) + validate_contract_shapes(contract, config.shapes) + self._validate_fixed_baselines(config) + self._operator_api_contract = contract + + def _validate_fixed_baselines(self, config: OptimizerConfig) -> None: + """Fail in PREPARE before creating repos or launching GPU workers.""" + missing: list[str] = [] + for shape_id, shape in config.shapes.items(): + try: + fixed_triton_graph_baseline(shape_id, shape.params) + except ValueError as exc: + missing.append(str(exc)) + if missing: + details = "\n- ".join(missing) + raise ValueError( + "fixed Triton Graph baseline coverage is incomplete; " + "measure and freeze every requested shape before starting " + f"GPU workers:\n- {details}" + ) + + def _prepare_worktrees( + self, config: OptimizerConfig, task_id: str + ) -> None: + """Create the seed git repo, using target_repo_path if set. + + ``config.target_repo_path`` is the concrete sibling kernel-repos + directory resolved from the New Task repository name. The agent + writes directly there; workspace_dir/main is only a symlink. + """ + self.workspace_dir.mkdir(parents=True, exist_ok=True) + seed = self.workspace_dir / "main" + contract = getattr( + self, + "_operator_api_contract", + resolve_operator_api(config.operator, config.dtype), + ) + + if config.target_repo_path is not None: + # User specified a kernel-repos path — work directly there. + kernel_dir = config.target_repo_path.resolve() + task_owns_existing_repo = ( + seed.is_symlink() and seed.resolve() == kernel_dir + ) + if ( + (kernel_dir / ".git").exists() + and not task_owns_existing_repo + ): + raise RuntimeError( + f"kernel repo {kernel_dir} already exists. A new " + "optimization task cannot reuse existing HIP code; " + "choose a new repository name. Existing repositories " + "will only be accepted by the explicit continuation " + "workflow." + ) + kernel_dir.mkdir(parents=True, exist_ok=True) + if not (kernel_dir / ".git").exists(): + if any(kernel_dir.iterdir()): + raise RuntimeError( + f"requested kernel repo {kernel_dir} is non-empty " + "but is not a git repository" + ) + _run(["git", "init"], cwd=kernel_dir) + _run(["git", "config", "user.name", "MetaInfer Agent"], cwd=kernel_dir) + _run([ + "git", "config", "user.email", "metainfer@localhost", + ], cwd=kernel_dir) + (kernel_dir / "README.md").write_text( + "# Auto-generated kernel repository\n\n" + "Prepared by the MetaInfer control plane. Kernel " + "implementations are owned by child worker branches.\n", + encoding="utf-8", + ) + _run(["git", "add", "README.md"], cwd=kernel_dir) + _run([ + "git", "commit", "-m", + "empty seed for kernel generation", + ], cwd=kernel_dir) + # Create symlink so workspace/main → kernel-repos/... dir. + if seed.is_symlink(): + if seed.resolve() != kernel_dir: + raise RuntimeError( + f"workspace main already points to {seed.resolve()}, " + f"not requested kernel repo {kernel_dir}" + ) + elif seed.exists(): + raise RuntimeError( + f"workspace main exists and is not a symlink to " + f"requested kernel repo {kernel_dir}" + ) + else: + # A relative link resolves correctly both inside the container + # (/workspace/...) and on the bind-mounted host checkout. + relative_target = os.path.relpath( + kernel_dir, start=seed.parent + ) + seed.symlink_to(relative_target, target_is_directory=True) + + elif not (seed / ".git").exists(): + seed.mkdir(parents=True, exist_ok=True) + _run(["git", "init"], cwd=seed) + _run(["git", "config", "user.name", "MetaInfer Agent"], cwd=seed) + _run([ + "git", "config", "user.email", "metainfer@localhost", + ], cwd=seed) + (seed / "README.md").write_text( + "# Auto-generated kernel repository\n\n" + "Prepared by the MetaInfer control plane. Kernel " + "implementations are owned by child worker branches.\n", + encoding="utf-8", + ) + _run(["git", "add", "README.md"], cwd=seed) + _run([ + "git", "commit", "-m", "empty seed for kernel generation", + ], cwd=seed) + + attributes_path = seed / _GIT_ATTRIBUTES_FILE + if attributes_path.exists(): + if attributes_path.read_text(encoding="utf-8") != _GIT_ATTRIBUTES: + raise RuntimeError( + "fresh task repository has unexpected .gitattributes" + ) + else: + attributes_path.write_text(_GIT_ATTRIBUTES, encoding="utf-8") + _run(["git", "add", _GIT_ATTRIBUTES_FILE], cwd=seed) + if _run( + ["git", "diff", "--cached", "--name-only"], cwd=seed + ).stdout.strip(): + _run([ + "git", "commit", "-m", + "enforce LF in generated Git worktrees", + ], cwd=seed) + + staged_contract = stage_operator_api(contract, seed) + staged_references = stage_operator_references(contract, seed) + staged_operator_files = [staged_contract, *staged_references] + staged_operator_relatives = [ + str(path.relative_to(seed)) for path in staged_operator_files + ] + if not _run( + ["git", "status", "--short", "--", *staged_operator_relatives], + cwd=seed, + ).stdout.strip(): + pass + else: + _run(["git", "add", *staged_operator_relatives], cwd=seed) + _run([ + "git", "commit", "-m", + "stage immutable operator API and optional references", + ], cwd=seed) + + scaffold_files = _stage_trusted_baseline(seed) + if scaffold_files: + _run(["git", "add", *scaffold_files], cwd=seed) + _run([ + "git", "commit", "-m", + "stage W8A8 build and loader scaffolding", + ], cwd=seed) + + scaffold_manifest = { + "schema_version": SCHEMA_VERSION, + "task_id": task_id, + "fresh_repository": True, + "continuation": False, + "implementation_inherited": False, + "control_plane_files": { + _GIT_ATTRIBUTES_FILE: file_digest(attributes_path), + staged_contract.name: file_digest(staged_contract), + **{ + str(path.relative_to(seed)): file_digest(path) + for path in staged_references + }, + **{ + relative: file_digest(seed / relative) + for relative in _SCAFFOLD_FILES + }, + }, + "initial_kernel": "pending_parallel_explore_child_generation", + "created_at": time.time(), + } + write_json(seed / "scaffold_manifest.json", scaffold_manifest) + _run(["git", "add", "scaffold_manifest.json"], cwd=seed) + if _run( + ["git", "diff", "--cached", "--name-only"], cwd=seed + ).stdout.strip(): + _run([ + "git", "commit", "-m", + "record fresh task scaffold provenance", + ], cwd=seed) + self._task_api_contract = _task_local_api_contract(contract, seed) + self.store.append_timeline( + "fresh_repository_created", + { + "repository": str(seed.resolve()), + "task_id": task_id, + "implementation_inherited": False, + "scaffold_manifest": str( + seed.resolve() / "scaffold_manifest.json" + ), + }, + ) + + if config.target_repo_path is not None: + _match_tree_owner( + seed.resolve(), config.target_repo_path.parent + ) + + # Shared directories — worker directories come later. + for name in ("shared_baseline", "final_validation", "skills"): + (self.workspace_dir / name).mkdir(parents=True, exist_ok=True) + + # ------------------------------------------------------------------ # + # Generate phase + # ------------------------------------------------------------------ # + + def _generate_kernel_repo( + self, config: OptimizerConfig + ) -> List[WorkerAssignment]: + """Resolve shape/GPU ownership without creating any HIP kernel.""" + seed = self.workspace_dir / "main" + agent_source_dir = seed.resolve() + harness_path = agent_source_dir / _HARNESS_FILENAME + origin_contract = getattr( + self, + "_operator_api_contract", + resolve_operator_api(config.operator, config.dtype), + ) + contract = getattr( + self, + "_task_api_contract", + _task_local_api_contract(origin_contract, seed), + ) + contract_path = agent_source_dir / contract.destination_name + trusted_contract_digest = file_digest(contract.source) + reference_paths = [ + agent_source_dir / "references" / source.name + for source in contract.reference_sources + ] + trusted_reference_digests = { + str(path.relative_to(agent_source_dir)): file_digest(path) + for path in reference_paths + } + + shapes_for_prompt: Dict[str, Dict[str, Any]] = { + sid: shape.params for sid, shape in config.shapes.items() + } + fixed_assignment: Dict[str, Dict[str, Any]] + if config.assignment_mode == "manual": + actual = { + item.worker_id: item.gpu for item in config.assignments + } + expected = { + f"worker_{item.gpu}": item.gpu + for item in config.assignments + } + if actual != expected: + raise ValueError( + "manual assignments must map worker_N to GPU N; " + f"got {actual}" + ) + fixed_assignment = { + item.worker_id: { + "gpu": item.gpu, + "shapes": list(item.shape_ids), + } + for item in config.assignments + } + self.store.append_timeline( + "generate_manual_assignment", + { + "role": "user", + "gpu_assignment": fixed_assignment, + "main_agent_review_required": True, + }, + ) + else: + fixed_assignment = shape_balanced_assignment( + shapes_for_prompt + ) + self.store.append_timeline( + "generate_automatic_assignment", + { + "role": "control_plane", + "strategy": "exact_shape_lpt_by_2mnk", + "gpu_assignment": fixed_assignment, + "main_agent_review_required": True, + }, + ) + + # Assignment validity is deterministic and must fail before launching + # an Agent. The coordinator reviews this source of truth; it does not + # invent a second assignment and burn retries on policy disagreements. + validated_assignment = validate_gpu_assignment( + config.shapes, fixed_assignment + ) + _require_valid_child_assignments(validated_assignment) + + preflight = _validate_generate_scaffold( + agent_source_dir, + contract_sha256=trusted_contract_digest, + shapes=shapes_for_prompt, + ) + preflight_path = seed / "generation_preflight.json" + write_json(preflight_path, preflight) + _run( + ["git", "add", "generation_preflight.json"], + cwd=seed, + ) + if _run( + ["git", "diff", "--cached", "--name-only"], cwd=seed + ).stdout.strip(): + _run([ + "git", "commit", "-m", + "validate trusted Generate scaffold", + ], cwd=seed) + preflight_digest = file_digest(preflight_path) + self.store.append_timeline( + "generate_scaffold_validated", + { + "path": str(preflight_path.resolve()), + "sha256": preflight_digest, + "harness_self_test_passed": True, + "gpu_probe_passed": True, + "cudagraph_available": True, + "python_graph_wrapper_staged": True, + "pmc_script_checked": True, + "hipprof_executable": True, + "kernel_source_created": False, + }, + ) + failure_reason: str | None = None + + for attempt in range(1, _MAX_GENERATE_RETRIES + 1): + self.store.append_timeline( + "generate_attempt", + {"attempt": attempt, "max": _MAX_GENERATE_RETRIES, + "prev_failure": failure_reason}, + ) + + prompt = generate_kernel_prompt( + operator=config.operator, + dtype=config.dtype, + shapes=shapes_for_prompt, + hardware=config.hardware, + kernel_language=config.kernel_language, + source_dir=agent_source_dir, + harness_path=harness_path, + api_contract_path=contract_path, + iteration=attempt, + prev_failure=failure_reason, + fixed_assignment=fixed_assignment, + ) + prompt_file = seed / ".metainfer-generate.prompt.txt" + prompt_file.write_text(prompt, encoding="utf-8") + + agent_name = f"kernel-coordinator-attempt{attempt}" + spec = AgentSpec( + name=agent_name, + role="kernel_coordinator", + prompt_file=prompt_file, + workdir=seed, + log_dir=seed / ".metainfer-generate-logs", + timeout_s=1800, + stuck_timeout_s=600, + max_retries=0, + extra_args=list(_COORDINATOR_AGENT_ARGS), + ) + self.store.append_timeline( + "agent_launch", + {"name": agent_name, "role": "kernel_coordinator", + "attempt": attempt}, + ) + self.manager.launch(spec) + agent_result = self.manager.result(agent_name) + + if agent_result is None or not agent_result.success: + failure_reason = ( + f"Generate agent attempt {attempt} failed: " + f"{agent_result.error if agent_result else 'no result'}" + ) + self.store.append_timeline( + "generate_agent_failed", + {"attempt": attempt, "error": failure_reason}, + ) + continue + + if file_digest(contract_path) != trusted_contract_digest: + _run([ + "git", "restore", "--source=HEAD", "--", + contract.destination_name, + ], cwd=seed) + failure_reason = ( + f"Generate agent attempt {attempt} modified or removed " + f"immutable API contract {contract_path.name}" + ) + self.store.append_timeline( + "generate_contract_modified", + {"attempt": attempt, "path": str(contract_path)}, + ) + continue + modified_reference = next( + ( + relative for relative, digest + in trusted_reference_digests.items() + if file_digest(agent_source_dir / relative) != digest + ), + None, + ) + if modified_reference is not None: + _run([ + "git", "restore", "--source=HEAD", "--", + modified_reference, + ], cwd=seed) + failure_reason = ( + f"Generate agent modified or removed immutable optional " + f"reference {modified_reference}" + ) + self.store.append_timeline( + "generate_reference_modified", + {"attempt": attempt, "path": modified_reference}, + ) + continue + + status_paths = [] + for line in _run( + [ + "git", "status", "--porcelain", + "--untracked-files=all", + ], + cwd=seed, + ).stdout.splitlines(): + path = line[3:].strip() + if path and not path.startswith(".metainfer-"): + status_paths.append(path) + unexpected = [ + path for path in status_paths if path != "proposal.json" + ] + if unexpected: + failure_reason = ( + f"Main coordinator exceeded its role and changed files " + f"other than proposal.json: {unexpected}" + ) + self.store.append_timeline( + "generate_role_violation", + {"attempt": attempt, "paths": unexpected}, + ) + continue + + # Parse and validate GPU assignment from proposal.json. + proposal_path = seed / "proposal.json" + if not proposal_path.is_file(): + failure_reason = ( + f"Generate agent did not write proposal.json" + ) + self.store.append_timeline( + "generate_missing_proposal", + {"attempt": attempt}, + ) + continue + + try: + proposal = json.loads( + proposal_path.read_text(encoding="utf-8") + ) + except (OSError, ValueError) as exc: + failure_reason = ( + f"Cannot parse proposal.json: {exc}" + ) + self.store.append_timeline( + "generate_bad_proposal", + {"attempt": attempt, "error": str(exc)}, + ) + continue + + raw_assignment = proposal.get("gpu_assignment") + if not isinstance(raw_assignment, dict) or not raw_assignment: + failure_reason = ( + "proposal.json is missing gpu_assignment" + ) + self.store.append_timeline( + "generate_missing_gpu_assignment", + {"attempt": attempt}, + ) + continue + + review = proposal.get("scaffold_review") + required_review = { + "preflight_file": "generation_preflight.json", + "preflight_status": "passed", + "harness_reference_self_test_passed": True, + "gpu_probe_passed": True, + "cudagraph_available": True, + "python_graph_wrapper_staged": True, + "pmc_script_checked": True, + "no_hip_implementation": True, + } + if not isinstance(review, dict) or any( + review.get(key) != value + for key, value in required_review.items() + ): + failure_reason = ( + "Generate agent did not confirm the trusted scaffold " + f"preflight: expected {required_review}, got {review}" + ) + self.store.append_timeline( + "generate_scaffold_review_rejected", + {"attempt": attempt, "error": failure_reason}, + ) + continue + + try: + assignments = validate_gpu_assignment( + config.shapes, raw_assignment + ) + _require_valid_child_assignments(assignments) + except ValueError as exc: + failure_reason = ( + f"Invalid gpu_assignment: {exc}" + ) + self.store.append_timeline( + "generate_invalid_gpu_assignment", + {"attempt": attempt, "error": str(exc)}, + ) + continue + if ( + fixed_assignment is not None + and raw_assignment != fixed_assignment + ): + failure_reason = ( + "Generate agent changed the authoritative GPU assignment" + ) + self.store.append_timeline( + "generate_assignment_modified", + { + "attempt": attempt, + "expected": fixed_assignment, + "actual": raw_assignment, + }, + ) + continue + + # Persist only the coordination decision. HIP implementation + # starts later, independently, on child worker branches. + proposal["profiling_plan"] = dict(_PMC_PLAN) + proposal["assignment_source"] = ( + "manual_new_task" + if config.assignment_mode == "manual" + else "control_plane_shape_balance" + ) + proposal["generation_preflight_sha256"] = preflight_digest + write_json(seed / "coordination_plan.json", proposal) + write_json(seed / "generation_review.json", { + "schema_version": SCHEMA_VERSION, + "task_id": str(self.req.get("task_id", "task")), + "attempt": attempt, + "role": "main_coordinator", + "preflight_sha256": preflight_digest, + "scaffold_review": review, + "gpu_assignment": raw_assignment, + "kernel_source_created": False, + "timestamp": time.time(), + }) + proposal_path.unlink(missing_ok=True) + _run([ + "git", "add", + "coordination_plan.json", + "generation_review.json", + ], cwd=seed) + staged = _run( + ["git", "diff", "--cached", "--name-only"], cwd=seed + ).stdout.strip() + if staged: + _run([ + "git", "commit", "-m", + f"Coordinate W8A8 shape/GPU assignment " + f"(attempt {attempt})", + ], cwd=seed) + commit = _run( + ["git", "rev-parse", "HEAD"], cwd=seed + ).stdout.strip() + + self.store.append_timeline( + "kernel_saved", + {"path": str(seed.resolve())}, + ) + self.store.append_timeline( + "generate_success", + { + "attempt": attempt, + "role": "main_coordinator", + "kernel_source_created": False, + "kernel_source": "pending_child_generation", + "api_contract": str(contract.source), + "api_contract_sha256": trusted_contract_digest, + "generation_preflight_sha256": preflight_digest, + "generation_review": str( + seed.resolve() / "generation_review.json" + ), + "harness_self_test_passed": True, + "gpu_probe_passed": True, + "cudagraph_available": True, + "python_graph_wrapper_staged": True, + "pmc_script_checked": True, + "commit": commit, + "gpu_assignment": raw_assignment, + }, + ) + return assignments + + raise RuntimeError( + f"Kernel coordination failed after {_MAX_GENERATE_RETRIES} " + f"attempts. Last failure: {failure_reason}" + ) + + # Worker worktree creation (called after GPU assignment is known) + # ------------------------------------------------------------------ # + + def _create_worker_worktrees( + self, config: OptimizerConfig, task_id: str + ) -> None: + """Create worker worktrees and expose them from the kernel repo. + + Workers must remain isolated because they all edit the same backend + filenames. ``candidates/`` is a stable directory containing + a live ``source`` link, compatibility links such as ``csrc``, and + immutable per-round snapshots added by the optimization loop. + """ + seed = self.workspace_dir / "main" + candidate_root = seed.resolve() / "candidates" + candidate_root.mkdir(parents=True, exist_ok=True) + exclude_path = seed.resolve() / ".git" / "info" / "exclude" + exclude_path.parent.mkdir(parents=True, exist_ok=True) + exclude_text = ( + exclude_path.read_text(encoding="utf-8") + if exclude_path.is_file() else "" + ) + if "candidates/" not in exclude_text.splitlines(): + with exclude_path.open("a", encoding="utf-8") as handle: + if exclude_text and not exclude_text.endswith("\n"): + handle.write("\n") + handle.write("candidates/\n") + + for assignment in config.assignments: + root = self.workspace_dir / "workers" / assignment.worker_id + # W8A8Runner forwards these cache paths to the host-side agent. + # Create them before ownership is matched; otherwise the + # container's root user creates them later as root:root and the + # host agent exits before emitting its first stream-json event. + for name in ( + "build", + "cache", + "cache/torch", + "cache/triton", + "cache/xdg", + "cache/tmp", + "logs", + "runs", + "artifacts", + ): + (root / name).mkdir(parents=True, exist_ok=True) + source = root / "source" + if not source.exists(): + branch = f"agent/{_safe(task_id)}/{assignment.worker_id}" + _run([ + "git", "worktree", "add", "-b", branch, + str(source), "HEAD", + ], cwd=seed) + candidate_dir = candidate_root / assignment.worker_id + relative_source = os.path.relpath( + source, start=candidate_dir + ) + if candidate_dir.is_symlink(): + if candidate_dir.resolve() != source.resolve(): + raise RuntimeError( + f"candidate view {candidate_dir} points outside its " + "managed worker source" + ) + candidate_dir.unlink() + elif candidate_dir.exists() and not candidate_dir.is_dir(): + raise RuntimeError( + f"candidate view {candidate_dir} exists and is not " + "a directory" + ) + candidate_dir.mkdir(parents=True, exist_ok=True) + + source_link = candidate_dir / "source" + if source_link.is_symlink(): + if os.readlink(source_link) != relative_source: + source_link.unlink() + elif source_link.exists(): + raise RuntimeError( + f"candidate source view {source_link} is not a symlink" + ) + if not source_link.exists(): + source_link.symlink_to( + relative_source, target_is_directory=True + ) + + # Preserve the old candidates/worker_N/csrc entry while making + # room beside it for iteration1, iteration2, ... + for relative in ( + "csrc", + W8A8_API_FILENAME, + W8A8_BACKEND_FILENAME, + "setup.py", + "profile_pmc.sh", + "README.md", + ): + compatibility_link = candidate_dir / relative + target = Path("source") / relative + if compatibility_link.is_symlink(): + if os.readlink(compatibility_link) != str(target): + compatibility_link.unlink() + elif compatibility_link.exists(): + raise RuntimeError( + f"candidate compatibility view " + f"{compatibility_link} is not a symlink" + ) + if ( + not compatibility_link.exists() + and not compatibility_link.is_symlink() + ): + compatibility_link.symlink_to( + target, + target_is_directory=(relative == "csrc"), + ) + if config.target_repo_path is not None: + _match_tree_owner( + root, config.target_repo_path.parent + ) + + def _bootstrap_worker_repos( + self, config: OptimizerConfig + ) -> Dict[str, Dict[str, Any]]: + """Let each child Agent create and validate its own initial HIP code.""" + + def generate_one( + assignment: WorkerAssignment, + ) -> tuple[str, Dict[str, Dict[str, Any]]]: + root = self.workspace_dir / "workers" / assignment.worker_id + source = root / "source" + if any(char.isspace() for char in str(source.resolve())): + raise RuntimeError( + "worker source path must not contain whitespace" + ) + runner = W8A8Runner(root, assignment.gpu) + shapes = { + shape_id: config.shapes[shape_id].params + for shape_id in assignment.shape_ids + } + fixed_targets = { + shape_id: fixed_triton_graph_baseline( + shape_id, shape + ) + for shape_id, shape in shapes.items() + } + progress_path = root / "bootstrap_progress.json" + proposal_path = source / "proposal.json" + kernel_path = source / _GENERATED_KERNEL_FILE + immutable_paths = [ + source / _GIT_ATTRIBUTES_FILE, + source / W8A8_API_FILENAME, + source / W8A8_VARIANTS_RELATIVE, + *[source / relative for relative in _SCAFFOLD_FILES], + source / "coordination_plan.json", + source / "scaffold_manifest.json", + ] + immutable_digests = { + str(path.relative_to(source)): file_digest(path) + for path in immutable_paths + if path.is_file() + } + previous_failure: str | None = None + + def persist( + attempt: int, + status: str, + *, + metrics: Dict[str, Dict[str, Any]] | None = None, + hypothesis: str | None = None, + error: str | None = None, + ) -> None: + payload = { + "schema_version": SCHEMA_VERSION, + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "attempt": attempt, + "status": status, + "hypothesis": hypothesis, + "shapes": list(shapes), + "metrics": metrics or {}, + "error": error, + "source": "child_agent_generated", + "timestamp": time.time(), + } + write_json(progress_path, payload) + write_json( + root / "iterations" / "bootstrap" + / f"iteration{attempt}" / "iteration.json", + payload, + ) + + for attempt in range(1, _MAX_BOOTSTRAP_RETRIES + 1): + metrics_by_shape: Dict[str, Dict[str, Any]] = {} + paired_fallback_metrics: Dict[str, Dict[str, Any]] = {} + verified: list[str] = [] + hypothesis = "Initial child-generated HIP kernel." + proposal_path.unlink(missing_ok=True) + if kernel_path.is_file() and not _run( + ["git", "ls-files", "--", _GENERATED_KERNEL_FILE], + cwd=source, + ).stdout.strip(): + kernel_path.unlink() + _status( + root, + assignment, + state="bootstrap_agent_running", + iteration=0, + shape_id=None, + ) + persist( + attempt, + "agent_running", + hypothesis=( + "Child Agent is generating its initial HIP kernel." + ), + ) + prompt = bootstrap_worker_prompt( + worker_id=assignment.worker_id, + gpu=assignment.gpu, + shapes=shapes, + hardware=config.hardware, + kernel_language=config.kernel_language, + source_dir=source, + harness_path=HARNESS_PATH, + api_contract_path=source / W8A8_API_FILENAME, + attempt=attempt, + prev_failure=previous_failure, + ) + prompt_file = ( + root / "logs" + / f"bootstrap-attempt-{attempt}.prompt.txt" + ) + prompt_file.write_text(prompt, encoding="utf-8") + agent_name = ( + f"{assignment.worker_id}-bootstrap-attempt{attempt}" + ) + self.store.append_timeline( + "worker_bootstrap_launch", + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "attempt": attempt, + "agent": agent_name, + "source": "child_agent_generated", + }, + ) + try: + self.manager.launch(AgentSpec( + name=agent_name, + role="dcu_w8a8_bootstrap_generator", + prompt_file=prompt_file, + workdir=source, + log_dir=root / "logs", + timeout_s=600, + stuck_timeout_s=600, + max_retries=0, + extra_args=list(_SOURCE_ONLY_AGENT_ARGS), + env_overrides=runner.env, + )) + agent_result = self.manager.result(agent_name) + if agent_result is None or not agent_result.success: + raise RuntimeError( + f"{agent_name} failed: " + f"{agent_result.error if agent_result else 'no result'}" + ) + if not proposal_path.is_file(): + raise RuntimeError( + f"{agent_name} did not write proposal.json" + ) + proposal = json.loads( + proposal_path.read_text(encoding="utf-8") + ) + if not isinstance(proposal, dict): + raise RuntimeError("proposal.json must be an object") + hypothesis = str( + proposal.get("hypothesis") + or "Initial child-generated HIP kernel." + ) + for relative, digest in immutable_digests.items(): + if file_digest(source / relative) != digest: + raise RuntimeError( + "bootstrap Agent modified immutable " + f"control-plane file {relative}" + ) + changed_paths = sorted( + line[3:].strip() + for line in _run( + [ + "git", "status", "--porcelain", + "--untracked-files=all", + ], + cwd=source, + ).stdout.splitlines() + if ( + line[3:].strip() != "proposal.json" + and not _is_control_plane_artifact( + line[3:].strip() + ) + ) + ) + unexpected = [ + path for path in changed_paths + if path != _GENERATED_KERNEL_FILE + ] + if unexpected: + raise RuntimeError( + "bootstrap Agent changed files outside its HIP " + f"ownership: {unexpected}" + ) + if not kernel_path.is_file(): + raise RuntimeError( + "bootstrap Agent did not generate " + f"{_GENERATED_KERNEL_FILE}" + ) + + _status( + root, + assignment, + state="bootstrap_validating", + iteration=0, + shape_id=None, + ) + for shape_id, params in shapes.items(): + metrics = runner.benchmark(params) + metrics_by_shape[shape_id] = metrics + persist( + attempt, + "validating", + metrics=metrics_by_shape, + hypothesis=hypothesis, + ) + self.store.append_timeline( + "worker_bootstrap_shape_measured", + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "attempt": attempt, + "shape_id": shape_id, + "passed": bool(metrics.get("passed")), + "graph_capture_passed": metrics.get( + "graph_capture_passed" + ), + "timing_mode": metrics.get("timing_mode"), + "median_us": metrics.get("median_us"), + "p90_us": metrics.get("p90_us"), + }, + ) + if ( + not metrics.get("passed") + or metrics.get("graph_capture_passed") is not True + ): + raise RuntimeError( + "child-generated kernel Graph/correctness " + "validation failed " + f"for {shape_id}: {json.dumps(metrics)}" + ) + if int(params.get("M", 0)) == 16: + paired_params = {**params, "M": 2} + paired_metrics = runner.benchmark(paired_params) + paired_fallback_metrics[shape_id] = paired_metrics + if ( + not paired_metrics.get("passed") + or paired_metrics.get( + "graph_capture_passed" + ) is not True + ): + raise RuntimeError( + "child-generated kernel paired M=2 " + "fallback failed " + f"for {shape_id}: " + f"{json.dumps(paired_metrics)}" + ) + verified.append(shape_id) + + iteration_dir = ( + root / "iterations" / "bootstrap" + / f"iteration{attempt}" + ) + archived_kernel = ( + iteration_dir / _GENERATED_KERNEL_FILE + ) + archived_kernel.parent.mkdir( + parents=True, exist_ok=True + ) + shutil.copyfile(kernel_path, archived_kernel) + proposal_path.unlink(missing_ok=True) + _run( + ["git", "add", _GENERATED_KERNEL_FILE], + cwd=source, + ) + _run([ + "git", "commit", "-m", + f"generate initial HIP kernel for " + f"{assignment.worker_id}", + ], cwd=source) + bootstrap_commit = _run( + ["git", "rev-parse", "HEAD"], cwd=source + ).stdout.strip() + for shape_id, params in shapes.items(): + snapshot_accepted_kernel_artifact( + worker_root=root, + shape_id=shape_id, + shape=params, + metrics=metrics_by_shape[shape_id], + commit=bootstrap_commit, + ) + fixed_baselines = { + shape_id: { + **fixed_targets[shape_id], + "bootstrap_metrics": metrics_by_shape[shape_id], + } + for shape_id in shapes + } + result = { + "schema_version": SCHEMA_VERSION, + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "attempt": attempt, + "status": "passed", + "passed": True, + "hypothesis": hypothesis, + "shapes": list(shapes), + "metrics": metrics_by_shape, + "paired_m2_fallback_metrics": ( + paired_fallback_metrics + ), + "comparison_baselines": fixed_baselines, + "shapes_verified": verified, + "generated_files": [_GENERATED_KERNEL_FILE], + "source": "child_agent_generated", + "timestamp": time.time(), + } + write_json(root / "bootstrap_result.json", result) + write_json( + iteration_dir / "iteration.json", result + ) + _status( + root, + assignment, + state="bootstrap_passed", + iteration=0, + shape_id=None, + ) + self.store.append_timeline( + "worker_bootstrap_success", + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "attempt": attempt, + "shapes_verified": verified, + "source": "child_agent_generated", + }, + ) + return assignment.worker_id, fixed_baselines + except Exception as exc: + previous_failure = str(exc) + proposal_path.unlink(missing_ok=True) + if kernel_path.is_file() and not _run( + [ + "git", "ls-files", "--", + _GENERATED_KERNEL_FILE, + ], + cwd=source, + ).stdout.strip(): + kernel_path.unlink() + persist( + attempt, + "failed", + metrics=metrics_by_shape, + hypothesis=hypothesis, + error=previous_failure, + ) + if immutable_digests: + _run( + [ + "git", "restore", "--", + *immutable_digests.keys(), + ], + cwd=source, + ) + self.store.append_timeline( + "worker_bootstrap_attempt_failed", + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "attempt": attempt, + "error": previous_failure, + }, + ) + + raise RuntimeError( + f"{assignment.worker_id} could not generate a valid initial " + f"HIP kernel after {_MAX_BOOTSTRAP_RETRIES} attempts: " + f"{previous_failure}" + ) + + completed: list[str] = [] + baseline: Dict[str, Dict[str, Any]] = {} + with ThreadPoolExecutor( + max_workers=len(config.assignments) + ) as pool: + futures = { + pool.submit(generate_one, assignment): assignment + for assignment in config.assignments + } + for future in as_completed(futures): + assignment = futures[future] + try: + worker_id, metrics = future.result() + except Exception as exc: + error = str(exc) + state = ( + "timed_out" + if any( + token in error.lower() + for token in ( + "timeout", "timed out", "stuck", "killed", + ) + ) + else "failed" + ) + failure = { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_ids": assignment.shape_ids, + "state": state, + "stage": "bootstrap", + "error": error, + "timestamp": time.time(), + } + self._worker_failures[assignment.worker_id] = failure + root = ( + self.workspace_dir / "workers" + / assignment.worker_id + ) + _status( + root, + assignment, + state=state, + iteration=0, + shape_id=None, + error=error, + ) + write_json(root / "failure.json", failure) + self.store.append_timeline("worker_failed", failure) + continue + completed.append(worker_id) + baseline.update(metrics) + if not completed: + raise RuntimeError( + "worker bootstrap has fewer than one successful GPU worker" + ) + self.store.append_timeline( + "worker_bootstrap_complete", + { + "workers": sorted(completed), + "ignored_workers": sorted(self._worker_failures), + "kernel_source": "child_agent_generated", + }, + ) + self.store.append_timeline( + "kernel_generation_complete", + { + "repository": str( + (self.workspace_dir / "main").resolve() + ), + "candidate_workers": sorted(completed), + "ignored_workers": sorted(self._worker_failures), + }, + ) + return baseline + + def _parallel_lane_lifecycles( + self, + config: OptimizerConfig, + ) -> tuple[Dict[str, Dict[str, Any]], Dict[str, Any]]: + """Generate, optimize for five rounds, and write Skill per child GPU. + + A slow or broken lane must not hold successful GPUs at a global + barrier. Aggregate only after every lane has reached a terminal state. + """ + baseline: Dict[str, Dict[str, Any]] = {} + workers: Dict[str, Any] = {} + + def run_lane( + assignment: WorkerAssignment, + ) -> tuple[ + str, + Dict[str, Dict[str, Any]], + Dict[str, Any], + ]: + lane_config = replace_assignments(config, [assignment]) + lane_baseline = self._bootstrap_worker_repos(lane_config) + lane_workers = self._parallel_agents( + lane_config, lane_baseline + ) + return assignment.worker_id, lane_baseline, lane_workers + + with ThreadPoolExecutor( + max_workers=len(config.assignments) + ) as pool: + futures = { + pool.submit(run_lane, assignment): assignment + for assignment in config.assignments + } + for future in as_completed(futures): + assignment = futures[future] + try: + _, lane_baseline, lane_workers = future.result() + except Exception as exc: + if assignment.worker_id not in self._worker_failures: + error = str(exc) + state = ( + "timed_out" + if any( + token in error.lower() + for token in ( + "timeout", "timed out", "stuck", "killed", + ) + ) + else "failed" + ) + failure = { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_ids": assignment.shape_ids, + "state": state, + "stage": "lane_lifecycle", + "error": error, + "timestamp": time.time(), + } + self._worker_failures[ + assignment.worker_id + ] = failure + root = ( + self.workspace_dir / "workers" + / assignment.worker_id + ) + _status( + root, + assignment, + state=state, + iteration=0, + shape_id=None, + error=error, + ) + write_json(root / "failure.json", failure) + self.store.append_timeline( + "worker_failed", failure + ) + continue + baseline.update(lane_baseline) + workers.update(lane_workers) + + write_json( + self.workspace_dir / "shared_baseline" / "results.json", + { + "schema_version": SCHEMA_VERSION, + "operator": "int8_w8a8_gemm", + "source": "user_supplied_fixed_triton_graph_table", + "shapes": { + shape_id: { + key: value + for key, value in record.items() + if key != "bootstrap_metrics" + } + for shape_id, record in baseline.items() + }, + }, + ) + minimum_completed = 1 + if len(workers) < minimum_completed: + raise RuntimeError( + "parallel explore has fewer than " + f"{minimum_completed} successful GPU workers " + f"({len(workers)}/{len(config.assignments)} completed)" + ) + self.store.append_timeline( + "parallel_lane_lifecycles_complete", + { + "completed_workers": sorted(workers), + "ignored_workers": sorted(self._worker_failures), + "policy": ( + "each child GPU generates its initial HIP kernel, runs " + "five optimization iterations, and writes a worker Skill; " + "one completed lane is enough for main Skill synthesis" + ), + }, + ) + return baseline, dict(sorted(workers.items())) + + def _synthesize_final_candidate( + self, + config: OptimizerConfig, + workers: Dict[str, Any], + initial_metrics: Dict[str, Dict[str, Any]], + task_id: str, + ) -> Dict[str, Any]: + """Link exact benchmarked worker objects behind trusted C++ dispatch.""" + seed = self.workspace_dir / "main" + root = self.workspace_dir / "final" + for name in ( + "cache", + "cache/torch", + "cache/triton", + "cache/xdg", + "cache/tmp", + "logs", + ): + (root / name).mkdir(parents=True, exist_ok=True) + source = root / "source" + if not source.exists(): + branch = f"agent/{_safe(task_id)}/final" + _run([ + "git", "worktree", "add", "-b", branch, str(source), "HEAD", + ], cwd=seed) + if config.target_repo_path is not None: + _match_tree_owner(root, config.target_repo_path.parent) + + # The final worktree's committed API is the immutable task contract. + # Never reload defaults through the live global API path: that path may + # legitimately change while a long-running optimization is in flight. + _run(["git", "restore", "."], cwd=source) + origin_contract = getattr( + self, + "_operator_api_contract", + resolve_operator_api(config.operator, config.dtype), + ) + frozen_contract = _task_local_api_contract(origin_contract, source) + + completed_assignments = [ + assignment for assignment in config.assignments + if assignment.worker_id in workers + ] + if not completed_assignments: + raise RuntimeError("serial validation has no completed worker") + assigned_gpus = {item.gpu for item in completed_assignments} + serial_gpu_raw = os.environ.get("METAINFER_SERIAL_VALIDATE_GPU") + try: + serial_gpu = ( + int(serial_gpu_raw) + if serial_gpu_raw is not None + else completed_assignments[0].gpu + ) + except ValueError as exc: + raise RuntimeError( + "METAINFER_SERIAL_VALIDATE_GPU must be an integer" + ) from exc + if serial_gpu not in assigned_gpus: + raise RuntimeError( + "serial validation GPU must be one of the completed worker " + f"GPUs {sorted(assigned_gpus)}, got {serial_gpu}" + ) + runner = W8A8Runner(root, serial_gpu) + best_by_shape = { + shape_id: workers[assignment.worker_id]["shapes"][shape_id][ + "metrics" + ] + for assignment in completed_assignments + for shape_id in assignment.shape_ids + } + optimized_ids = set(best_by_shape) + api_shapes = default_optimization_shapes(frozen_contract) + fallback_shapes = [ + shape for shape in api_shapes + if str(shape["id"]) not in optimized_ids + ] + for relative in ( + "prebuilt", + "accepted_sources", + "csrc/w8a8_dispatch.cpp", + "artifact_manifest.json", + ): + path = source / relative + if path.is_dir(): + shutil.rmtree(path) + else: + path.unlink(missing_ok=True) + + artifacts: list[Dict[str, Any]] = [] + for assignment in completed_assignments: + worker_root = ( + self.workspace_dir / "workers" / assignment.worker_id + ) + for shape_id in assignment.shape_ids: + if shape_id not in workers[assignment.worker_id]["shapes"]: + continue + shape_result = workers[assignment.worker_id]["shapes"][ + shape_id + ] + artifact = shape_result.get("artifact") + if not isinstance(artifact, dict): + raise RuntimeError( + f"{assignment.worker_id}/{shape_id} did not hand off " + "an accepted compiled artifact" + ) + source_object = worker_root / str(artifact["object"]) + source_hip = worker_root / str(artifact["source"]) + if ( + _sha256_file(source_object) + != artifact.get("object_sha256") + ): + raise RuntimeError( + f"accepted object digest changed for {shape_id}" + ) + if ( + _sha256_file(source_hip) + != artifact.get("source_sha256") + ): + raise RuntimeError( + f"accepted HIP digest changed for {shape_id}" + ) + destination_object = ( + source / "prebuilt" / f"{_safe(shape_id)}.o" + ) + namespaced = _namespace_prebuilt_object( + source_object, destination_object, shape_id + ) + archived_source = ( + source / "accepted_sources" + / f"{_safe(shape_id)}.hip" + ) + archived_source.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source_hip, archived_source) + artifacts.append({ + **artifact, + **namespaced, + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_id": shape_id, + "shape": config.shapes[shape_id].params, + "final_object": str( + destination_object.relative_to(source) + ), + "audited_hip": str( + archived_source.relative_to(source) + ), + }) + artifacts.sort(key=lambda item: str(item["shape_id"])) + dispatch_path = source / "csrc" / "w8a8_dispatch.cpp" + dispatch_path.write_text( + _render_prebuilt_dispatch(artifacts), encoding="utf-8" + ) + proposal = { + "schema_version": SCHEMA_VERSION, + "strategy": "link_exact_accepted_worker_objects", + "hip_recompiled_by_main": False, + "hip_rewritten_by_main": False, + "dispatch_language": "C++", + "artifacts": artifacts, + } + write_json(source / "artifact_manifest.json", proposal) + self.store.append_timeline( + "final_artifact_link_launch", + { + "artifacts": [ + { + "shape_id": item["shape_id"], + "worker_id": item["worker_id"], + "object_sha256": item["object_sha256"], + } + for item in artifacts + ], + "gpu": serial_gpu, + }, + ) + + validation: Dict[str, Dict[str, Any]] = {} + validation_shapes = [ + {"id": shape.id, **shape.params} + for shape in config.shapes.values() + if shape.id in optimized_ids + ] + fallback_shapes + try: + for shape in validation_shapes: + shape_id = str(shape["id"]) + params = { + key: value for key, value in shape.items() + if key != "id" + } + if shape_id in optimized_ids: + metrics = runner.benchmark(params) + else: + metrics = runner.benchmark( + params, warmups=2, samples=3 + ) + validation[shape_id] = metrics + if not metrics.get("passed"): + raise RuntimeError( + f"correctness failed for {shape_id}: " + f"{json.dumps(metrics)}" + ) + if shape_id not in optimized_ids: + continue + best_median = float( + best_by_shape[shape_id].get("median_us") + or float("inf") + ) + metrics = _final_performance_gate( + shape_id=shape_id, + best_median=best_median, + metrics=metrics, + benchmark=lambda: runner.benchmark(params), + max_retries=_PERF_GATE_MAX_RETRIES, + retry_interval_s=_PERF_GATE_RETRY_INTERVAL_S, + store=self.store, + ) + validation[shape_id] = metrics + except Exception as exc: + self.store.append_timeline( + "final_artifact_link_rejected", + {"error": str(exc)}, + ) + raise RuntimeError( + "Final prebuilt-object W8A8 validation failed: " + f"{exc}" + ) from exc + + changed_paths = sorted( + line[3:].strip() + for line in _run( + ["git", "status", "--porcelain", "--untracked-files=all"], + cwd=source, + ).stdout.splitlines() + if line[3:].strip() + ) + _run(["git", "add", *changed_paths], cwd=source) + _run([ + "git", "commit", "-m", + "link accepted per-shape W8A8 objects", + ], cwd=source) + commit = _run( + ["git", "rev-parse", "HEAD"], cwd=source + ).stdout.strip() + _run(["git", "cherry-pick", commit], cwd=seed) + published_commit = _run( + ["git", "rev-parse", "HEAD"], cwd=seed + ).stdout.strip() + result = { + "schema_version": SCHEMA_VERSION, + "attempt": 1, + "candidate_commit": commit, + "commit": published_commit, + "published_repo": str(seed.resolve()), + "proposal": proposal, + "validation": validation, + "initial_metrics": initial_metrics, + "workers": [ + assignment.worker_id + for assignment in completed_assignments + ], + "all_shapes_in_one_extension": True, + "optimized_shapes": sorted(optimized_ids), + "fallback_regression_shapes": [ + str(shape["id"]) for shape in fallback_shapes + ], + "all_api_shapes_validated": True, + "hip_recompiled_by_main": False, + "serial_validation_gpu": serial_gpu, + } + write_json(root / "result.json", result) + self.store.append_timeline( + "final_artifact_link_success", + {"attempt": 1, "commit": published_commit}, + ) + return result + + # ------------------------------------------------------------------ # + # Main pipeline + # ------------------------------------------------------------------ # + + def run(self, *, dry_run: bool = False) -> Dict[str, Any]: + """Full generate-then-optimize pipeline. + + 1. PREPARE: parse config (may lack GPU assignments), create seed repo. + 2. GENERATE: control plane assigns shapes/GPUs; main Agent reviews. + 3. Create worker worktrees. + 4. EXPLORE: child Agents create HIP, then run five optimization rounds. + 5. SYNTHESIZE → VALIDATE → REPORT. + """ + task_id = str(self.req.get("task_id", "task")) + self.store.init_or_resume(task_id) + self.store.update_run( + finished=False, + final_status=None, + last_outcome=None, + last_transition_label=None, + notes=[], + ) + started = time.time() + try: + # ---- PREPARE -------------------------------------------------- # + self._phase(phases.PREPARE) + config = load_config(self.req) + self._validate_contract(config) + self._prepare_worktrees(config, task_id) + plan = self._plan(config) + write_json(self.workspace_dir / "plan.json", plan) + if dry_run: + return plan + + # ---- GENERATE ------------------------------------------------- # + self._phase(phases.GENERATE) + assignments = self._generate_kernel_repo(config) + config = replace_assignments(config, assignments) + self._create_worker_worktrees(config, task_id) + plan = self._plan(config) + write_json(self.workspace_dir / "plan.json", plan) + + # ---- EXPLORE -------------------------------------------------- # + self._phase(phases.EXPLORE, workers=len(config.assignments)) + initial_metrics, workers = self._parallel_lane_lifecycles(config) + + # ---- SYNTHESIZE ----------------------------------------------- # + self._phase( + phases.SYNTHESIZE, + action="merge_completed_worker_skills", + ) + completed_assignments = [ + item for item in config.assignments + if item.worker_id in workers + ] + merged_skill = self._author_merged_skill( + config, completed_assignments + ) + + # ---- VALIDATE ------------------------------------------------- # + self._phase( + phases.VALIDATE, + action="merge_and_validate_final_kernel", + ) + worker_validation = { + shape_id: { + "passed": bool( + shape_result.get("metrics", {}).get("passed") + ), + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "candidate": shape_result.get("candidate"), + "metrics": shape_result.get("metrics") or {}, + "artifact": shape_result.get("artifact") or {}, + "source": "child_accepted_compiled_artifact", + "rerun_by_main": False, + } + for assignment in config.assignments + if assignment.worker_id in workers + for shape_id, shape_result in workers[ + assignment.worker_id + ].get("shapes", {}).items() + } + synthesis = self._synthesize_final_candidate( + config, workers, initial_metrics, task_id + ) + validation = synthesis["validation"] + + # ---- REPORT --------------------------------------------------- # + self._phase(phases.REPORT) + final_target = evaluate_final_target( + baseline=initial_metrics, + validation=validation, + target_improvement_percent=( + config.minimum_improvement_percent + ), + ) + report = { + "schema_version": SCHEMA_VERSION, + "task_id": task_id, + "task_type": "dcu-kernel-auto-opt", + "mode": "generate-and-optimize", + "started_at": started, + "finished_at": time.time(), + "duration_s": round(time.time() - started, 4), + "config": plan, + "initial_metrics": initial_metrics, + "workers": workers, + "worker_failures": self._worker_failures, + "worker_validation": worker_validation, + "synthesis": synthesis, + "merged_skill": merged_skill, + "final_validation": validation, + "final_target": final_target, + "real_gpu_used": True, + "kernel_generated": True, + "gpu_assignment_agent_decided": False, + "target_repo_modified": True, + "status": ( + "partial_success" if self._worker_failures else "success" + ), + } + write_json(self.workspace_dir / "final_report.json", report) + self.store.update_run( + current_iteration=config.mock_iterations, + current_phase=phases.FINISHED, + finished=True, + final_status="success", + last_outcome="ok", + last_transition_label="generate + optimize complete", + ) + self.store.append_timeline( + "orchestrator_success", + { + "workers": len(workers), + "real_gpu_used": True, + "kernel_generated": True, + "gpu_assignment_agent_decided": False, + "operator": "int8_w8a8_gemm", + }, + ) + return report + except Exception as exc: + self.store.append_timeline( + "orchestrator_error", {"error": repr(exc)} + ) + self.store.update_run( + current_phase=self._current_phase, + finished=True, + final_status="stopped", + last_outcome="infra_fail", + notes=[str(exc)], + ) + raise + + @staticmethod + def _plan(config: OptimizerConfig) -> Dict[str, Any]: + return { + "schema_version": SCHEMA_VERSION, + "execution_mode": GEN_AND_OPT_MODE, + "operator": "INT8 W8A8 GEMM", + "dtype": config.dtype, + "hardware": config.hardware, + "kernel_language": config.kernel_language, + "claude_model": config.claude_model, + "max_iterations": config.mock_iterations, + "minimum_improvement_percent": ( + config.minimum_improvement_percent + ), + "minimum_improvement_semantics": ( + "final validated result versus fixed baseline" + ), + "round_acceptance_improvement_percent": ( + ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + ), + "shape_scope": config.shape_scope, + "assignment_mode": config.assignment_mode, + "harness": "trusted PyTorch FP32-accumulate W8A8 reference", + "kernel_source": "generated by assigned child agents", + "kernel_repo": ( + str(config.target_repo_path) + if config.target_repo_path is not None else None + ), + "gpu_assignment_source": ( + "control-plane-shape-balanced" + if config.assignment_mode == "ai" + else "manual-new-task" + ), + "main_agent_role": "coordination only", + "contract": { + "A": "int8[M,K], row-major contiguous", + "B": "int8[K,N], row-major contiguous", + "x_scale": "float32[M,1]", + "weight_scale": "float32[N,1]", + "output": "bfloat16[M,N]", + }, + "shapes": [ + {"id": shape.id, **shape.params} + for shape in config.shapes.values() + ], + "assignments": [ + { + "worker_id": item.worker_id, + "gpu": item.gpu, + "shapes": item.shape_ids, + } + for item in config.assignments + ], + "real_gpu_used": True, + } diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_binding.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_binding.py new file mode 100644 index 00000000..cef27317 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_binding.py @@ -0,0 +1,28 @@ +"""GPU visibility policy for worker29. + +DTK accepts either HIP_VISIBLE_DEVICES or ROCR_VISIBLE_DEVICES. Setting both +to the same non-zero physical index applies filtering twice and can hide the +device, so this task standardizes on HIP_VISIBLE_DEVICES only. +""" + +from __future__ import annotations + +from typing import Dict, MutableMapping + + +def bind_worker_gpu( + env: MutableMapping[str, str], physical_gpu: int +) -> Dict[str, str]: + if physical_gpu not in range(4): + raise ValueError("physical_gpu must be one of 0,1,2,3") + env.pop("ROCR_VISIBLE_DEVICES", None) + env["HIP_VISIBLE_DEVICES"] = str(physical_gpu) + return {"HIP_VISIBLE_DEVICES": str(physical_gpu)} + + +def hide_gpus_from_control_plane( + env: MutableMapping[str, str], +) -> Dict[str, str]: + env.pop("ROCR_VISIBLE_DEVICES", None) + env["HIP_VISIBLE_DEVICES"] = "" + return {"HIP_VISIBLE_DEVICES": ""} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/guidance.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/guidance.py new file mode 100644 index 00000000..63c56104 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/guidance.py @@ -0,0 +1,102 @@ +"""Durable, per-worker optimization guidance queue.""" + +from __future__ import annotations + +import fcntl +import json +import os +import time +import uuid +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Dict, Iterator + + +WORKER_IDS = tuple(f"worker_{index}" for index in range(4)) + + +def _validate_worker(worker_id: str) -> None: + if worker_id not in WORKER_IDS: + raise ValueError(f"unknown worker: {worker_id}") + + +def _path(root: Path, worker_id: str) -> Path: + return root / f"{worker_id}.json" + + +@contextmanager +def _locked(root: Path, worker_id: str) -> Iterator[None]: + root.mkdir(parents=True, exist_ok=True) + lock_path = root / f".{worker_id}.lock" + with lock_path.open("a+", encoding="utf-8") as handle: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + + +def _read(path: Path) -> list[Dict[str, Any]]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return [] + return value if isinstance(value, list) else [] + + +def _write(path: Path, entries: list[Dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + temporary.write_text( + json.dumps(entries, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + os.replace(temporary, path) + + +def list_guidance(root: Path, worker_id: str) -> list[Dict[str, Any]]: + _validate_worker(worker_id) + with _locked(root, worker_id): + return _read(_path(root, worker_id)) + + +def add_guidance(root: Path, worker_id: str, text: str) -> Dict[str, Any]: + _validate_worker(worker_id) + normalized = text.strip() + if not normalized: + raise ValueError("guidance cannot be empty") + if len(normalized) > 4000: + raise ValueError("guidance must be at most 4000 characters") + entry: Dict[str, Any] = { + "id": uuid.uuid4().hex, + "worker_id": worker_id, + "text": normalized, + "source": "manual", + "status": "pending", + "created_at": time.time(), + } + with _locked(root, worker_id): + path = _path(root, worker_id) + entries = _read(path) + entries.append(entry) + _write(path, entries) + return entry + + +def claim_next_guidance( + root: Path, worker_id: str, iteration: int +) -> Dict[str, Any] | None: + """Atomically consume the oldest pending instruction for one worker.""" + _validate_worker(worker_id) + with _locked(root, worker_id): + path = _path(root, worker_id) + entries = _read(path) + for entry in entries: + if entry.get("status") != "pending": + continue + entry["status"] = "consumed" + entry["consumed_iteration"] = iteration + entry["consumed_at"] = time.time() + _write(path, entries) + return dict(entry) + return None diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/integrate_restarted_worker.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/integrate_restarted_worker.py new file mode 100644 index 00000000..a0cd7f54 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/integrate_restarted_worker.py @@ -0,0 +1,178 @@ +"""Integrate a completed restarted lane after the original orchestrator exits.""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from pathlib import Path + +from metainfer.orchestrator._bootstrap import make_subagent_manager +from metainfer.orchestrator.state import StateStore + +from .config import load_config, resolve_claude_bin +from .gen_and_opt_pipeline import GenAndOptPipeline +from .result_store import write_json +from .w8a8_pipeline import evaluate_final_target + + +def _load(path: Path) -> dict: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return value if isinstance(value, dict) else {} + + +def _orchestrator_running(state_dir: Path) -> bool: + record = _load(state_dir / "orchestrator.pid") + try: + pid = int(record.get("pid") or 0) + except (TypeError, ValueError): + return False + if pid <= 0: + return False + try: + os.kill(pid, 0) + except OSError: + return False + return True + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("requirements", type=Path) + parser.add_argument("--state-dir", type=Path, required=True) + parser.add_argument("--workspace-dir", type=Path, required=True) + parser.add_argument("--worker-id", required=True) + parser.add_argument( + "--claude-bin", + default=None, + help=( + "Agent binary override; defaults resolved from agent_framework " + "(ccb -> METAINFER_CLAUDE_BIN or 'ccb', dsh -> " + "bridge/dsh/dsh_agent.py)." + ), + ) + parser.add_argument("--timeout", type=int, default=43200) + args = parser.parse_args() + + req = _load(args.requirements) + config = load_config(req) + claude_bin = resolve_claude_bin( + config.agent_framework, explicit=args.claude_bin + ) + assignment = next( + item for item in config.assignments + if item.worker_id == args.worker_id + ) + worker_result_path = ( + args.workspace_dir / "workers" / args.worker_id + / "restart_result.json" + ) + final_report_path = args.workspace_dir / "final_report.json" + deadline = time.time() + args.timeout + while time.time() < deadline: + worker_result = _load(worker_result_path) + report = _load(final_report_path) + if ( + worker_result.get("status") == "completed" + and report + and not _orchestrator_running(args.state_dir) + ): + break + time.sleep(5) + else: + raise TimeoutError("timed out waiting for worker and original task") + + workers = dict(report.get("workers") or {}) + workers.update(worker_result.get("workers") or {}) + initial_metrics = dict(report.get("initial_metrics") or {}) + initial_metrics.update(worker_result.get("baseline") or {}) + failures = dict(report.get("worker_failures") or {}) + failures.pop(args.worker_id, None) + + os.environ["METAINFER_TASK_ID"] = str(req.get("task_id") or "task") + os.environ["METAINFER_SERIAL_VALIDATE_GPU"] = str(assignment.gpu) + manager = make_subagent_manager( + claude_bin=claude_bin, + model=config.claude_model, + permission_mode="bypassPermissions", + effort="max", + extra_add_dirs=[args.workspace_dir], + snapshot_file=( + args.workspace_dir / "workers" / args.worker_id + / "restart_integration_agents.json" + ), + max_concurrent=1, + ) + store = StateStore(args.state_dir) + pipeline = GenAndOptPipeline( + req=req, + state_dir=args.state_dir, + workspace_dir=args.workspace_dir, + store=store, + manager=manager, + ) + pipeline._worker_failures = failures + try: + completed_assignments = [ + item for item in config.assignments + if item.worker_id in workers + ] + merged_skill = pipeline._author_merged_skill( + config, completed_assignments + ) + synthesis = pipeline._synthesize_final_candidate( + config, + workers, + initial_metrics, + str(req.get("task_id") or "task"), + ) + validation = synthesis["validation"] + worker_validation = dict(report.get("worker_validation") or {}) + for shape_id, shape_result in workers[args.worker_id].get( + "shapes", {} + ).items(): + worker_validation[shape_id] = { + "passed": bool(shape_result.get("metrics", {}).get("passed")), + "worker_id": args.worker_id, + "physical_gpu": assignment.gpu, + "candidate": shape_result.get("candidate"), + "metrics": shape_result.get("metrics") or {}, + "artifact": shape_result.get("artifact") or {}, + "source": "restarted_child_accepted_compiled_artifact", + "rerun_by_main": False, + } + report.update({ + "workers": workers, + "worker_failures": failures, + "initial_metrics": initial_metrics, + "worker_validation": worker_validation, + "synthesis": synthesis, + "merged_skill": merged_skill, + "final_validation": validation, + "final_target": evaluate_final_target( + baseline=initial_metrics, + validation=validation, + target_improvement_percent=config.minimum_improvement_percent, + ), + "status": "partial_success" if failures else "success", + "restart_integrated_worker": args.worker_id, + "restart_integrated_at": time.time(), + }) + write_json(final_report_path, report) + store.append_timeline("worker_restart_integrated", { + "worker_id": args.worker_id, + "physical_gpu": assignment.gpu, + "serial_validation_gpu": assignment.gpu, + "workers": sorted(workers), + }) + return 0 + finally: + manager.shutdown() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/isa_analysis.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/isa_analysis.py new file mode 100644 index 00000000..e3d4451e --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/isa_analysis.py @@ -0,0 +1,434 @@ +"""Trusted gfx928 code-object extraction and compact ISA evidence.""" + +from __future__ import annotations + +import hashlib +import os +import re +import shutil +import subprocess +import tempfile +from collections import Counter +from pathlib import Path +from typing import Any, Dict, Sequence + + +_DTK_LLVM_CANDIDATES = ( + Path("/opt/dtk/aillvm/bin"), + Path("/opt/dtk/llvm/bin"), + Path("/opt/dtk/dcc/bin"), +) +_DTK_LLVM_TOOLS = ( + "llvm-objcopy", + "clang-offload-bundler", + "llvm-readobj", + "llvm-objdump", +) +_KEY_OPCODE_PREFIXES = ( + "s_load_", "global_load", "global_store", "buffer_load", + "buffer_store", "flat_load", "flat_store", "ds_read", "ds_write", + "s_waitcnt", "s_barrier", "v_mmac", "v_fma", "v_fmac", "v_mad", + "v_mac", "v_dot", "v_pk_", +) +_CATEGORY_PATTERNS = { + "smem_load": r"^s_load_", + "global_load": r"^global_load", + "global_store": r"^global_store", + "buffer_load": r"^buffer_load", + "buffer_store": r"^buffer_store", + "flat_load": r"^flat_load", + "flat_store": r"^flat_store", + "ds_read": r"^ds_read", + "ds_write": r"^ds_write", + "waitcnt": r"^s_waitcnt$", + "barrier": r"^s_barrier$", + "mmac": r"^v_mmac", + "fma_mac": r"^v_(?:fma|fmac|mad|mac)", + "dot": r"^v_dot", + "packed_alu": r"^v_pk_", + "dpp": r"dpp", +} + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _run(args: list[str], *, timeout: int = 60) -> str: + result = subprocess.run( + args, + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=timeout, + ) + return result.stdout + + +def resolve_dtk_llvm_bin(explicit: Path | None = None) -> Path: + """Find one complete DTK LLVM toolchain instead of assuming one layout.""" + configured = os.environ.get("METAINFER_DTK_LLVM_BIN") + candidates = [ + candidate for candidate in ( + explicit, + Path(configured) if configured else None, + *_DTK_LLVM_CANDIDATES, + ) + if candidate is not None + ] + checked: list[str] = [] + for candidate in candidates: + missing = [ + name for name in _DTK_LLVM_TOOLS + if not (candidate / name).is_file() + ] + if not missing: + return candidate + checked.append(f"{candidate}: missing {missing}") + raise FileNotFoundError( + "no complete DTK LLVM toolchain found; checked " + "; ".join(checked) + ) + + +def _opcodes(disassembly: str) -> list[tuple[str, str]]: + result: list[tuple[str, str]] = [] + for raw in disassembly.splitlines(): + stripped = raw.strip() + if not stripped or stripped.endswith(":") or "file format" in stripped: + continue + instruction = stripped.split("//", 1)[0].strip() + opcode = instruction.split(None, 1)[0].lower() if instruction else "" + if opcode and not opcode.startswith("0"): + result.append((opcode, stripped)) + return result + + +def summarize_isa_text(disassembly: str) -> Dict[str, Any]: + """Summarize observable instructions without inferring performance.""" + instructions = _opcodes(disassembly) + counts = { + name: sum(bool(re.search(pattern, opcode)) for opcode, _ in instructions) + for name, pattern in _CATEGORY_PATTERNS.items() + } + waits: Counter[str] = Counter() + key_instructions: list[tuple[str, str]] = [] + for opcode, line in instructions: + if opcode == "s_waitcnt": + expression = line.split("s_waitcnt", 1)[1].split("//", 1)[0].strip() + waits[expression] += 1 + if opcode.startswith(_KEY_OPCODE_PREFIXES) or "dpp" in opcode: + key_instructions.append((opcode, line)) + + serialized_chains = 0 + for index, (opcode, _) in enumerate(key_instructions): + if not opcode.startswith(("global_load", "buffer_load", "flat_load")): + continue + window = key_instructions[index + 1:index + 5] + wait_index = next( + ( + offset for offset, (candidate, line) in enumerate(window) + if candidate == "s_waitcnt" and "vmcnt(0)" in line + ), + None, + ) + if wait_index is None: + continue + if any( + candidate.startswith("ds_write") + for candidate, _ in window[wait_index + 1:] + ): + serialized_chains += 1 + + return { + "instruction_counts": counts, + "waitcnt_expressions": dict(sorted(waits.items())), + "load_wait0_ds_write_windows": serialized_chains, + "key_instruction_excerpt": [ + line for _, line in key_instructions[:96] + ], + "interpretation_guard": ( + "Counts and ordering come from the exact gfx928 code object. " + "They identify audit targets, not performance conclusions." + ), + } + + +def _metadata_resources(metadata: str) -> Dict[str, Any]: + fields = { + "vgpr_count": r"\.vgpr_count:\s+(\d+)", + "sgpr_count": r"\.sgpr_count:\s+(\d+)", + "lds_bytes": r"\.group_segment_fixed_size:\s+(\d+)", + "scratch_bytes": r"\.private_segment_fixed_size:\s+(\d+)", + "vgpr_spill_count": r"\.vgpr_spill_count:\s+(\d+)", + "sgpr_spill_count": r"\.sgpr_spill_count:\s+(\d+)", + "wavefront_size": r"\.wavefront_size:\s+(\d+)", + } + resources: Dict[str, Any] = {} + for name, pattern in fields.items(): + values = [int(value) for value in re.findall(pattern, metadata)] + resources[name] = max(values) if values else None + return resources + + +_FUNCTION_HEADER = re.compile( + r"(?m)^[0-9a-fA-F]+ <([^>]+)>:\s*$" +) + + +def _disassembly_functions(disassembly: str) -> Dict[str, str]: + matches = list(_FUNCTION_HEADER.finditer(disassembly)) + return { + match.group(1): disassembly[ + match.start():matches[index + 1].start() + if index + 1 < len(matches) else len(disassembly) + ] + for index, match in enumerate(matches) + } + + +def _metadata_kernels(metadata: str) -> Dict[str, str]: + starts = [ + match.start() for match in re.finditer(r"(?m)^ - \.args:\s*$", metadata) + ] + result: Dict[str, str] = {} + for index, start in enumerate(starts): + end = starts[index + 1] if index + 1 < len(starts) else len(metadata) + section = metadata[start:end] + name = re.search(r"(?m)^ \.name:\s+(\S+)\s*$", section) + if name: + result[name.group(1)] = section + return result + + +def _is_auxiliary_symbol(name: str) -> bool: + lowered = name.lower() + return any(token in lowered for token in ( + "combine", "reduce", "finalize", "epilogue", "copy", "pack", + )) + + +def summarize_kernel_symbols( + disassembly: str, + metadata: str, + *, + kernel_names: Sequence[str] | None = None, + primary_kernel_name: str | None = None, +) -> Dict[str, Any]: + """Bind instruction/resource evidence to exact gfx928 kernel symbols.""" + functions = _disassembly_functions(disassembly) + metadata_sections = _metadata_kernels(metadata) + if kernel_names: + requested = list(dict.fromkeys(str(name) for name in kernel_names)) + else: + requested = [name for name in functions if "w8a8" in name.lower()] + + kernels: list[Dict[str, Any]] = [] + unmatched: list[str] = [] + for name in requested: + function = functions.get(name) + if function is None: + unmatched.append(name) + continue + item = summarize_isa_text(function) + item.update({ + "kernel_name": name, + "resources": _metadata_resources(metadata_sections.get(name, "")), + "resource_semantics": ( + "Code-object metadata for this exact symbol. LDS is static " + "group-segment size; use the matching PMC launch record for " + "dynamic LDS." + ), + }) + kernels.append(item) + + if not kernels: + summary = summarize_isa_text(disassembly) + summary.update({ + "kernel_name": None, + "resources": _metadata_resources(metadata), + "profiled_kernels": [], + "kernel_ownership": { + "requested": requested, + "matched": [], + "unmatched": unmatched, + }, + }) + return summary + + by_name = {item["kernel_name"]: item for item in kernels} + primary = by_name.get(str(primary_kernel_name or "")) + if primary is None: + non_auxiliary = [ + item for item in kernels + if not _is_auxiliary_symbol(item["kernel_name"]) + ] + candidates = non_auxiliary or kernels + primary = max( + candidates, + key=lambda item: ( + item["instruction_counts"].get("mmac", 0), + item["instruction_counts"].get("global_load", 0), + ), + ) + + summary = { + key: value for key, value in primary.items() + if key not in {"kernel_name", "resources"} + } + summary.update({ + "kernel_name": primary["kernel_name"], + "resources": primary["resources"], + "profiled_kernels": kernels, + "kernel_ownership": { + "requested": requested, + "matched": [item["kernel_name"] for item in kernels], + "unmatched": unmatched, + "primary": primary["kernel_name"], + }, + }) + return summary + + +def inspect_gfx928_object( + object_path: Path, + output_dir: Path, + *, + toolchain_dir: Path | None = None, + kernel_names: Sequence[str] | None = None, + primary_kernel_name: str | None = None, +) -> Dict[str, Any]: + """Extract the gfx928 bundle and archive disassembly plus metadata.""" + object_path = object_path.resolve() + output_dir.mkdir(parents=True, exist_ok=True) + llvm_bin = resolve_dtk_llvm_bin(toolchain_dir) + objcopy = llvm_bin / "llvm-objcopy" + bundler = llvm_bin / "clang-offload-bundler" + readobj = llvm_bin / "llvm-readobj" + objdump = llvm_bin / "llvm-objdump" + + code_object = output_dir / "gfx928.co" + metadata_path = output_dir / "metadata.txt" + disassembly_path = output_dir / "isa.txt" + with tempfile.TemporaryDirectory(prefix="metainfer-isa-") as temp: + temp_dir = Path(temp) + copied_object = temp_dir / "kernel.o" + fatbin = temp_dir / "hip.fatbin" + shutil.copy2(object_path, copied_object) + _run([ + str(objcopy), "--dump-section", + f".hip_fatbin={fatbin}", str(copied_object), + ]) + targets = [ + line.strip() for line in _run([ + str(bundler), "--list", "--type=o", f"--input={fatbin}", + ]).splitlines() + if line.strip() + ] + gfx928_targets = [target for target in targets if "gfx928" in target] + if len(gfx928_targets) != 1: + raise ValueError( + f"expected one gfx928 offload bundle, found {gfx928_targets}" + ) + target = gfx928_targets[0] + _run([ + str(bundler), "--unbundle", "--type=o", f"--targets={target}", + f"--input={fatbin}", f"--output={code_object}", + ]) + + metadata = _run([str(readobj), "--notes", str(code_object)]) + disassembly = _run([ + str(objdump), "-d", "--mcpu=gfx928", str(code_object), + ]) + metadata_path.write_text(metadata, encoding="utf-8") + disassembly_path.write_text(disassembly, encoding="utf-8") + summary = summarize_kernel_symbols( + disassembly, + metadata, + kernel_names=kernel_names, + primary_kernel_name=primary_kernel_name, + ) + summary.update({ + "available": True, + "bundle_target": target, + "host_object_sha256": _sha256(object_path), + "code_object_sha256": _sha256(code_object), + "artifact_paths": { + "code_object": str(code_object), + "metadata": str(metadata_path), + "disassembly": str(disassembly_path), + }, + }) + return summary + + +_ASM_PATTERN = re.compile( + r"\b(?:asm|__asm|__asm__)\s*" + r"(?:(?:volatile|__volatile|__volatile__)\s*)?\(\s*" + r'"((?:[^"\\]|\\.)*)"', + re.DOTALL, +) + + +def analyze_inline_asm_source(source: str) -> Dict[str, Any]: + """Distinguish compiler barriers from non-empty instruction asm.""" + bodies = [match.group(1) for match in _ASM_PATTERN.finditer(source)] + normalized = [re.sub(r"\\[ntr]", " ", body).strip() for body in bodies] + raw = sorted({body for body in normalized if body}) + return { + "asm_block_count": len(bodies), + "compiler_barrier_count": sum(not body for body in normalized), + "raw_instruction_asm_count": sum(bool(body) for body in normalized), + "raw_instruction_fingerprints": raw, + } + + +def evaluate_inline_asm_gate( + *, + before: Dict[str, Any], + after: Dict[str, Any], + proposal: Dict[str, Any], + isa_evidence: Dict[str, Any], + raw_inline_asm_allowed: bool | None = None, + verified_target_instructions: list[str] | None = None, +) -> Dict[str, Any]: + """Require explicit intent and trusted ISA evidence for new raw asm.""" + old = set(before.get("raw_instruction_fingerprints") or []) + new = sorted( + set(after.get("raw_instruction_fingerprints") or []) - old + ) + required = bool(new) + reasons: list[str] = [] + isa_plan = proposal.get("isa_optimization") + if required: + if raw_inline_asm_allowed is False: + reasons.append( + "control-plane phase does not allow raw inline asm" + ) + if not isinstance(isa_plan, dict): + reasons.append("proposal.json is missing isa_optimization") + else: + if isa_plan.get("strategy") != "inline_asm": + reasons.append("isa_optimization.strategy must be inline_asm") + targets = isa_plan.get("target_instructions") + if not isinstance(targets, list) or not targets: + reasons.append("target_instructions must be a non-empty list") + elif verified_target_instructions is not None and not set( + targets + ).issubset(set(verified_target_instructions)): + reasons.append( + "target_instructions exceed the compiler limitation " + "verified by the preceding ISA-guided round" + ) + if isa_evidence.get("available") is not True: + reasons.append("trusted candidate gfx928 ISA audit is unavailable") + return { + "required": required, + "passed": not reasons, + "new_raw_instruction_asm": new, + "reasons": reasons, + } diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/orchestrator.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/orchestrator.py new file mode 100644 index 00000000..8d74875c --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/orchestrator.py @@ -0,0 +1,108 @@ +"""Lifecycle wrapper for the mock optimization pipeline.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any, Dict, Optional + +from metainfer.orchestrator._bootstrap import ( + clear_pid_file, + install_subagent_shutdown_handlers, + make_subagent_manager, + set_process_name, + write_pid_file, +) +from metainfer.orchestrator.state import StateStore + +from .config import ( + GEN_AND_OPT_MODE, + LEGACY_SMOKE_MODE, + SMOKE_MODE, + W8A8_MODE, + load_config, +) +from .gen_and_opt_pipeline import GenAndOptPipeline +from .pipeline import MockOptimizationPipeline +from .real_pipeline import RealSmokeOptimizationPipeline +from .w8a8_pipeline import RealW8A8OptimizationPipeline + + +def run_with_requirements( + requirements_path: Path, + *, + state_dir: Optional[Path] = None, + workspace_dir: Optional[Path] = None, + dry_run: bool = False, + claude_bin: str = "ccb", +) -> int: + if not requirements_path.is_file(): + raise FileNotFoundError(requirements_path) + req: Dict[str, Any] = json.loads( + requirements_path.read_text(encoding="utf-8") + ) + task_id = str(req.get("task_id", "task")) + # The container-side bridge client forwards this ownership tag to the + # host bridge so task deletion can terminate only this task's agents. + os.environ["METAINFER_TASK_ID"] = task_id + if state_dir is None or workspace_dir is None: + from metainfer.server import paths as web_paths + state_dir = state_dir or web_paths.task_dir(task_id) + workspace_dir = workspace_dir or web_paths.workspace_dir(task_id) + state_dir.mkdir(parents=True, exist_ok=True) + workspace_dir.mkdir(parents=True, exist_ok=True) + target_req = state_dir / "requirements.json" + if requirements_path.resolve() != target_req.resolve(): + target_req.write_text( + requirements_path.read_text(encoding="utf-8"), encoding="utf-8" + ) + + set_process_name("metainfer-dkao") + pid_file = state_dir / "orchestrator.pid" + write_pid_file(pid_file, task_id) + answers = req.get("answers") + mode_source = answers if isinstance(answers, dict) else req + mode = str(mode_source.get("execution_mode", "Mock (no GPU)")) + config = load_config(req) + manager = make_subagent_manager( + claude_bin=claude_bin, + model=config.claude_model, + permission_mode="bypassPermissions", + effort=( + "low" + if mode in {LEGACY_SMOKE_MODE, SMOKE_MODE} + else "max" + ), + extra_add_dirs=[workspace_dir], + snapshot_file=state_dir / "agents.json", + max_concurrent=4, + ) + restore = install_subagent_shutdown_handlers(manager, pid_file=pid_file) + try: + common = { + "req": req, + "state_dir": state_dir, + "workspace_dir": workspace_dir, + "store": StateStore(state_dir), + } + if mode in {LEGACY_SMOKE_MODE, SMOKE_MODE}: + pipeline = RealSmokeOptimizationPipeline( + manager=manager, **common + ) + elif mode == W8A8_MODE: + pipeline = RealW8A8OptimizationPipeline( + manager=manager, **common + ) + elif mode == GEN_AND_OPT_MODE: + pipeline = GenAndOptPipeline( + manager=manager, **common + ) + else: + pipeline = MockOptimizationPipeline(**common) + pipeline.run(dry_run=dry_run) + return 0 + finally: + manager.shutdown() + restore() + clear_pid_file(pid_file) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/phases.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/phases.py new file mode 100644 index 00000000..163a0cbe --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/phases.py @@ -0,0 +1,75 @@ +"""Small top-level state machine for the multi-worker optimizer.""" + +from __future__ import annotations + +from typing import Any, Dict, List + + +PREPARE = "prepare" +GENERATE = "generate_kernel_repo" +BASELINE = "baseline" +EXPLORE = "parallel_explore" +SYNTHESIZE = "skill_synthesis" +VALIDATE = "serial_validate" +REPORT = "report" +FINISHED = "finished" + +_ORDER = [ + PREPARE, GENERATE, BASELINE, EXPLORE, SYNTHESIZE, VALIDATE, REPORT, FINISHED +] +_LABELS = { + PREPARE: "Prepare", + GENERATE: "Generate kernel repo", + BASELINE: "Baseline", + EXPLORE: "Parallel explore", + SYNTHESIZE: "Skill synthesis", + VALIDATE: "Serial validate", + REPORT: "Report", + FINISHED: "Finished", +} + + +def graph_payload( + current: str, + last_outcome: str | None = None, + last_label: str | None = None, + *, + include_baseline: bool = True, +) -> Dict[str, Any]: + order = ( + _ORDER + if include_baseline + else [phase for phase in _ORDER if phase != BASELINE] + ) + nodes: List[Dict[str, Any]] = [ + { + "id": phase, + "label": _LABELS[phase], + "description": _LABELS[phase], + "is_terminal": phase == FINISHED, + } + for phase in order + ] + edges = [ + { + "from": order[i], + "to": order[i + 1], + "label": "ok", + "outcomes": ["ok"], + } + for i in range(len(order) - 1) + ] + active_edge = None + if current in order and current != PREPARE: + idx = order.index(current) + active_edge = f"{order[idx - 1]} / {current}" + return { + "current": current, + "nodes": nodes, + "edges": edges, + "active_edge": active_edge, + "last_outcome": last_outcome, + "last_transition_label": last_label, + "terminal_nodes": [FINISHED], + "outcome_legend": [{"id": "ok", "label": "OK"}], + } diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/pipeline.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/pipeline.py new file mode 100644 index 00000000..e4e876a4 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/pipeline.py @@ -0,0 +1,299 @@ +"""MVP control plane: baseline, parallel mock workers, serial validation.""" + +from __future__ import annotations + +import json +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import Any, Dict + +from metainfer.orchestrator.state import StateStore + +from . import phases +from .adapters.mock import MockKernelAdapter +from .config import ( + ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT, + OptimizerConfig, + load_config, +) +from .result_store import SCHEMA_VERSION, write_json +from .skill_store import generate_merged_skill, generate_worker_skill +from .worker import run_mock_worker + + +class MockOptimizationPipeline: + def __init__( + self, + *, + req: Dict[str, Any], + state_dir: Path, + workspace_dir: Path, + store: StateStore, + ) -> None: + self.req = req + self.state_dir = state_dir + self.workspace_dir = workspace_dir + self.store = store + + def _phase(self, phase: str, **payload: Any) -> None: + self.store.update_run(current_phase=phase) + self.store.append_timeline( + "phase_start", {"phase": phase, **payload} + ) + + def run(self, *, dry_run: bool = False) -> Dict[str, Any]: + task_id = str(self.req.get("task_id", "task")) + self.store.init_or_resume(task_id) + self.store.update_run(finished=False, final_status=None, notes=[]) + started = time.time() + try: + self._phase(phases.PREPARE) + config = load_config(self.req) + self._create_layout(config) + plan = self._plan_payload(config) + write_json(self.workspace_dir / "plan.json", plan) + if dry_run: + result = {**plan, "dry_run": True, "status": "success"} + write_json(self.workspace_dir / "final_report.json", result) + self.store.update_run( + current_phase=phases.FINISHED, + finished=True, + final_status="success", + last_outcome="ok", + ) + return result + + self._phase(phases.BASELINE) + baseline = self._baseline(config) + + self._phase(phases.EXPLORE, workers=len(config.assignments)) + worker_results = self._parallel_workers(config, baseline) + + self._phase(phases.SYNTHESIZE) + merged_skill = generate_merged_skill( + config=config, + assignments=config.assignments, + workspace_dir=self.workspace_dir, + ) + + self._phase(phases.VALIDATE) + validation = self._serial_validate(config, worker_results) + + self._phase(phases.REPORT) + report = { + "schema_version": SCHEMA_VERSION, + "task_id": task_id, + "task_type": "dcu-kernel-auto-opt", + "mode": "mock", + "started_at": started, + "finished_at": time.time(), + "duration_s": round(time.time() - started, 4), + "config": plan, + "baseline": baseline, + "workers": worker_results, + "merged_skill": merged_skill, + "final_validation": validation, + "real_gpu_used": False, + "target_repo_modified": False, + "status": "success", + } + write_json(self.workspace_dir / "final_report.json", report) + (self.workspace_dir / "report.md").write_text( + self._markdown_report(report), encoding="utf-8" + ) + self.store.write_iteration(1, { + "iteration": 1, + "status": "success", + "goal": "validate multi-worker mock orchestration", + "started_at": started, + "ended_at": report["finished_at"], + "duration_s": report["duration_s"], + "perf": {}, + "workers": worker_results, + }) + self.store.update_run( + current_iteration=1, + current_phase=phases.FINISHED, + finished=True, + final_status="success", + last_outcome="ok", + last_transition_label="report complete", + ) + self.store.append_timeline( + "orchestrator_success", + {"workers": len(worker_results), "real_gpu_used": False}, + ) + return report + except Exception as exc: + self.store.append_timeline( + "orchestrator_error", {"error": repr(exc)} + ) + self.store.update_run( + current_phase=phases.FINISHED, + finished=True, + final_status="stopped", + last_outcome="infra_fail", + notes=[str(exc)], + ) + raise + + def _create_layout(self, config: OptimizerConfig) -> None: + self.workspace_dir.mkdir(parents=True, exist_ok=True) + for name in ("main", "shared_baseline", "final_validation", "workers"): + (self.workspace_dir / name).mkdir(parents=True, exist_ok=True) + for assignment in config.assignments: + (self.workspace_dir / "workers" / assignment.worker_id).mkdir( + parents=True, exist_ok=True + ) + + @staticmethod + def _plan_payload(config: OptimizerConfig) -> Dict[str, Any]: + return { + "schema_version": SCHEMA_VERSION, + "execution_mode": config.execution_mode, + "operator": config.operator, + "dtype": config.dtype, + "hardware": config.hardware, + "kernel_language": config.kernel_language, + "claude_model": config.claude_model, + "target_repo_path": ( + str(config.target_repo_path) if config.target_repo_path else None + ), + "mock_iterations": config.mock_iterations, + "minimum_improvement_percent": config.minimum_improvement_percent, + "minimum_improvement_semantics": ( + "final validated result versus fixed baseline" + ), + "round_acceptance_improvement_percent": ( + ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + ), + "shapes": [ + {"id": shape.id, **shape.params} + for shape in config.shapes.values() + ], + "assignments": [ + { + "worker_id": a.worker_id, + "gpu": a.gpu, + "shapes": a.shape_ids, + } + for a in config.assignments + ], + "real_gpu_used": False, + } + + def _baseline( + self, config: OptimizerConfig + ) -> Dict[str, Dict[str, float]]: + adapter = MockKernelAdapter() + adapter.prepare(self.workspace_dir / "shared_baseline") + baseline: Dict[str, Dict[str, float]] = {} + for shape in config.shapes.values(): + correct = adapter.correctness( + self.workspace_dir / "shared_baseline", shape + ) + if not correct.success: + raise RuntimeError(f"baseline correctness failed: {shape.id}") + result = adapter.benchmark( + self.workspace_dir / "shared_baseline", shape, iteration=0 + ) + if not result.success: + raise RuntimeError(f"baseline benchmark failed: {shape.id}") + baseline[shape.id] = result.metrics + write_json( + self.workspace_dir / "shared_baseline" / "results.json", + {"schema_version": SCHEMA_VERSION, "shapes": baseline}, + ) + return baseline + + def _parallel_workers( + self, + config: OptimizerConfig, + baseline: Dict[str, Dict[str, float]], + ) -> Dict[str, Any]: + out: Dict[str, Any] = {} + with ThreadPoolExecutor( + max_workers=len(config.assignments), + thread_name_prefix="mock-gpu-worker", + ) as pool: + futures = { + pool.submit( + run_mock_worker, + assignment=assignment, + config=config, + baseline=baseline, + worker_root=self.workspace_dir / "workers" / assignment.worker_id, + guidance_root=self.state_dir / "guidance", + adapter_factory=MockKernelAdapter, + ): assignment.worker_id + for assignment in config.assignments + } + for future in as_completed(futures): + worker_id = futures[future] + out[worker_id] = future.result() + assignment = next( + item for item in config.assignments + if item.worker_id == worker_id + ) + out[worker_id]["skill"] = generate_worker_skill( + config=config, + assignment=assignment, + workspace_dir=self.workspace_dir, + ) + self.store.append_timeline( + "worker_complete", { + "worker_id": worker_id, + "skill": out[worker_id]["skill"]["name"], + } + ) + return dict(sorted(out.items())) + + def _serial_validate( + self, config: OptimizerConfig, workers: Dict[str, Any] + ) -> Dict[str, Any]: + adapter = MockKernelAdapter() + root = self.workspace_dir / "final_validation" + results: Dict[str, Any] = {} + for assignment in config.assignments: + for shape_id in assignment.shape_ids: + shape = config.shapes[shape_id] + best = workers[assignment.worker_id]["shapes"][shape_id] + correctness = adapter.correctness(root, shape) + benchmark = adapter.benchmark( + root, shape, iteration=int(best["iteration"]) + ) + passed = correctness.success and benchmark.success + results[shape_id] = { + "passed": passed, + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "candidate": best["mock_candidate"], + "metrics": benchmark.metrics, + "serial": True, + } + if not passed: + raise RuntimeError(f"final validation failed: {shape_id}") + write_json(root / "results.json", { + "schema_version": SCHEMA_VERSION, + "shapes": results, + "real_gpu_used": False, + }) + return results + + @staticmethod + def _markdown_report(report: Dict[str, Any]) -> str: + lines = [ + "# DCU Kernel Auto-Optimization — Mock MVP", + "", + "No GPU was used and no target repository was modified.", + "", + "| Shape | Worker | Candidate | Median (us) |", + "|---|---|---|---:|", + ] + for shape_id, result in report["final_validation"].items(): + lines.append( + f"| {shape_id} | {result['worker_id']} | " + f"{result['candidate']} | {result['metrics']['median_us']} |" + ) + return "\n".join(lines) + "\n" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/plugin.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/plugin.py new file mode 100644 index 00000000..382b6a49 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/plugin.py @@ -0,0 +1,11 @@ +"""Task plugin descriptor.""" + +from metainfer.orchestrator.tasks.base import TaskPlugin + + +PLUGIN = TaskPlugin( + task_type="dcu-kernel-auto-opt", + cli_module="metainfer.tasks.dcu_kernel_auto_opt.orchestrator.cli", + phases_module="metainfer.tasks.dcu_kernel_auto_opt.orchestrator.phases", + diagnostic_globs=("*.json", "*.jsonl", "*.log"), +) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/pmc_profile.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/pmc_profile.py new file mode 100644 index 00000000..2fce2c33 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/pmc_profile.py @@ -0,0 +1,343 @@ +"""Trusted parsing and summarization for hipprof PMC CSV output.""" + +from __future__ import annotations + +import csv +from pathlib import Path +from typing import Any, Dict, Iterable + + +_DIRECT_COUNTERS = { + "grbm_count": "GRBM_COUNT", + "grbm_gui_active": "GRBM_GUI_ACTIVE", + "valu_active_instructions": "SQ_ACTIVE_INST_VALU", + "flat_lds_instructions": "SQ_INSTS_FLAT_LDS_ONLY", + "lds_instructions": "SQ_INSTS_LDS", + "valu_instructions": "SQ_INSTS_VALU", + "vmem_read_instructions": "SQ_INSTS_VMEM_RD", + "vmem_write_instructions": "SQ_INSTS_VMEM_WR", + "lds_bank_conflicts": "SQ_LDS_BANK_CONFLICT", + "lds_wait_instructions": "SQ_WAIT_INST_LDS", +} + + +def _integer(value: object) -> int: + try: + return int(str(value or "0"), 0) + except ValueError: + return 0 + + +def _sum_prefix(row: Dict[str, str], prefix: str) -> int: + return sum( + _integer(value) + for key, value in row.items() + if key.startswith(prefix) + ) + + +def _candidate_rows( + rows: Iterable[Dict[str, str]], +) -> list[Dict[str, str]]: + return [ + row for row in rows + if "w8a8" in str(row.get("KernelName", "")).lower() + ] + + +def _last_candidate_group(path: Path) -> list[Dict[str, str]]: + with path.open(newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + matching_indices = [ + index for index, row in enumerate(rows) + if "w8a8" in str(row.get("KernelName", "")).lower() + ] + if not matching_indices: + raise ValueError(f"no W8A8 kernel dispatch found in {path}") + # The harness emits one or more consecutive W8A8 dispatches per operator + # replay. Keep the final contiguous group so split-K partial + combine (or + # another multi-kernel implementation) remains one attributable replay. + start = end = matching_indices[-1] + matching = set(matching_indices) + while start - 1 in matching: + start -= 1 + return rows[start:end + 1] + + +def _last_candidate(path: Path) -> Dict[str, str]: + """Backward-compatible helper for callers that need the final dispatch.""" + return _last_candidate_group(path)[-1] + + +def _is_auxiliary_kernel(name: object) -> bool: + lowered = str(name or "").lower() + return any(token in lowered for token in ( + "combine", "reduce", "finalize", "epilogue", "copy", "pack", + )) + + +def _summarize_pmc_row(row: Dict[str, str]) -> Dict[str, Any]: + grid_workitems = _integer(row.get("grd")) + workgroup_size = _integer(row.get("wgr")) + begin_ns = _integer(row.get("BeginNs")) + end_ns = _integer(row.get("EndNs")) + l2_hits = _sum_prefix(row, "TCC_HIT[") + l2_misses = _sum_prefix(row, "TCC_MISS[") + grbm_count = _integer(row.get("GRBM_COUNT")) + grbm_gui_active = _integer(row.get("GRBM_GUI_ACTIVE")) + return { + "kernel_name": row.get("KernelName", ""), + "gpu_id": _integer(row.get("gpu-id")), + "grid_workitems": grid_workitems, + "workgroup_size": workgroup_size, + "grid_blocks": ( + grid_workitems // workgroup_size if workgroup_size else None + ), + "lds_bytes": _integer(row.get("lds")), + "scratch_bytes": _integer(row.get("scr")), + "arch_vgpr": _integer(row.get("arch_vgpr")), + "accum_vgpr": _integer(row.get("accum_vgpr")), + "sgpr": _integer(row.get("sgpr")), + "wave_size": _integer(row.get("wave_size")), + "profiled_duration_us": ( + round((end_ns - begin_ns) / 1000.0, 3) + if end_ns >= begin_ns else None + ), + "counters": { + name: _integer(row.get(column)) + for name, column in _DIRECT_COUNTERS.items() + }, + "l2_hits": l2_hits, + "l2_misses": l2_misses, + "l2_hit_rate_percent": ( + round(100.0 * l2_hits / (l2_hits + l2_misses), 3) + if l2_hits + l2_misses else None + ), + "gpu_active_percent": ( + round(100.0 * grbm_gui_active / grbm_count, 3) + if grbm_count else None + ), + } + + +def parse_pmc_csv(path: Path) -> Dict[str, Any]: + """Return per-kernel and operator evidence for the last W8A8 replay.""" + # The harness invokes the candidate once for correctness and again for its + # timed sample. The last matching dispatch is the closest representation + # of the code the next optimization round will inherit. + kernels = [_summarize_pmc_row(row) for row in _last_candidate_group(path)] + primary = next( + (item for item in kernels if not _is_auxiliary_kernel(item["kernel_name"])), + kernels[0], + ) + aggregate_hits = sum(item["l2_hits"] for item in kernels) + aggregate_misses = sum(item["l2_misses"] for item in kernels) + aggregate_counters = { + name: sum(item["counters"][name] for item in kernels) + for name in _DIRECT_COUNTERS + } + aggregate_duration = sum( + float(item["profiled_duration_us"] or 0.0) for item in kernels + ) + + evidence: Dict[str, Any] = { + "available": True, + "profile_mode": "hipprof --pmc --pmc-type 3", + **primary, + "primary_kernel_name": primary["kernel_name"], + "profiled_kernels": kernels, + "operator_aggregate": { + "kernel_count": len(kernels), + "kernel_names": [item["kernel_name"] for item in kernels], + "profiled_duration_us": round(aggregate_duration, 3), + "counters": aggregate_counters, + "l2_hits": aggregate_hits, + "l2_misses": aggregate_misses, + "l2_hit_rate_percent": ( + round( + 100.0 * aggregate_hits / (aggregate_hits + aggregate_misses), + 3, + ) + if aggregate_hits + aggregate_misses else None + ), + }, + "interpretation_guard": ( + "PMC perturbs latency. Top-level launch/resources identify the " + "primary GEMM kernel; profiled_kernels and operator_aggregate keep " + "all dispatches in the final operator replay. Use the normal " + "benchmark median/P90 for acceptance. " + "algorithmic_bandwidth_gb_s is theoretical minimum bytes divided " + "by time, not measured HBM traffic." + ), + } + return evidence + + +def parse_memory_traffic_csv( + read_path: Path, + write_path: Path, +) -> Dict[str, Any]: + """Apply DTK's documented TCC/EA request-size formulas.""" + reads = _last_candidate_group(read_path) + writes = _last_candidate_group(write_path) + read_names = [row.get("KernelName", "") for row in reads] + write_names = [row.get("KernelName", "") for row in writes] + if read_names != write_names: + raise ValueError( + "hipprof read/write replays contain different W8A8 dispatches: " + f"read={read_names}, write={write_names}" + ) + + kernels: list[Dict[str, Any]] = [] + for read, write in zip(reads, writes): + kernels.append(_memory_traffic_for_rows(read, write)) + read_bytes = sum(item["read_bytes"] for item in kernels) + write_bytes = sum(item["write_bytes"] for item in kernels) + total_bytes = read_bytes + write_bytes + reconciliation = [ + adjustment + for item in kernels + for adjustment in item.get("counter_reconciliation", []) + ] + return { + "source": "hipprof TCC/EA request counters", + "read_bytes_per_operator_replay": read_bytes, + "write_bytes_per_operator_replay": write_bytes, + "total_bytes_per_operator_replay": total_bytes, + # Compatibility aliases. These now represent the complete final + # operator replay, not whichever W8A8 kernel happened to run last. + "read_bytes_per_dispatch": read_bytes, + "write_bytes_per_dispatch": write_bytes, + "total_bytes_per_dispatch": total_bytes, + "kernel_count": len(kernels), + "kernel_names": read_names, + "kernels": kernels, + "counter_reconciliation": reconciliation, + "formula": "DTK derived_counters.xml FETCH_SIZE + WRITE_SIZE", + "read_kernel_name": read_names[0] if len(read_names) == 1 else None, + "write_kernel_name": write_names[0] if len(write_names) == 1 else None, + } + + +def _memory_traffic_for_rows( + read: Dict[str, str], + write: Dict[str, str], +) -> Dict[str, Any]: + """Calculate traffic for one named kernel across read/write replays.""" + + rd = _sum_prefix(read, "TCC_EA_RDREQ[") + rd32 = _sum_prefix(read, "TCC_EA_RDREQ_32B[") + rd1 = _sum_prefix(read, "TCC_EA1_RDREQ[") + rd1_32 = _sum_prefix(read, "TCC_EA1_RDREQ_32B[") + wr = _sum_prefix(write, "TCC_EA_WRREQ[") + wr64 = _sum_prefix(write, "TCC_EA_WRREQ_64B[") + wr1 = _sum_prefix(write, "TCC_EA1_WRREQ[") + wr1_64 = _sum_prefix(write, "TCC_EA1_WRREQ_64B[") + + reconciliation: list[Dict[str, Any]] = [] + rd = _reconcile_request_subcounter( + rd, rd32, "TCC_EA_RDREQ", "TCC_EA_RDREQ_32B", reconciliation + ) + rd1 = _reconcile_request_subcounter( + rd1, + rd1_32, + "TCC_EA1_RDREQ", + "TCC_EA1_RDREQ_32B", + reconciliation, + ) + wr = _reconcile_request_subcounter( + wr, wr64, "TCC_EA_WRREQ", "TCC_EA_WRREQ_64B", reconciliation + ) + wr1 = _reconcile_request_subcounter( + wr1, + wr1_64, + "TCC_EA1_WRREQ", + "TCC_EA1_WRREQ_64B", + reconciliation, + ) + + read_bytes = ( + rd32 * 32 + (rd - rd32) * 64 + + rd1_32 * 32 + (rd1 - rd1_32) * 64 + ) + write_bytes = ( + (wr - wr64) * 32 + wr64 * 64 + + (wr1 - wr1_64) * 32 + wr1_64 * 64 + ) + return { + "kernel_name": read.get("KernelName", ""), + "read_bytes": read_bytes, + "write_bytes": write_bytes, + "total_bytes": read_bytes + write_bytes, + "counter_reconciliation": reconciliation, + } + + +def _reconcile_request_subcounter( + total: int, + subset: int, + total_name: str, + subset_name: str, + reconciliation: list[Dict[str, Any]], +) -> int: + """Reconcile only tiny cross-replay counter skew, never large mismatch. + + hipprof may collect a request total and its size subcounter on separate + deterministic replays. A few requests of skew are possible around replay + boundaries. Promote the total to the observed subset only when the delta + is no more than eight requests or 0.01%, whichever is larger, and retain + explicit evidence. Larger inconsistencies remain hard failures. + """ + if subset <= total: + return total + delta = subset - total + tolerance = max(8, (total + 9_999) // 10_000) + if delta > tolerance: + raise ValueError( + "invalid hipprof memory request counters: " + f"{subset_name}={subset} exceeds {total_name}={total} " + f"by {delta} (tolerance={tolerance})" + ) + reconciliation.append( + { + "total_counter": total_name, + "subset_counter": subset_name, + "observed_total": total, + "observed_subset": subset, + "promoted_total": subset, + "delta_requests": delta, + "tolerance_requests": tolerance, + } + ) + return subset + + +def add_unprofiled_bandwidth( + evidence: Dict[str, Any], + median_us: object, +) -> None: + """Combine counter-derived traffic with an unprofiled benchmark latency.""" + traffic = evidence.get("memory_traffic") + if not isinstance(traffic, dict): + return + try: + latency = float(median_us) + total_bytes = int(traffic["total_bytes_per_operator_replay"]) + except (KeyError, TypeError, ValueError): + return + if latency <= 0: + return + traffic["unprofiled_median_us"] = latency + bandwidth = round( + total_bytes / latency / 1000.0, 6 + ) + traffic["counter_derived_operator_hbm_bandwidth_gb_s"] = bandwidth + traffic["counter_derived_hbm_bandwidth_gb_s"] = bandwidth + traffic["counter_derived_video_memory_bandwidth_gb_s"] = bandwidth + traffic["bandwidth_semantics"] = ( + "Sum of DTK FETCH_SIZE/WRITE_SIZE-equivalent TCC/EA request bytes " + "for every kernel in one marked operator replay, " + "divided by the separate unprofiled whole-operator median for " + "the identical source. Per-kernel bytes are not divided by the whole-" + "operator latency." + ) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/prompts.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/prompts.py new file mode 100644 index 00000000..fe80a0c7 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/prompts.py @@ -0,0 +1,793 @@ +"""Prompt templates for coordination and kernel implementation agents.""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import Any, Dict, Optional + + +HARNESS_PATH = ( + Path(__file__).resolve().parent.parent / "assets" / "w8a8_bench.py" +) + + +def split_k_candidate_set( + shape: Dict[str, Any], + *, + cu_count: int, + max_split: int = 16, +) -> list[int]: + """Suggest occupancy probes; callers must still benchmark outside them.""" + try: + n = int(shape["N"]) + k = int(shape["K"]) + cus = int(cu_count) + except (KeyError, TypeError, ValueError): + return [] + if n <= 0 or k < 64 or cus <= 0 or max_split < 2: + return [] + stage_count = k // 64 + candidates: set[int] = {2} + for waves_per_block in (1, 2, 3, 4, 6, 8): + n_blocks = math.ceil(n / (16 * waves_per_block)) + for resident_batches in (1, 2): + ideal = cus * resident_batches / n_blocks + candidates.update((math.floor(ideal), math.ceil(ideal))) + return sorted( + split for split in candidates + if 2 <= split <= min(max_split, stage_count) + ) + + +def _shapes_table(shapes: Dict[str, Dict[str, Any]]) -> str: + return "\n".join( + "| {sid} | {m} | {n} | {k} |".format( + sid=sid, + m=params.get("M", "?"), + n=params.get("N", "?"), + k=params.get("K", "?"), + ) + for sid, params in shapes.items() + ) + + +def shape_balanced_assignment( + shapes: Dict[str, Dict[str, Any]], +) -> Dict[str, Dict[str, Any]]: + """Deterministically balance exact shapes across up to four workers.""" + shape_work: Dict[str, int] = {} + for shape_id, params in shapes.items(): + try: + work = ( + 2 + * int(params["M"]) + * int(params["N"]) + * int(params["K"]) + ) + except (KeyError, TypeError, ValueError): + work = 1 + shape_work[shape_id] = work + + worker_count = min(4, len(shapes)) + bins: list[tuple[int, list[str]]] = [ + (0, []) for _ in range(worker_count) + ] + for shape_id in sorted( + shapes, key=lambda name: (-shape_work[name], name) + ): + index = min( + range(worker_count), key=lambda item: (bins[item][0], item) + ) + current_work, current_shapes = bins[index] + bins[index] = ( + current_work + shape_work[shape_id], + current_shapes + [shape_id], + ) + return { + f"worker_{gpu}": {"gpu": gpu, "shapes": bins[gpu][1]} + for gpu in range(worker_count) + if bins[gpu][1] + } + + +def w8a8_strategy_guidance( + shapes: Dict[str, Dict[str, Any]], +) -> str: + """Return gfx928 strategy guidance tailored to the assigned M values.""" + m_values = { + int(params["M"]) + for params in shapes.values() + if "M" in params + } + if m_values and max(m_values) <= 4: + return """This is a small-M decode lane (M <= 4). Start with a +bandwidth-conscious scalar/vector path, not DUMMA by reflex: +- map adjacent wavefront lanes to adjacent N columns so B[k, n] loads are + coalesced at each K step; +- use blockDim 128 or 256 (a multiple of the gfx928 wavefront size 64); +- stage/reuse the tiny A tile in LDS or registers while accumulating int32; +- if staging A in LDS, stage the whole tiny A matrix once and use at most one + synchronization before compute; never add a barrier for every K tile; +- let one thread or lane compute one or a small fixed number of output columns; +- do not use split-K in bootstrap. Consider it later only when measured grid + parallelism is insufficient and include reduction/combine cost.""" + if m_values == {16}: + return """This is an M=16 lane. The primary candidate is native gfx928 +DUMMA INT8 m16n16k32: +- include `` and use the exact installed DTK API below: + `du::dumma::DUFragment` + for A, the analogous `matrix_b` fragment for B, and + `du::dumma::DUFragment` for C; +- call `du::dumma::du_fill_fragment`, `du::dumma::du_load_matrix_sync`, + `du::dumma::du_mma_sync`, and `du::dumma::du_store_matrix_sync`; +- `row_major` is a fragment layout type, while the store layout argument is + `du::dumma::mem_row_major`; do not use CUDA-like unprefixed function names; +- one 64-thread wavefront owns each independent DUMMA fragment/tile; +- K is divisible by 32 and N by 16 for every assigned shape; +- use int8 A/B fragments and int32 accumulation, then apply x_scale[m] and + packed_weight_scale[n] before bf16 store; +- choose 1, 2, or 4 wavefronts per block only after checking N-grid + parallelism, LDS use, and register pressure; +- use coalesced global-to-LDS loads and a layout accepted by the installed + du_mma.h. Do not invent CUDA WMMA/PTX APIs.""" + if m_values and min(m_values) >= 128: + return """This is a large-prefill lane (M >= 128). Treat it as a +throughput GEMM, not a decode kernel: +- use native gfx928 INT8 DUMMA m16n16k32 and build larger 2-D macro-tiles from + multiple wavefront-owned fragments; +- sweep M/N block tiles and waves per block; start from 64x64, 64x128 and + 128x64 output tiles, then select using measured VGPR, LDS and occupancy; +- cooperatively vector-load A and packed B into LDS with padded/swizzled + layouts, reuse both operands across adjacent fragments, and compare single + versus double buffering using measured memory stalls; +- keep int32 accumulation resident through the full K loop, fuse scale and + bf16 conversion into the final store, and avoid a separate epilogue; +- do not use split-K by default: the MxN grid already exposes substantial + parallelism. Test it only for an evidenced K-pipeline imbalance and include + workspace reduction cost; +- report achieved TOPS and HBM traffic together with median/P90 so compute, + LDS, and memory bottlenecks are distinguished.""" + return """This lane contains different M regimes. Implement explicit +shape/M dispatch: use a coalesced scalar/LDS path for M <= 4 and evaluate +native gfx928 DUMMA INT8 m16n16k32 for M=16. For M>=128, use a throughput +DUMMA macro-tile with 2-D A/B reuse and tune tile/LDS/occupancy. Do not force +one launch geometry onto all regimes.""" + + +def w8a8_round_strategy( + shape: Dict[str, Any], + iteration: int, + history: list[Dict[str, Any]] | None = None, + pmc_evidence: Dict[str, Any] | None = None, + max_iterations: int = 10, + isa_policy: Dict[str, Any] | None = None, +) -> str: + """Choose an architecture-first round plan from measured evidence.""" + history = history or [] + pmc_evidence = pmc_evidence or {} + isa_policy = isa_policy or {} + + faster_wrong = [ + record for record in history + if record.get("correctness_passed") is False + and isinstance(record.get("speedup"), (int, float)) + and float(record["speedup"]) > 1.0 + ] + if faster_wrong: + candidate = max( + faster_wrong, key=lambda record: float(record["speedup"]) + ) + return ( + "Highest priority: repair the faster but incorrect candidate from " + f"iteration {candidate.get('iteration')} (measured speedup " + f"{candidate.get('speedup')}x). Read its archived source at " + f"`{candidate.get('artifact_dir')}` and preserve the fast mapping. " + "Fix only the smallest correctness defect: signed int8 unpacking, " + "tail bounds, scale indexing, bf16 conversion, or a race. Do not " + "replace it with an unrelated architecture." + ) + + if history: + latest = history[-1] + failure = str(latest.get("failure_reason") or "").lower() + if any(token in failure for token in ( + "timeout", "timed out", "killed", "no result", "exit 143", + )): + return ( + "The preceding attempt failed in the agent infrastructure. " + "Return to the accepted best source and start a new bounded " + "experiment; do not repair or replay a partially written " + "candidate. This failure does not count as a completed " + "optimization round." + ) + if latest.get("build_success") is False: + return ( + "Repair the immediately preceding candidate from iteration " + f"{latest.get('iteration')} at `{latest.get('artifact_dir')}`. " + "Keep its strategy and make only the minimum compile/API/syntax " + "correction; do not start another redesign this round." + ) + + phase = str(isa_policy.get("phase") or "hip_only") + if phase == "isa_guided_hip": + completed = int(isa_policy.get("valid_isa_guided_rounds") or 0) + return ( + "ISA-guided HIP round. This is successful ISA experiment " + f"{completed + 1} of at least 2. Select one measured memory or " + "compute bottleneck, compare the exact primary-kernel ISA with " + "the preceding code object, and make one HIP/DUMMA/intrinsic " + "code-shaping change. Raw inline asm remains forbidden. Record " + "a compiler limitation only when the before/after binary proves " + "it and name the exact target instructions." + ) + if phase == "conditional_inline_asm": + return ( + "Conditional inline-asm experiment. Target only the compiler " + "limitation and exact instructions verified by the immediately " + "preceding ISA-guided HIP round. Keep the asm block minimal, " + "preserve complete constraints/clobbers, and reject it unless " + "the candidate ISA, exact correctness, median, P90, and resources " + "all validate. Do not write raw global/buffer/flat loads or MMAC." + ) + + small_m = { + 1: "Vectorize contiguous K loads with exact signed-int8 semantics.", + 2: ( + "Increase instruction-level parallelism with independent int32 " + "accumulators or adjacent N outputs; avoid per-K-tile LDS barriers." + ), + 3: ( + "Stage all of tiny A once with at most one barrier, or tune unroll " + "one step if whole-A staging is not cheaper." + ), + 4: ( + "Change one launch variable only: waves per block, N columns per " + "wave, or unroll factor." + ), + 5: ( + "HIP-only memory round: change one vector-load width, contiguous " + "N mapping, or whole-A reuse decision. Raw inline asm is forbidden." + ), + 6: ( + "HIP-only pipeline round: reduce one dependency chain or barrier " + "using ordinary HIP/intrinsics. Raw inline asm is forbidden." + ), + 7: ( + "HIP-only resource round: tune one block size, unroll factor, or " + "live range while preserving coalescing. Raw inline asm is forbidden." + ), + 8: ( + "HIP-only consolidation round: revisit the fastest correct archived " + "mapping and make one final architecture/codegen improvement. Raw " + "inline asm is forbidden." + ), + } + m16 = { + 1: ( + "Establish a minimal 16x16x32 DUMMA tile with the exact API below, " + "one wave per output tile and explicit int32 accumulation. If the " + "seed already has a correct DUMMA kernel, preserve it and instead " + "test the smallest one-wave-per-block, one-N-tile geometry with no " + "cross-wave barrier; do not spend the round reimplementing it." + ), + 2: ( + "Architecture round: measure grid parallelism before polishing. " + "Explore one complete launch geometry among 1/2/4 waves per block " + "and 1/2/4 adjacent N tiles. Prefer enough independent blocks to " + "cover at least all device CUs; report grid_blocks, waves_per_block " + "and estimated_active_cus in proposal.json." + ), + 3: ( + "Architecture round: if the unsplit grid has fewer than two " + "blocks per device CU and K >= 1024, implement and measure " + "split-K=2 plus at least one CU-aligned candidate (which may be " + "non-power-of-two), or test a one-wave zero-barrier geometry that " + "reaches the same parallelism. Write int32 partials into the " + "caller workspace and include the combine+scale kernel in the " + "timed Graph." + ), + 4: ( + "Architecture/pipeline round: explore one of multi-N-tile reuse, " + "A-only staging, or bounded register/LDS prefetch. Retain enough " + "blocks to cover all CUs, state which A/B bytes are reused, and " + "measure whether the change improves normal median/P90." + ), + 5: ( + "HIP-only packed-weight/staging round: compare one packed layout, " + "A-only staging, or B-only staging design. Raw inline asm is forbidden." + ), + 6: ( + "Pipeline round: choose exactly one staging family from direct, " + "A-only LDS, B-only LDS, or A+B LDS using L2/VMEM evidence. Use " + "coalesced 8- or 16-byte cooperative loads and report HBM/LDS byte " + "changes; do not claim asynchronous overlap without evidence." + ), + 7: ( + "Pipeline round: compare single buffering with double buffering " + "only when K>=1024, L2 hit rate is below 70%, and the doubled LDS " + "budget stays below 48 KiB. Count barriers per K step." + ), + 8: ( + "HIP-only resource round: tune one occupancy limiter using actual " + "PMC evidence: waves per block, VGPR live range, LDS footprint, " + "or spill removal. Do not trade repeated HBM reads for occupancy." + ), + 9: ( + "Late ISA-diagnosis round. Only if the control-plane plateau gate " + "is open, use one selected ISA Skill and trusted disassembly to " + "shape compiler output through HIP/DUMMA/intrinsics. Raw inline " + "asm remains forbidden. Otherwise continue HIP-only exploration." + ), + 10: ( + "Final conditional inline-asm round. Raw asm is allowed only when " + "the control plane confirms a HIP plateau and the prior ISA-guided " + "round recorded one concrete compiler limitation plus target " + "instructions. Otherwise make one HIP-only consolidation change." + ), + } + large_m = { + 1: ( + "Establish a correct DUMMA throughput baseline using a 2-D " + "macro-tile. Benchmark 64x64, 64x128, and 128x64 block tiles; " + "record waves per block, VGPRs, LDS bytes, occupancy and TOPS." + ), + 2: ( + "Operand-reuse round: compare direct loads with cooperative " + "A+B LDS staging. Quantify A/B reuse per macro-tile and use " + "vectorized coalesced global loads with a bank-safe LDS layout." + ), + 3: ( + "Pipeline round: compare single and double buffering across K " + "tiles. Retain double buffering only when ISA/PMC evidence shows " + "reduced VMEM stalls without harmful LDS or occupancy growth." + ), + 4: ( + "Tile-shape round: tune M-tile versus N-tile aspect ratio for " + "this exact M/N/K, balancing B reuse, A reuse and enough blocks " + "to occupy every CU. Do not inherit decode launch geometry." + ), + 5: ( + "Packing round: test one weight packing/swizzle that makes each " + "DUMMA B tile vector-loadable and LDS-bank-safe. Include packing " + "outside timing and validate the graph-stable packed layout." + ), + 6: ( + "Epilogue round: fuse per-row and per-column scales, bf16 " + "conversion and the final coalesced store into the compute " + "kernel; remove any unnecessary workspace/combine pass." + ), + 7: ( + "Compute-pipeline round: tune DUMMA issue grouping, prefetch " + "distance and accumulator independence using ISA stall evidence. " + "Raw inline asm remains forbidden." + ), + 8: ( + "Resource round: tune waves per block, VGPR live ranges and LDS " + "footprint from measured occupancy. Recheck the best tile family " + "with normal median/P90 measurements." + ), + 9: m16[9], + 10: m16[10], + } + m = int(shape.get("M", 0)) + if m >= 128: + portfolio = large_m + elif m < 16: + portfolio = small_m + else: + portfolio = m16 + late_start = max(9, max_iterations - 1) + if m < 16: + if iteration < late_start: + return portfolio.get(iteration, portfolio[8]) + return portfolio[9] if iteration < max_iterations else portfolio[10] + if iteration < late_start: + decision = portfolio.get(iteration, portfolio[8]) + else: + decision = portfolio[9] if iteration < max_iterations else portfolio[10] + grid_blocks = pmc_evidence.get("grid_blocks") + cu_count = pmc_evidence.get("device_cu_count") + split_candidates = split_k_candidate_set( + shape, + cu_count=int(cu_count) if isinstance(cu_count, (int, float)) else 0, + ) + if iteration >= 2 and split_candidates: + decision += ( + " Trusted occupancy-probe split candidates for the measured CU " + f"count are {split_candidates}. They include non-power-of-two " + "values where useful, are not a whitelist, and must fit the " + "workspace and stage-alignment constraints. Explore outside this " + "set when evidence supports it." + ) + if ( + iteration >= 2 + and isinstance(grid_blocks, (int, float)) + and isinstance(cu_count, (int, float)) + and grid_blocks < 2 * cu_count + ): + decision += ( + f" Trusted control-plane warning: current grid has {grid_blocks} " + f"blocks for {cu_count} CUs, below the two-blocks-per-CU latency-" + "hiding target. Before micro-optimization, benchmark a finer " + "one-wave zero-barrier grid or multiple legal split-K candidates " + "including combine cost." + ) + return decision + + +def generate_kernel_prompt( + *, + operator: str, + dtype: str, + shapes: Dict[str, Dict[str, Any]], + hardware: str, + kernel_language: str, + source_dir: Path, + harness_path: Path, + api_contract_path: Path | None = None, + iteration: int = 0, + prev_failure: Optional[str] = None, + fixed_assignment: Dict[str, Dict[str, Any]] | None = None, +) -> str: + """Prompt the main Agent to coordinate, never implement, kernels.""" + contract_path = api_contract_path or ( + source_dir / "int8_w8a8_gemm_api.py" + ) + failure_block = "" + if prev_failure: + failure_block = f""" +## Previous attempt failed + +The last coordination attempt failed with: + +{prev_failure} + +Correct only `proposal.json`. Do not respond by writing implementation code. +""" + assignment = ( + fixed_assignment + if fixed_assignment is not None + else shape_balanced_assignment(shapes) + ) + assignment_example = json.dumps(assignment, indent=2) + assignment_rule = ( + "The control plane has already selected the authoritative " + "`gpu_assignment` shown in the deliverable. Preserve it exactly; " + "do not rebalance, regroup, add, or remove shapes." + if fixed_assignment is not None + else ( + "Use the deterministic shape-balanced `gpu_assignment` shown " + "in the deliverable." + ) + ) + return f"""You are the MAIN COORDINATOR for a DCU kernel optimization task. + +## Your responsibilities + +1. Read the immutable Python API contract at `{contract_path}`. +2. Read the trusted Generate preflight evidence at + `{source_dir / "generation_preflight.json"}`. +3. Inspect the staged correctness harness at `{harness_path}`, the build + scaffold, `profile_pmc.sh`, hardware, and complete shape list below. +4. Confirm that the repository contains no HIP implementation. +5. Review or preserve the GPU assignment as instructed below. +6. Write only `{source_dir / "proposal.json"}`. + +## Hard role boundary + +You are not a kernel implementation agent. Do not write, edit, rename, or +delete any HIP, C++, CUDA, Python backend, setup/build, test, benchmark, or +public API file. In particular, do not create `.hip`, `.cu`, `.cpp`, +`w8a8_backend.py`, or `setup.py`. Do not compile or run the correctness +harness. The control plane has staged only the immutable API and build/loader +scaffolding. After assignment, child implementation Agents create their +initial HIP kernels from scratch during Parallel explore, and the control +plane validates them. + +The API contract is user-owned and immutable. The repository scaffold and +interface have already been prepared by the control plane. +The fixed public call is `w8a8_gemm_out(...)`; its backend operation is +`torch.ops.zth_w8a8.gemm_out`. +The trusted scaffold contains `w8a8_bench.py` and `profile_pmc.sh`. Do not +modify or execute them. The control plane has already run the harness's CPU +PyTorch-reference self-test, probed GPU visibility, checked PMC script syntax, +verified the real hipprof command, and recorded those facts in +`generation_preflight.json`. Actual kernel correctness and PMC collection are +deferred until a child Agent creates a kernel during Parallel explore. + +## GPU assignment policy + +## Task context + +- Operator: {operator} +- Dtype: {dtype} +- Hardware: {hardware} +- Kernel language for child agents: {kernel_language} +- Child workers: between 1 and 4 non-empty workers +- Mapping: worker_N→physical GPU N; one to four non-empty workers +- Assignment instruction: {assignment_rule} + +| ID | M | N | K | +|---|---|---|---| +{_shapes_table(shapes)} + +The assignment unit is one exact shape ID. Different M variants of the same +logical operator may be assigned to different workers; for example, +`m4_wqkv_a` and `m16_wqkv_a` may use different shape-specific HIP kernels. +The fixed public API remains one interface, and final synthesis dispatches to +the selected kernel by shape. Every target shape must appear exactly once. +{failure_block} +## Deliverable + +Write strict JSON to `{source_dir / "proposal.json"}`: + +```json +{{ + "iteration": {iteration}, + "generated": false, + "hypothesis": "brief shape grouping and load-balance rationale", + "profile_evidence": {{ + "hardware": "{hardware}", + "coordination_only": true + }}, + "profiling_plan": {{ + "script": "profile_pmc.sh", + "mode": "hipprof_pmc_csv", + "trigger": "usable DUMMA bootstrap, accepted best, or late ISA decision", + "reuse_when_source_digest_matches": true, + "skip_scalar_bootstrap": true, + "acceptance_timing": "unprofiled_cuda_graph_replay_median_p90" + }}, + "scaffold_review": {{ + "preflight_file": "generation_preflight.json", + "preflight_status": "passed", + "harness_reference_self_test_passed": true, + "gpu_probe_passed": true, + "cudagraph_available": true, + "python_graph_wrapper_staged": true, + "pmc_script_checked": true, + "no_hip_implementation": true + }}, + "gpu_assignment": {assignment_example} +}} +``` + +Every target shape must appear exactly once. Every emitted worker must receive +at least one shape. Use up to four GPUs; a subset task may use fewer when its +operator-family grouping does not provide four independent units of work. +Do not change any other file. +""" + + +def bootstrap_worker_prompt( + *, + worker_id: str, + gpu: int, + shapes: Dict[str, Dict[str, Any]], + hardware: str, + kernel_language: str, + source_dir: Path, + harness_path: Path, + api_contract_path: Path, + attempt: int, + prev_failure: Optional[str] = None, +) -> str: + """Prompt one child to generate a new initial HIP implementation.""" + del harness_path + failure_block = "" + if prev_failure: + failure_block = f""" +## Previous attempt failed + +{prev_failure} + +Make the smallest source-only correction for the reported issue before +returning it to the trusted control plane. Do not inspect the machine, +toolchain, network, or unrelated files. +""" + has_large_prefill = any( + int(shape.get("M", 0)) >= 128 for shape in shapes.values() + ) + if has_large_prefill: + bootstrap_strategy = """ +Bootstrap is iteration 0 and remains correctness-first, but a large-Prefill +scalar K loop is not a usable profiling baseline. For every assigned shape +with M >= 128, create a simple native INT8 DUMMA m16n16k32 tiled kernel with +int32 accumulation. Use enough M/N tiles to expose CU parallelism; keep the +first implementation direct or single-buffered and avoid split-K, raw asm, +or speculative deep pipelines. + +When variant reference code is present, you may read and adapt measured +implementations as reference. Variants live in a staged tree under +`references/` organized as `////.hip` +(for example `references/int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip`); the +legacy flat file `references/w8a8_gemm_variants.hip` is also available. +Navigate the tree with `ls`/`find` to the directory matching your exact +operator/dtype, model, TP, and M family, then read only the specific operator +file you need. Variants are neither a whitelist nor a restriction: take only +the minimum code needed for a correct starting point and continue exploring +freely in later rounds. Do not claim their measurements until the trusted +control plane revalidates this source. Reusing a variant never locks you into +it: adapt freely, combine ideas, or ignore it and start fresh — every +accepted kernel must still pass the full correctness/Graph/median-P90/ +resource validation. + +Keep one simple scalar int8/int32 fallback for unmatched shapes and small-M +API cases. Shapes with M < 128 may use that scalar fallback for bootstrap. +""" + bootstrap_path = "dumma_prefill_with_scalar_fallback" + else: + bootstrap_strategy = """ +Bootstrap is iteration 0, not a performance round. For every assigned shape, +including M=16, implement one simple scalar int8 dot-product kernel: + +- map one thread to one output element or one/few adjacent N columns; +- use blockDim 128 or 256 and a straightforward grid over M*N; +- compute the complete K loop exactly in int32, then apply the two float + scales and store bf16; +- prioritize code that compiles and passes exact correctness. + +Do not use DUMMA, split-K, double buffering, complicated LDS layouts, inline +assembly, or speculative performance machinery for this small-M bootstrap. +Those belong to measured optimization iterations after correctness passes. +""" + bootstrap_path = "scalar_correctness" + return f"""You are CHILD kernel implementation agent `{worker_id}` in +Parallel explore. You own physical GPU {gpu} and only these assigned shapes: + +| ID | M | N | K | +|---|---|---|---| +{_shapes_table(shapes)} + +The main coordinator has already fixed the API, repository, hardware, and +shape/GPU assignment. This is a new task, not continuation mode: no HIP +implementation has been supplied. Create the initial HIP kernel for the +assigned shapes from scratch. Only the staged optional variant file may be +adapted as a reference under the rules below; do not copy another task repo. + +## Immutable interface + +Read `{api_contract_path}` and preserve it byte-for-byte. Implement its private +backend contract without creating a second public API. + +## Required implementation + +The control plane owns these staged scaffold files under `{source_dir}`: + +- `w8a8_backend.py` +- `w8a8_graph.py` +- `setup.py` +- `csrc/bindings.cpp` + +You own and must create `csrc/w8a8_gemm_hip.hip`. Do not assume an existing +compiled kernel is present. + +The trusted binding already registers both `gemm_out` and the optional +out-of-timed-region `pack_weight` operation. Your HIP file must provide these +two stable host launch symbols: + +- `launch_w8a8_gemm(..., void* workspace, int64_t workspace_bytes, + int m, int n, int k, hipStream_t stream)`; +- `launch_pack_w8a8_weight(raw_weight, weight_scale, packed_weight, + packed_weight_scale, int k, int n, hipStream_t stream)`. + +Bootstrap `launch_pack_w8a8_weight` as an identity device-to-device copy. +Later Parallel explore rounds may change only its HIP implementation and the +matching GEMM interpretation to test packed layouts. The main coordinator and +Generate phase never write the HIP implementation. + +Register `torch.ops.zth_w8a8.gemm_out`, return the exact caller-provided `out` +tensor, use PyTorch's current HIP stream, and preserve graph-safe behavior. +The trusted harness rejects the bootstrap unless this fixed Python API can be +captured on a non-default stream with `torch.cuda.CUDAGraph`, replayed for +exact correctness, and called from Python through `w8a8_graph.py`. +The mathematical reference is: +`(A.float() @ B.float()) * x_scale * weight_scale.T`, converted to bfloat16. + +The timed operator must perform no allocation, compilation, autotuning, +weight packing, host synchronization, device synchronization, or default- +stream launch. It may use only the caller-provided `out` and `workspace`. +Weight packing, if implemented, belongs in the optional `pack_weight` op +outside the timed region. + +`w8a8_backend.py` must keep the control-plane entry point +`load_extension()`. It uses `torch.utils.cpp_extension.load(..., +is_python_module=False)` because the operator is registered through +`TORCH_LIBRARY`. Never replace it with an import of a prebuilt `.so` or an +installed package. + +## Hardware + +- Target: {hardware} +- Language: {kernel_language} +- Compile target: gfx928 +- Visible physical GPU: {gpu} +- Native wavefront size: 64 +- LDS capacity: 64 KiB per CU +- INT8 DUMMA support is m16n16k32 with int32 accumulation +{failure_block} + +## Correctness-first bootstrap strategy + +{bootstrap_strategy} + +Include ``, use `hip_bfloat16`, and use its supported +float conversion; alternatively use the existing proven manual uint16 store +without switching type families. + +For every path: + +1. Include the installed headers in this known-good order: + ``, then ``, then ``. + This DTK's `du_mma.h` is not self-contained when included before the HIP + runtime headers. Do not include or invent `du_mma_common.h`; it is not an + installed public header. +2. Keep adjacent lanes on adjacent addresses in the fastest-changing N + dimension. Audit every hard-coded warp value: gfx928 wavefront is 64, not + 32; blockDim must be a multiple of 64. +3. Accumulate the integer dot product in int32. The maximum assigned K keeps + the exact int8 dot within int32 range. Convert to float only for + `dot * x_scale[m] * packed_weight_scale[n]`, then store bf16. +4. Use `hipLaunchKernelGGL` on PyTorch's current HIP stream. Never use + stream 0, `hipDeviceSynchronize`, `hipStreamSynchronize`, or a temporary + allocation in `gemm_out`. +5. Never use NVIDIA `wmma`, `mma.sync`, PTX, warp=32 masks, FP8, or INT4 on + gfx928. +6. Keep all barriers on paths reached by every thread in the block. Budget + LDS before choosing single/double buffering; bootstrap should prefer one + correct, understandable buffer over speculative complexity. +7. Support every assigned shape through explicit dispatch where launch + geometry differs. Do not specialize for only the first shape. +8. Keep a scalar generic fallback for every unmatched `(m,n,k)`. Later + optimization rounds will specialize only an exact shape and the accepted + object will be linked directly into the final extension. In particular, + an M=16 specialization must remain guarded so the paired M=2 API shape + with the same `(N,K)` still reaches the scalar fallback correctly. Keep + `launch_pack_w8a8_weight` valid for unmatched `(K,N)` as well; identity + device-to-device packing is its generic fallback. + +## Execution boundary + +Do not run the benchmark or correctness harness. Do not invoke Docker, SSH, +Skill tools, pip, apt, conda, network access, package installation, +environment activation, filesystem-wide searches such as `find /`, or +PyTorch/CUDA environment probes. Read only the immutable API, named scaffold +files, and optional `references/w8a8_gemm_variants.hip`. The trusted control plane already owns the correct +DTK/PyTorch environment and will compile, run exact correctness, and measure +median/P90 after you return. Your job is source implementation only. + +You may inspect the immutable API and scaffold source read-only. Preserve +`w8a8_backend.py`, `setup.py`, and `csrc/bindings.cpp`. You may change only +`csrc/w8a8_gemm_hip.hip` plus `proposal.json`. Do not edit tests, the harness, +cache/build artifacts, or files outside this worktree. + +Then write `{source_dir / "proposal.json"}` as strict JSON: + +```json +{{ + "iteration": {attempt}, + "generated": true, + "hypothesis": "specific initial kernel strategy chosen for these M/N/K shapes", + "profile_evidence": {{ + "worker_id": "{worker_id}", + "physical_gpu": {gpu}, + "shapes_targeted": {json.dumps(list(shapes))}, + "path": "{bootstrap_path}", + "block_threads": 128, + "tile_m": 0, + "tile_n": 0, + "tile_k": 0, + "lds_bytes": 0, + "split_k": 1, + "expected_bottleneck": "memory, compute, launch, or occupancy", + "risk": "main correctness or performance risk", + "validation_owner": "trusted_control_plane" + }}, + "files_changed": [ + "csrc/w8a8_gemm_hip.hip" + ] +}} +``` +""" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/real_pipeline.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/real_pipeline.py new file mode 100644 index 00000000..dc36db91 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/real_pipeline.py @@ -0,0 +1,652 @@ +"""Real Claude-agent + real DCU smoke pipeline. + +This validates the production isolation and lifecycle without pretending that +the built-in vector kernel is the user's future operator adapter. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import Any, Dict + +from metainfer.orchestrator.state import StateStore +from metainfer.orchestrator.subagent_manager import AgentSpec, SubAgentManager + +from . import phases +from .config import ( + ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT, + OptimizerConfig, + WorkerAssignment, + load_config, +) +from .gpu_binding import bind_worker_gpu +from .guidance import claim_next_guidance +from .result_store import SCHEMA_VERSION, append_jsonl, write_json +from .skill_store import generate_merged_skill, generate_worker_skill + + +ASSET = Path(__file__).resolve().parent.parent / "assets" / "smoke_harness.cpp" + + +def _run( + command: list[str], + *, + cwd: Path, + env: Dict[str, str] | None = None, + timeout: int = 180, +) -> subprocess.CompletedProcess[str]: + effective_command = list(command) + if effective_command and effective_command[0] == "git": + # Repositories mounted from the host are host-user-owned, while the + # orchestrator runs as root. Trust only this exact repository (and its + # main worktree when cwd is a linked worktree), never a global wildcard. + safe_directories = [cwd.resolve()] + git_marker = cwd / ".git" + if git_marker.is_file(): + marker = git_marker.read_text(encoding="utf-8").strip() + if marker.startswith("gitdir:"): + git_dir = Path(marker.removeprefix("gitdir:").strip()) + common_repo, separator, _ = str(git_dir).partition("/.git/worktrees/") + if separator: + safe_directories.append(Path(common_repo).resolve()) + + effective_command = ["git"] + for safe_directory in dict.fromkeys(safe_directories): + effective_command.extend(["-c", f"safe.directory={safe_directory}"]) + effective_command.extend(command[1:]) + + result = subprocess.run( + effective_command, cwd=cwd, env=env, text=True, capture_output=True, + timeout=timeout, check=False, + ) + if result.returncode != 0: + raise RuntimeError( + f"{' '.join(command)} failed ({result.returncode}): " + f"{result.stderr[-2000:] or result.stdout[-2000:]}" + ) + return result + + +def _last_json(text: str) -> Dict[str, Any]: + for line in reversed(text.splitlines()): + try: + value = json.loads(line) + except ValueError: + continue + if isinstance(value, dict): + return value + raise ValueError(f"no JSON object in output: {text[-1000:]}") + + +def _elements(shape: Dict[str, Any]) -> int: + try: + value = int(shape.get("M", 1)) * int(shape.get("N", 1)) + except (TypeError, ValueError): + value = 1 << 20 + value = max(1 << 20, min(value, 16 << 20)) + return (value + 3) // 4 * 4 + + +class SmokeRunner: + def __init__(self, worker_root: Path, gpu: int) -> None: + self.worker_root = worker_root + self.source = worker_root / "source" + self.binary = worker_root / "build" / "smoke_harness" + self.env = dict(os.environ) + bind_worker_gpu(self.env, gpu) + self.env.update({ + "TORCH_EXTENSIONS_DIR": str(worker_root / "cache" / "torch"), + "TRITON_CACHE_DIR": str(worker_root / "cache" / "triton"), + "XDG_CACHE_HOME": str(worker_root / "cache" / "xdg"), + "TMPDIR": str(worker_root / "cache" / "tmp"), + }) + for key in ( + "TORCH_EXTENSIONS_DIR", "TRITON_CACHE_DIR", "XDG_CACHE_HOME", + "TMPDIR", + ): + Path(self.env[key]).mkdir(parents=True, exist_ok=True) + + def build(self) -> None: + self.binary.parent.mkdir(parents=True, exist_ok=True) + _run([ + "/opt/dtk/bin/hipcc", "-O3", "--offload-arch=gfx928", + str(self.source / "smoke_harness.cpp"), "-o", str(self.binary), + ], cwd=self.source, env=self.env, timeout=240) + + def probe(self) -> Dict[str, Any]: + return _last_json( + _run( + [str(self.binary), "--probe"], cwd=self.source, env=self.env + ).stdout + ) + + def benchmark(self, shape: Dict[str, Any], variant: str) -> Dict[str, Any]: + if variant not in {"scalar", "vector4"}: + raise ValueError(f"unsupported smoke variant: {variant}") + result = _run( + [str(self.binary), str(_elements(shape)), variant], + cwd=self.source, env=self.env, timeout=180, + ) + return _last_json(result.stdout) + + +def _status( + worker_root: Path, + assignment: WorkerAssignment, + *, + state: str, + iteration: int, + shape_id: str | None, + probe: Dict[str, Any] | None = None, + **details: Any, +) -> None: + payload = { + "schema_version": SCHEMA_VERSION, + "worker_id": assignment.worker_id, + "state": state, + "iteration": iteration, + "shape_id": shape_id, + "pid": os.getpid(), + "physical_gpu": assignment.gpu, + "logical_gpu": 0, + "gpu_binding": { + "HIP_VISIBLE_DEVICES": str(assignment.gpu), + "strategy": "HIP_VISIBLE_DEVICES-only", + "enforced": True, + "visible_devices": (probe or {}).get("visible_devices"), + "device_name": (probe or {}).get("device_name"), + }, + "last_update": time.time(), + } + payload.update(details) + write_json(worker_root / "status.json", payload) + + +class RealSmokeOptimizationPipeline: + def __init__( + self, + *, + req: Dict[str, Any], + state_dir: Path, + workspace_dir: Path, + store: StateStore, + manager: SubAgentManager, + ) -> None: + self.req = req + self.state_dir = state_dir + self.workspace_dir = workspace_dir + self.store = store + self.manager = manager + self._current_phase = phases.PREPARE + self._progress_lock = threading.Lock() + self._reported_iteration = 0 + + def _phase(self, phase: str, **payload: Any) -> None: + self._current_phase = phase + self.store.update_run(current_phase=phase) + self.store.append_timeline("phase_start", {"phase": phase, **payload}) + + def run(self, *, dry_run: bool = False) -> Dict[str, Any]: + task_id = str(self.req.get("task_id", "task")) + self.store.init_or_resume(task_id) + self.store.update_run( + finished=False, + final_status=None, + last_outcome=None, + last_transition_label=None, + notes=[], + ) + started = time.time() + try: + self._phase(phases.PREPARE) + config = load_config(self.req) + self._prepare_worktrees(config, task_id) + plan = self._plan(config) + write_json(self.workspace_dir / "plan.json", plan) + if dry_run: + return plan + + self._phase(phases.BASELINE) + baseline = self._parallel_baseline(config) + + self._phase(phases.EXPLORE, workers=len(config.assignments)) + workers = self._parallel_agents(config, baseline) + + self._phase(phases.SYNTHESIZE) + merged_skill = generate_merged_skill( + config=config, assignments=config.assignments, + workspace_dir=self.workspace_dir, + ) + + self._phase(phases.VALIDATE) + validation = self._serial_validate(config, workers) + + self._phase(phases.REPORT) + report = { + "schema_version": SCHEMA_VERSION, + "task_id": task_id, + "task_type": "dcu-kernel-auto-opt", + "mode": "real-agent-dcu-smoke", + "started_at": started, + "finished_at": time.time(), + "duration_s": round(time.time() - started, 4), + "config": plan, + "baseline": baseline, + "workers": workers, + "merged_skill": merged_skill, + "final_validation": validation, + "real_gpu_used": True, + "target_repo_modified": False, + "status": "success", + } + write_json(self.workspace_dir / "final_report.json", report) + self.store.update_run( + current_iteration=config.mock_iterations, + current_phase=phases.FINISHED, + finished=True, + final_status="success", + last_outcome="ok", + last_transition_label="real smoke complete", + ) + self.store.append_timeline( + "orchestrator_success", + {"workers": len(workers), "real_gpu_used": True}, + ) + return report + except Exception as exc: + self.store.append_timeline( + "orchestrator_error", {"error": repr(exc)} + ) + self.store.update_run( + # Keep the state machine on the phase that failed. A stopped + # task must not look like a successfully completed workflow. + current_phase=self._current_phase, finished=True, + final_status="stopped", last_outcome="infra_fail", + notes=[str(exc)], + ) + raise + + def _prepare_worktrees( + self, config: OptimizerConfig, task_id: str + ) -> None: + self.workspace_dir.mkdir(parents=True, exist_ok=True) + seed = self.workspace_dir / "main" + if not (seed / ".git").exists(): + seed.mkdir(parents=True, exist_ok=True) + shutil.copy2(ASSET, seed / "smoke_harness.cpp") + _run(["git", "init"], cwd=seed) + _run(["git", "config", "user.name", "MetaInfer Agent"], cwd=seed) + _run([ + "git", "config", "user.email", "metainfer@localhost" + ], cwd=seed) + _run(["git", "add", "smoke_harness.cpp"], cwd=seed) + _run(["git", "commit", "-m", "seed trusted DCU smoke harness"], cwd=seed) + for assignment in config.assignments: + root = self.workspace_dir / "workers" / assignment.worker_id + for name in ("build", "cache", "logs", "runs", "artifacts"): + (root / name).mkdir(parents=True, exist_ok=True) + source = root / "source" + if not source.exists(): + branch = f"agent/{_safe(task_id)}/{assignment.worker_id}" + _run([ + "git", "worktree", "add", "-b", branch, + str(source), "HEAD", + ], cwd=seed) + for name in ("shared_baseline", "final_validation", "skills"): + (self.workspace_dir / name).mkdir(parents=True, exist_ok=True) + + @staticmethod + def _plan(config: OptimizerConfig) -> Dict[str, Any]: + return { + "schema_version": SCHEMA_VERSION, + "execution_mode": config.execution_mode, + "operator": config.operator, + "dtype": config.dtype, + "hardware": config.hardware, + "kernel_language": config.kernel_language, + "claude_model": config.claude_model, + "mock_iterations": config.mock_iterations, + "minimum_improvement_percent": config.minimum_improvement_percent, + "minimum_improvement_semantics": ( + "final validated result versus fixed baseline" + ), + "round_acceptance_improvement_percent": ( + ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + ), + "harness": "trusted built-in DCU vector smoke harness", + "shapes": [ + {"id": shape.id, **shape.params} + for shape in config.shapes.values() + ], + "assignments": [ + {"worker_id": item.worker_id, "gpu": item.gpu, + "shapes": item.shape_ids} + for item in config.assignments + ], + "real_gpu_used": True, + } + + def _parallel_baseline( + self, config: OptimizerConfig + ) -> Dict[str, Dict[str, Any]]: + output: Dict[str, Dict[str, Any]] = {} + + def run_one(assignment: WorkerAssignment) -> Dict[str, Dict[str, Any]]: + root = self.workspace_dir / "workers" / assignment.worker_id + runner = SmokeRunner(root, assignment.gpu) + _status(root, assignment, state="building", iteration=0, shape_id=None) + runner.build() + probe = runner.probe() + if probe.get("visible_devices") != 1: + raise RuntimeError( + f"{assignment.worker_id} sees " + f"{probe.get('visible_devices')} GPUs" + ) + _status( + root, assignment, state="baseline", iteration=0, + shape_id=None, probe=probe, + ) + return { + shape_id: runner.benchmark( + config.shapes[shape_id].params, "scalar" + ) + for shape_id in assignment.shape_ids + } + + with ThreadPoolExecutor(max_workers=len(config.assignments)) as pool: + futures = { + pool.submit(run_one, item): item for item in config.assignments + } + for future in as_completed(futures): + output.update(future.result()) + write_json( + self.workspace_dir / "shared_baseline" / "results.json", + {"schema_version": SCHEMA_VERSION, "shapes": output}, + ) + return output + + def _parallel_agents( + self, + config: OptimizerConfig, + baseline: Dict[str, Dict[str, Any]], + ) -> Dict[str, Any]: + output: Dict[str, Any] = {} + with ThreadPoolExecutor(max_workers=len(config.assignments)) as pool: + futures = { + pool.submit( + self._run_worker, config, assignment, baseline + ): assignment + for assignment in config.assignments + } + for future in as_completed(futures): + assignment = futures[future] + result = future.result() + result["skill"] = generate_worker_skill( + config=config, assignment=assignment, + workspace_dir=self.workspace_dir, + ) + output[assignment.worker_id] = result + self.store.append_timeline( + "worker_complete", + {"worker_id": assignment.worker_id, + "skill": result["skill"]["name"]}, + ) + return dict(sorted(output.items())) + + def _run_worker( + self, + config: OptimizerConfig, + assignment: WorkerAssignment, + baseline: Dict[str, Dict[str, Any]], + ) -> Dict[str, Any]: + root = self.workspace_dir / "workers" / assignment.worker_id + runner = SmokeRunner(root, assignment.gpu) + probe = runner.probe() + result: Dict[str, Any] = { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "branch": f"agent/{_safe(self.req.get('task_id', 'task'))}/" + f"{assignment.worker_id}", + "worktree_created": True, + "mode": "real-agent-dcu-smoke", + "gpu_probe": probe, + "shapes": {}, + } + for shape_id in assignment.shape_ids: + shape = config.shapes[shape_id] + best_variant = "scalar" + best_metrics = baseline[shape_id] + experiments_path = root / "runs" / shape_id / "experiments.jsonl" + for iteration in range(1, config.mock_iterations + 1): + guidance = claim_next_guidance( + self.state_dir / "guidance", + assignment.worker_id, + iteration, + ) + _status( + root, assignment, state="agent_running", + iteration=iteration, shape_id=shape_id, probe=probe, + ) + proposal_path = root / "source" / "proposal.json" + try: + proposal_path.unlink() + except FileNotFoundError: + pass + prompt = self._worker_prompt( + assignment, shape_id, shape.params, baseline[shape_id], + root, iteration, guidance, + ) + prompt_file = root / "logs" / ( + f"{shape_id}-iteration-{iteration}.prompt.txt" + ) + prompt_file.write_text(prompt, encoding="utf-8") + agent_name = ( + f"{assignment.worker_id}-{shape_id}-iter{iteration}" + ) + spec = AgentSpec( + name=agent_name, + role="dcu_kernel_worker", + prompt_file=prompt_file, + workdir=root / "source", + log_dir=root / "logs", + timeout_s=600, + stuck_timeout_s=240, + max_retries=0, + env_overrides=runner.env, + ) + self.store.append_timeline( + "agent_launch", + {"name": agent_name, "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu}, + ) + self.manager.launch(spec) + agent_result = self.manager.result(agent_name) + if agent_result is None or not agent_result.success: + raise RuntimeError( + f"{agent_name} failed: " + f"{agent_result.error if agent_result else 'no result'}" + ) + proposal = json.loads(proposal_path.read_text(encoding="utf-8")) + variant = str(proposal.get("variant")) + metrics = runner.benchmark(shape.params, variant) + speedup = ( + float(baseline[shape_id]["median_us"]) + / float(metrics["median_us"]) + ) + round_improvement = ( + float(best_metrics["median_us"]) + / float(metrics["median_us"]) + - 1.0 + ) * 100.0 + accepted = ( + bool(metrics.get("passed")) + and float(metrics["median_us"]) + < float(best_metrics["median_us"]) + and round_improvement + >= ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + ) + experiment = { + "schema_version": SCHEMA_VERSION, + "worker_id": assignment.worker_id, + "iteration": iteration, + "shape_id": shape_id, + "shape": shape.params, + "hypothesis": proposal.get("hypothesis"), + "changes": [f"select trusted smoke variant: {variant}"], + "profile_evidence": proposal.get("profile_evidence") or {}, + "build_success": True, + "correctness_passed": bool(metrics.get("passed")), + "metrics": { + key: metrics[key] for key in ( + "median_us", "p90_us", "min_us", "max_us", + "tflops", "bandwidth_gb_s", + ) + }, + "baseline_us": baseline[shape_id]["median_us"], + "speedup": round(speedup, 6), + "round_improvement_percent": round( + round_improvement, 6 + ), + "accepted": accepted, + "commit": None, + "failure_reason": None, + "manual_guidance": ( + guidance["text"] if guidance else None + ), + "guidance_id": guidance["id"] if guidance else None, + "agent_session_id": agent_result.session_id, + "timestamp": time.time(), + } + if accepted: + _run(["git", "add", "proposal.json"], cwd=root / "source") + _run([ + "git", "commit", "-m", + f"{shape_id}: select {variant} in iteration {iteration}", + ], cwd=root / "source") + commit = _run( + ["git", "rev-parse", "HEAD"], cwd=root / "source" + ).stdout.strip() + experiment["commit"] = commit + best_variant = variant + best_metrics = metrics + else: + tracked = _run( + ["git", "ls-files", "proposal.json"], + cwd=root / "source", + ).stdout.strip() + if tracked: + _run( + ["git", "restore", "proposal.json"], + cwd=root / "source", + ) + else: + proposal_path.unlink(missing_ok=True) + append_jsonl(experiments_path, experiment) + # Surface live progress in the task header while workers run + # concurrently. StateStore serializes updates within this + # orchestrator process; never move the displayed round back. + with self._progress_lock: + if iteration > self._reported_iteration: + self.store.update_run(current_iteration=iteration) + self._reported_iteration = iteration + result["shapes"][shape_id] = { + "shape_id": shape_id, + "variant": best_variant, + "metrics": best_metrics, + } + _status( + root, assignment, state="completed", + iteration=config.mock_iterations, shape_id=None, probe=probe, + ) + write_json(root / "result.json", result) + return result + + @staticmethod + def _worker_prompt( + assignment: WorkerAssignment, + shape_id: str, + shape: Dict[str, Any], + baseline: Dict[str, Any], + root: Path, + iteration: int, + guidance: Dict[str, Any] | None, + ) -> str: + guidance_text = ( + guidance["text"] if guidance else "(none; decide independently)" + ) + return f"""You are {assignment.worker_id}, an autonomous DCU smoke-tuning worker. +You are bound to physical GPU {assignment.gpu}; inside this process it must be +the only visible GPU and is logical device 0. + +This is a real infrastructure smoke run, not the final operator integration. +Do not edit smoke_harness.cpp. Inspect it and make an evidence-based choice +between its scalar and vector4 candidates for shape {shape_id}: {json.dumps(shape)}. +Baseline: {json.dumps(baseline)}. +Human guidance for this round: {guidance_text} + +First run `{root / 'build' / 'smoke_harness'} --probe`; stop if visible_devices +is not exactly 1. Then benchmark BOTH candidates with: +`{root / 'build' / 'smoke_harness'} {_elements(shape)} scalar` +`{root / 'build' / 'smoke_harness'} {_elements(shape)} vector4` + +Write `{root / 'source' / 'proposal.json'}` as strict JSON: +{{ + "iteration": {iteration}, + "variant": "scalar or vector4", + "hypothesis": "short evidence-based explanation", + "profile_evidence": {{ + "scalar_median_us": 0.0, + "vector4_median_us": 0.0, + "visible_devices": 1, + "device_name": "actual device" + }} +}} +Do not choose before measuring both. Do not change correctness coverage or work. +""" + + def _serial_validate( + self, config: OptimizerConfig, workers: Dict[str, Any] + ) -> Dict[str, Any]: + results: Dict[str, Any] = {} + for assignment in config.assignments: + runner = SmokeRunner( + self.workspace_dir / "workers" / assignment.worker_id, + assignment.gpu, + ) + for shape_id in assignment.shape_ids: + winner = workers[assignment.worker_id]["shapes"][shape_id] + metrics = runner.benchmark( + config.shapes[shape_id].params, winner["variant"] + ) + if not metrics.get("passed"): + raise RuntimeError( + f"serial validation failed: {shape_id}" + ) + results[shape_id] = { + "passed": True, + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "candidate": winner["variant"], + "metrics": metrics, + "serial": True, + } + write_json( + self.workspace_dir / "final_validation" / "results.json", + {"schema_version": SCHEMA_VERSION, "shapes": results, + "real_gpu_used": True}, + ) + return results + + +def _safe(value: Any) -> str: + return "".join( + char if char.isalnum() or char in "-_" else "-" + for char in str(value) + )[:48] diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/recover_serial_validate.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/recover_serial_validate.py new file mode 100644 index 00000000..0d9f0645 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/recover_serial_validate.py @@ -0,0 +1,239 @@ +"""Recovery driver: re-run final synthesis + serial validation + report. + +Use when a generate-and-optimize task stopped in ``serial_validate`` (e.g. an +infrastructure failure inside the validation bench) after the parallel lanes +already produced accepted artifacts. The driver reconstructs the per-worker +results from ``workers//accepted//manifest.json``, re-runs +``_synthesize_final_candidate`` (which rebuilds ``final/source`` and performs +the full serial validation against every API shape), writes ``final_report.json`` +and marks the task finished/success. + +The accepted artifacts are the single source of truth; the in-memory +``workers``/``initial_metrics`` of the crashed orchestrator are rebuilt from +disk (accepted manifests + ``shared_baseline/results.json``). +""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from pathlib import Path + +from metainfer.orchestrator._bootstrap import make_subagent_manager +from metainfer.orchestrator.state import StateStore + +from . import phases +from .config import load_config, resolve_claude_bin +from .gen_and_opt_pipeline import GenAndOptPipeline +from .result_store import SCHEMA_VERSION, write_json +from .w8a8_pipeline import evaluate_final_target + + +def _load_json(path: Path) -> dict: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return value if isinstance(value, dict) else {} + + +def reconstruct_workers( + workspace_dir: Path, config +) -> dict[str, dict]: + """Rebuild ``workers[worker_id] = {"shapes": {shape_id: {...}}}`` from the + accepted artifact manifests written by the (possibly restarted) lanes.""" + workspace_dir = Path(workspace_dir) + workers: dict[str, dict] = {} + for assignment in config.assignments: + root = workspace_dir / "workers" / assignment.worker_id + lane_shapes: dict[str, dict] = {} + for shape_id in assignment.shape_ids: + manifest_path = root / "accepted" / shape_id / "manifest.json" + if not manifest_path.is_file(): + continue + manifest = _load_json(manifest_path) + if not manifest.get("source") or not manifest.get("object"): + continue + lane_shapes[shape_id] = { + "metrics": manifest.get("metrics") or {}, + "artifact": { + "source": manifest["source"], + "object": manifest["object"], + "source_sha256": manifest.get("source_sha256"), + "object_sha256": manifest.get("object_sha256"), + "commit": manifest.get("commit"), + }, + "candidate": manifest.get("commit"), + } + if lane_shapes: + workers[assignment.worker_id] = {"shapes": lane_shapes} + return workers + + +def main() -> int: + parser = argparse.ArgumentParser( + prog="dcu-kernel-auto-opt-recover-serial-validate" + ) + parser.add_argument("requirements", type=Path) + parser.add_argument("--state-dir", type=Path, required=True) + parser.add_argument("--workspace-dir", type=Path, required=True) + parser.add_argument( + "--claude-bin", + default=None, + help=( + "Agent binary override for the skill-synthesis agent; defaults " + "resolved from agent_framework." + ), + ) + args = parser.parse_args() + + req = _load_json(args.requirements) + config = load_config(req) + task_id = str(req.get("task_id") or "task") + claude_bin = resolve_claude_bin( + config.agent_framework, explicit=args.claude_bin + ) + os.environ["METAINFER_TASK_ID"] = task_id + # Serial validation is GPU-contention sensitive (µs-scale decode kernels; + # see bridge/dsh/README.md). Use an env override so recovery can target an + # idle GPU instead of always hammering GPU 0; default to a non-zero card. + os.environ.setdefault("METAINFER_SERIAL_VALIDATE_GPU", "3") + + workspace = args.workspace_dir + baseline = _load_json(workspace / "shared_baseline" / "results.json") + initial_metrics = dict(baseline.get("shapes") or {}) + workers = reconstruct_workers(workspace, config) + if not workers: + raise RuntimeError( + "no accepted artifacts found under workers/*/accepted; " + "cannot recover serial validation" + ) + completed_assignments = [ + assignment + for assignment in config.assignments + if assignment.worker_id in workers + ] + print( + f"recovered workers: {sorted(workers)} " + f"({sum(len(w['shapes']) for w in workers.values())} shapes)", + flush=True, + ) + + manager = make_subagent_manager( + claude_bin=claude_bin, + model=config.claude_model, + permission_mode="bypassPermissions", + effort="max", + extra_add_dirs=[workspace], + snapshot_file=workspace / "recovery_agents.json", + max_concurrent=1, + ) + store = StateStore(args.state_dir) + pipeline = GenAndOptPipeline( + req=req, + state_dir=args.state_dir, + workspace_dir=workspace, + store=store, + manager=manager, + ) + pipeline._worker_failures = {} + try: + pipeline._phase( + phases.VALIDATE, + action="merge_and_validate_final_kernel", + ) + try: + merged_skill = pipeline._author_merged_skill( + config, completed_assignments + ) + except Exception as exc: # noqa: BLE001 - skill is documentation only + print( + f"skill synthesis agent failed ({exc!r}); " + "continuing with an empty merged skill", + flush=True, + ) + merged_skill = { + "schema_version": 1, + "task_id": task_id, + "status": "skipped", + "reason": "recovery: skill synthesis agent unavailable", + } + synthesis = pipeline._synthesize_final_candidate( + config, workers, initial_metrics, task_id + ) + validation = synthesis["validation"] + worker_validation = { + shape_id: { + "passed": bool(shape_result.get("metrics", {}).get("passed")), + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "candidate": shape_result.get("candidate"), + "metrics": shape_result.get("metrics") or {}, + "artifact": shape_result.get("artifact") or {}, + "source": "accepted_compiled_artifact", + "rerun_by_main": False, + } + for assignment in completed_assignments + for shape_id, shape_result in workers[ + assignment.worker_id + ].get("shapes", {}).items() + } + pipeline._phase(phases.REPORT) + report = { + "schema_version": SCHEMA_VERSION, + "task_id": task_id, + "task_type": "dcu-kernel-auto-opt", + "mode": "generate-and-optimize", + "started_at": _load_json( + workspace / "plan.json" + ).get("created_at", time.time()), + "finished_at": time.time(), + "duration_s": None, + "config": _load_json(workspace / "plan.json"), + "initial_metrics": initial_metrics, + "workers": workers, + "worker_failures": {}, + "worker_validation": worker_validation, + "synthesis": synthesis, + "merged_skill": merged_skill, + "final_validation": validation, + "final_target": evaluate_final_target( + baseline=initial_metrics, + validation=validation, + target_improvement_percent=config.minimum_improvement_percent, + ), + "real_gpu_used": True, + "kernel_generated": True, + "gpu_assignment_agent_decided": False, + "target_repo_modified": True, + "status": "success", + "recovered_serial_validate": True, + "recovered_at": time.time(), + } + write_json(workspace / "final_report.json", report) + store.update_run( + current_iteration=config.mock_iterations, + current_phase=phases.FINISHED, + finished=True, + final_status="success", + last_outcome="ok", + last_transition_label="recovered serial validate complete", + ) + store.append_timeline( + "orchestrator_success", + { + "task_id": task_id, + "status": "success", + "recovered_serial_validate": True, + }, + ) + print(f"recovery complete: {workspace / 'final_report.json'}", flush=True) + return 0 + finally: + manager.shutdown() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/rename_kernel_repo.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/rename_kernel_repo.py new file mode 100644 index 00000000..f0e3effa --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/rename_kernel_repo.py @@ -0,0 +1,276 @@ +"""Rename the kernel repository of one task and repair every reference. + +Mutates, in order: + +* ``kernel-repos/`` -> ``kernel-repos/`` (the git repo directory); +* git worktree administrative files (the repo moved under its linked + worktrees; ``git worktree repair`` rewrites the ``.git`` pointer files); +* ``/main`` symlink (relative, same depth as before); +* task ``requirements.json``: ``label`` + ``target_repo_path``; +* workspace ``plan.json``: ``kernel_repo`` absolute path; +* state ``timeline.jsonl``: one ``kernel_repo_renamed`` event. + +The task id (state/workspace directory names) never changes: the WebUI and +the orchestrator key on it, and the operator shapes inside the repo are +name-agnostic. Refuses to rename a task whose orchestrator is still running +(``run.json`` not finished or a live ``orchestrator.pid``). + +The kernel-repos root is resolved exactly like ``config._kernel_repos_root()`` +(``METAINFER_KERNEL_REPOS`` override, else sibling of the MetaInfer root), so +a rename stays consistent with how tasks resolve repositories. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import shutil +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict + +from .config import _kernel_repos_root + +# Safe repository directory names: no path separators, no traversal, and the +# same charset the New Task form accepts for repository names. +_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") + + +def _load_json(path: Path, default: Any = None) -> Any: + try: + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return default + + +def _write_json(path: Path, value: Any) -> None: + path.write_text( + json.dumps(value, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + + +def _append_timeline(state_dir: Path, event_type: str, payload: Dict[str, Any]) -> None: + path = state_dir / "timeline.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write( + json.dumps( + {"ts": time.time(), "type": event_type, "payload": payload}, + ensure_ascii=False, + ) + + "\n" + ) + + +def _task_running(state_dir: Path) -> bool: + """True when the orchestrator may still be executing this task.""" + run = _load_json(state_dir / "run.json", None) + if run is None: + # No run state yet: only a live orchestrator PID blocks a rename. + pid_path = state_dir / "orchestrator.pid" + try: + pid = int(pid_path.read_text(encoding="utf-8").strip()) + except (OSError, ValueError): + pid = 0 + if pid > 0: + try: + os.kill(pid, 0) + return True + except OSError: + pass + return False + if run.get("finished") is True: + return False + # Started but not marked finished: the orchestrator may still be live, + # so renaming the repo under it is unsafe. + return True + + +def _repair_worktrees(repo: Path) -> None: + """Fix git worktree admin files after the repo directory moved. + + ``git worktree repair`` (run in the moved repo) rewrites every linked + worktree's ``.git`` pointer file, which embeds the old repo path. If git + is unavailable or misses one, fall back to a deterministic rewrite: each + linked worktree's ``.git`` file must contain + ``gitdir: /.git/worktrees/``. + """ + try: + subprocess.run( + ["git", "-C", str(repo), "worktree", "repair"], + check=False, + capture_output=True, + text=True, + ) + except (OSError, subprocess.SubprocessError): + pass + _repair_worktrees_manual(repo) + + +def _repair_worktrees_manual(repo: Path) -> None: + """Deterministic fallback: rewrite every linked worktree ``.git`` file.""" + worktrees_dir = repo / ".git" / "worktrees" + if not worktrees_dir.is_dir(): + return + for entry in worktrees_dir.iterdir(): + gitdir_file = entry / "gitdir" + if not gitdir_file.is_file(): + continue + worktree_git = Path(gitdir_file.read_text(encoding="utf-8").strip()) + if not worktree_git.is_file(): + continue + expected = f"gitdir: {repo / '.git' / 'worktrees' / entry.name}" + if worktree_git.read_text(encoding="utf-8").strip() == expected: + continue + worktree_git.write_text(expected + "\n", encoding="utf-8") + + +def rename_kernel_repo( + workspace_dir: Path | str, + new_name: str, + *, + state_dir: Path | str | None = None, + kernel_repos_root: Path | str | None = None, +) -> Dict[str, Any]: + """Rename the kernel repository behind one task workspace. + + Returns a summary dict; raises ``ValueError`` on invalid input and + ``RuntimeError`` when the task is still running or the repo is busy. + """ + workspace = Path(workspace_dir) + main = workspace / "main" + if not main.is_symlink(): + raise ValueError( + f"workspace main is not a symlink to a kernel repo: {main}" + ) + old_repo = main.resolve() + if not old_repo.is_dir(): + raise ValueError(f"kernel repo does not exist: {old_repo}") + + name = (new_name or "").strip() + if not _NAME_RE.fullmatch(name): + raise ValueError( + "repository name must match [A-Za-z0-9][A-Za-z0-9._-]* " + f"(no separators or traversal), got {new_name!r}" + ) + + root = ( + Path(kernel_repos_root).expanduser().resolve() + if kernel_repos_root is not None + else _kernel_repos_root() + ) + new_repo = root / name + if new_repo == old_repo: + raise ValueError(f"repository is already named {name!r}") + if new_repo.exists(): + raise ValueError(f"kernel repo already exists: {new_repo}") + + state: Path | None = Path(state_dir) if state_dir is not None else None + if state is not None and _task_running(state): + raise RuntimeError( + "task is still running; stop or wait for it before renaming " + "its kernel repository" + ) + + old_path = str(old_repo) + new_path = str(new_repo) + + # 1. Move the repository directory (same filesystem). + shutil.move(old_path, new_path) + + # 2. Repair git worktree administrative files. + _repair_worktrees(new_repo) + + # 3. Re-point the workspace main symlink (relative, same depth). + main.unlink(missing_ok=True) + relative_target = os.path.relpath(new_repo, start=main.parent) + main.symlink_to(relative_target, target_is_directory=True) + + updated: list[str] = [] + + # 4. Task references: requirements.json (label + target_repo_path). + if state is not None: + requirements_path = state / "requirements.json" + requirements = _load_json(requirements_path, {}) or {} + changed = False + if requirements.get("target_repo_path") == old_repo.name: + requirements["target_repo_path"] = name + changed = True + if requirements.get("label") == old_repo.name: + requirements["label"] = name + changed = True + if changed: + _write_json(requirements_path, requirements) + updated.append(str(requirements_path)) + + # 5. Workspace plan.json: kernel_repo absolute path. + plan_path = workspace / "plan.json" + plan = _load_json(plan_path, {}) or {} + if plan.get("kernel_repo") == old_path: + plan["kernel_repo"] = new_path + _write_json(plan_path, plan) + updated.append(str(plan_path)) + + # 6. Timeline event for the WebUI. + if state is not None: + _append_timeline(state, "kernel_repo_renamed", { + "task_id": state.name, + "old_name": old_repo.name, + "new_name": name, + "old_repo": old_path, + "new_repo": new_path, + }) + + return { + "renamed": True, + "old_name": old_repo.name, + "new_name": name, + "old_repo": old_path, + "new_repo": new_path, + "updated_references": updated, + } + + +def main() -> int: + parser = argparse.ArgumentParser( + prog="dcu-kernel-auto-opt-rename-kernel-repo", + description=( + "Rename the kernel repository of one dcu-kernel-auto-opt task " + "and repair every reference (worktrees, main symlink, " + "requirements.json, plan.json)." + ), + ) + parser.add_argument("state_dir", type=Path) + parser.add_argument("workspace_dir", type=Path) + parser.add_argument("new_name", type=str) + parser.add_argument( + "--kernel-repos-root", + type=Path, + default=None, + help=( + "Override the kernel-repos root (defaults to " + "METAINFER_KERNEL_REPOS or the sibling of the MetaInfer root)." + ), + ) + args = parser.parse_args() + try: + result = rename_kernel_repo( + args.workspace_dir, + args.new_name, + state_dir=args.state_dir, + kernel_repos_root=args.kernel_repos_root, + ) + except (ValueError, RuntimeError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + print(json.dumps(result, indent=2, ensure_ascii=False)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/restart_worker.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/restart_worker.py new file mode 100644 index 00000000..e8ec9b6a --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/restart_worker.py @@ -0,0 +1,142 @@ +"""Restart one failed DCU auto-opt lane without stopping sibling lanes.""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import time +from pathlib import Path + +from metainfer.orchestrator._bootstrap import make_subagent_manager +from metainfer.orchestrator.state import StateStore + +from .config import ( + load_config, + replace_assignments, + resolve_claude_bin, +) +from .gen_and_opt_pipeline import GenAndOptPipeline +from .result_store import write_json + + +def main() -> int: + parser = argparse.ArgumentParser(prog="dcu-kernel-auto-opt-restart-worker") + parser.add_argument("requirements", type=Path) + parser.add_argument("--state-dir", type=Path, required=True) + parser.add_argument("--workspace-dir", type=Path, required=True) + parser.add_argument("--worker-id", required=True) + parser.add_argument( + "--claude-bin", + default=None, + help=( + "Agent binary override; defaults resolved from agent_framework " + "(ccb -> METAINFER_CLAUDE_BIN or 'ccb', dsh -> " + "bridge/dsh/dsh_agent.py)." + ), + ) + args = parser.parse_args() + + req = json.loads(args.requirements.read_text(encoding="utf-8")) + config = load_config(req) + claude_bin = resolve_claude_bin( + config.agent_framework, explicit=args.claude_bin + ) + matches = [ + assignment + for assignment in config.assignments + if assignment.worker_id == args.worker_id + ] + if len(matches) != 1: + raise RuntimeError( + f"expected one assignment for {args.worker_id}, got {len(matches)}" + ) + assignment = matches[0] + lane_config = replace_assignments(config, [assignment]) + task_id = str(req.get("task_id", "task")) + os.environ["METAINFER_TASK_ID"] = task_id + + worker_root = args.workspace_dir / "workers" / args.worker_id + restart_root = worker_root / "restarts" / str(int(time.time())) + restart_root.mkdir(parents=True, exist_ok=False) + prior_failure = worker_root / "failure.json" + if prior_failure.is_file(): + # shutil.move, not Path.replace: on overlayfs the archive dir is + # freshly created (upper layer) while failure.json may live in a + # lower layer, so os.rename raises EXDEV. + shutil.move(str(prior_failure), str(restart_root / "prior_failure.json")) + + store = StateStore(args.state_dir) + store.append_timeline( + "worker_restart_started", + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_ids": assignment.shape_ids, + "mode": "isolated_lane_sidecar", + "restart_dir": str(restart_root), + }, + ) + manager = make_subagent_manager( + claude_bin=claude_bin, + model=config.claude_model, + permission_mode="bypassPermissions", + effort="max", + extra_add_dirs=[args.workspace_dir], + snapshot_file=restart_root / "agents.json", + max_concurrent=1, + ) + pipeline = GenAndOptPipeline( + req=req, + state_dir=args.state_dir, + workspace_dir=args.workspace_dir, + store=store, + manager=manager, + ) + try: + baseline = pipeline._bootstrap_worker_repos(lane_config) + workers = pipeline._parallel_agents(lane_config, baseline) + result = { + "schema_version": 1, + "task_id": task_id, + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_ids": assignment.shape_ids, + "status": "completed", + "baseline": baseline, + "workers": workers, + "finished_at": time.time(), + } + write_json(restart_root / "result.json", result) + write_json(worker_root / "restart_result.json", result) + store.append_timeline( + "worker_restart_complete", + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_ids": assignment.shape_ids, + "restart_dir": str(restart_root), + }, + ) + return 0 + except Exception as exc: + failure = { + "schema_version": 1, + "task_id": task_id, + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_ids": assignment.shape_ids, + "status": "failed", + "error": repr(exc), + "finished_at": time.time(), + } + write_json(restart_root / "result.json", failure) + store.append_timeline("worker_restart_failed", failure) + raise + finally: + manager.shutdown() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/result_store.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/result_store.py new file mode 100644 index 00000000..0cfbeec9 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/result_store.py @@ -0,0 +1,26 @@ +"""Atomic worker result storage.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any, Dict + + +SCHEMA_VERSION = 1 + + +def write_json(path: Path, payload: Dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + ".tmp") + tmp.write_text(json.dumps(payload, indent=2), encoding="utf-8") + os.replace(tmp, path) + + +def append_jsonl(path: Path, payload: Dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "a", encoding="utf-8") as f: + f.write(json.dumps(payload) + "\n") + f.flush() + os.fsync(f.fileno()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/skill_store.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/skill_store.py new file mode 100644 index 00000000..99a38082 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/skill_store.py @@ -0,0 +1,743 @@ +"""Generate, inspect, and explicitly publish optimization skills.""" + +from __future__ import annotations + +import difflib +import json +import os +import re +import shutil +import time +from pathlib import Path +from typing import Any, Dict, Iterable + +from .config import OptimizerConfig, WorkerAssignment, resolve_claude_bin + + +def _slug(value: str) -> str: + value = re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-") + return value[:63] or "kernel-optimization" + + +def _frontmatter(name: str, description: str) -> str: + return ( + "---\n" + f"name: {name}\n" + f"description: {description}\n" + "---\n\n" + ) + + +def _write_candidate( + pending_root: Path, + *, + name: str, + kind: str, + source: str, + content: str, +) -> Dict[str, Any]: + target = pending_root / name + target.mkdir(parents=True, exist_ok=True) + (target / "SKILL.md").write_text(content, encoding="utf-8") + manifest = { + "name": name, + "kind": kind, + "source": source, + "status": "pending", + "created_at": time.time(), + } + (target / "manifest.json").write_text( + json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + return manifest + + +def _read_experiments(worker_root: Path) -> list[Dict[str, Any]]: + experiments: list[Dict[str, Any]] = [] + for path in sorted((worker_root / "runs").glob("*/experiments.jsonl")): + for line in path.read_text( + encoding="utf-8", errors="replace" + ).splitlines(): + try: + item = json.loads(line) + except ValueError: + continue + if isinstance(item, dict): + experiments.append(item) + return experiments + + +def generate_worker_skill( + *, + config: OptimizerConfig, + assignment: WorkerAssignment, + workspace_dir: Path, + agent_draft: str | None = None, +) -> Dict[str, Any]: + """Write one skill after a worker has finished all assigned shapes.""" + worker_root = workspace_dir / "workers" / assignment.worker_id + experiments = _read_experiments(worker_root) + shape_label = "-".join(assignment.shape_ids) + name = _slug( + f"{config.dtype}-{config.operator}-{assignment.worker_id}-{shape_label}" + ) + description = ( + f"Optimize {config.dtype} {config.operator} on {config.hardware} for " + f"the shape group {', '.join(assignment.shape_ids)}. Use when tuning " + "a matching kernel shape or transferring measured worker findings." + ) + if agent_draft and agent_draft.strip(): + content = _frontmatter(name, description) + agent_draft.strip() + "\n" + return _write_candidate( + workspace_dir / "skills" / "pending", + name=name, + kind="worker", + source=assignment.worker_id, + content=content, + ) + + lines = [ + _frontmatter(name, description), + f"# {config.dtype} {config.operator}: {assignment.worker_id}", + "", + "## Scope", + "", + f"- Hardware: {config.hardware}", + f"- Stack: {config.kernel_language}", + f"- GPU worker: {assignment.worker_id} / physical GPU {assignment.gpu}", + f"- Shapes: {', '.join(assignment.shape_ids)}", + "", + "## Measured results", + "", + "| Shape | Round | Plan | Median (us) | INT8 TOPS | Algorithmic BW (GB/s) | Accepted |", + "|---|---:|---|---:|---:|---:|---|", + ] + for item in experiments: + metrics = item.get("metrics") or {} + plan = str(item.get("hypothesis") or "").replace("|", "/") + lines.append( + f"| {item.get('shape_id', '—')} | {item.get('iteration', '—')} | " + f"{plan} | {metrics.get('median_us', '—')} | " + f"{metrics.get('logical_tops', metrics.get('tflops', '—'))} | " + f"{metrics.get('algorithmic_bandwidth_gb_s', metrics.get('bandwidth_gb_s', '—'))} | " + f"{'yes' if item.get('accepted') else 'no'} |" + ) + manual = [ + str(item["manual_guidance"]) + for item in experiments if item.get("manual_guidance") + ] + lines.extend([ + "", + "## Optimization procedure", + "", + "1. Match the incoming dimensions to the measured shape group.", + "2. Start from the best accepted candidate and preserve the trusted reference.", + "3. Change one optimization variable per round and record the hypothesis.", + "4. Reject correctness failures and unstable median or P90 regressions.", + "5. Re-run the winner in serial validation before integration.", + ]) + if manual: + lines.extend([ + "", + "## Human guidance applied", + "", + *[f"- {text}" for text in dict.fromkeys(manual)], + ]) + lines.extend([ + "", + "Treat these results as shape-specific evidence; remeasure before " + "applying them outside the listed dimensions.", + "", + ]) + return _write_candidate( + workspace_dir / "skills" / "pending", + name=name, + kind="worker", + source=assignment.worker_id, + content="\n".join(lines), + ) + + +def generate_merged_skill( + *, + config: OptimizerConfig, + assignments: Iterable[WorkerAssignment], + workspace_dir: Path, + agent_draft: str | None = None, + failed_workers: Dict[str, Any] | None = None, +) -> Dict[str, Any]: + """Create the main-agent synthesis after every worker skill exists.""" + assignments = list(assignments) + name = _slug(f"{config.dtype}-{config.operator}-optimization") + description = ( + f"Optimize {config.dtype} {config.operator} across multiple shape " + f"regimes on {config.hardware}. Use when selecting, tuning, validating, " + "or integrating a kernel covered by the measured shape families." + ) + if agent_draft and agent_draft.strip(): + content = _frontmatter(name, description) + agent_draft.strip() + "\n" + return _write_candidate( + workspace_dir / "skills" / "pending", + name=name, + kind="merged", + source="main_agent", + content=content, + ) + + lines = [ + _frontmatter(name, description), + f"# {config.dtype} {config.operator} optimization", + "", + "## Shape routing", + "", + "| Worker evidence | Physical GPU | Shapes |", + "|---|---:|---|", + ] + for assignment in assignments: + lines.append( + f"| {assignment.worker_id} | {assignment.gpu} | " + f"{', '.join(assignment.shape_ids)} |" + ) + if failed_workers: + lines.extend([ + "", + "## Unavailable worker evidence", + "", + *[ + f"- `{worker_id}`: {failure.get('error', 'worker failed')}" + for worker_id, failure in sorted(failed_workers.items()) + ], + ]) + lines.extend([ + "", + "## Workflow", + "", + "1. Route the target shape to the closest measured shape family.", + "2. Read that worker skill before choosing the first candidate.", + "3. Preserve quantization, scale, reference, build, and call contracts.", + "4. Compare stable median, P90, TFLOPS, and bandwidth every round.", + "5. Feed human guidance only to the selected worker's next round.", + "6. Validate winners serially and keep per-shape fallbacks.", + "", + "## Evidence files", + "", + ]) + for assignment in assignments: + worker_name = _slug( + f"{config.dtype}-{config.operator}-{assignment.worker_id}-" + f"{'-'.join(assignment.shape_ids)}" + ) + lines.append(f"- `{worker_name}/SKILL.md`") + lines.extend([ + "", + "Do not infer an unmeasured shape rule solely from a nearby result. " + "Benchmark the new shape and append evidence before generalizing.", + "", + ]) + return _write_candidate( + workspace_dir / "skills" / "pending", + name=name, + kind="merged", + source="main_agent", + content="\n".join(lines), + ) + + +def dsh_skills_root() -> Path: + """Authoritative skill library (DeepSeek Harness). + + dsh is the canonical library: the fusion agent scans it, and new/updated + skills land here first. ``sync_skill_libraries()`` then mirrors them into + the ccb (Claude Code) library so both agent frameworks share one set of + skills. + """ + for env in ("DSH_SKILLS_DIR", "METAINFER_SKILLS_DIR"): + configured = os.environ.get(env) + if configured: + return Path(configured).expanduser() + return Path.home() / ".dsh" / "skills" + + +def ccb_skills_root() -> Path: + """Claude Code skill library — the mirror target for dsh skills.""" + configured = os.environ.get("METAINFER_CLAUDE_SKILLS_DIR") + return ( + Path(configured).expanduser() + if configured else Path.home() / ".claude" / "skills" + ) + + +def existing_skills_root() -> Path: + """The library publish/fuse write into (dsh, the authoritative library).""" + return dsh_skills_root() + + +def _skill_entries(root: Path, *, status: str) -> list[Dict[str, Any]]: + entries = [] + if not root.exists(): + return entries + for path in sorted(root.glob("*/SKILL.md")): + manifest_path = path.parent / "manifest.json" + manifest: Dict[str, Any] = {} + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, ValueError): + pass + entries.append({ + "name": path.parent.name, + "status": status, + "kind": manifest.get("kind", "existing"), + "source": manifest.get("source", "claude"), + "path": str(path), + "content": path.read_text(encoding="utf-8", errors="replace"), + }) + return entries + + +def _frontmatter_description(content: str) -> str: + """Pull the one-line description out of a SKILL.md frontmatter block.""" + match = re.match(r"^---\n(.*?)\n---\n", content, flags=re.DOTALL) + if not match: + return "" + desc = re.search(r"(?m)^description:\s*(.+)$", match.group(1)) + if not desc: + return "" + return desc.group(1).strip().strip('"').strip("'") + + +def _copy_skill_dir(src: Path, dst: Path) -> None: + """Copy one skill directory, skipping manifests and Windows Zone junk.""" + dst.mkdir(parents=True, exist_ok=True) + for item in src.iterdir(): + if item.name == "manifest.json" or "Zone.Identifier" in item.name: + continue + if item.is_dir(): + shutil.copytree(item, dst / item.name, dirs_exist_ok=True) + else: + shutil.copy2(item, dst / item.name) + + +def _backup_skill(skill_dir: Path) -> Path | None: + """Snapshot SKILL.md as ``SKILL.md.bak-`` for rollback.""" + skill = skill_dir / "SKILL.md" + if not skill.is_file(): + return None + backup = skill_dir / f"SKILL.md.bak-{int(time.time() * 1000)}" + shutil.copy2(skill, backup) + return backup + + +def _skill_name_set(root: Path) -> set[str]: + if not root.exists(): + return set() + return {path.parent.name for path in root.glob("*/SKILL.md")} + + +# --------------------------------------------------------------------------- # +# dsh -> ccb one-way mirror +# --------------------------------------------------------------------------- # + +def _sync_summary_path(workspace_dir: Path) -> Path: + return workspace_dir / "skills" / "sync.json" + + +def read_sync_summary(workspace_dir: Path) -> Dict[str, Any]: + try: + return json.loads(_sync_summary_path(workspace_dir).read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + + +def _write_sync_summary(workspace_dir: Path, summary: Dict[str, Any]) -> None: + _sync_summary_path(workspace_dir).parent.mkdir(parents=True, exist_ok=True) + _sync_summary_path(workspace_dir).write_text( + json.dumps(summary, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + + +def sync_skill_libraries(*, workspace_dir: Path | None = None) -> Dict[str, Any]: + """Mirror the authoritative dsh library into the ccb library (one-way). + + Copies skills that are new or changed in dsh into ccb; never deletes + ccb-only skills. A ccb SKILL.md is backed up before being overwritten so + a bad mirror can be rolled back. Idempotent and safe to call repeatedly. + """ + source = dsh_skills_root() + target = ccb_skills_root() + added: list[str] = [] + updated: list[str] = [] + skipped: list[str] = [] + if source.exists(): + target.mkdir(parents=True, exist_ok=True) + for src_dir in sorted(source.glob("*/")): + src_skill = src_dir / "SKILL.md" + if not src_skill.is_file(): + continue + name = src_dir.name + dst = target / name + dst_skill = dst / "SKILL.md" + if not dst_skill.is_file(): + _copy_skill_dir(src_dir, dst) + added.append(name) + elif dst_skill.read_bytes() != src_skill.read_bytes(): + _backup_skill(dst) + _copy_skill_dir(src_dir, dst) + updated.append(name) + else: + skipped.append(name) + summary = { + "source": str(source), + "target": str(target), + "added": added, + "updated": updated, + "skipped": skipped, + "ccb_only": sorted(_skill_name_set(target) - _skill_name_set(source)), + "ts": time.time(), + } + if workspace_dir is not None: + _write_sync_summary(workspace_dir, summary) + return summary + + +def rollback_skill(name: str, *, workspace_dir: Path) -> Dict[str, Any]: + """Restore the latest SKILL.md backup in the dsh library, then re-mirror.""" + skill_dir = dsh_skills_root() / name + if not skill_dir.is_dir(): + raise FileNotFoundError(f"skill not found in dsh library: {name}") + backups = sorted(skill_dir.glob("SKILL.md.bak-*")) + if not backups: + raise FileNotFoundError(f"no backup available for skill {name}") + latest = backups[-1] + shutil.copy2(latest, skill_dir / "SKILL.md") + latest.unlink() + sync = sync_skill_libraries(workspace_dir=workspace_dir) + return { + "name": name, + "restored_from": str(latest), + "synced_to_ccb": sync, + } + + +# --------------------------------------------------------------------------- # +# Main-agent skill fusion (UI-triggered) +# --------------------------------------------------------------------------- # + +def _fuse_status_path(workspace_dir: Path) -> Path: + return workspace_dir / "skills" / "fuse" / "status.json" + + +def read_fuse_status(workspace_dir: Path) -> Dict[str, Any]: + try: + return json.loads(_fuse_status_path(workspace_dir).read_text(encoding="utf-8")) + except (OSError, ValueError): + return {"status": "idle"} + + +def mark_fuse_running(workspace_dir: Path, skill_name: str) -> None: + _fuse_status_path(workspace_dir).parent.mkdir(parents=True, exist_ok=True) + _fuse_status_path(workspace_dir).write_text( + json.dumps({ + "status": "running", + "skill_name": skill_name, + "ts": time.time(), + }, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + + +def _write_fuse_status(workspace_dir: Path, payload: Dict[str, Any]) -> None: + _fuse_status_path(workspace_dir).parent.mkdir(parents=True, exist_ok=True) + _fuse_status_path(workspace_dir).write_text( + json.dumps(payload, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + + +def _build_fuse_prompt( + existing: list[Dict[str, Any]], + pending_content: str, + decision_path: Path, +) -> str: + lines = [ + "You are the skill-fusion agent for DCU kernel optimization.", + "", + "Fold the task's newly written optimization skill into the existing", + "DeepSeek Harness skill library so future agents reuse it. The library", + "is mirrored to Claude Code automatically, so this one decision applies", + "to both frameworks.", + "", + "## Existing skills in the library (name, description, body excerpt)", + "", + ] + if not existing: + lines.append("(the library is empty — add a new skill)") + for item in existing: + lines.append(f"### {item.get('name')}") + lines.append(f"description: {item.get('description') or ''}") + body = item.get("content") or "" + if body.count("---\n") >= 2: + first = body.find("---\n") + body = body[body.find("---\n", first + 1) + 4:] + body = body.strip() + if len(body) > 1200: + body = body[:1200] + "\n... (truncated; read the full file with tools if needed)" + lines.append(body) + lines.append("") + lines.extend([ + "## New skill to fuse", + "", + pending_content, + "", + "## Decision contract", + "", + "Decide whether the new skill should be ADDED as a brand-new skill file or", + "MERGED into one of the existing skills above (only when the new material is", + "a natural extension of an existing skill's scope).", + "", + f"Write ONE JSON object to `{decision_path}`. No markdown fences, no prose", + "around it. Schema:", + '{"action": "new" | "merge", "name": "", "description": "", "content": ""}', + "", + '- action "new": name must be a short kebab-case id (<=63 chars) that does', + " not collide with an existing skill; content is the complete body.", + '- action "merge": name must be EXACTLY an existing skill\'s name; content is', + " the complete UPDATED body after folding the new material in.", + "- Preserve all existing working knowledge; add the new evidence/rules where", + " they fit.", + "- The body must be concise, actionable Markdown for a kernel-optimization", + " agent that runs on DCU (gfx928).", + ]) + return "\n".join(lines) + "\n" + + +def _parse_fuse_decision(text: str) -> Dict[str, Any]: + cleaned = text.strip() + if cleaned.startswith("```"): + cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned) + cleaned = re.sub(r"\s*```$", "", cleaned) + try: + decision = json.loads(cleaned) + except ValueError: + start = cleaned.find("{") + end = cleaned.rfind("}") + if start < 0 or end <= start: + raise ValueError( + "skill fusion agent did not return a JSON decision" + ) from None + try: + decision = json.loads(cleaned[start:end + 1]) + except ValueError as exc: + raise ValueError( + "skill fusion agent returned unparseable JSON" + ) from exc + if not isinstance(decision, dict): + raise ValueError("skill fusion decision must be a JSON object") + action = str(decision.get("action") or "").strip().lower() + if action not in {"new", "merge"}: + raise ValueError( + f"skill fusion decision action must be new or merge, got {action!r}" + ) + name = _slug(str(decision.get("name") or "")) + if not name: + raise ValueError("skill fusion decision missing name") + content = str(decision.get("content") or "").strip() + if not content: + raise ValueError("skill fusion decision missing content") + description = str(decision.get("description") or "").strip() + return { + "action": action, + "name": name, + "description": description, + "content": content, + } + + +def _apply_fuse_decision(decision: Dict[str, Any]) -> Dict[str, Any]: + """Apply a parsed fusion decision into the dsh library.""" + root = dsh_skills_root() + root.mkdir(parents=True, exist_ok=True) + action = decision["action"] + name = decision["name"] + content = decision["content"] + description = decision["description"] + + if action == "new": + target = root / name + if target.exists(): + raise FileExistsError( + f"skill {name} already exists in the dsh library; " + "choose merge or a new name" + ) + target.mkdir(parents=True, exist_ok=True) + body = _frontmatter( + name, description or f"{name} optimization skill" + ) + content + "\n" + (target / "SKILL.md").write_text(body, encoding="utf-8") + return {"action": "new", "name": name, "path": str(target / "SKILL.md")} + + target = root / name + if not (target / "SKILL.md").is_file(): + raise FileNotFoundError( + f"merge target skill {name} does not exist in the dsh library" + ) + previous = (target / "SKILL.md").read_text(encoding="utf-8", errors="replace") + backup = _backup_skill(target) + # Keep the existing frontmatter (the dir name is the skill identity); + # only the body is replaced. + fm = re.match(r"^(---\n.*?\n---\n)", previous, flags=re.DOTALL) + body = (fm.group(1) if fm else _frontmatter(name, description or f"{name} optimization skill")) + content + "\n" + (target / "SKILL.md").write_text(body, encoding="utf-8") + diff = "".join(difflib.unified_diff( + previous.splitlines(keepends=True), + body.splitlines(keepends=True), + fromfile=f"{name}/SKILL.md (before)", + tofile=f"{name}/SKILL.md (after)", + )) + return { + "action": "merge", + "name": name, + "path": str(target / "SKILL.md"), + "backup": str(backup) if backup else None, + "diff": diff, + } + + +def fuse_skill( + *, + config: OptimizerConfig, + workspace_dir: Path, + state_dir: Path, + skill_name: str, +) -> Dict[str, Any]: + """Fuse one pending skill into the dsh library via the task's main agent. + + Runs the agent with the same framework/model as the task (ccb or dsh). + The agent scans the existing dsh library, writes a JSON decision + (new file vs merge into an existing skill), the decision is applied to + the dsh library, and the result is mirrored to the ccb library. + """ + pending = workspace_dir / "skills" / "pending" / skill_name + if not (pending / "SKILL.md").is_file(): + raise FileNotFoundError(f"pending skill not found: {skill_name}") + pending_content = (pending / "SKILL.md").read_text( + encoding="utf-8", errors="replace" + ) + + existing = [ + { + "name": item["name"], + "description": _frontmatter_description(item["content"]), + "content": item["content"], + } + for item in _skill_entries(dsh_skills_root(), status="existing") + ] + + fuse_dir = workspace_dir / "skills" / "fuse" + fuse_dir.mkdir(parents=True, exist_ok=True) + decision_path = fuse_dir / "decision.json" + decision_path.unlink(missing_ok=True) + prompt_file = fuse_dir / "prompt.txt" + prompt_file.write_text( + _build_fuse_prompt(existing, pending_content, decision_path), + encoding="utf-8", + ) + + from metainfer.orchestrator._bootstrap import make_subagent_manager + from metainfer.orchestrator.subagent_manager import AgentSpec + + manager = make_subagent_manager( + claude_bin=resolve_claude_bin(config.agent_framework), + model=config.claude_model, + permission_mode="bypassPermissions", + effort="max", + extra_add_dirs=[workspace_dir, dsh_skills_root()], + snapshot_file=state_dir / "agents.json", + max_concurrent=1, + ) + try: + manager.launch(AgentSpec( + name="skill-fusion", + role="dcu_skill_fusion", + prompt_file=prompt_file, + workdir=fuse_dir, + log_dir=fuse_dir / "logs", + timeout_s=600, + stuck_timeout_s=300, + max_retries=1, + )) + agent_result = manager.result("skill-fusion") + if agent_result is None or not agent_result.success: + raise RuntimeError( + agent_result.error if agent_result and agent_result.error + else "skill fusion agent failed" + ) + if not decision_path.is_file(): + raise RuntimeError( + "skill fusion agent did not write a decision file" + ) + decision = _parse_fuse_decision( + decision_path.read_text(encoding="utf-8", errors="replace") + ) + applied = _apply_fuse_decision(decision) + sync = sync_skill_libraries(workspace_dir=workspace_dir) + outcome = { + **applied, + "synced_to_ccb": { + "added": sync["added"], + "updated": sync["updated"], + }, + } + _write_fuse_status(workspace_dir, { + "status": "done", "ok": True, **outcome, "ts": time.time(), + }) + return outcome + except Exception as exc: # noqa: BLE001 - surfaced via the status file + _write_fuse_status(workspace_dir, { + "status": "error", "ok": False, + "error": str(exc), "ts": time.time(), + }) + raise + finally: + manager.close() + + +def list_skill_library(workspace_dir: Path) -> Dict[str, Any]: + pending_root = workspace_dir / "skills" / "pending" + return { + "existing_root": str(dsh_skills_root()), + "ccb_mirror_root": str(ccb_skills_root()), + "existing": _skill_entries(dsh_skills_root(), status="existing"), + "pending": _skill_entries(pending_root, status="pending"), + "fuse_status": read_fuse_status(workspace_dir), + "sync": read_sync_summary(workspace_dir), + } + + +def publish_skill(workspace_dir: Path, name: str) -> Dict[str, Any]: + """Publish one exact pending directory into the dsh library, then mirror + the result to the ccb library. Never overwrites an existing skill.""" + if _slug(name) != name: + raise ValueError("invalid skill name") + source = workspace_dir / "skills" / "pending" / name + if not (source / "SKILL.md").is_file(): + raise FileNotFoundError(f"pending skill not found: {name}") + root = dsh_skills_root() + root.mkdir(parents=True, exist_ok=True) + target = root / name + if target.exists(): + raise FileExistsError(f"skill already exists: {name}") + shutil.copytree( + source, target, ignore=shutil.ignore_patterns("manifest.json") + ) + published = workspace_dir / "skills" / "published" / name + published.parent.mkdir(parents=True, exist_ok=True) + shutil.move(str(source), str(published)) + sync = sync_skill_libraries(workspace_dir=workspace_dir) + return { + "name": name, + "path": str(target), + "status": "existing", + "synced_to_ccb": {"added": sync["added"], "updated": sync["updated"]}, + } diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_store.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_store.py new file mode 100644 index 00000000..d1774e03 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_store.py @@ -0,0 +1,279 @@ +"""Fine-grained implementation variants for agent reference. + +Variants are accepted kernels stored as a directory tree that mirrors the +form taxonomy: + + variant/////.hip + +e.g. ``variant/int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip``. + +Each leaf file carries a machine-parseable header: + + // @@variant shape= commit= added= + // median_us=.. p90_us=.. tops=.. bandwidth_gb_s=.. speedup=.. baseline_us=.. + // source= + + +Semantics (kept in sync with prompts.py): variants are evidence, not policy — +agents may read and adapt the file matching their (operator/dtype, model, TP, +M, operator) family, but must not be locked into it, and any reuse requires +full re-validation. The tree is staged read-only into generated kernel +repositories under ``references/variants/`` and is not part of the build. +""" + +from __future__ import annotations + +import re +import shutil +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +_HEADER_RE = re.compile(r"(?m)^// @@variant shape=(\S+)(.*?)\n(?:.*?\n)*?// @@end") +# Captures the whole header block (first line tail + ``// key=value`` lines) +# up to ``// @@end`` for structured parsing of metrics/source/commit. +_HEADER_BLOCK_RE = re.compile( + r"^// @@variant shape=(\S+)(.*?)\n// @@end", + re.DOTALL | re.MULTILINE, +) + +MODEL_SLUGS = { + "deepseek v4 flash": "deepseek-v4", + "deepseek-v4": "deepseek-v4", + "hy3 (hunyuan 3)": "hy3", + "hy3": "hy3", + "minimax m3": "minimax-m3", + "minimax-m3": "minimax-m3", + "glm5.2": "glm52", + "glm52": "glm52", +} + +OPERATOR_TYPE_SLUGS = { + "quantized gemm": "gemm", + "attention": "attention", + "rmsnorm / layernorm": "rmsnorm", + "rope": "rope", + "custom operator": "custom", +} + + +def _slug(value: str) -> str: + # "INT8 W8A8" -> "int8w8a8", "Quantized GEMM" -> "gemm" + return re.sub(r"[^a-z0-9]+", "", value.strip().lower()) + + +def operator_type_slug(operator: str) -> str: + key = operator.strip().lower() + return OPERATOR_TYPE_SLUGS.get(key, re.sub(r"[^a-z0-9]+", "-", key).strip("-") or "custom") + + +def dtype_slug(dtype: str) -> str: + return _slug(dtype) or "other" + + +def model_slug(label: str) -> str: + key = label.strip().lower() + return MODEL_SLUGS.get(key, re.sub(r"[^a-z0-9]+", "-", key).strip("-") or "unknown-model") + + +def parse_shape_meta(shape_id: str) -> Dict[str, Any]: + """Split a shape id like ``hy3_tp4_o_proj_m4096`` into tp/operator/m.""" + tp_match = re.search(r"(?:^|_)(tp\d+)_", shape_id) + tp = int(tp_match.group(1)[2:]) if tp_match else None + m_match = re.search(r"_m(\d+)$", shape_id) + m = int(m_match.group(1)) if m_match else None + operator = shape_id + if m_match: + operator = operator[: m_match.start()] + if tp_match: + operator = operator[tp_match.end():] + operator = operator.strip("_") or shape_id + return {"tp": tp, "operator": operator, "m": m} + + +def derive_variant_meta(answers: Dict[str, Any], shape_id: str) -> Dict[str, Any]: + """Derive the fine-grained taxonomy for one shape from the form answers.""" + parsed = parse_shape_meta(shape_id) + operator = str(answers.get("operator") or "") + dtype = str(answers.get("dtype") or "") + return { + "shape": shape_id, + "operator": operator, + "dtype": dtype, + "family": f"{dtype_slug(dtype)}-{operator_type_slug(operator)}", + "model": model_slug(str(answers.get("model") or "")), + "tp": parsed["tp"], + "operator_name": parsed["operator"], + "m": parsed["m"], + } + + +def variant_root() -> Path: + return Path(__file__).resolve().parents[1] / "variant" + + +def variant_path(meta: Dict[str, Any]) -> Path: + """The leaf file path for one specific operator's variant.""" + tp = f"TP{meta['tp']}" if meta.get("tp") is not None else "TP?" + m = f"M{meta['m']}" if meta.get("m") is not None else "M?" + return ( + variant_root() / meta["family"] / meta["model"] / tp / m + / f"{meta['operator_name']}.hip" + ) + + +def section_header( + meta: Dict[str, Any], + *, + commit: str, + metrics: Dict[str, Any], + source_task: str = "", +) -> str: + lines = [ + f"// @@variant shape={meta['shape']} commit={commit or '?'} " + f"added={time.strftime('%Y-%m-%d')}", + ] + metric_parts = [] + for key, label in ( + ("median_us", "median_us"), + ("p90_us", "p90_us"), + ("logical_tops", "tops"), + ("algorithmic_bandwidth_gb_s", "bandwidth_gb_s"), + ("speedup", "speedup"), + ("baseline_us", "baseline_us"), + ): + value = metrics.get(key) + if value is not None: + metric_parts.append( + f"{label}={value:.4g}" if isinstance(value, float) else f"{label}={value}" + ) + if metric_parts: + lines.append("// " + " ".join(metric_parts)) + if source_task: + lines.append(f"// source={source_task}") + lines.append("") + return "\n".join(lines) + + +def add_variant( + *, + meta: Dict[str, Any], + kernel_source: str, + commit: str, + metrics: Dict[str, Any], + source_task: str = "", + backup: bool = True, + reject_slower_than_existing: bool = False, +) -> Dict[str, Any]: + """Write one specific operator's accepted kernel into the variant tree. + + Replaces an existing file for the same leaf (backing it up first). When + ``reject_slower_than_existing`` is set and both the existing file's header + and ``metrics`` carry a ``median_us``, a candidate that is strictly slower + than the existing variant is rejected with ``ValueError`` (equal or faster + may replace). + """ + target = variant_path(meta) + target.parent.mkdir(parents=True, exist_ok=True) + if target.exists() and reject_slower_than_existing: + _reject_slower_replacement(target, metrics) + header = section_header( + meta, commit=commit, metrics=metrics, source_task=source_task + ) + content = header + kernel_source.rstrip() + "\n// @@end\n" + action = "updated" if target.exists() else "added" + backup_path = None + if target.exists() and backup: + backup_path = target.with_name(f"{target.name}.bak-{int(time.time() * 1000)}") + shutil.copy2(target, backup_path) + target.write_text(content, encoding="utf-8") + return { + "action": action, + "shape": meta["shape"], + "path": str(target), + "backup": str(backup_path) if backup_path else None, + "meta": meta, + } + + +def _reject_slower_replacement(target: Path, metrics: Dict[str, Any]) -> None: + """Raise ValueError when the candidate is strictly slower than the + existing variant's recorded median (both must be known).""" + existing = _parse_variant_header(target.read_text(encoding="utf-8", errors="replace")) + existing_median = existing.get("median_us") + new_median = metrics.get("median_us") + if existing_median is None or new_median is None: + # Cannot compare -> allow (guard is best-effort on known medians). + return + if float(new_median) > float(existing_median): + raise ValueError( + "rejecting slower variant replacement: existing median_us=" + f"{float(existing_median):.4g} vs candidate median_us=" + f"{float(new_median):.4g} " + "(only an equal-or-faster candidate may replace)" + ) + + +def list_variant_index() -> List[Dict[str, Any]]: + """Walk the variant tree and return the section index (path-derived + taxonomy + header fields).""" + out: List[Dict[str, Any]] = [] + root = variant_root() + if not root.exists(): + return out + # skip the legacy single-file variant (kept for historical reference) + for file in sorted(root.rglob("*.hip")): + if file.name.startswith("w8a8_gemm_variants"): + continue + text = file.read_text(encoding="utf-8", errors="replace") + header = _parse_variant_header(text) + shape = header.get("shape") or file.stem + rel = file.relative_to(root).parts + out.append({ + "path": str(file.relative_to(root)), + "shape": shape, + "family": rel[0] if len(rel) > 0 else "", + "model": rel[1] if len(rel) > 1 else "", + "tp": rel[2] if len(rel) > 2 else "", + "m": rel[3] if len(rel) > 3 else "", + "operator": file.stem, + "commit": header.get("commit", ""), + "added": header.get("added", ""), + "source": header.get("source", ""), + "median_us": header.get("median_us"), + "p90_us": header.get("p90_us"), + "speedup": header.get("speedup"), + "baseline_us": header.get("baseline_us"), + }) + return out + + +# Header fields that carry machine-readable performance evidence. Only these +# are surfaced in the variant index (and the WebUI comparison dialog); the +# rest of the header/body is opaque to the index. +_HEADER_NUMERIC_FIELDS = frozenset({ + "median_us", "p90_us", "tops", "bandwidth_gb_s", "speedup", "baseline_us", +}) +_HEADER_TEXT_FIELDS = frozenset({"commit", "added", "source"}) + + +def _parse_variant_header(text: str) -> Dict[str, Any]: + """Parse the machine-readable variant header of one kernel file. + + Returns ``{"shape": ..., "commit": ..., "median_us": ..., ...}`` with + numeric metric fields converted to float; unknown ``key=value`` tokens + inside the header block are ignored. + """ + block = _HEADER_BLOCK_RE.search(text) + if block is None: + return {} + fields: Dict[str, Any] = {"shape": block.group(1)} + for key, value in re.findall(r"(\w+)=(\S+)", block.group(2)): + if key in _HEADER_NUMERIC_FIELDS: + try: + fields[key] = float(value) + except ValueError: + fields[key] = value + elif key in _HEADER_TEXT_FIELDS: + fields[key] = value + return fields diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_baselines.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_baselines.py new file mode 100644 index 00000000..c5f8df79 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_baselines.py @@ -0,0 +1,319 @@ +"""Control-plane-owned W8A8 Triton Graph comparison baselines.""" + +from __future__ import annotations + +from typing import Any, Dict, Mapping + + +# User-supplied, fixed decode/prefill Graph measurements. These values are +# comparison targets, not measurements of the child-generated bootstrap HIP +# source. Latencies are stored in microseconds. +_TRITON_GRAPH_BASELINES_US = { + (4, 2, 1536, 4096): ("wqkv_a", 66.924), + (4, 16, 1536, 4096): ("wqkv_a", 66.597), + (4, 2, 8192, 1024): ("wq_b", 47.181), + (4, 16, 8192, 1024): ("wq_b", 48.049), + (4, 2, 4096, 2048): ("wo_b", 54.453), + (4, 16, 4096, 2048): ("wo_b", 54.617), + (4, 2, 1024, 4096): ("shared_gate_up", 60.389), + (4, 16, 1024, 4096): ("shared_gate_up", 60.965), + (4, 2, 4096, 512): ("shared_down", 18.975), + (4, 16, 4096, 512): ("shared_down", 21.372), + # M=3072 prefill baselines supplied on 2026-08-03. + (4, 3072, 1536, 4096): ("wqkv_a", 9488.0), + (4, 3072, 8192, 1024): ("wq_b", 15437.0), + (4, 3072, 4096, 2048): ("wo_b", 14680.0), + (4, 3072, 1024, 4096): ("shared_gate_up", 5985.0), + (4, 3072, 4096, 512): ("shared_down", 4199.0), + # M=4096 prefill baselines measured on 2026-08-06 (worker29, gfx928) with + # lmslim int8_utils.matmul_kernel (the W8A8 Triton baseline), GPU events, + # CUDA Graph replay, hot cache, warmups=10, samples=20, launches/sample=5. + # The (4096, 8192, 1024) entry also covers indexer.wq_b, which shares the + # same logical (K, N) shape as wq_b. + (4, 4096, 1536, 4096): ("wqkv_a", 13247.253), + (4, 4096, 8192, 1024): ("wq_b", 20590.225), + (4, 4096, 4096, 2048): ("wo_b", 19881.949), + (4, 4096, 1024, 4096): ("shared_gate_up", 8790.192), + (4, 4096, 4096, 512): ("shared_down", 5545.821), + # Model-catalog baselines measured on 2026-08-12 (worker29, gfx928) with + # lmslim int8_utils.matmul_kernel (the W8A8 Triton baseline), GPU events, + # CUDA Graph replay, hot cache, warmups=10, samples=20, launches/sample=5. + # Config follows int8_utils.matmul_int8 defaults per M (decode M<=32 uses + # BM16/BN32/BK256; M>1024 uses BM256/BN256/BK64). Covers the DeepSeek TP1 + # and TP8 M=2 additions plus the Hy3 (Hunyuan 3), MiniMax M3 and GLM5.2 + # catalogs. Keys are shared across models when (tp, M, N, K) coincide, + # e.g. (1, *, 8192, 4096) is DeepSeek wo_b / Hy3 o_proj and (4, *, 8192, + # 1024) covers indexer.wq_b. + (1, 2, 1536, 4096): ("wqkv_a", 78.046), + (1, 2, 2624, 6144): ("fused_qkv_a_proj", 96.623), + (1, 2, 3072, 4096): ("shared_gate_up_proj", 81.854), + (1, 2, 4096, 1536): ("shared_down_proj", 47.839), + (1, 2, 4096, 2048): ("shared_down_proj", 62.351), + (1, 2, 4096, 4096): ("shared_gate_up_proj", 109.998), + (1, 2, 4096, 6144): ("shared_gate_up_proj", 151.045), + (1, 2, 4096, 8192): ("wo_b", 194.748), + (1, 2, 6144, 2048): ("shared_down_proj", 66.098), + (1, 2, 6144, 3072): ("shared_down_proj", 92.174), + (1, 2, 6144, 6144): ("shared_gate_up_proj", 155.341), + (1, 2, 6144, 8192): ("o_proj", 190.044), + (1, 2, 6144, 16384): ("o_proj", 371.539), + (1, 2, 8192, 1024): ("wq_b", 55.599), + (1, 2, 9216, 6144): ("qkv_proj", 232.987), + (1, 2, 9856, 6144): ("qkv_proj_and_indexer_qk", 216.684), + (1, 2, 10240, 4096): ("qkv_proj", 175.532), + (1, 2, 16384, 2048): ("q_b_proj", 145.087), + (1, 2, 28672, 512): ("kv_b_proj", 82.624), + (1, 2, 32768, 1024): ("wq_b", 155.197), + (1, 16, 1536, 4096): ("wqkv_a", 80.574), + (1, 16, 2624, 6144): ("fused_qkv_a_proj", 98.463), + (1, 16, 3072, 4096): ("shared_gate_up_proj", 82.638), + (1, 16, 4096, 1536): ("shared_down_proj", 49.439), + (1, 16, 4096, 2048): ("shared_down_proj", 65.567), + (1, 16, 4096, 4096): ("shared_gate_up_proj", 110.606), + (1, 16, 4096, 6144): ("shared_gate_up_proj", 151.685), + (1, 16, 4096, 8192): ("wo_b", 205.644), + (1, 16, 6144, 2048): ("shared_down_proj", 67.346), + (1, 16, 6144, 3072): ("shared_down_proj", 91.838), + (1, 16, 6144, 6144): ("shared_gate_up_proj", 156.653), + (1, 16, 6144, 8192): ("o_proj", 192.412), + (1, 16, 6144, 16384): ("o_proj", 358.435), + (1, 16, 8192, 1024): ("wq_b", 56.847), + (1, 16, 9216, 6144): ("qkv_proj", 235.083), + (1, 16, 9856, 6144): ("qkv_proj_and_indexer_qk", 219.564), + (1, 16, 10240, 4096): ("qkv_proj", 171.069), + (1, 16, 16384, 2048): ("q_b_proj", 146.719), + (1, 16, 28672, 512): ("kv_b_proj", 83.504), + (1, 16, 32768, 1024): ("wq_b", 158.989), + (1, 3072, 1536, 4096): ("wqkv_a", 9795.766), + (1, 3072, 2624, 6144): ("fused_qkv_a_proj", 28944.543), + (1, 3072, 3072, 4096): ("shared_gate_up_proj", 21047.242), + (1, 3072, 4096, 1536): ("shared_down_proj", 11363.852), + (1, 3072, 4096, 2048): ("shared_down_proj", 14921.346), + (1, 3072, 4096, 4096): ("shared_gate_up_proj", 29325.629), + (1, 3072, 4096, 6144): ("shared_gate_up_proj", 43830.049), + (1, 3072, 4096, 8192): ("wo_b", 58484.25), + (1, 3072, 6144, 2048): ("shared_down_proj", 22320.748), + (1, 3072, 6144, 3072): ("shared_down_proj", 33673.799), + (1, 3072, 6144, 6144): ("shared_gate_up_proj", 67255.209), + (1, 3072, 6144, 8192): ("o_proj", 90545.355), + (1, 3072, 6144, 16384): ("o_proj", 182662.598), + (1, 3072, 8192, 1024): ("wq_b", 15305.526), + (1, 3072, 9216, 6144): ("qkv_proj", 103111.029), + (1, 3072, 9856, 6144): ("qkv_proj_and_indexer_qk", 110813.373), + (1, 3072, 10240, 4096): ("qkv_proj", 74301.431), + (1, 3072, 16384, 2048): ("q_b_proj", 59796.738), + (1, 3072, 28672, 512): ("kv_b_proj", 29183.929), + (1, 3072, 32768, 1024): ("wq_b", 62264.438), + (4, 2, 768, 4096): ("shared_gate_up_proj", 71.711), + (4, 2, 1024, 6144): ("shared_gate_up_proj", 108.278), + (4, 2, 1536, 6144): ("shared_gate_up_proj", 109.822), + (4, 2, 2304, 6144): ("qkv_proj", 113.134), + (4, 2, 2560, 4096): ("qkv_proj", 77.854), + (4, 2, 2560, 6144): ("qkv_proj_and_indexer_qk", 111.79), + (4, 2, 2624, 6144): ("fused_qkv_a_proj", 98.756), + (4, 2, 4096, 384): ("shared_down_proj", 21.664), + (4, 2, 6144, 512): ("shared_down_proj", 23.553), + (4, 2, 6144, 768): ("shared_down_proj", 33.407), + (4, 2, 6144, 2048): ("o_proj", 65.135), + (4, 2, 6144, 4096): ("o_proj", 110.501), + (4, 2, 7168, 512): ("kv_b_proj", 25.009), + (4, 16, 768, 4096): ("shared_gate_up_proj", 73.262), + (4, 16, 1024, 6144): ("shared_gate_up_proj", 108.422), + (4, 16, 1536, 6144): ("shared_gate_up_proj", 113.006), + (4, 16, 2304, 6144): ("qkv_proj", 113.422), + (4, 16, 2560, 4096): ("qkv_proj", 80.19), + (4, 16, 2560, 6144): ("qkv_proj_and_indexer_qk", 112.27), + (4, 16, 2624, 6144): ("fused_qkv_a_proj", 98.708), + (4, 16, 4096, 384): ("shared_down_proj", 23.216), + (4, 16, 6144, 512): ("shared_down_proj", 24.977), + (4, 16, 6144, 768): ("shared_down_proj", 34.063), + (4, 16, 6144, 2048): ("o_proj", 67.535), + (4, 16, 6144, 4096): ("o_proj", 112.198), + (4, 16, 7168, 512): ("kv_b_proj", 27.793), + (4, 3072, 768, 4096): ("shared_gate_up_proj", 4988.476), + (4, 3072, 1024, 6144): ("shared_gate_up_proj", 9708.601), + (4, 3072, 1536, 6144): ("shared_gate_up_proj", 14650.233), + (4, 3072, 2304, 6144): ("qkv_proj", 22634.012), + (4, 3072, 2560, 4096): ("qkv_proj", 16984.091), + (4, 3072, 2560, 6144): ("qkv_proj_and_indexer_qk", 25396.649), + (4, 3072, 2624, 6144): ("fused_qkv_a_proj", 28947.94), + (4, 3072, 4096, 384): ("shared_down_proj", 3366.413), + (4, 3072, 6144, 512): ("shared_down_proj", 6270.013), + (4, 3072, 6144, 768): ("shared_down_proj", 8900.2), + (4, 3072, 6144, 2048): ("o_proj", 22307.655), + (4, 3072, 6144, 4096): ("o_proj", 44517.157), + (4, 3072, 7168, 512): ("kv_b_proj", 7308.89), + (4, 4096, 768, 4096): ("shared_gate_up_proj", 6678.282), + (4, 4096, 1024, 6144): ("shared_gate_up_proj", 13036.336), + (4, 4096, 1536, 6144): ("shared_gate_up_proj", 19776.488), + (4, 4096, 2304, 6144): ("qkv_proj", 31556.192), + (4, 4096, 2560, 4096): ("qkv_proj", 24045.12), + (4, 4096, 2560, 6144): ("qkv_proj_and_indexer_qk", 35918.768), + (4, 4096, 2624, 6144): ("fused_qkv_a_proj", 40458.051), + (4, 4096, 4096, 384): ("shared_down_proj", 4381.289), + (4, 4096, 6144, 512): ("shared_down_proj", 8243.047), + (4, 4096, 6144, 768): ("shared_down_proj", 11759.473), + (4, 4096, 6144, 2048): ("o_proj", 29702.473), + (4, 4096, 6144, 4096): ("o_proj", 58712.341), + (4, 4096, 7168, 512): ("kv_b_proj", 9702.386), + (8, 2, 384, 4096): ("shared_gate_up_proj", 64.511), + (8, 2, 512, 4096): ("shared_gate_up_proj", 72.687), + (8, 2, 512, 6144): ("shared_gate_up_proj", 97.606), + (8, 2, 768, 6144): ("shared_gate_up_proj", 104.351), + (8, 2, 1280, 4096): ("qkv_proj", 73.535), + (8, 2, 1280, 6144): ("qkv_proj", 110.638), + (8, 2, 1536, 4096): ("wqkv_a", 74.767), + (8, 2, 1536, 6144): ("qkv_proj_and_indexer_qk", 110.19), + (8, 2, 2048, 2048): ("q_b_proj", 40.371), + (8, 2, 2624, 6144): ("fused_qkv_a_proj", 95.43), + (8, 2, 3584, 512): ("kv_b_proj", 17.873), + (8, 2, 4096, 192): ("shared_down_proj", 16.144), + (8, 2, 4096, 256): ("shared_down_proj", 19.344), + (8, 2, 4096, 1024): ("wq_b_or_wo_b", 36.271), + (8, 2, 6144, 256): ("shared_down_proj", 16.881), + (8, 2, 6144, 384): ("shared_down_proj", 21.968), + (8, 2, 6144, 1024): ("o_proj", 37.407), + (8, 2, 6144, 2048): ("o_proj", 64.948), + (8, 2, 8192, 1024): ("wq_b", 56.239), + (8, 16, 384, 4096): ("shared_gate_up_proj", 66.527), + (8, 16, 512, 6144): ("shared_gate_up_proj", 99.318), + (8, 16, 768, 6144): ("shared_gate_up_proj", 104.527), + (8, 16, 1280, 4096): ("qkv_proj", 74.782), + (8, 16, 1280, 6144): ("qkv_proj", 110.478), + (8, 16, 1536, 6144): ("qkv_proj_and_indexer_qk", 111.486), + (8, 16, 2048, 2048): ("q_b_proj", 42.035), + (8, 16, 2624, 6144): ("fused_qkv_a_proj", 97.99), + (8, 16, 3584, 512): ("kv_b_proj", 18.833), + (8, 16, 4096, 192): ("shared_down_proj", 17.312), + (8, 16, 6144, 256): ("shared_down_proj", 18.177), + (8, 16, 6144, 384): ("shared_down_proj", 23.152), + (8, 16, 6144, 1024): ("o_proj", 38.655), + (8, 16, 6144, 2048): ("o_proj", 66.98), + (8, 3072, 384, 4096): ("shared_gate_up_proj", 3009.348), + (8, 3072, 512, 6144): ("shared_gate_up_proj", 4355.572), + (8, 3072, 768, 6144): ("shared_gate_up_proj", 7087.451), + (8, 3072, 1280, 4096): ("qkv_proj", 8489.847), + (8, 3072, 1280, 6144): ("qkv_proj", 12553.842), + (8, 3072, 1536, 6144): ("qkv_proj_and_indexer_qk", 14614.569), + (8, 3072, 2048, 2048): ("q_b_proj", 6821.235), + (8, 3072, 2624, 6144): ("fused_qkv_a_proj", 29185.017), + (8, 3072, 3584, 512): ("kv_b_proj", 3685.827), + (8, 3072, 4096, 192): ("shared_down_proj", 2043.847), + (8, 3072, 6144, 256): ("shared_down_proj", 3593.156), + (8, 3072, 6144, 384): ("shared_down_proj", 4925.667), + (8, 3072, 6144, 1024): ("o_proj", 11664.713), + (8, 3072, 6144, 2048): ("o_proj", 22305.402), + # Model-catalog TP8 M=4096 large-prefill baselines measured on 2026-08-27 + # (worker29, gfx928) with the same protocol as the 2026-08-12 catalog: + # lmslim int8_utils.matmul_kernel (the W8A8 Triton baseline), GPU events, + # CUDA Graph replay, hot cache, warmups=10, samples=20, launches/sample=5, + # BM256/BN256/BK64 (matmul_int8 default for M>1024). Covers the Hy3 / + # MiniMax M3 / GLM5.2 TP8 operators at M=4096 (see + # MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES in the operator API contract). + (8, 4096, 1280, 4096): ("qkv_proj", 19381.383), + (8, 4096, 4096, 1024): ("o_proj", 14151.792), + (8, 4096, 384, 4096): ("shared_gate_up_proj", 16712.231), + (8, 4096, 4096, 192): ("shared_down_proj", 10762.616), + (8, 4096, 1280, 6144): ("qkv_proj", 41000.618), + (8, 4096, 1536, 6144): ("qkv_proj_and_indexer_qk", 35604.698), + (8, 4096, 6144, 1024): ("o_proj", 19897.517), + (8, 4096, 768, 6144): ("shared_gate_up_proj", 31882.257), + (8, 4096, 6144, 384): ("shared_down_proj", 22389.162), + (8, 4096, 2624, 6144): ("fused_qkv_a_proj", 68661.417), + (8, 4096, 2048, 2048): ("q_b_proj", 14385.571), + (8, 4096, 3584, 512): ("kv_b_proj", 6939.939), + (8, 4096, 6144, 2048): ("o_proj", 54775.797), + (8, 4096, 512, 6144): ("shared_gate_up_proj", 21254.589), + (8, 4096, 6144, 256): ("shared_down_proj", 17700.262), +} + + +# TP8 hot-cache measurements supplied in milliseconds and normalized here to +# microseconds. Graph median is the optimization comparison target; eager and +# both P90 values are retained as measurement evidence. +_TRITON_TP8_BASELINES_US = { + (8, 16, 1536, 4096): ("wqkv_a", 97.660, 98.404, 66.591, 67.347), + (8, 16, 4096, 1024): ("wq_b_or_wo_b", 95.284, 96.092, 32.137, 35.969), + (8, 16, 8192, 1024): ("indexer.wq_b", 96.204, 96.796, 48.274, 50.794), + (8, 16, 512, 4096): ("shared_gate_up", 96.056, 97.788, 59.210, 59.962), + (8, 16, 4096, 256): ("shared_down", 94.540, 95.740, 14.849, 18.001), + (8, 3072, 1536, 4096): ("wqkv_a", 9423.243, 9446.877, 9422.539, 9486.942), + (8, 3072, 4096, 1024): ("wq_b_or_wo_b", 7637.209, 8033.128, 7617.880, 7686.395), + (8, 3072, 8192, 1024): ("indexer.wq_b", 15461.000, 15528.571, 15466.648, 15958.569), + (8, 3072, 512, 4096): ("shared_gate_up", 2775.526, 2780.134, 2740.757, 2751.685), + (8, 3072, 4096, 256): ("shared_down", 2414.377, 2424.569, 2370.583, 2376.343), +} + + +def fixed_triton_graph_baseline( + shape_id: str, + shape: Mapping[str, Any], + *, + bootstrap_metrics: Mapping[str, Any] | None = None, +) -> Dict[str, Any]: + """Return the immutable Graph baseline for one exact M/N/K shape. + + ``bootstrap_metrics`` remains separate because it describes the current + child HIP source used for PMC, rollback and iterative improvement. + """ + try: + # Legacy callers predate TP-aware shape metadata and are TP4-only. + tp_size = int(shape.get("tp_size", 4)) + key = ( + tp_size, + int(shape["M"]), + int(shape["N"]), + int(shape["K"]), + ) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError( + f"invalid W8A8 shape for fixed baseline: {shape_id}: {shape}" + ) from exc + tp8_measurement = _TRITON_TP8_BASELINES_US.get(key) + fixed_measurement = _TRITON_GRAPH_BASELINES_US.get(key) + if tp8_measurement is None and fixed_measurement is None: + raise ValueError( + "no fixed Triton Graph baseline for " + f"{shape_id} with TP={key[0]}, M={key[1]}, " + f"N={key[2]}, K={key[3]}" + ) + + if tp8_measurement is not None: + case, eager_median_us, eager_p90_us, latency_us, graph_p90_us = ( + tp8_measurement + ) + else: + assert fixed_measurement is not None + case, latency_us = fixed_measurement + + record: Dict[str, Any] = { + "median_us": latency_us, + "baseline_us": latency_us, + "baseline_kind": "triton_graph", + "case": case, + "tp_size": tp_size, + "shape": {"M": key[1], "N": key[2], "K": key[3]}, + "source": "user_supplied_fixed_table", + "timing_scope": ( + "prefill_graph_replay" if key[1] > 16 + else "decode_graph_replay" + ), + "distribution_stats_available": tp8_measurement is not None, + } + if tp8_measurement is not None: + record.update( + { + "p90_us": graph_p90_us, + "eager_median_us": eager_median_us, + "eager_p90_us": eager_p90_us, + "cache_state": "hot", + "measurement_protocol": { + "warmups": 50 if key[1] == 16 else 10, + "samples": 50 if key[1] == 16 else 20, + "launches_per_sample": 20 if key[1] == 16 else 5, + }, + } + ) + if bootstrap_metrics is not None: + record["bootstrap_metrics"] = dict(bootstrap_metrics) + return record diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_pipeline.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_pipeline.py new file mode 100644 index 00000000..41d3d61a --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_pipeline.py @@ -0,0 +1,3128 @@ +"""Real multi-agent optimization pipeline for the extracted W8A8 GEMM.""" + +from __future__ import annotations + +import json +import hashlib +import os +import shutil +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import Any, Dict + +from metainfer.orchestrator.state import StateStore +from metainfer.orchestrator.subagent_manager import AgentSpec, SubAgentManager + +from . import phases +from .api_contracts import W8A8_API_FILENAME, file_digest +from .config import ( + ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT, + OptimizerConfig, + WorkerAssignment, + load_config, +) +from .gpu_binding import bind_worker_gpu +from .experience_store import load_verified_experience +from .guidance import claim_next_guidance +from .isa_analysis import ( + analyze_inline_asm_source, + evaluate_inline_asm_gate, + inspect_gfx928_object, +) +from .prompts import w8a8_round_strategy, w8a8_strategy_guidance +from .pmc_profile import ( + add_unprofiled_bandwidth, + parse_memory_traffic_csv, + parse_pmc_csv, +) +from .real_pipeline import _last_json, _run, _safe, _status +from .result_store import SCHEMA_VERSION, append_jsonl, write_json +from .skill_store import generate_merged_skill, generate_worker_skill +from .w8a8_baselines import fixed_triton_graph_baseline + + +HARNESS = ( + Path(__file__).resolve().parent.parent / "assets" / "w8a8_bench.py" +) + +_REQUIRED_FILE_SETS = ( + # Current generated-repository contract. + ( + "int8_w8a8_gemm_api.py", + "w8a8_backend.py", + "w8a8_bench.py", + "setup.py", + "csrc/bindings.cpp", + "csrc/w8a8_gemm_hip.hip", + ), + # Legacy optimize-existing-repository contract. + ( + "w8a8_gemm.py", + "csrc/bindings.cpp", + "csrc/w8a8_gemm_hip.hip", + ), +) +_SOURCE_ONLY_AGENT_ARGS = [ + "--tools", + "Read,Glob,Grep,Write,Edit", +] +_ISA_AGENT_ARGS = [ + "--tools", + "Read,Glob,Grep,Write,Edit,Skill", +] +_MAX_IN_ROUND_REPAIRS = 4 +_MAX_INFRASTRUCTURE_RECOVERY_ROUNDS = 5 +_MAX_PHASE_EXTENSION_ROUNDS = 8 +_REQUIRED_VALID_ISA_GUIDED_ROUNDS = 2 +_PLATEAU_MAX_REGRESSION_PERCENT = 2.0 +_SHADOW_MIN_IMPROVEMENT_PERCENT = 0.3 +_AGENT_TIMEOUT_S = 900 +_AGENT_STUCK_TIMEOUT_S = 600 +# Trusted harness subprocess budget for one correctness-checked benchmark or +# PMC profile run. With the reference auto-seeded (_REFERENCE_PREPARE_TIMEOUT_S) +# this only needs to cover hipcc compile (~30s) + graph capture + timing; +# 1200s adds margin for slow compiles or a busy host (was 900s, which M=4096 +# correctness runs blew through before reference pre-seeding existed). +_BENCHMARK_TIMEOUT_S = 1200 +# The CPU int64 exact reference for M>=3072 is the dominant cost of a +# correctness-checked benchmark (measured ~0.2 GFLOPS for torch.mm int64, so +# M=4096/K=6144 needs ~10+ minutes solo and more under worker contention). +# It is prepared once per shape outside the benchmark subprocess budget. +_REFERENCE_PREPARE_TIMEOUT_S = 3600 +_UNSUPPORTED_SKILL_CLAIMS = ( + "scalar is optimal", + "dumma is unavailable", + "does not support int8 dumma", + "lacks int8", + "hiplaunchkernelggl cannot", + "template kernels cannot", +) + + +def isa_round_policy( + *, + iteration: int, + max_iterations: int, + history: list[Dict[str, Any]], +) -> Dict[str, Any]: + """Gate ISA Skills and raw asm behind completed HIP-only exploration.""" + required_hip_rounds = max(1, max_iterations - 2) + base = { + "phase": "hip_only", + "skill_allowed": False, + "raw_inline_asm_allowed": False, + "plateau": False, + "max_iterations": max_iterations, + "required_valid_hip_rounds": required_hip_rounds, + "required_valid_isa_guided_rounds": _REQUIRED_VALID_ISA_GUIDED_ROUNDS, + "reason": "At least eight HIP-only rounds are required.", + } + valid = [ + record for record in history + if record.get("build_success") is True + and record.get("correctness_passed") is True + and (record.get("metrics") or {}).get("graph_capture_passed") is True + ] + valid_hip = [ + record for record in valid + if (record.get("isa_policy") or {}).get("phase", "hip_only") + == "hip_only" + ] + if len(valid_hip) < required_hip_rounds: + return { + **base, + "valid_hip_rounds": len(valid_hip), + "reason": ( + f"Require {required_hip_rounds} completed correct HIP-only " + "experiments; infrastructure and invalid-candidate failures " + f"do not count. Current valid HIP rounds: {len(valid_hip)}." + ), + } + + recent = valid_hip[-3:] + improvements = [ + float((record.get("acceptance") or {}).get( + "improvement_percent", float("inf") + )) + for record in recent + ] + plateau = ( + len(recent) == 3 + and all( + -_PLATEAU_MAX_REGRESSION_PERCENT <= value < 2.0 + for value in improvements + ) + ) + if not plateau: + return { + **base, + "reason": ( + "HIP plateau is not proven: require three recent valid " + "HIP candidates within [-2%, +2%) of the then-current best, " + "after the required HIP exploration. Large regressions do " + "not prove a plateau. Continue HIP-only work." + ), + "recent_valid_improvements_percent": improvements, + } + + policy = { + **base, + "phase": "isa_guided_hip", + "skill_allowed": True, + "plateau": True, + "reason": ( + "HIP plateau proven. Complete two valid ISA-guided HIP rounds. " + "One selected ISA Skill may guide each HIP/DUMMA/intrinsic " + "code-shaping change." + ), + "recent_valid_improvements_percent": improvements, + } + valid_isa = [ + record for record in valid + if (record.get("isa_policy") or {}).get("phase") + == "isa_guided_hip" + ] + policy["valid_isa_guided_rounds"] = len(valid_isa) + if len(valid_isa) < _REQUIRED_VALID_ISA_GUIDED_ROUNDS: + return policy + + previous = valid_isa[-1] + isa_plan = previous.get("isa_optimization") or {} + limitation_confirmed = ( + isa_plan.get("compiler_limitation_confirmed") is True + and isinstance(isa_plan.get("target_instructions"), list) + and bool(isa_plan.get("target_instructions")) + and (previous.get("candidate_isa") or {}).get("available") is True + ) + if not limitation_confirmed: + return { + **policy, + "reason": ( + "HIP plateau is proven, but the prior ISA-guided round did " + "not confirm a compiler limitation with target instructions " + "and trusted candidate ISA. The two-round ISA requirement is " + "complete, but raw asm remains forbidden." + ), + } + return { + **policy, + "phase": "conditional_inline_asm", + "raw_inline_asm_allowed": True, + "verified_target_instructions": list( + isa_plan.get("target_instructions") or [] + ), + "reason": ( + "Final-round raw asm gate is open for one minimal block targeting " + "the compiler limitation verified in the preceding round." + ), + } + + +def phase_extension_reason( + *, + max_iterations: int, + history: list[Dict[str, Any]], +) -> str | None: + """Return the successful phase still owed after the nominal budget.""" + valid = [ + record for record in history + if record.get("build_success") is True + and record.get("correctness_passed") is True + and (record.get("metrics") or {}).get("graph_capture_passed") is True + ] + required_hip = max(1, max_iterations - 2) + valid_hip = [ + record for record in valid + if (record.get("isa_policy") or {}).get("phase", "hip_only") + == "hip_only" + ] + if len(valid_hip) < required_hip: + return ( + f"need {required_hip - len(valid_hip)} more valid HIP-only " + "experiment(s)" + ) + + policy = isa_round_policy( + iteration=len(history) + 1, + max_iterations=max_iterations, + history=history, + ) + if policy.get("plateau") is not True: + return "need a bounded three-experiment HIP plateau" + + valid_isa = [ + record for record in valid + if (record.get("isa_policy") or {}).get("phase") + == "isa_guided_hip" + ] + if len(valid_isa) < _REQUIRED_VALID_ISA_GUIDED_ROUNDS: + return ( + "need " + f"{_REQUIRED_VALID_ISA_GUIDED_ROUNDS - len(valid_isa)} more valid " + "ISA-guided HIP experiment(s)" + ) + + previous_plan = valid_isa[-1].get("isa_optimization") or {} + limitation_confirmed = ( + previous_plan.get("compiler_limitation_confirmed") is True + and isinstance(previous_plan.get("target_instructions"), list) + and bool(previous_plan.get("target_instructions")) + and (valid_isa[-1].get("candidate_isa") or {}).get("available") is True + ) + valid_inline = [ + record for record in valid + if (record.get("isa_policy") or {}).get("phase") + == "conditional_inline_asm" + ] + if limitation_confirmed and not valid_inline: + return "need one conditional inline-asm experiment" + return None + + +def pmc_profile_decision( + *, + iteration: int, + history: list[Dict[str, Any]], + source_uses_dumma: bool, + isa_policy: Dict[str, Any], +) -> Dict[str, Any]: + """Profile only when counters can change the next optimization decision.""" + if isa_policy.get("skill_allowed"): + return { + "profile": True, + "reason": "late-round ISA or plateau decision requires fresh PMC", + } + if iteration == 1: + if source_uses_dumma: + return { + "profile": True, + "reason": "usable DUMMA bootstrap needs one initial profile", + } + return { + "profile": False, + "reason": ( + "skip PMC for scalar bootstrap; first establish a validated " + "DUMMA/current-best implementation" + ), + } + if history and history[-1].get("accepted") is True: + return { + "profile": True, + "reason": "the preceding round established a new official best", + } + return { + "profile": False, + "reason": ( + "current best is unchanged; reuse matching prior counters and " + "reserve fresh PMC for accepted-best or late ISA decisions" + ), + } + + +def _cached_pmc_evidence( + profile_root: Path, + source_digest: str, +) -> Dict[str, Any] | None: + candidates = sorted( + profile_root.glob("iteration*/pmc.json"), + key=lambda path: path.stat().st_mtime_ns, + reverse=True, + ) + for path in candidates: + try: + evidence = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + continue + if ( + evidence.get("available") is not False + and evidence.get("source_hip_digest") == source_digest + ): + return { + **evidence, + "reused": True, + "reused_from": str(path), + } + return None + + +def _compact_metrics_for_prompt(metrics: Dict[str, Any]) -> Dict[str, Any]: + """Keep decision signals while leaving full samples in the fact ledger.""" + keys = ( + "median_us", "p90_us", "min_us", "max_us", "passed", + "graph_capture_passed", "mismatch_count", "max_abs_error", + "logical_tops", "algorithmic_bandwidth_gb_s", "shape", + "official_best_median_us", "official_best_p90_us", + "shadow_candidate_active", "correctness_passed_in_precheck", + ) + compact = {key: metrics.get(key) for key in keys if key in metrics} + fallback = metrics.get("paired_m2_fallback_validation") + if isinstance(fallback, dict): + compact["paired_m2_fallback_validation"] = { + key: fallback.get(key) + for key in ( + "passed", "graph_capture_passed", "mismatch_count", + "median_us", "p90_us", "shape", + ) + if key in fallback + } + return compact + + +def _compact_pmc_for_prompt(evidence: Dict[str, Any]) -> Dict[str, Any]: + """Retain actionable PMC/ISA facts without repeating raw artifacts.""" + compact = { + key: evidence.get(key) + for key in ( + "available", "reused", "reused_from", "skipped", "skip_reason", + "error", "shape", "source_hip_digest", "compile_cache_key", + "primary_kernel_name", "profiled_duration_us", "grid_blocks", + "workgroup_size", "lds_bytes", "scratch_bytes", "arch_vgpr", + "accum_vgpr", "sgpr", "wave_size", "counters", "l2_hits", + "l2_misses", "l2_hit_rate_percent", "device_cu_count", + "interpretation_guard", + ) + if key in evidence + } + memory = evidence.get("memory_traffic") + if isinstance(memory, dict): + compact["memory_traffic"] = { + key: memory.get(key) + for key in ( + "available", "read_bytes_per_operator_replay", + "write_bytes_per_operator_replay", + "total_bytes_per_operator_replay", + "counter_derived_operator_hbm_bandwidth_gb_s", "reason", + ) + if key in memory + } + isa = evidence.get("isa") + if isinstance(isa, dict): + compact["isa"] = { + key: isa.get(key) + for key in ( + "available", "kernel_name", "instruction_counts", + "waitcnt_expressions", "resources", + "profiled_launch_resources", "key_instruction_excerpt", + "error", "interpretation_guard", + ) + if key in isa + } + return compact + + +def is_infrastructure_failure(reason: str | None) -> bool: + normalized = str(reason or "").lower() + return any(token in normalized for token in ( + "timeout", "timed out", "killed", "no result", "exit 143", + )) + + +def evaluate_candidate_acceptance( + *, + passed: bool, + metrics: Dict[str, Any], + best_metrics: Dict[str, Any], + minimum_improvement_percent: float, + shadow_metrics: Dict[str, Any] | None = None, +) -> Dict[str, Any]: + """Apply stable median and no-P90-regression acceptance gates.""" + candidate_us = float(metrics.get("median_us") or float("inf")) + best_us = float(best_metrics["median_us"]) + candidate_p90 = float(metrics.get("p90_us") or candidate_us) + best_p90 = float(best_metrics.get("p90_us") or best_us) + improvement = (best_us / candidate_us - 1.0) * 100.0 + p90_guard_passed = candidate_p90 <= best_p90 + accepted = ( + passed + and candidate_us < best_us + and improvement >= minimum_improvement_percent + and p90_guard_passed + ) + shadow_us = float( + (shadow_metrics or {}).get("median_us") or float("inf") + ) + shadow_p90 = float( + (shadow_metrics or {}).get("p90_us") or shadow_us + ) + improves_shadow = ( + candidate_us < shadow_us and candidate_p90 <= shadow_p90 + ) + shadow_eligible = ( + passed + and not accepted + and _SHADOW_MIN_IMPROVEMENT_PERCENT <= improvement + < minimum_improvement_percent + and p90_guard_passed + and improves_shadow + ) + return { + "accepted": accepted, + "candidate_us": candidate_us, + "best_us": best_us, + "candidate_p90_us": candidate_p90, + "best_p90_us": best_p90, + "improvement_percent": improvement, + "p90_guard_passed": p90_guard_passed, + "shadow_eligible": shadow_eligible, + "shadow_base_us": ( + shadow_us if shadow_us != float("inf") else None + ), + "improves_shadow": improves_shadow, + } + + +def evaluate_final_target( + *, + baseline: Dict[str, Any], + validation: Dict[str, Any], + target_improvement_percent: float, +) -> Dict[str, Any]: + """Evaluate the user target only against the fixed baseline at report time.""" + validation_shapes = validation.get("shapes", validation) + shapes: Dict[str, Any] = {} + for shape_id, baseline_record in baseline.items(): + final_record = validation_shapes.get(shape_id) or {} + final_metrics = final_record.get("metrics", final_record) + baseline_value = baseline_record.get("median_us") + final_value = final_metrics.get("median_us") + baseline_us = float(baseline_value) if baseline_value else None + final_us = float(final_value) if final_value else None + improvement = ( + (baseline_us / final_us - 1.0) * 100.0 + if baseline_us is not None and final_us is not None else None + ) + shapes[shape_id] = { + "baseline_us": baseline_us, + "final_us": final_us, + "improvement_percent": improvement, + "target_met": ( + final_record.get("passed") is True + and improvement is not None + and improvement >= target_improvement_percent + ), + } + return { + "target_improvement_percent": target_improvement_percent, + "semantics": "final validated result versus fixed baseline", + "all_shapes_met": bool(shapes) and all( + record["target_met"] for record in shapes.values() + ), + "shapes": shapes, + } + + +def experiment_fact_ledger(worker_root: Path) -> list[Dict[str, Any]]: + """Return control-plane facts suitable for Skill authoring.""" + facts: list[Dict[str, Any]] = [] + for path in sorted((worker_root / "runs").glob("*/experiments.jsonl")): + for line in path.read_text(encoding="utf-8").splitlines(): + try: + record = json.loads(line) + except ValueError: + continue + if not isinstance(record, dict): + continue + facts.append({ + "shape_id": record.get("shape_id"), + "iteration": record.get("iteration"), + "accepted": record.get("accepted"), + "build_success": record.get("build_success"), + "correctness_passed": record.get("correctness_passed"), + "metrics": record.get("metrics") or {}, + "architecture": record.get("architecture") or {}, + "baseline_us": record.get("baseline_us"), + "speedup": record.get("speedup"), + "p90_guard_passed": record.get("p90_guard_passed"), + "failure_category": ( + "compile_or_agent_failure" + if record.get("build_success") is False + else None + ), + }) + return facts + + +def validate_skill_draft( + content: str, + facts: list[Dict[str, Any]], +) -> None: + """Reject common unsupported conclusions before they become memory.""" + lowered = content.lower() + violations = [ + claim for claim in _UNSUPPORTED_SKILL_CLAIMS + if claim in lowered + ] + if violations: + raise ValueError( + "skill draft contains unsupported conclusions: " + f"{violations}" + ) + if facts and not any(item.get("accepted") for item in facts): + markers = ( + "no accepted", + "none accepted", + "zero accepted", + "0 accepted", + "no candidate was accepted", + "no optimization was accepted", + ) + if not any(marker in lowered for marker in markers): + raise ValueError( + "skill draft does not clearly state that no candidate was " + "accepted" + ) + + +def archive_iteration_candidate( + source: Path, + iteration_dir: Path, + extra_paths: list[str] | None = None, +) -> list[str]: + """Keep an immutable-looking copy of implementation code for one round. + + The operator API contract deliberately remains shared in ``source``. + Candidate implementation and HIP files are copied before a rejected + working tree is restored or a later agent overwrites it. + """ + relative_paths = { + "w8a8_gemm.py", + "w8a8_backend.py", + "setup.py", + "csrc/bindings.cpp", + *(extra_paths or []), + } + for path in source.rglob("*.hip"): + if ".git" not in path.parts: + relative_paths.add(path.relative_to(source).as_posix()) + + archived = [] + for relative in sorted(relative_paths): + if ( + relative in {W8A8_API_FILENAME, "proposal.json"} + or relative.startswith("references/") + ): + continue + candidate = source / relative + if not candidate.is_file(): + continue + destination = iteration_dir / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(candidate, destination) + archived.append(relative) + return archived + + +def candidate_iteration_destination( + workspace_dir: Path, + assignment: WorkerAssignment, + shape_id: str, + iteration: int, +) -> Path | None: + """Return the stable kernel-repo path for one immutable round snapshot.""" + candidate_dir = ( + workspace_dir / "main" / "candidates" / assignment.worker_id + ) + if not candidate_dir.is_dir() or candidate_dir.is_symlink(): + return None + if len(assignment.shape_ids) == 1: + return candidate_dir / f"iteration{iteration}" + return candidate_dir / shape_id / f"iteration{iteration}" + + +def publish_iteration_candidate( + iteration_dir: Path, + destination: Path, +) -> None: + """Copy one archived round into the user-visible kernel repository.""" + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copytree(iteration_dir, destination, dirs_exist_ok=True) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def compiled_kernel_object(worker_root: Path) -> Path: + """Return the exact HIP object produced by the trusted worker benchmark.""" + current_build = worker_root / "cache" / "current_build.json" + if current_build.is_file(): + try: + metadata = json.loads(current_build.read_text(encoding="utf-8")) + extension_name = str(metadata["extension_name"]) + expected = ( + worker_root / "cache" / "torch" / extension_name + / "w8a8_gemm_hip.cuda.o" + ) + if expected.is_file(): + return expected + except (KeyError, TypeError, ValueError, OSError): + pass + expected = ( + worker_root / "cache" / "torch" / "metainfer_w8a8_backend" + / "w8a8_gemm_hip.cuda.o" + ) + if expected.is_file(): + return expected + matches = sorted( + (worker_root / "cache" / "torch").glob( + "**/w8a8_gemm_hip.cuda.o" + ), + key=lambda path: path.stat().st_mtime_ns, + reverse=True, + ) + if not matches: + raise FileNotFoundError( + "trusted benchmark did not leave a compiled " + "w8a8_gemm_hip.cuda.o" + ) + return matches[0] + + +def snapshot_accepted_kernel_artifact( + *, + worker_root: Path, + shape_id: str, + shape: Dict[str, Any], + metrics: Dict[str, Any], + commit: str, + isa_evidence: Dict[str, Any] | None = None, + isa_dir: Path | None = None, +) -> Dict[str, Any]: + """Freeze the benchmarked HIP source and its exact compiled object.""" + source_hip = worker_root / "source" / "csrc" / "w8a8_gemm_hip.hip" + object_file = compiled_kernel_object(worker_root) + destination = worker_root / "accepted" / _safe(shape_id) + if destination.exists(): + shutil.rmtree(destination) + destination.mkdir(parents=True) + archived_hip = destination / "kernel.hip" + archived_object = destination / "kernel.cuda.o" + shutil.copy2(source_hip, archived_hip) + shutil.copy2(object_file, archived_object) + build_ninja = object_file.parent / "build.ninja" + if build_ninja.is_file(): + shutil.copy2(build_ninja, destination / "build.ninja") + archived_isa: list[str] = [] + if isa_dir is not None: + for name in ("gfx928.co", "metadata.txt", "isa.txt"): + candidate = isa_dir / name + if candidate.is_file(): + shutil.copy2(candidate, destination / name) + archived_isa.append(name) + manifest = { + "schema_version": SCHEMA_VERSION, + "shape_id": shape_id, + "shape": { + key: int(shape[key]) for key in ("M", "N", "K") + }, + "commit": commit, + "source": str(archived_hip.relative_to(worker_root)), + "object": str(archived_object.relative_to(worker_root)), + "source_sha256": _sha256_file(archived_hip), + "object_sha256": _sha256_file(archived_object), + "compile_target": "gfx928", + "compiler_contract": ( + "trusted torch cpp_extension worker build; exact accepted object" + ), + "launch_abi": "launch_w8a8_gemm", + "pack_abi": "launch_pack_w8a8_weight", + "metrics": { + key: metrics[key] + for key in ( + "median_us", "p90_us", "min_us", "max_us", + "graph_capture_passed", "mismatch_count", "max_abs_error", + ) + if key in metrics + }, + "isa_artifacts": archived_isa, + "isa_evidence": { + key: value for key, value in (isa_evidence or {}).items() + if key != "artifact_paths" + }, + } + write_json(destination / "manifest.json", manifest) + return manifest + + +def _check_required_files(repo: Path) -> bool: + """Return True for either supported complete kernel-repo layout.""" + return any( + all((repo / name).is_file() for name in required_files) + for required_files in _REQUIRED_FILE_SETS + ) + + +class W8A8Runner: + def __init__(self, worker_root: Path, gpu: int) -> None: + self.worker_root = worker_root + self.source = worker_root / "source" + staged_harness = self.source / "w8a8_bench.py" + self.harness = ( + staged_harness if staged_harness.is_file() else HARNESS + ) + self.env = dict(os.environ) + bind_worker_gpu(self.env, gpu) + self.env.update({ + "MAX_JOBS": "2", + "PYTHONDONTWRITEBYTECODE": "1", + "PYTORCH_ROCM_ARCH": "gfx928", + "TORCH_EXTENSIONS_DIR": str(worker_root / "cache" / "torch"), + "TRITON_CACHE_DIR": str(worker_root / "cache" / "triton"), + "XDG_CACHE_HOME": str(worker_root / "cache" / "xdg"), + "TMPDIR": str(worker_root / "cache" / "tmp"), + }) + for key in ( + "TORCH_EXTENSIONS_DIR", "TRITON_CACHE_DIR", "XDG_CACHE_HOME", + "TMPDIR", + ): + Path(self.env[key]).mkdir(parents=True, exist_ok=True) + + def _prepare_compile_cache(self) -> Dict[str, Any]: + """Stage immutable content-addressed compiler inputs for Ninja reuse.""" + inputs = [self.source / "csrc" / "bindings.cpp"] + prebuilt = sorted((self.source / "prebuilt").glob("*.o")) + dispatch = self.source / "csrc" / "w8a8_dispatch.cpp" + if prebuilt: + inputs.extend([dispatch, *prebuilt]) + else: + inputs.append(self.source / "csrc" / "w8a8_gemm_hip.hip") + missing = [str(path) for path in inputs if not path.is_file()] + if missing: + raise FileNotFoundError( + f"compile inputs are missing: {missing}" + ) + digest = hashlib.sha256() + for path in inputs: + relative = path.relative_to(self.source) + digest.update(str(relative).encode("utf-8")) + digest.update(b"\0") + digest.update(bytes.fromhex(_sha256_file(path))) + build_key = digest.hexdigest() + cache_source = ( + self.worker_root / "cache" / "compile_sources" / build_key + ) + if not cache_source.is_dir(): + temporary = cache_source.with_name( + f"{cache_source.name}.tmp-{os.getpid()}" + ) + if temporary.exists(): + shutil.rmtree(temporary) + for path in inputs: + destination = temporary / path.relative_to(self.source) + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(path, destination) + write_json(temporary / "manifest.json", { + "build_key": build_key, + "inputs": { + str(path.relative_to(self.source)): _sha256_file(path) + for path in inputs + }, + "arch": "gfx928", + "flags": ["-O3", "--offload-arch=gfx928"], + }) + temporary.replace(cache_source) + extension_name = f"metainfer_w8a8_backend_{build_key[:24]}" + self.env["METAINFER_W8A8_COMPILE_SOURCE_DIR"] = str(cache_source) + self.env["METAINFER_W8A8_BUILD_KEY"] = build_key + metadata = { + "build_key": build_key, + "extension_name": extension_name, + "compile_source_dir": str(cache_source), + } + write_json(self.worker_root / "cache" / "current_build.json", metadata) + return metadata + + def probe(self) -> Dict[str, Any]: + result = _run([ + "python3", str(self.harness), "--source", str(self.source), + "--m", "1", "--n", "1", "--k", "1", "--probe", + ], cwd=self.source, env=self.env, timeout=60) + return _last_json(result.stdout) + + def benchmark( + self, + shape: Dict[str, Any], + *, + warmups: int | None = None, + samples: int | None = None, + replays_per_sample: int | None = None, + check_correctness: bool = True, + ) -> Dict[str, Any]: + try: + m = int(shape["M"]) + n = int(shape["N"]) + k = int(shape["K"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("W8A8 shapes require integer M, N and K") from exc + if check_correctness: + # The CPU int64 reference is the slow part of a + # correctness-checked benchmark for M>=3072; prepare and cache it + # once per shape outside the 900s benchmark subprocess budget so + # the timed run below always hits the cache. + self._ensure_reference_cache(m, n, k) + command = [ + "python3", str(self.harness), "--source", str(self.source), + "--m", str(m), "--n", str(n), "--k", str(k), + "--reference-cache-dir", + str(self.worker_root / "cache" / "references"), + ] + if warmups is not None: + command.extend(["--warmups", str(warmups)]) + if samples is not None: + command.extend(["--samples", str(samples)]) + if replays_per_sample is not None: + command.extend([ + "--replays-per-sample", str(replays_per_sample) + ]) + if not check_correctness: + command.append("--skip-correctness") + build = self._prepare_compile_cache() + result = _run( + command, cwd=self.source, env=self.env, + timeout=_BENCHMARK_TIMEOUT_S, + ) + metrics = _last_json(result.stdout) + metrics["compile_cache_key"] = build["build_key"] + return metrics + + def _reference_cache_path(self, m: int, n: int, k: int) -> Path: + """Path of the exact int64 reference cache for one (M, N, K).""" + return ( + self.worker_root / "cache" / "references" + / f"exact-int64-v1-m{m}-n{n}-k{k}.pt" + ) + + def _ensure_reference_cache(self, m: int, n: int, k: int) -> None: + """Compute and cache the exact CPU int64 reference once per shape. + + ``w8a8_bench.py --prepare-reference`` regenerates the exact + deterministic inputs the timed run uses (fixed seed, CUDA RNG) and + persists the reference, so a cache hit inside the benchmark is + bit-identical to computing it there. M>=3072 references can take + ~10+ minutes (CPU int64 GEMM at ~0.2 GFLOPS), which is why this runs + under its own generous timeout instead of the benchmark's 900s. + """ + reference_path = self._reference_cache_path(m, n, k) + if reference_path.is_file(): + return + cache_dir = reference_path.parent + cache_dir.mkdir(parents=True, exist_ok=True) + command = [ + "python3", str(self.harness), "--source", str(self.source), + "--m", str(m), "--n", str(n), "--k", str(k), + "--reference-cache-dir", str(cache_dir), + "--prepare-reference", + ] + result = _run( + command, cwd=self.source, env=self.env, + timeout=_REFERENCE_PREPARE_TIMEOUT_S, + ) + payload = _last_json(result.stdout) + if ( + not payload.get("reference_prepared") + or not reference_path.is_file() + ): + raise RuntimeError( + "reference cache preparation failed for " + f"M={m}, N={n}, K={k}: {result.stdout[-2000:]}" + ) + + def profile_pmc( + self, + shape: Dict[str, Any], + output_dir: Path, + ) -> Dict[str, Any]: + """Profile the current trusted source once and return compact PMC data.""" + try: + m = int(shape["M"]) + n = int(shape["N"]) + k = int(shape["K"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("W8A8 shapes require integer M, N and K") from exc + script = self.source / "profile_pmc.sh" + if not script.is_file(): + raise FileNotFoundError( + f"trusted PMC script is missing: {script}" + ) + output_dir.mkdir(parents=True, exist_ok=True) + build = self._prepare_compile_cache() + _run( + [ + "bash", + str(script), + str(self.harness), + str(self.source), + str(m), + str(n), + str(k), + str(output_dir), + ], + cwd=self.source, + env=self.env, + timeout=_BENCHMARK_TIMEOUT_S, + ) + evidence = parse_pmc_csv(output_dir / "pmc.csv") + read_csv = output_dir / "pmc-read.csv" + write_csv = output_dir / "pmc-write.csv" + if read_csv.is_file() and write_csv.is_file(): + evidence["memory_traffic"] = parse_memory_traffic_csv( + read_csv, write_csv + ) + else: + evidence["memory_traffic"] = { + "available": False, + "reason": ( + "This repository predates trusted HBM traffic " + "profiling. General PMC remains valid; do not make an " + "HBM bandwidth claim." + ), + } + evidence["shape"] = {"M": m, "N": n, "K": k} + evidence["source_hip_digest"] = file_digest( + self.source / "csrc" / "w8a8_gemm_hip.hip" + ) + evidence["compile_cache_key"] = build["build_key"] + evidence["csv_path"] = str(output_dir / "pmc.csv") + evidence["memory_csv_paths"] = { + "read": str(read_csv), + "write": str(write_csv), + } + return evidence + + def inspect_isa( + self, + output_dir: Path, + *, + kernel_names: list[str] | None = None, + primary_kernel_name: str | None = None, + ) -> Dict[str, Any]: + """Audit the exact object most recently built by the trusted runner.""" + evidence = inspect_gfx928_object( + compiled_kernel_object(self.worker_root), + output_dir, + kernel_names=kernel_names, + primary_kernel_name=primary_kernel_name, + ) + evidence["source_hip_digest"] = file_digest( + self.source / "csrc" / "w8a8_gemm_hip.hip" + ) + return evidence + + +class RealW8A8OptimizationPipeline: + def __init__( + self, + *, + req: Dict[str, Any], + state_dir: Path, + workspace_dir: Path, + store: StateStore, + manager: SubAgentManager, + ) -> None: + self.req = req + self.state_dir = state_dir + self.workspace_dir = workspace_dir + self.store = store + self.manager = manager + self._current_phase = phases.PREPARE + self._progress_lock = threading.Lock() + self._reported_iteration = 0 + self._worker_failures: Dict[str, Dict[str, Any]] = {} + + def _phase(self, phase: str, **payload: Any) -> None: + self._current_phase = phase + self.store.update_run(current_phase=phase) + self.store.append_timeline("phase_start", {"phase": phase, **payload}) + + def run(self, *, dry_run: bool = False) -> Dict[str, Any]: + task_id = str(self.req.get("task_id", "task")) + self.store.init_or_resume(task_id) + self.store.update_run( + current_iteration=0, + finished=False, + final_status=None, + last_outcome=None, + last_transition_label=None, + notes=[], + ) + started = time.time() + try: + self._phase(phases.PREPARE) + config = load_config(self.req) + self._validate_contract(config) + self._prepare_worktrees(config, task_id) + plan = self._plan(config) + write_json(self.workspace_dir / "plan.json", plan) + if dry_run: + return plan + + self._phase(phases.BASELINE) + baseline = self._parallel_baseline(config) + + self._phase(phases.EXPLORE, workers=len(config.assignments)) + workers = self._parallel_agents(config, baseline) + + self._phase(phases.SYNTHESIZE) + completed_assignments = [ + item for item in config.assignments + if item.worker_id in workers + ] + merged_skill = self._author_merged_skill( + config, completed_assignments + ) + + self._phase(phases.VALIDATE) + validation = self._serial_validate(config, workers) + + self._phase(phases.REPORT) + final_target = evaluate_final_target( + baseline=baseline, + validation=validation, + target_improvement_percent=( + config.minimum_improvement_percent + ), + ) + report = { + "schema_version": SCHEMA_VERSION, + "task_id": task_id, + "task_type": "dcu-kernel-auto-opt", + "mode": "real-int8-w8a8-gemm", + "started_at": started, + "finished_at": time.time(), + "duration_s": round(time.time() - started, 4), + "config": plan, + "baseline": baseline, + "workers": workers, + "worker_failures": self._worker_failures, + "merged_skill": merged_skill, + "final_validation": validation, + "final_target": final_target, + "real_gpu_used": True, + "target_repo_modified": False, + "status": ( + "partial_success" if self._worker_failures else "success" + ), + } + write_json(self.workspace_dir / "final_report.json", report) + self.store.update_run( + current_iteration=config.mock_iterations, + current_phase=phases.FINISHED, + finished=True, + final_status="success", + last_outcome="ok", + last_transition_label="real W8A8 optimization complete", + ) + self.store.append_timeline( + "orchestrator_success", + { + "workers": len(workers), + "real_gpu_used": True, + "operator": "int8_w8a8_gemm", + }, + ) + return report + except Exception as exc: + draft.unlink(missing_ok=True) + self.store.append_timeline( + "orchestrator_error", {"error": repr(exc)} + ) + self.store.update_run( + current_phase=self._current_phase, + finished=True, + final_status="stopped", + last_outcome="infra_fail", + notes=[str(exc)], + ) + raise + + @staticmethod + def _validate_contract(config: OptimizerConfig) -> None: + if config.operator != "Quantized GEMM": + raise ValueError( + "Real INT8 W8A8 GEMM mode requires operator=Quantized GEMM" + ) + if config.dtype != "INT8 W8A8": + raise ValueError( + "Real INT8 W8A8 GEMM mode requires dtype=INT8 W8A8" + ) + for shape in config.shapes.values(): + for key in ("M", "N", "K"): + if key not in shape.params: + raise ValueError(f"{shape.id} is missing {key}") + + def _prepare_worktrees( + self, config: OptimizerConfig, task_id: str + ) -> None: + self.workspace_dir.mkdir(parents=True, exist_ok=True) + seed = self.workspace_dir / "main" + if not (seed / ".git").exists(): + source = config.target_repo_path + repo_exists = ( + source is not None + and source.is_dir() + and _check_required_files(source) + ) + if repo_exists: + assert source is not None + shutil.copytree( + source, + seed, + dirs_exist_ok=True, + ignore=shutil.ignore_patterns( + ".git", "__pycache__", "*.pyc", "profiles", + "build", ".pytest_cache", + ), + ) + _run(["git", "init"], cwd=seed) + _run(["git", "config", "user.name", "MetaInfer Agent"], cwd=seed) + _run([ + "git", "config", "user.email", "metainfer@localhost", + ], cwd=seed) + _run(["git", "add", "."], cwd=seed) + _run([ + "git", "commit", "-m", "seed extracted INT8 W8A8 GEMM", + ], cwd=seed) + else: + raise RuntimeError( + "The optimize-existing-repo mode requires a complete " + "kernel repository. Use Generate + Optimize mode when " + "child agents should create kernels from the fixed API." + ) + for assignment in config.assignments: + root = self.workspace_dir / "workers" / assignment.worker_id + for name in ("build", "cache", "logs", "runs", "artifacts"): + (root / name).mkdir(parents=True, exist_ok=True) + source = root / "source" + if not source.exists(): + branch = f"agent/{_safe(task_id)}/{assignment.worker_id}" + _run([ + "git", "worktree", "add", "-b", branch, + str(source), "HEAD", + ], cwd=seed) + for name in ("shared_baseline", "final_validation", "skills"): + (self.workspace_dir / name).mkdir(parents=True, exist_ok=True) + + @staticmethod + def _plan(config: OptimizerConfig) -> Dict[str, Any]: + return { + "schema_version": SCHEMA_VERSION, + "execution_mode": config.execution_mode, + "operator": "INT8 W8A8 GEMM", + "dtype": config.dtype, + "hardware": config.hardware, + "kernel_language": config.kernel_language, + "claude_model": config.claude_model, + "max_iterations": config.mock_iterations, + "minimum_improvement_percent": ( + config.minimum_improvement_percent + ), + "minimum_improvement_semantics": ( + "final validated result versus fixed baseline" + ), + "round_acceptance_improvement_percent": ( + ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + ), + "shape_scope": config.shape_scope, + "assignment_mode": config.assignment_mode, + "target_repo_path": ( + str(config.target_repo_path) if config.target_repo_path else None + ), + "harness": "trusted PyTorch FP32-accumulate W8A8 reference", + "contract": { + "A": "int8[M,K], row-major contiguous", + "B": "int8[K,N], row-major contiguous", + "x_scale": "float32[M,1]", + "weight_scale": "float32[N,1]", + "output": "bfloat16[M,N]", + }, + "shapes": [ + {"id": shape.id, **shape.params} + for shape in config.shapes.values() + ], + "assignments": [ + { + "worker_id": item.worker_id, + "gpu": item.gpu, + "shapes": item.shape_ids, + } + for item in config.assignments + ], + "real_gpu_used": True, + } + + def _parallel_baseline( + self, config: OptimizerConfig + ) -> Dict[str, Dict[str, Any]]: + output: Dict[str, Dict[str, Any]] = {} + + def run_one( + assignment: WorkerAssignment, + ) -> Dict[str, Dict[str, Any]]: + root = self.workspace_dir / "workers" / assignment.worker_id + runner = W8A8Runner(root, assignment.gpu) + _status( + root, assignment, state="building", + iteration=0, shape_id=None, + ) + probe = runner.probe() + if probe.get("visible_devices") != 1: + raise RuntimeError( + f"{assignment.worker_id} sees " + f"{probe.get('visible_devices')} GPUs" + ) + _status( + root, assignment, state="baseline", + iteration=0, shape_id=None, probe=probe, + ) + result: Dict[str, Dict[str, Any]] = {} + for shape_id in assignment.shape_ids: + shape = config.shapes[shape_id].params + metrics = runner.benchmark(shape) + if not metrics.get("passed"): + raise RuntimeError( + f"baseline correctness failed: {shape_id}" + ) + result[shape_id] = fixed_triton_graph_baseline( + shape_id, + shape, + bootstrap_metrics=metrics, + ) + return result + + with ThreadPoolExecutor(max_workers=len(config.assignments)) as pool: + futures = { + pool.submit(run_one, item): item + for item in config.assignments + } + for future in as_completed(futures): + assignment = futures[future] + try: + output.update(future.result()) + except Exception as exc: + error = str(exc) + state = ( + "timed_out" + if any( + token in error.lower() + for token in ( + "timeout", "timed out", "stuck", "killed", + ) + ) + else "failed" + ) + failure = { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_ids": assignment.shape_ids, + "state": state, + "stage": "baseline", + "error": error, + "timestamp": time.time(), + } + self._worker_failures[assignment.worker_id] = failure + root = ( + self.workspace_dir / "workers" + / assignment.worker_id + ) + _status( + root, assignment, state=state, + iteration=0, shape_id=None, error=error, + ) + write_json(root / "failure.json", failure) + self.store.append_timeline("worker_failed", failure) + completed_workers = { + assignment.worker_id + for assignment in config.assignments + if assignment.worker_id not in self._worker_failures + } + minimum_completed = 1 + if len(completed_workers) < minimum_completed: + raise RuntimeError( + "parallel baseline has fewer than " + f"{minimum_completed} successful GPU workers " + f"({len(completed_workers)}/{len(config.assignments)} completed)" + ) + write_json( + self.workspace_dir / "shared_baseline" / "results.json", + {"schema_version": SCHEMA_VERSION, "shapes": output}, + ) + return output + + def _parallel_agents( + self, + config: OptimizerConfig, + baseline: Dict[str, Dict[str, Any]], + ) -> Dict[str, Any]: + output: Dict[str, Any] = {} + active_assignments = [ + item for item in config.assignments + if item.worker_id not in self._worker_failures + ] + with ThreadPoolExecutor( + max_workers=max(1, len(active_assignments)) + ) as pool: + futures = { + pool.submit( + self._run_worker, config, assignment, baseline + ): assignment + for assignment in active_assignments + } + for future in as_completed(futures): + assignment = futures[future] + root = self.workspace_dir / "workers" / assignment.worker_id + try: + result = future.result() + except Exception as exc: + error = str(exc) + state = ( + "timed_out" + if any( + token in error.lower() + for token in ("timeout", "timed out", "stuck", "killed") + ) + else "failed" + ) + failure = { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_ids": assignment.shape_ids, + "state": state, + "error": error, + "timestamp": time.time(), + } + self._worker_failures[assignment.worker_id] = failure + _status( + root, assignment, state=state, + iteration=0, shape_id=None, error=error, + ) + write_json(root / "failure.json", failure) + self.store.append_timeline("worker_failed", failure) + continue + _status( + root, assignment, state="skill_writing", + iteration=config.mock_iterations, shape_id=None, + ) + result["skill"] = self._author_worker_skill(config, assignment) + _status( + root, assignment, state="completed", + iteration=config.mock_iterations, shape_id=None, + ) + output[assignment.worker_id] = result + self.store.append_timeline( + "worker_complete", + { + "worker_id": assignment.worker_id, + "skill": result["skill"]["name"], + }, + ) + minimum_completed = 1 + if len(output) < minimum_completed: + raise RuntimeError( + "parallel explore has fewer than " + f"{minimum_completed} successful GPU workers " + f"({len(output)}/{len(config.assignments)} completed)" + ) + if self._worker_failures: + self.store.append_timeline( + "parallel_explore_degraded", + { + "completed_workers": sorted(output), + "ignored_workers": sorted(self._worker_failures), + "policy": "continue when at least one worker completes", + }, + ) + return dict(sorted(output.items())) + + def _author_worker_skill( + self, + config: OptimizerConfig, + assignment: WorkerAssignment, + ) -> Dict[str, Any]: + """Ask the finished GPU worker to distill its five-round evidence.""" + root = self.workspace_dir / "workers" / assignment.worker_id + draft = root / "source" / "worker_skill_draft.md" + draft.unlink(missing_ok=True) + fact_ledger = experiment_fact_ledger(root) + prompt = f"""You are the evidence writer for {assignment.worker_id}. +The optimization rounds are complete. Read `{root / 'result.json'}` and every +`experiments.jsonl` below `{root / 'runs'}`. Write a concise reusable Markdown +skill body to `{draft}`. + +The following control-plane fact ledger is authoritative: + +```json +{json.dumps(fact_ledger, ensure_ascii=False, indent=2)} +``` + +Do not add YAML frontmatter and do not modify any other file. Include: +- exact hardware, dtype, operator and shape scope; +- accepted and rejected changes with measured evidence; +- the best reusable optimization rules; +- correctness/performance pitfalls and non-generalizable conclusions; +- how a later agent should validate a transferred idea. + +Separate measured facts from hypotheses. Missing or failed measurements must +not be presented as successful optimizations. A compile error proves only that +candidate failed to compile; it does not prove a dtype/API/hardware feature is +unsupported. Do not call a baseline optimal unless an accepted-candidate +search establishes that fact. Do not claim template HIP kernels cannot be +launched; `HIP_KERNEL_NAME(...)` is valid with `hipLaunchKernelGGL`. +""" + prompt_file = root / "logs" / "worker-skill.prompt.txt" + prompt_file.write_text(prompt, encoding="utf-8") + name = f"{assignment.worker_id}-skill" + self.store.append_timeline( + "worker_skill_launch", {"worker_id": assignment.worker_id} + ) + try: + self.manager.launch(AgentSpec( + name=name, + role="dcu_worker_skill_writer", + prompt_file=prompt_file, + workdir=root / "source", + log_dir=root / "logs", + timeout_s=300, + stuck_timeout_s=300, + max_retries=0, + extra_args=list(_SOURCE_ONLY_AGENT_ARGS), + )) + agent_result = self.manager.result(name) + if ( + agent_result is None + or not agent_result.success + or not draft.is_file() + ): + raise RuntimeError( + agent_result.error + if agent_result is not None and agent_result.error + else "skill writer did not produce a draft" + ) + content = draft.read_text(encoding="utf-8", errors="replace") + if not content.strip(): + raise RuntimeError("skill writer produced an empty draft") + validate_skill_draft(content, fact_ledger) + draft.unlink(missing_ok=True) + self.store.append_timeline( + "worker_skill_success", {"worker_id": assignment.worker_id} + ) + return generate_worker_skill( + config=config, + assignment=assignment, + workspace_dir=self.workspace_dir, + agent_draft=content, + ) + except Exception as exc: + draft.unlink(missing_ok=True) + self.store.append_timeline( + "worker_skill_fallback", + {"worker_id": assignment.worker_id, "error": str(exc)}, + ) + return generate_worker_skill( + config=config, + assignment=assignment, + workspace_dir=self.workspace_dir, + ) + + def _author_merged_skill( + self, + config: OptimizerConfig, + assignments: list[WorkerAssignment], + ) -> Dict[str, Any]: + """Let the main agent summarize only the workers with usable evidence.""" + workdir = ( + self.workspace_dir / "workers" + / assignments[0].worker_id / "source" + ) + draft = workdir / "main_agent_skill_draft.md" + draft.unlink(missing_ok=True) + worker_skills = sorted( + (self.workspace_dir / "skills" / "pending").glob("*/SKILL.md") + ) + merged_facts = [ + fact + for assignment in assignments + for fact in experiment_fact_ledger( + self.workspace_dir / "workers" / assignment.worker_id + ) + ] + prompt = f"""You are the main DCU optimization synthesis agent. +Read these completed worker skills: +{json.dumps([str(path) for path in worker_skills], ensure_ascii=False, indent=2)} + +Unavailable GPU workers: +{json.dumps(self._worker_failures, ensure_ascii=False, indent=2)} + +Authoritative control-plane fact ledger: +```json +{json.dumps(merged_facts, ensure_ascii=False, indent=2)} +``` + +Write one concise Markdown skill body to `{draft}`. Do not add YAML frontmatter +and do not modify any other file. Summarize shape routing, accepted evidence, +rejected ideas, integration/validation rules, and fallback behavior. Clearly +mark unavailable shapes as unoptimized; never infer success from a timed-out +or failed worker. Compiler errors are candidate failures, not proof that +DUMMA, a dtype, or a HIP API is unsupported. Never call the baseline optimal +when no candidate was accepted. If the task spans both phases, keep decode +(M<=32) and prefill (M>32) in clearly separated sections and mirror the +canonical skill family structure (int8-w8a8-gemm-decode / +int8-w8a8-gemm-prefill + int8-w8a8-gemm-foundations): shared contract/layout/ +benchmark rules belong to the foundations section, phase-specific recipes and +acceptance rules to the phase section. +""" + prompt_file = self.workspace_dir / "skills" / "main-agent.prompt.txt" + prompt_file.write_text(prompt, encoding="utf-8") + self.store.append_timeline( + "main_skill_launch", + { + "completed_workers": [item.worker_id for item in assignments], + "ignored_workers": sorted(self._worker_failures), + }, + ) + try: + self.manager.launch(AgentSpec( + name="main-skill-synthesis", + role="dcu_main_skill_synthesizer", + prompt_file=prompt_file, + workdir=workdir, + log_dir=self.workspace_dir / "skills" / "logs", + timeout_s=420, + stuck_timeout_s=420, + max_retries=0, + extra_args=list(_SOURCE_ONLY_AGENT_ARGS), + )) + agent_result = self.manager.result("main-skill-synthesis") + if ( + agent_result is None + or not agent_result.success + or not draft.is_file() + ): + raise RuntimeError( + agent_result.error + if agent_result is not None and agent_result.error + else "main skill synthesizer did not produce a draft" + ) + content = draft.read_text(encoding="utf-8", errors="replace") + if not content.strip(): + raise RuntimeError("main skill synthesizer produced an empty draft") + validate_skill_draft(content, merged_facts) + draft.unlink(missing_ok=True) + self.store.append_timeline("main_skill_success", {}) + return generate_merged_skill( + config=config, + assignments=assignments, + workspace_dir=self.workspace_dir, + agent_draft=content, + failed_workers=self._worker_failures, + ) + except Exception as exc: + self.store.append_timeline( + "main_skill_fallback", {"error": str(exc)} + ) + return generate_merged_skill( + config=config, + assignments=assignments, + workspace_dir=self.workspace_dir, + failed_workers=self._worker_failures, + ) + + def _run_worker( + self, + config: OptimizerConfig, + assignment: WorkerAssignment, + baseline: Dict[str, Dict[str, Any]], + ) -> Dict[str, Any]: + root = self.workspace_dir / "workers" / assignment.worker_id + source = root / "source" + runner = W8A8Runner(root, assignment.gpu) + probe = runner.probe() + result: Dict[str, Any] = { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "branch": ( + f"agent/{_safe(self.req.get('task_id', 'task'))}/" + f"{assignment.worker_id}" + ), + "worktree_created": True, + "mode": "real-int8-w8a8-gemm", + "gpu_probe": probe, + "shapes": {}, + } + for shape_id in assignment.shape_ids: + shape = config.shapes[shape_id] + verified_experience = load_verified_experience( + ( + config.target_repo_path.parent + if config.target_repo_path is not None else None + ), + shape.params, + exclude_repo=config.target_repo_path, + ) + comparison_baseline = baseline[shape_id] + best_metrics = comparison_baseline.get( + "bootstrap_metrics", comparison_baseline + ) + comparison_target = { + key: value + for key, value in comparison_baseline.items() + if key != "bootstrap_metrics" + } + best_commit = _run( + ["git", "rev-parse", "HEAD"], cwd=source + ).stdout.strip() + artifact_manifest_path = ( + root / "accepted" / _safe(shape_id) / "manifest.json" + ) + if not artifact_manifest_path.is_file(): + snapshot_accepted_kernel_artifact( + worker_root=root, + shape_id=shape_id, + shape=shape.params, + metrics=best_metrics, + commit=best_commit, + ) + experiments_path = root / "runs" / shape_id / "experiments.jsonl" + iteration = 1 + attempt_limit = config.mock_iterations + infrastructure_recoveries = 0 + phase_extensions = 0 + shadow_metrics: Dict[str, Any] | None = None + shadow_path = ( + root / "shadow" / _safe(shape_id) / "w8a8_gemm_hip.hip" + ) + session_path = ( + root / "sessions" / f"{_safe(shape_id)}.json" + ) + shape_session_id: str | None = None + if session_path.is_file(): + try: + stored_session = json.loads( + session_path.read_text(encoding="utf-8") + ) + candidate_session = stored_session.get("session_id") + if isinstance(candidate_session, str) and candidate_session: + shape_session_id = candidate_session + except (OSError, ValueError, TypeError): + shape_session_id = None + while iteration <= attempt_limit: + working_metrics = shadow_metrics or best_metrics + guidance = claim_next_guidance( + self.state_dir / "guidance", + assignment.worker_id, + iteration, + ) + source_hip_path = source / "csrc" / "w8a8_gemm_hip.hip" + round_source_digest = file_digest(source_hip_path) + source_text = source_hip_path.read_text(encoding="utf-8") + history: list[Dict[str, Any]] = [] + if experiments_path.is_file(): + for line in experiments_path.read_text( + encoding="utf-8" + ).splitlines(): + try: + record = json.loads(line) + except json.JSONDecodeError: + continue + if isinstance(record, dict): + history.append(record) + isa_policy = isa_round_policy( + iteration=iteration, + max_iterations=config.mock_iterations, + history=history, + ) + pmc_decision = pmc_profile_decision( + iteration=iteration, + history=history, + source_uses_dumma=( + "du_mma_sync" in source_text + or "DUFragment" in source_text + ), + isa_policy=isa_policy, + ) + profile_dir = ( + root / "profiles" / shape_id / f"iteration{iteration}" + ) + if pmc_decision["profile"]: + _status( + root, assignment, state="profiling_current_best", + iteration=iteration, shape_id=shape_id, probe=probe, + ) + try: + pmc_evidence = runner.profile_pmc( + shape.params, profile_dir + ) + pmc_evidence["device_cu_count"] = probe.get( + "multi_processor_count", 0 + ) + add_unprofiled_bandwidth( + pmc_evidence, working_metrics.get("median_us") + ) + except Exception as exc: + pmc_evidence = { + "available": False, + "error": str(exc), + "shape": shape.params, + "source_hip_digest": round_source_digest, + "interpretation_guard": ( + "PMC collection failed. Do not infer a " + "bottleneck from algorithmic_bandwidth_gb_s; " + "it is not measured HBM." + ), + } + try: + isa_evidence = runner.inspect_isa( + profile_dir / "current-best-isa", + kernel_names=[ + str(item.get("kernel_name") or "") + for item in pmc_evidence.get( + "profiled_kernels", [] + ) + if item.get("kernel_name") + ], + primary_kernel_name=str( + pmc_evidence.get("primary_kernel_name") or "" + ) or None, + ) + launches = { + str(item.get("kernel_name") or ""): item + for item in pmc_evidence.get( + "profiled_kernels", [] + ) + } + for item in isa_evidence.get( + "profiled_kernels", [] + ): + launch = launches.get( + str(item.get("kernel_name") or "") + ) + if launch is not None: + item["profiled_launch_resources"] = { + key: launch.get(key) for key in ( + "grid_blocks", "workgroup_size", + "lds_bytes", "scratch_bytes", + "arch_vgpr", "accum_vgpr", "sgpr", + "wave_size", + ) + } + primary_launch = launches.get( + str(isa_evidence.get("kernel_name") or "") + ) + if primary_launch is not None: + isa_evidence["profiled_launch_resources"] = { + key: primary_launch.get(key) for key in ( + "grid_blocks", "workgroup_size", + "lds_bytes", "scratch_bytes", "arch_vgpr", + "accum_vgpr", "sgpr", "wave_size", + ) + } + pmc_evidence["isa"] = isa_evidence + except Exception as exc: + pmc_evidence["isa"] = { + "available": False, + "error": str(exc), + "interpretation_guard": ( + "Current-best ISA extraction failed. Do not " + "invent instruction-level claims." + ), + } + else: + cached = _cached_pmc_evidence( + root / "profiles" / shape_id, + round_source_digest, + ) + pmc_evidence = cached or { + "available": False, + "shape": shape.params, + "source_hip_digest": round_source_digest, + "skipped": True, + "skip_reason": pmc_decision["reason"], + "interpretation_guard": ( + "Fresh PMC was intentionally skipped. Use only " + "normal benchmark metrics until an accepted best " + "or late ISA decision triggers profiling." + ), + } + write_json(profile_dir / "pmc.json", pmc_evidence) + self.store.append_timeline( + ( + "worker_pmc_profile" + if pmc_decision["profile"] + else "worker_pmc_skipped" + ), + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_id": shape_id, + "iteration": iteration, + "available": pmc_evidence.get("available", False), + "path": str(profile_dir / "pmc.json"), + "error": pmc_evidence.get("error"), + "reason": pmc_decision["reason"], + "reused": bool(pmc_evidence.get("reused")), + }, + ) + _status( + root, assignment, state="agent_running", + iteration=iteration, shape_id=shape_id, probe=probe, + pmc_available=pmc_evidence.get("available", False), + ) + proposal_path = source / "proposal.json" + proposal_path.unlink(missing_ok=True) + baseline_inline_asm = analyze_inline_asm_source( + source_text + ) + contract_path = source / W8A8_API_FILENAME + contract_digest = file_digest(contract_path) + prompt_evidence = dict(pmc_evidence) + if not isa_policy["skill_allowed"]: + prompt_evidence.pop("isa", None) + prompt_metrics = dict(working_metrics) + prompt_metrics["official_best_median_us"] = best_metrics.get( + "median_us" + ) + prompt_metrics["official_best_p90_us"] = best_metrics.get( + "p90_us" + ) + prompt_metrics["shadow_candidate_active"] = ( + shadow_metrics is not None + ) + prompt = self._worker_prompt( + assignment, shape_id, shape.params, prompt_metrics, + root, iteration, guidance, history, prompt_evidence, + verified_experience, comparison_target, + isa_policy=isa_policy, + continuation=shape_session_id is not None, + ) + prompt_file = root / "logs" / ( + f"{shape_id}-iteration-{iteration}.prompt.txt" + ) + prompt_file.write_text(prompt, encoding="utf-8") + agent_name = ( + f"{assignment.worker_id}-{shape_id}-iter{iteration}" + ) + spec = AgentSpec( + name=agent_name, + role="dcu_w8a8_gemm_worker", + prompt_file=prompt_file, + workdir=source, + log_dir=root / "logs", + timeout_s=_AGENT_TIMEOUT_S, + stuck_timeout_s=_AGENT_STUCK_TIMEOUT_S, + max_retries=0, + extra_args=list( + _ISA_AGENT_ARGS + if isa_policy["skill_allowed"] + else _SOURCE_ONLY_AGENT_ARGS + ), + env_overrides=runner.env, + resume_session_id=shape_session_id, + ) + self.store.append_timeline( + "agent_launch", + { + "name": agent_name, + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "operator": "int8_w8a8_gemm", + "session_continuation": shape_session_id is not None, + }, + ) + agent_result = None + proposal: Dict[str, Any] = {} + failure_reason: str | None = None + compiled = False + agent_ready = False + repair_records: list[Dict[str, Any]] = [] + try: + self.manager.launch(spec) + agent_result = self.manager.result(agent_name) + if agent_result is not None and agent_result.session_id: + shape_session_id = agent_result.session_id + write_json(session_path, { + "worker_id": assignment.worker_id, + "shape_id": shape_id, + "session_id": shape_session_id, + "last_iteration": iteration, + "updated_at": time.time(), + }) + if agent_result is None or not agent_result.success: + raise RuntimeError( + f"{agent_name} failed: " + f"{agent_result.error if agent_result else 'no result'}" + ) + if not proposal_path.is_file(): + raise RuntimeError( + f"{agent_name} did not write proposal.json" + ) + proposal = json.loads( + proposal_path.read_text(encoding="utf-8") + ) + if not isinstance(proposal, dict): + raise RuntimeError( + f"{agent_name} proposal.json must be an object" + ) + agent_ready = True + except Exception as exc: + failure_reason = str(exc) + + def changed_paths() -> list[str]: + return [ + line for line in _run( + ["git", "diff", "--name-only"], cwd=source + ).stdout.splitlines() + if line and line != "proposal.json" + ] + + def preflight() -> list[str]: + if ( + contract_digest is not None + and file_digest(contract_path) != contract_digest + ): + raise RuntimeError( + f"agent modified immutable API contract " + f"{W8A8_API_FILENAME}" + ) + paths = changed_paths() + if file_digest(source_hip_path) == round_source_digest: + raise RuntimeError( + "agent made no new kernel change relative to the " + "current official/shadow working source" + ) + if not paths: + raise RuntimeError( + "agent made no tracked kernel change" + ) + unexpected = sorted( + path for path in paths + if path != "csrc/w8a8_gemm_hip.hip" + ) + if unexpected: + raise RuntimeError( + "agent changed control-plane-owned extension " + f"infrastructure: {unexpected}; only " + "csrc/w8a8_gemm_hip.hip may change" + ) + candidate_inline_asm = analyze_inline_asm_source( + source_hip_path.read_text(encoding="utf-8") + ) + old_asm = set( + baseline_inline_asm.get( + "raw_instruction_fingerprints" + ) or [] + ) + new_asm = set( + candidate_inline_asm.get( + "raw_instruction_fingerprints" + ) or [] + ) - old_asm + if new_asm: + if not isa_policy["raw_inline_asm_allowed"]: + raise RuntimeError( + "raw inline asm is forbidden by this round's " + f"ISA policy: {isa_policy['reason']}" + ) + isa_plan = proposal.get("isa_optimization") + if not isinstance(isa_plan, dict): + raise RuntimeError( + "new raw inline asm requires an " + "isa_optimization object in proposal.json" + ) + if isa_plan.get("strategy") != "inline_asm": + raise RuntimeError( + "new raw inline asm requires " + "isa_optimization.strategy=inline_asm" + ) + targets = isa_plan.get("target_instructions") + if not isinstance(targets, list) or not targets: + raise RuntimeError( + "new raw inline asm requires non-empty " + "isa_optimization.target_instructions" + ) + if int(shape.params.get("M", 0)) == 16: + architecture = proposal.get("architecture") + required_architecture_fields = { + "family", + "grid_blocks", + "waves_per_block", + "tiles_per_block", + "split_k", + "estimated_active_cus", + "packed_layout", + "staging", + "vector_load_bytes", + "barriers_per_k_step", + } + if not isinstance(architecture, dict): + raise RuntimeError( + "M=16 proposal.json must contain an " + "architecture object" + ) + missing_architecture = sorted( + required_architecture_fields - architecture.keys() + ) + if missing_architecture: + raise RuntimeError( + "M=16 proposal architecture is missing " + f"fields: {missing_architecture}" + ) + return paths + + metrics: Dict[str, Any] = {} + session_id = ( + getattr(agent_result, "session_id", None) + if agent_result is not None else None + ) + for repair_index in range(_MAX_IN_ROUND_REPAIRS + 1): + if not agent_ready: + break + if failure_reason is None: + try: + preflight() + _status( + root, + assignment, + state="validating_candidate", + iteration=iteration, + shape_id=shape_id, + probe=probe, + repair=repair_index, + ) + validation_metrics = runner.benchmark( + shape.params, + warmups=0, + samples=1, + replays_per_sample=1, + ) + compiled = True + metrics = validation_metrics + if validation_metrics.get("passed"): + # The exact same source and deterministic + # inputs just passed the CPU-int64 reference. + # The stable timing pass must not repeat that + # expensive reference, especially for large M. + metrics = runner.benchmark( + shape.params, + check_correctness=False, + ) + if ( + metrics.get("compile_cache_key") + != validation_metrics.get( + "compile_cache_key" + ) + ): + raise RuntimeError( + "timing source changed after exact " + "correctness precheck" + ) + metrics["correctness_passed_in_precheck"] = True + if int(shape.params.get("M", 0)) == 16: + paired_fallback = { + **shape.params, + "M": 2, + } + fallback_metrics = runner.benchmark( + paired_fallback, + warmups=0, + samples=1, + replays_per_sample=1, + ) + metrics[ + "paired_m2_fallback_validation" + ] = fallback_metrics + if ( + not fallback_metrics.get("passed") + or fallback_metrics.get( + "graph_capture_passed" + ) is not True + ): + raise RuntimeError( + "paired M=2 fallback failed for " + f"{shape_id}: " + f"{json.dumps(fallback_metrics)}" + ) + failure_reason = None + break + failure_reason = ( + "exact correctness failed: " + f"{json.dumps(validation_metrics)}" + ) + except Exception as exc: + failure_reason = str(exc) + + if repair_index >= _MAX_IN_ROUND_REPAIRS: + break + repair_number = repair_index + 1 + _status( + root, + assignment, + state="repairing_candidate", + iteration=iteration, + shape_id=shape_id, + probe=probe, + repair=repair_number, + max_repairs=_MAX_IN_ROUND_REPAIRS, + ) + proposal_path.unlink(missing_ok=True) + repair_prompt = self._repair_prompt( + assignment=assignment, + shape_id=shape_id, + shape=shape.params, + root=root, + iteration=iteration, + repair=repair_number, + error=failure_reason or "unknown validation failure", + metrics=metrics, + pmc_evidence=prompt_evidence, + isa_policy=isa_policy, + ) + repair_prompt_file = root / "logs" / ( + f"{shape_id}-iteration-{iteration}-" + f"repair-{repair_number}.prompt.txt" + ) + repair_prompt_file.write_text( + repair_prompt, encoding="utf-8" + ) + repair_name = ( + f"{assignment.worker_id}-{shape_id}-iter{iteration}-" + f"repair{repair_number}" + ) + self.store.append_timeline( + "worker_repair_launch", + { + "name": repair_name, + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_id": shape_id, + "iteration": iteration, + "repair": repair_number, + "max_repairs": _MAX_IN_ROUND_REPAIRS, + }, + ) + repair_failure = failure_reason + try: + self.manager.launch(AgentSpec( + name=repair_name, + role="dcu_w8a8_gemm_repair", + prompt_file=repair_prompt_file, + workdir=source, + log_dir=root / "logs", + timeout_s=_AGENT_TIMEOUT_S, + stuck_timeout_s=_AGENT_STUCK_TIMEOUT_S, + max_retries=0, + extra_args=list( + _ISA_AGENT_ARGS + if isa_policy["skill_allowed"] + else _SOURCE_ONLY_AGENT_ARGS + ), + env_overrides=runner.env, + resume_session_id=session_id, + )) + repair_result = self.manager.result(repair_name) + if ( + repair_result is None + or not repair_result.success + ): + raise RuntimeError( + f"{repair_name} failed: " + f"{repair_result.error if repair_result else 'no result'}" + ) + if repair_result.session_id: + session_id = repair_result.session_id + shape_session_id = session_id + write_json(session_path, { + "worker_id": assignment.worker_id, + "shape_id": shape_id, + "session_id": shape_session_id, + "last_iteration": iteration, + "last_repair": repair_number, + "updated_at": time.time(), + }) + if not proposal_path.is_file(): + raise RuntimeError( + f"{repair_name} did not write proposal.json" + ) + repair_proposal = json.loads( + proposal_path.read_text(encoding="utf-8") + ) + if not isinstance(repair_proposal, dict): + raise RuntimeError( + f"{repair_name} proposal.json must be an object" + ) + if repair_proposal.get("hypothesis"): + proposal["last_repair_hypothesis"] = ( + repair_proposal["hypothesis"] + ) + if isinstance( + repair_proposal.get("architecture"), dict + ): + proposal["architecture"] = ( + repair_proposal["architecture"] + ) + if isinstance( + repair_proposal.get("isa_optimization"), dict + ): + proposal["isa_optimization"] = ( + repair_proposal["isa_optimization"] + ) + failure_reason = None + repair_agent_ready = True + except Exception as exc: + failure_reason = str(exc) + repair_agent_ready = False + repair_records.append({ + "repair": repair_number, + "input_failure": repair_failure, + "agent_failure": failure_reason, + "session_id": session_id, + }) + if not repair_agent_ready: + break + + changed = changed_paths() + + iteration_dir = ( + root / "iterations" / shape_id + / f"iteration{iteration}" + ) + candidate_inline_asm = analyze_inline_asm_source( + source_hip_path.read_text(encoding="utf-8") + ) + candidate_isa: Dict[str, Any] = { + "available": False, + "error": "candidate did not compile", + } + if compiled: + try: + candidate_isa = runner.inspect_isa( + iteration_dir / "isa" + ) + except Exception as exc: + candidate_isa = { + "available": False, + "error": str(exc), + } + inline_asm_gate = evaluate_inline_asm_gate( + before=baseline_inline_asm, + after=candidate_inline_asm, + proposal=proposal, + isa_evidence=candidate_isa, + raw_inline_asm_allowed=bool( + isa_policy.get("raw_inline_asm_allowed") + ), + verified_target_instructions=( + list(isa_policy.get("verified_target_instructions") or []) + if isa_policy.get("raw_inline_asm_allowed") + else None + ), + ) + if not inline_asm_gate["passed"]: + gate_error = ( + "inline asm acceptance gate failed: " + + "; ".join(inline_asm_gate["reasons"]) + ) + failure_reason = ( + f"{failure_reason}; {gate_error}" + if failure_reason else gate_error + ) + + _status( + root, assignment, state="recording_result", + iteration=iteration, shape_id=shape_id, probe=probe, + ) + passed = ( + bool(metrics.get("passed")) + and metrics.get("graph_capture_passed") is True + and inline_asm_gate["passed"] + ) + acceptance = evaluate_candidate_acceptance( + passed=passed, + metrics=metrics, + best_metrics=best_metrics, + minimum_improvement_percent=( + ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + ), + shadow_metrics=shadow_metrics, + ) + candidate_us = acceptance["candidate_us"] + accepted = acceptance["accepted"] + shadow_eligible = acceptance["shadow_eligible"] + p90_guard_passed = acceptance["p90_guard_passed"] + experiment = { + "schema_version": SCHEMA_VERSION, + "worker_id": assignment.worker_id, + "iteration": iteration, + "shape_id": shape_id, + "shape": shape.params, + "hypothesis": ( + proposal.get("hypothesis") + or f"Round {iteration} failed before a valid proposal." + ), + "changes": changed, + "profile_evidence": ( + proposal.get("profile_evidence") or {} + ), + "architecture": proposal.get("architecture") or {}, + "isa_optimization": ( + proposal.get("isa_optimization") or {} + ), + "isa_policy": isa_policy, + "candidate_isa": candidate_isa, + "inline_asm_source": candidate_inline_asm, + "inline_asm_gate": inline_asm_gate, + "build_success": compiled, + "correctness_passed": passed, + "metrics": { + key: metrics[key] for key in ( + "median_us", "p90_us", "min_us", "max_us", + "logical_tops", "algorithmic_bytes", + "algorithmic_bandwidth_gb_s", + "latency_samples_us", "max_abs_error", + "mismatch_count", "first_mismatch", + "graph_capture_passed", "timing_mode", + "python_callable", "python_graph_api", + "graph_error", + ) if key in metrics + }, + "pmc_evidence": pmc_evidence, + "in_round_repairs": repair_records, + "baseline_us": baseline[shape_id]["median_us"], + "baseline_kind": baseline[shape_id].get( + "baseline_kind", "measured_bootstrap" + ), + "speedup": ( + round( + float(baseline[shape_id]["median_us"]) + / candidate_us, + 6, + ) + if candidate_us != float("inf") else 0.0 + ), + "accepted": accepted, + "shadow_eligible": shadow_eligible, + "shadow_base_active": shadow_metrics is not None, + "p90_guard_passed": p90_guard_passed, + "acceptance": acceptance, + "beats_baseline": ( + passed + and candidate_us + < float(baseline[shape_id]["median_us"]) + ), + "commit": None, + "failure_reason": failure_reason, + "manual_guidance": ( + guidance["text"] if guidance else None + ), + "guidance_id": guidance["id"] if guidance else None, + "agent_session_id": ( + getattr(agent_result, "session_id", None) + ), + "timestamp": time.time(), + } + candidate_files = archive_iteration_candidate( + source, iteration_dir, changed + ) + experiment["artifact_dir"] = str( + iteration_dir.relative_to(root) + ) + experiment["candidate_files"] = candidate_files + candidate_destination = candidate_iteration_destination( + self.workspace_dir, + assignment, + shape_id, + iteration, + ) + if candidate_destination is not None: + experiment["candidate_repo_dir"] = str( + candidate_destination + ) + proposal_path.unlink(missing_ok=True) + if accepted: + experiment["shadow_promoted"] = shadow_metrics is not None + _run(["git", "add", "-u"], cwd=source) + _run([ + "git", "commit", "-m", + f"{shape_id}: accept iteration {iteration}", + ], cwd=source) + best_commit = _run( + ["git", "rev-parse", "HEAD"], cwd=source + ).stdout.strip() + experiment["commit"] = best_commit + best_metrics = metrics + shadow_metrics = None + snapshot_accepted_kernel_artifact( + worker_root=root, + shape_id=shape_id, + shape=shape.params, + metrics=best_metrics, + commit=best_commit, + isa_evidence=candidate_isa, + isa_dir=iteration_dir / "isa", + ) + elif shadow_eligible: + shadow_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source_hip_path, shadow_path) + shadow_metrics = metrics + experiment["shadow_candidate"] = { + "source": str(shadow_path), + "median_us": metrics.get("median_us"), + "p90_us": metrics.get("p90_us"), + "official_best_median_us": best_metrics.get( + "median_us" + ), + "policy": ( + "experimental base only; official best is unchanged " + "until cumulative improvement reaches the normal " + "acceptance threshold" + ), + } + else: + _run(["git", "restore", "."], cwd=source) + if shadow_metrics is not None and shadow_path.is_file(): + shutil.copy2(shadow_path, source_hip_path) + write_json( + iteration_dir / "iteration.json", experiment + ) + append_jsonl(experiments_path, experiment) + if candidate_destination is not None: + publish_iteration_candidate( + iteration_dir, candidate_destination + ) + if failure_reason: + self.store.append_timeline( + "worker_iteration_failed", + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_id": shape_id, + "iteration": iteration, + "error": failure_reason, + "next_iteration": ( + iteration + 1 + if iteration < attempt_limit + else None + ), + "policy": ( + "record failed round, restore best commit, " + "and continue" + ), + }, + ) + if ( + is_infrastructure_failure(failure_reason) + and infrastructure_recoveries + < _MAX_INFRASTRUCTURE_RECOVERY_ROUNDS + ): + infrastructure_recoveries += 1 + attempt_limit += 1 + self.store.append_timeline( + "worker_iteration_recovered", + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_id": shape_id, + "failed_iteration": iteration, + "replacement_iteration": attempt_limit, + "recovery": infrastructure_recoveries, + "max_recoveries": ( + _MAX_INFRASTRUCTURE_RECOVERY_ROUNDS + ), + }, + ) + elif ( + iteration >= attempt_limit + and not is_infrastructure_failure(failure_reason) + ): + extension_reason = phase_extension_reason( + max_iterations=config.mock_iterations, + history=[*history, experiment], + ) + if ( + extension_reason is not None + and phase_extensions < _MAX_PHASE_EXTENSION_ROUNDS + ): + phase_extensions += 1 + attempt_limit += 1 + self.store.append_timeline( + "worker_phase_extended", + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_id": shape_id, + "completed_iteration": iteration, + "replacement_iteration": attempt_limit, + "reason": extension_reason, + "extension": phase_extensions, + "max_extensions": _MAX_PHASE_EXTENSION_ROUNDS, + }, + ) + with self._progress_lock: + if iteration > self._reported_iteration: + self.store.update_run(current_iteration=iteration) + self._reported_iteration = iteration + iteration += 1 + discarded_shadow = None + if shadow_metrics is not None: + discarded_shadow = { + "median_us": shadow_metrics.get("median_us"), + "p90_us": shadow_metrics.get("p90_us"), + "source": str(shadow_path), + "reason": ( + "cumulative gain did not reach the official acceptance " + "threshold before the phase budget ended" + ), + } + _run(["git", "restore", "."], cwd=source) + result["shapes"][shape_id] = { + "shape_id": shape_id, + "candidate": best_commit, + "metrics": best_metrics, + "discarded_shadow": discarded_shadow, + "artifact": json.loads( + artifact_manifest_path.read_text(encoding="utf-8") + ), + } + _status( + root, assignment, state="optimization_complete", + iteration=iteration - 1, shape_id=None, probe=probe, + ) + write_json(root / "result.json", result) + return result + + @staticmethod + def _repair_prompt( + *, + assignment: WorkerAssignment, + shape_id: str, + shape: Dict[str, Any], + root: Path, + iteration: int, + repair: int, + error: str, + metrics: Dict[str, Any], + pmc_evidence: Dict[str, Any], + isa_policy: Dict[str, Any] | None = None, + ) -> str: + isa_policy = isa_policy or { + "skill_allowed": False, + "raw_inline_asm_allowed": False, + "reason": "HIP-only repair round", + } + isa_repair = ( + "If the candidate introduced raw inline asm, use the installed " + "`hygon-gfx928-memory-isa` or `hygon-gfx928-compute-isa` Skill " + "as appropriate, and repair the smallest constraint, clobber, " + "wait, EXEC, lane-layout, or operand defect. It is valid to " + "replace unsafe asm with HIP/DUMMA code." + if isa_policy.get("skill_allowed") + else "Do not use ISA Skills or add raw inline asm in this repair." + ) + isa_repair_schema = ( + ''' "isa_optimization": { + "skill": "hygon-gfx928-memory-isa or hygon-gfx928-compute-isa", + "evidence_level": "binary or experiment", + "bottleneck": "memory, compute, or mixed", + "strategy": "hip_codegen, intrinsic, or inline_asm", + "compiler_limitation_confirmed": false, + "current_isa_findings": [], + "target_instructions": [], + "expected_isa_change": "specific observable code-object change", + "constraint_and_clobber_risks": [] + }, +''' + if isa_policy.get("skill_allowed") else "" + ) + return f"""Continue the same gfx928 W8A8 optimization round for +{assignment.worker_id}, shape {shape_id}: {json.dumps(shape)}. + +This is in-round repair {repair}/{_MAX_IN_ROUND_REPAIRS} for iteration +{iteration}. The trusted control plane compiled or checked your current HIP +code and returned: + +```text +{error[-6000:]} +``` + +Structured validation data: + +```json +{json.dumps(metrics, indent=2, sort_keys=True)} +``` + +The pre-edit PMC evidence remains: + +```json +{json.dumps(pmc_evidence, indent=2, sort_keys=True)} +``` + +Repair the current `csrc/w8a8_gemm_hip.hip` in place. Preserve the proposed +mapping and performance idea; make the smallest change that fixes the reported +compile or exact-correctness defect. Use `mismatch_count` and +`first_mismatch` to distinguish tail/index/scale/signed-unpack/race errors. +{isa_repair} +Do not redesign the kernel, change infrastructure, add files, run commands, or +run the harness. The trusted control plane will revalidate after you return. +If `graph_capture_passed` is false, use `graph_error` to remove allocation, +synchronization, host callbacks, or default-stream launches from the timed +operator while preserving the current HIP computation. + +Write strict JSON to `{root / 'source' / 'proposal.json'}` with: + +```json +{{ + "iteration": {iteration}, + "repair": {repair}, + "hypothesis": "specific root cause and minimal repair", + "architecture": {{ + "family": "preserve the current family", + "grid_blocks": 0, + "waves_per_block": 1, + "tiles_per_block": 1, + "split_k": 1, + "estimated_active_cus": 0, + "packed_layout": "identity or exact layout name", + "staging": "direct, A_only, B_only, or A_and_B", + "vector_load_bytes": 0, + "barriers_per_k_step": 0 + }}, +{isa_repair_schema} + "files_changed": ["csrc/w8a8_gemm_hip.hip"] +}} +``` +""" + + @staticmethod + def _worker_prompt( + assignment: WorkerAssignment, + shape_id: str, + shape: Dict[str, Any], + best: Dict[str, Any], + root: Path, + iteration: int, + guidance: Dict[str, Any] | None, + history: list[Dict[str, Any]] | None = None, + pmc_evidence: Dict[str, Any] | None = None, + verified_experience: list[Dict[str, Any]] | None = None, + comparison_baseline: Dict[str, Any] | None = None, + isa_policy: Dict[str, Any] | None = None, + continuation: bool = False, + ) -> str: + history = history or [] + verified_experience = verified_experience or [] + pmc_evidence = pmc_evidence or { + "available": False, + "error": "PMC evidence was not supplied", + } + comparison_baseline = comparison_baseline or best + isa_policy = isa_policy or isa_round_policy( + iteration=iteration, + max_iterations=10, + history=history, + ) + if not isa_policy["skill_allowed"]: + pmc_evidence = dict(pmc_evidence) + pmc_evidence.pop("isa", None) + guidance_text = ( + guidance["text"] if guidance else "(none; decide independently)" + ) + history_summary = [ + { + "iteration": record.get("iteration"), + "accepted": record.get("accepted"), + "build_success": record.get("build_success"), + "correctness_passed": record.get("correctness_passed"), + "metrics": _compact_metrics_for_prompt( + record.get("metrics") or {} + ), + "baseline_us": record.get("baseline_us"), + "speedup": record.get("speedup"), + "hypothesis": record.get("hypothesis"), + "isa_optimization": ( + record.get("isa_optimization") or {} + ), + "candidate_isa": { + "available": (record.get("candidate_isa") or {}).get( + "available" + ), + "instruction_counts": ( + record.get("candidate_isa") or {} + ).get("instruction_counts"), + "resources": (record.get("candidate_isa") or {}).get( + "resources" + ), + }, + "inline_asm_gate": record.get("inline_asm_gate") or {}, + "failure_reason": str( + record.get("failure_reason") or "" + )[-1200:], + "artifact_dir": record.get("artifact_dir"), + } + for record in history[-5:] + ] + if not isa_policy["skill_allowed"]: + for record in history_summary: + record.pop("isa_optimization", None) + record.pop("candidate_isa", None) + record.pop("inline_asm_gate", None) + isa_section = "" + isa_discipline = ( + "6. ISA Skills and raw inline asm are disabled for this round. " + "Use HIP, DUMMA APIs, and ordinary compiler intrinsics only." + ) + skill_boundary = ( + "The Skill tool is disabled for this HIP-only round." + ) + isa_schema = "" + change_dimensions = ( + "memory_mapping, tile, block, lds, registers, dumma, or split_k" + ) + if isa_policy["skill_allowed"]: + raw_rule = ( + "One minimal raw inline-asm block is permitted by the " + "control-plane gate." + if isa_policy["raw_inline_asm_allowed"] + else "Raw inline asm remains forbidden in this round." + ) + isa_section = f""" +## Late-round ISA evidence and selected Skills + +The nested `isa` object in PMC evidence is generated by the trusted control +plane from the exact current-best gfx928 object. It is audit evidence, not +proof of a bottleneck. The read-only Skill tool is allowed only for one of: +- `hygon-gfx928-memory-isa` for VMEM/LDS/waitcnt/barrier evidence; +- `hygon-gfx928-compute-isa` for VALU/DPP/MMAC evidence. + +Select only the skill matching one measured bottleneck. Do not use either +skill to answer DUMMA C++ API questions. {raw_rule} +""" + isa_discipline = ( + "6. Inspect only the trusted ISA excerpt relevant to one " + "bottleneck. Prefer a HIP/DUMMA/intrinsic code-shaping change. " + + raw_rule + " Raw global/buffer/flat loads and raw MMAC remain " + "forbidden. Record whether a compiler limitation is actually " + "confirmed; do not infer lane mappings or clobbers." + ) + skill_boundary = ( + "The read-only Skill tool may be used only for the selected " + "memory or compute ISA skill named above." + ) + isa_schema = ''' "isa_optimization": { + "skill": "hygon-gfx928-memory-isa or hygon-gfx928-compute-isa", + "evidence_level": "binary or experiment", + "bottleneck": "memory, compute, or mixed", + "strategy": "hip_codegen, intrinsic, decline, or inline_asm", + "compiler_limitation_confirmed": false, + "current_isa_findings": ["facts visible in trusted ISA only"], + "target_instructions": [], + "expected_isa_change": "observable candidate disassembly change", + "constraint_and_clobber_risks": [], + "decline_reason": "why instruction-level work is not justified" + }, +''' + change_dimensions += ", isa_memory, isa_compute, or inline_asm" + round_strategy = w8a8_round_strategy( + shape, iteration, history, pmc_evidence, + max_iterations=int(isa_policy.get("max_iterations") or 10), + isa_policy=isa_policy, + ) + prompt_best = _compact_metrics_for_prompt(best) + prompt_baseline = _compact_metrics_for_prompt(comparison_baseline) + prompt_pmc = _compact_pmc_for_prompt(pmc_evidence) + if continuation: + return f"""Continue the existing shape-specialized gfx928 W8A8 +optimization session for {assignment.worker_id}, shape {shape_id}: +{json.dumps(shape, sort_keys=True)}. This is iteration {iteration}. + +The immutable API, graph/current-stream contract, exact int32 accumulation, +wavefront=64 rules, source ownership, optional reference freedom, acceptance +threshold, and proposal schema from the first turn remain unchanged. Do not +reread unchanged scaffold files. Inspect only the current HIP diff/sections +needed for this decision and the fact-ledger paths below. + +Current official/shadow metrics: +```json +{json.dumps(prompt_best, indent=2, sort_keys=True)} +``` + +Fixed Triton comparison baseline: +```json +{json.dumps(prompt_baseline, indent=2, sort_keys=True)} +``` + +Mandatory decision for this round: +{round_strategy} + +ISA policy: +```json +{json.dumps(isa_policy, indent=2, sort_keys=True)} +``` + +Current-best PMC/ISA summary (fresh or exact-source cached): +```json +{json.dumps(prompt_pmc, indent=2, sort_keys=True)} +``` + +Recent trusted experiments: +```json +{json.dumps(history_summary[-3:], indent=2, sort_keys=True)} +``` + +Full evidence remains available without being repeated in this prompt: +- `{root / 'runs' / shape_id / 'experiments.jsonl'}` +- `{root / 'profiles' / shape_id}` +- `{root / 'iterations' / shape_id}` +- `{root / 'source' / 'csrc' / 'w8a8_gemm_hip.hip'}` + +Human guidance: {guidance_text} + +Make one falsifiable, focused HIP change. Do not run the harness, profiler, +Docker, SSH, package tools, or broad searches. Preserve exact shape guards and +the generic fallback. Raw asm and ISA Skills follow only the policy above. +Write strict JSON to `{root / 'source' / 'proposal.json'}` with the unchanged +first-turn schema, including iteration={iteration}, hypothesis, +profile_evidence, architecture, optional isa_optimization when required, and +files_changed=["csrc/w8a8_gemm_hip.hip"]. +""" + return f"""You are {assignment.worker_id}, a shape-specialized native +HIP optimization worker for gfx928 INT8 W8A8 GEMM. +You own physical GPU {assignment.gpu} and the branch in {root / 'source'}. +Optimize only shape {shape_id}: {json.dumps(shape)}. + +## Immutable timed contract + +The immutable operator contract is defined by `int8_w8a8_gemm_api.py`: +- A: contiguous row-major int8 [M,K] +- B: contiguous row-major int8 [K,N] +- x_scale: float32 [M,1] +- weight_scale: float32 [N,1] +- output: bfloat16 [M,N] +- logical dot: int32 accumulation, then float scaling and bf16 store +- timed API: w8a8_gemm_out(..., out, workspace), returning the same out storage +- backend op: torch.ops.zth_w8a8.gemm_out +- no allocation, compilation, autotuning, packing, host/device sync, or + default-stream launch inside the timed API +- use PyTorch's current HIP stream and caller-owned out/workspace only +- the trusted binding passes `workspace.data_ptr()` and workspace byte count + to `launch_w8a8_gemm`; split-K partials and their combine kernel must use + only this storage and their total Graph replay time is the candidate time +- the trusted binding calls `launch_pack_w8a8_weight` before capture; + packing may be optimized in HIP but must never occur in `gemm_out` +- `csrc/w8a8_gemm_hip.hip` must preserve the binding-owned launch ABI: + `launch_w8a8_gemm(..., void* workspace, int64_t workspace_bytes, + int m, int n, int k, hipStream_t stream)` and + `launch_pack_w8a8_weight(raw_weight, weight_scale, packed_weight, + packed_weight_scale, int k, int n, hipStream_t stream)` +- every candidate is captured on a non-default stream and timed exclusively + through `torch.cuda.CUDAGraph.replay()`; Graph capture failure is a + correctness failure, not a performance result + +## Optional implementation reference + +`references/w8a8_gemm_variants.hip`, when present, is a read-only collection +of previously measured implementation variants. It is not compiled by the +default build, not a seed, and not a strategy whitelist. You may adapt, +combine, or ignore it; continue exploring any legal HIP/DUMMA architecture +that improves this exact shape. Never edit the reference file, and never +claim its measurements for the current source without revalidation. + +Current best HIP metrics: {json.dumps(prompt_best)}. +If `shadow_candidate_active` is true, the checked-out HIP source is a +provisional sub-1% improvement. Build on it, but treat +`official_best_median_us`/`official_best_p90_us` as the acceptance baseline. +The control plane will promote the accumulated source only after the normal +>=1% median threshold and P90 guard pass; otherwise it restores the official +best at the end. +Fixed comparison baseline: {json.dumps(prompt_baseline)}. +The baseline is the user-supplied Triton decode Graph replay latency. It is +not the bootstrap HIP source and it is not PMC evidence. Improve the current +HIP source iteratively; report speedup against this fixed Graph baseline. +Human guidance for this round: {guidance_text} +Accepted/rejected history, when present: +`{root / 'runs' / shape_id / 'experiments.jsonl'}`. + +## Mandatory decision for this round + +{round_strategy} + +Control-plane ISA policy: + +```json +{json.dumps(isa_policy, indent=2, sort_keys=True)} +``` + +## Trusted prior-round evidence + +```json +{json.dumps(history_summary, indent=2, sort_keys=True)} +``` + +## Trusted PMC evidence for the current best source + +The control plane supplies fresh counters only for a usable DUMMA bootstrap, +a newly accepted best, or a late plateau/ISA decision. Otherwise it reuses +evidence only when the exact source digest matches, or marks PMC skipped. +Use available counters and resource fields below to choose your change. +Profiled latency is perturbed and must not be used as the acceptance timing. + +```json +{json.dumps(prompt_pmc, indent=2, sort_keys=True)} +``` + +{isa_section} + +## Verified exact-shape experience from earlier tasks + +This ledger is generated from trusted `iteration.json` records, never from +free-form Skill prose. Treat measured fields and classification as facts. +Treat `proposed_change_not_verified_fact` only as a hypothesis. + +```json +{json.dumps(verified_experience, indent=2, sort_keys=True)} +``` + +## Shape-specific starting point + +{w8a8_strategy_guidance({shape_id: shape})} + +## One-round optimization discipline + +1. Read the current source and prior experiment history. You may read archived + snapshots below `{root / 'iterations' / shape_id}` as read-only evidence. + Follow the mandatory decision above and do not repeat an already rejected change. +2. State one falsifiable bottleneck hypothesis from the current metrics and + launch geometry. Architecture rounds may implement one complete architecture + family (including its required main/combine or pack/GEMM pair); polish rounds + change one conceptual factor only. Keep all HIP implementation and required + launch dispatch inside + `csrc/w8a8_gemm_hip.hip`. +3. Preserve gfx928 invariants: + - wavefront=64; blockDim is 64/128/256, never warp-32 logic; + - DUMMA INT8 support is m16n16k32 with int32 accumulation; + - installed DTK uses `` and namespace `du::dumma`; + - preserve the current compiling bf16 representation. When a typed bf16 + pointer is needed, include `` and use + `hip_bfloat16`; do not switch between bf16 type families speculatively; + - no NVIDIA WMMA, `mma.sync`, PTX, FP8, or INT4; + - all `__syncthreads()` calls are reached by every block thread; + - coalesced lanes follow contiguous N addresses; + - single-buffer LDS must fit 64 KiB; use double buffering only when the + two buffers fit about 48 KiB, leaving occupancy headroom. +4. Do not choose split-K merely because K is large. Require too few N/M tiles + to occupy the measured device CUs. Scan legal CU-aligned candidates, + including non-power-of-two splits, within the caller workspace capacity; + the suggested set is not a whitelist. Count the combine kernel in total + latency and accept only measured median/P90 improvements. +5. Do not force occupancy with `launch_bounds` unless the source already has + resource evidence. Extra VGPR spill, barriers, or repeated HBM reads can + erase any occupancy gain. + Treat `algorithmic_bandwidth_gb_s` from the normal benchmark as an + effective minimum-bytes rate, not measured HBM traffic. Use + `memory_traffic.counter_derived_operator_hbm_bandwidth_gb_s` for HBM traffic + claims; it sums every kernel in the operator replay before combining the + request counters with the separate unprofiled whole-operator median. + Never use `profiled_duration_us` for performance acceptance. +{isa_discipline} +7. Preserve exact correctness, current-stream behavior, graph safety, and all + assigned-shape dispatch paths. Keep every optimized launch behind an exact + `(m,n,k)` guard and preserve the scalar generic fallback for unmatched + shapes, especially the paired M=2 shape with the same `(N,K)`. The trusted + control plane archives and links the exact accepted object, so do not + replace the fallback with an unconditional shape-specialized launch. + Guard packed-weight specializations by exact `(k,n)` and retain identity + packing for unmatched pairs. + Never modify or rename + `int8_w8a8_gemm_api.py`. The control plane also owns + `w8a8_backend.py`, `setup.py`, `profile_pmc.sh`, and + `csrc/bindings.cpp`; do not edit them. + +## Execution boundary + +Do not run the harness, benchmark, profiler, Docker, SSH, pip, apt, conda, +package installation, network access, or filesystem-wide searches +such as `find /`, or environment probes. Use the authoritative DTK facts in +this prompt. {skill_boundary} The trusted control plane will compile, validate exact output, +measure median/P90, and accept or +restore your diff after you return. Do not modify tests, caches, build output, +or files outside this worktree. Do not add files. + +Inspect `README.md`, `w8a8_backend.py`, `csrc/bindings.cpp`, and +`csrc/w8a8_gemm_hip.hip`, make the focused tracked-source change, inspect the +diff, and then write `{root / 'source' / 'proposal.json'}` as strict JSON: + +```json +{{ + "iteration": {iteration}, + "hypothesis": "one falsifiable bottleneck hypothesis and the focused change", + "profile_evidence": {{ + "observed_best": {json.dumps(prompt_best)}, + "path": "scalar_lds or dumma_m16n16k32", + "change_dimension": "{change_dimensions}", + "block_threads": 64, + "tile_m": 0, + "tile_n": 0, + "tile_k": 0, + "lds_bytes": 0, + "double_buffer": false, + "split_k": 1, + "expected_effect": "which measured metric should improve and why", + "risk": "correctness, occupancy, bank conflict, or combine-overhead risk", + "validation_owner": "trusted_control_plane" + }}, + "architecture": {{ + "family": "direct, split_k, multi_n_tile, persistent, packed_weight, or staged", + "grid_blocks": 0, + "waves_per_block": 1, + "tiles_per_block": 1, + "split_k": 1, + "estimated_active_cus": 0, + "packed_layout": "identity or exact layout name", + "staging": "direct, A_only, B_only, or A_and_B", + "vector_load_bytes": 0, + "barriers_per_k_step": 0 + }}, +{isa_schema} + "files_changed": ["actual tracked source paths changed this round"] +}} +``` +""" + + def _serial_validate( + self, config: OptimizerConfig, workers: Dict[str, Any] + ) -> Dict[str, Any]: + results: Dict[str, Any] = {} + for assignment in config.assignments: + if assignment.worker_id not in workers: + continue + runner = W8A8Runner( + self.workspace_dir / "workers" / assignment.worker_id, + assignment.gpu, + ) + for shape_id in assignment.shape_ids: + metrics = runner.benchmark(config.shapes[shape_id].params) + if not metrics.get("passed"): + raise RuntimeError( + f"serial validation failed: {shape_id}" + ) + winner = workers[assignment.worker_id]["shapes"][shape_id] + results[shape_id] = { + "passed": True, + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "candidate": winner["candidate"], + "metrics": metrics, + "serial": True, + } + write_json( + self.workspace_dir / "final_validation" / "results.json", + { + "schema_version": SCHEMA_VERSION, + "shapes": results, + "real_gpu_used": True, + }, + ) + return results diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/worker.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/worker.py new file mode 100644 index 00000000..8ee4970e --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/worker.py @@ -0,0 +1,180 @@ +"""Independent mock worker loop with durable experiment records.""" + +from __future__ import annotations + +import time +from pathlib import Path +from typing import Any, Callable, Dict + +from .adapters.base import KernelAdapter +from .config import ( + ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT, + OptimizerConfig, + WorkerAssignment, +) +from .guidance import claim_next_guidance +from .result_store import SCHEMA_VERSION, append_jsonl, write_json + + +def _status( + path: Path, + *, + assignment: WorkerAssignment, + state: str, + iteration: int, + shape_id: str | None, + best: Dict[str, Any] | None, +) -> None: + write_json(path, { + "schema_version": SCHEMA_VERSION, + "worker_id": assignment.worker_id, + "state": state, + "iteration": iteration, + "shape_id": shape_id, + "pid": None, + "physical_gpu": assignment.gpu, + "logical_gpu": 0, + "gpu_binding": { + "HIP_VISIBLE_DEVICES": str(assignment.gpu), + "ROCR_VISIBLE_DEVICES": None, + "strategy": "HIP_VISIBLE_DEVICES-only", + "enforced": False, + "reason": "mock mode does not launch a GPU subprocess", + }, + "best": best, + "last_update": time.time(), + }) + + +def run_mock_worker( + *, + assignment: WorkerAssignment, + config: OptimizerConfig, + baseline: Dict[str, Dict[str, float]], + worker_root: Path, + guidance_root: Path, + adapter_factory: Callable[[], KernelAdapter], +) -> Dict[str, Any]: + """Run one deterministic worker. Each caller owns a disjoint root.""" + adapter = adapter_factory() + for name in ("source", "build", "cache", "logs", "runs", "artifacts"): + (worker_root / name).mkdir(parents=True, exist_ok=True) + status_path = worker_root / "status.json" + _status( + status_path, assignment=assignment, state="starting", + iteration=0, shape_id=None, best=None, + ) + + worker_result: Dict[str, Any] = { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "branch": f"metainfer/mock/{assignment.worker_id}", + "worktree_created": False, + "mode": "mock", + "shapes": {}, + } + + for shape_id in assignment.shape_ids: + shape = config.shapes[shape_id] + run_dir = worker_root / "runs" / shape_id + experiments_path = run_dir / "experiments.jsonl" + baseline_metrics = baseline[shape_id] + baseline_us = baseline_metrics["median_us"] + best: Dict[str, Any] = { + "shape_id": shape_id, + "iteration": 0, + "median_us": baseline_us, + "p90_us": baseline_metrics["p90_us"], + "speedup": 1.0, + "commit": None, + "mock_candidate": "baseline", + } + write_json(run_dir / "best.json", best) + + for iteration in range(1, config.mock_iterations + 1): + guidance = claim_next_guidance( + guidance_root, assignment.worker_id, iteration + ) + _status( + status_path, assignment=assignment, state="profiling", + iteration=iteration, shape_id=shape_id, best=best, + ) + profile = adapter.profile(worker_root, shape) + build = adapter.build(worker_root) + correct = adapter.correctness(worker_root, shape) + bench = ( + adapter.benchmark(worker_root, shape, iteration=iteration) + if build.success and correct.success else None + ) + median_us = ( + bench.metrics.get("median_us") if bench is not None else None + ) + speedup = ( + baseline_us / median_us + if median_us is not None and median_us > 0 else 0.0 + ) + improvement = ( + best["median_us"] / median_us - 1.0 + if median_us is not None and median_us > 0 else 0.0 + ) * 100.0 + accepted = bool( + build.success + and correct.success + and bench is not None + and bench.success + and median_us < best["median_us"] + and improvement >= ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + ) + experiment = { + "schema_version": SCHEMA_VERSION, + "worker_id": assignment.worker_id, + "iteration": iteration, + "shape_id": shape_id, + "shape": shape.params, + "hypothesis": ( + f"Apply manual guidance: {guidance['text']}" + if guidance else f"mock-hypothesis-{iteration}" + ), + "changes": [ + ( + f"manual plan: {guidance['text']}" + if guidance else f"synthetic candidate {iteration}" + ) + ], + "manual_guidance": guidance["text"] if guidance else None, + "guidance_id": guidance["id"] if guidance else None, + "profile_evidence": profile.evidence, + "build_success": build.success, + "correctness_passed": correct.success, + "metrics": bench.metrics if bench else {}, + "baseline_us": baseline_us, + "speedup": round(speedup, 6), + "accepted": accepted, + "commit": None, + "failure_reason": None, + "timestamp": time.time(), + } + append_jsonl(experiments_path, experiment) + if accepted: + best = { + "shape_id": shape_id, + "iteration": iteration, + "median_us": median_us, + "p90_us": bench.metrics["p90_us"], + "speedup": round(speedup, 6), + "commit": None, + "mock_candidate": f"candidate-{iteration}", + } + write_json(run_dir / "best.json", best) + _status( + status_path, assignment=assignment, state="benchmarking", + iteration=iteration, shape_id=shape_id, best=best, + ) + worker_result["shapes"][shape_id] = best + + _status( + status_path, assignment=assignment, state="completed", + iteration=config.mock_iterations, shape_id=None, best=None, + ) + write_json(worker_root / "result.json", worker_result) + return worker_result diff --git a/metainfer/tasks/dcu_kernel_auto_opt/server/__init__.py b/metainfer/tasks/dcu_kernel_auto_opt/server/__init__.py new file mode 100644 index 00000000..1a05360e --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/server/__init__.py @@ -0,0 +1 @@ +"""Web integration for dcu-kernel-auto-opt.""" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/server/plugin.py b/metainfer/tasks/dcu_kernel_auto_opt/server/plugin.py new file mode 100644 index 00000000..7b4cf8aa --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/server/plugin.py @@ -0,0 +1,26 @@ +"""Web plugin registration.""" + +from pathlib import Path + +from metainfer.server.registry import WebPlugin, register + +from .routes import build_router + + +PLUGIN_TYPE = "dcu-kernel-auto-opt" +_FRONTEND_DIR = Path(__file__).resolve().parent.parent / "static" + +plugin = WebPlugin( + type=PLUGIN_TYPE, + label="DCU kernel auto-opt", + description=( + "Run isolated multi-agent, multi-GPU kernel optimization with " + "trusted per-operator harnesses, including real gfx928 INT8 W8A8 GEMM." + ), + build_router=build_router, + detail_view_module="app/dkao-detail", + frontend_dir=_FRONTEND_DIR, + extra_stylesheets=["dkao.css"], +) + +register(plugin) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py b/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py new file mode 100644 index 00000000..6f3b845e --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py @@ -0,0 +1,844 @@ +"""Read-only task-specific Web routes.""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys +import threading +import time +from pathlib import Path +from typing import Any, Dict + +from fastapi import APIRouter, HTTPException, Request + +from metainfer.server._helpers import ( + require_task_type, + state_dir_for, + task_or_404, + workspace_dir_for, +) +from metainfer.server.state_reader import read_requirements, read_run + +from ..orchestrator import phases +from ..orchestrator.config import ( + GEN_AND_OPT_MODE, + LEGACY_SMOKE_MODE, + SMOKE_MODE, + load_config, +) +from ..orchestrator.guidance import add_guidance, list_guidance +from ..orchestrator.rename_kernel_repo import rename_kernel_repo +from ..orchestrator.skill_store import ( + fuse_skill, + list_skill_library, + mark_fuse_running, + publish_skill, + rollback_skill, + sync_skill_libraries, +) +from ..orchestrator.variant_store import ( + add_variant, + derive_variant_meta, + list_variant_index, +) + + +PLUGIN_TYPE = "dcu-kernel-auto-opt" + + +def _load(path: Path, default: Any) -> Any: + try: + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return default + + +def _read_jsonl(path: Path) -> list[Dict[str, Any]]: + if not path.exists(): + return [] + rows = [] + for line in path.read_text(encoding="utf-8", errors="replace").splitlines(): + try: + item = json.loads(line) + except ValueError: + continue + if isinstance(item, dict): + rows.append(item) + return rows + + +def _running_pid(path: Path) -> int | None: + record = _load(path, {}) or {} + try: + pid = int(record.get("pid") or 0) + except (TypeError, ValueError): + return None + if pid <= 0: + return None + try: + os.kill(pid, 0) + except (OSError, ProcessLookupError): + return None + return pid + + +def _restart_worker_commands( + requirements: Path, + state_dir: Path, + workspace_dir: Path, + worker_id: str, +) -> tuple[list[str], list[str]]: + """Build the restart + integration commands for one failed worker lane. + + Neither command passes an explicit ``--claude-bin``: the lane binaries + resolve the agent binary from the task's ``agent_framework`` themselves + (ccb -> METAINFER_CLAUDE_BIN / ``ccb``, dsh -> ``bridge/dsh/dsh_agent.py``), + matching how the original run launched its agents. + """ + common = [ + "--state-dir", str(state_dir), + "--workspace-dir", str(workspace_dir), + "--worker-id", worker_id, + ] + command = [ + sys.executable, + "-m", + "metainfer.tasks.dcu_kernel_auto_opt.orchestrator.restart_worker", + str(requirements), + *common, + ] + integration_command = [ + sys.executable, + "-m", + ( + "metainfer.tasks.dcu_kernel_auto_opt.orchestrator." + "integrate_restarted_worker" + ), + str(requirements), + *common, + ] + return command, integration_command + + +def _read_bootstrap_attempts( + worker_root: Path, + state_dir: Path | None, + worker_id: str, + assigned_shapes: list[str], +) -> list[Dict[str, Any]]: + if state_dir is None: + return [] + snapshot = _load(state_dir / "agents.json", {}) or {} + agents = snapshot.get("agents") or [] + snapshot_ts = float(snapshot.get("ts") or 0) + snapshot_stale = snapshot_ts > 0 and time.time() - snapshot_ts > 30 + pattern = re.compile( + rf"^{re.escape(worker_id)}-bootstrap-attempt(\d+)$" + ) + matched: list[tuple[int, Dict[str, Any]]] = [] + for agent in agents: + if not isinstance(agent, dict): + continue + match = pattern.match(str(agent.get("name") or "")) + if match: + matched.append((int(match.group(1)), agent)) + bootstrap_progress = _load( + worker_root / "bootstrap_progress.json", {} + ) or {} + progress_attempt = int(bootstrap_progress.get("attempt") or 0) + if not matched and progress_attempt: + progress_status = str( + bootstrap_progress.get("status") or "pending" + ) + matched.append((progress_attempt, { + "name": f"{worker_id}-bootstrap-attempt{progress_attempt}", + "status": ( + "running" + if progress_status in {"agent_running", "validating"} + else progress_status + ), + "success": ( + True if progress_status == "passed" + else False if progress_status == "failed" + else None + ), + })) + matched.sort(key=lambda item: item[0]) + if not matched: + return [] + + successful_workers = { + str(event.get("payload", {}).get("worker_id")) + for event in _read_jsonl(state_dir / "timeline.jsonl") + if event.get("type") == "worker_bootstrap_success" + } + source = worker_root / "source" + generated_files = [ + relative + for relative in ( + "w8a8_backend.py", + "setup.py", + "csrc/bindings.cpp", + "csrc/w8a8_gemm_hip.hip", + ) + if (source / relative).is_file() + ] + proposal = _load(source / "proposal.json", {}) or {} + bootstrap_result = _load(worker_root / "bootstrap_result.json", {}) or {} + latest_attempt = matched[-1][0] + attempts = [] + for attempt, agent in matched: + raw_status = str( + agent.get("status") or agent.get("phase") or "pending" + ) + success = agent.get("success") + error = agent.get("error") + persisted_success = ( + int(bootstrap_result.get("attempt") or 0) == attempt + and bool(bootstrap_result.get("passed")) + ) + progress_matches = progress_attempt == attempt + iteration_record = _load( + worker_root / "iterations" / "bootstrap" + / f"iteration{attempt}" / "iteration.json", + {}, + ) or {} + attempt_record = ( + iteration_record + if iteration_record else + bootstrap_progress + if progress_matches else {} + ) + progress_status = str(attempt_record.get("status") or "") + if persisted_success: + status = "passed" + elif progress_status == "failed": + status = "failed" + error = attempt_record.get("error") or error + elif progress_status == "passed": + status = "passed" + elif progress_status == "validating": + status = "validating" + elif progress_status == "agent_running" and raw_status == "running": + status = "running" + elif ( + raw_status == "running" + and snapshot_stale + ): + status = "orphaned" + error = error or ( + "The orchestrator stopped updating this agent; its last " + "recorded process state is stale." + ) + elif agent.get("killed") or raw_status == "failed" or success is False: + status = "failed" + elif ( + attempt == latest_attempt + and worker_id in successful_workers + ): + status = "passed" + elif raw_status == "running": + status = "running" + elif attempt < latest_attempt: + status = "retrying" + if not error: + error = ( + "Agent finished, but trusted validation requested " + "another bootstrap attempt." + ) + elif success is True: + status = "validating" + else: + status = raw_status + hypothesis = ( + bootstrap_result.get("hypothesis") + if persisted_success else proposal.get("hypothesis") + if attempt == latest_attempt and isinstance(proposal, dict) + else None + ) + if ( + not persisted_success + and attempt_record.get("hypothesis") + ): + hypothesis = attempt_record["hypothesis"] + attempt_metrics = ( + bootstrap_result.get("metrics") or {} + if persisted_success else + attempt_record.get("metrics") or {} + ) + attempts.append({ + "kind": "bootstrap", + "attempt": attempt, + "status": status, + "hypothesis": hypothesis or ( + "Create and validate the initial HIP implementation for " + f"{', '.join(assigned_shapes) or 'assigned shapes'}." + ), + "generated_files": generated_files, + "metrics": attempt_metrics, + "artifact_dir": attempt_record.get("artifact_dir"), + "candidate_files": attempt_record.get("candidate_files") or [], + "error": error, + "elapsed_s": agent.get("elapsed_s"), + "last_output_age_s": agent.get("last_output_age_s"), + "started_at": agent.get("started_at"), + }) + return attempts + + +def read_worker_lanes( + workspace_dir: Path, state_dir: Path | None = None +) -> Dict[str, Any]: + """Return exactly four worker lanes with their full iteration history.""" + plan = _load(workspace_dir / "plan.json", {}) or {} + max_iterations = int(plan.get("max_iterations") or 0) + assignment_by_worker = { + str(item.get("worker_id")): item + for item in (plan.get("assignments") or []) + if isinstance(item, dict) and item.get("worker_id") + } + agent_by_worker: Dict[str, Dict[str, Any]] = {} + if state_dir is not None: + snapshot = _load(state_dir / "agents.json", {}) or {} + for agent in snapshot.get("agents") or []: + if not isinstance(agent, dict): + continue + name = str(agent.get("name") or "") + match = re.match( + r"^(worker_[0-3])-(?:.+-iter\d+|skill|bootstrap-attempt\d+)$", + name, + ) + if not match: + continue + worker_id = match.group(1) + current = agent_by_worker.get(worker_id) + if ( + current is None + or float(agent.get("started_at") or 0) + >= float(current.get("started_at") or 0) + ): + agent_by_worker[worker_id] = agent + lanes = [] + for index in range(4): + worker_id = f"worker_{index}" + worker_root = workspace_dir / "workers" / worker_id + assignment = assignment_by_worker.get(worker_id, {}) + assigned_shapes = assignment.get("shapes") or [] + status = _load(worker_root / "status.json", {}) or {} + bootstrap_attempts = _read_bootstrap_attempts( + worker_root, state_dir, worker_id, assigned_shapes + ) + experiments = [] + runs_root = worker_root / "runs" + if runs_root.exists(): + for path in sorted(runs_root.glob("*/experiments.jsonl")): + experiments.extend(_read_jsonl(path)) + experiments.sort( + key=lambda item: ( + float(item.get("timestamp") or 0), + int(item.get("iteration") or 0), + ) + ) + lane_state = status.get("state") + if not lane_state and bootstrap_attempts: + lane_state = f"bootstrap_{bootstrap_attempts[-1]['status']}" + current_agent = agent_by_worker.get(worker_id, {}) + agent_status = str(current_agent.get("status") or "") + last_output_age = current_agent.get("last_output_age_s") + if lane_state == "building": + step = "Building trusted baseline" + elif lane_state == "baseline": + step = "Benchmarking trusted baseline" + elif str(lane_state).startswith("bootstrap_"): + bootstrap_state = str(lane_state).removeprefix("bootstrap_") + step = { + "running": "Bootstrap Agent generating initial kernel", + "agent_running": "Child Agent generating initial HIP kernel", + "validating": "Validating initial kernel", + "retrying": "Retrying initial kernel generation", + "passed": "Initial kernel validated", + "failed": "Initial kernel unavailable; lane will be skipped", + "orphaned": "Bootstrap Agent stopped responding", + }.get(bootstrap_state, "Preparing initial kernel") + elif lane_state == "agent_running": + step = "Agent planning and editing kernel" + elif lane_state == "profiling_current_best": + step = "Profiling current best kernel with PMC" + elif lane_state == "repairing_candidate": + repair = int(status.get("repair") or 0) + max_repairs = int(status.get("max_repairs") or 4) + step = ( + "Repairing compile/correctness failure " + f"({repair}/{max_repairs})" + ) + elif lane_state == "validating_candidate": + step = "Compiling and validating candidate" + elif lane_state == "recording_result": + step = "Recording round result" + elif lane_state == "optimization_complete": + step = "Optimization rounds complete" + elif lane_state == "skill_writing": + step = "Writing worker optimization skill" + elif lane_state == "completed": + step = "Worker skill ready" + elif lane_state in {"failed", "timed_out", "skipped"}: + step = "Lane unavailable; main Agent will skip it" + elif not assignment: + step = "No shapes assigned" + else: + step = "Waiting to start" + if agent_status in {"failed", "orphaned"} and lane_state == "agent_running": + step = "Agent failed; waiting for lane fallback" + completed_rounds = len(experiments) + target_rounds = max_iterations * max(1, len(assigned_shapes)) + current_iteration = int(status.get("iteration") or 0) + current_shape = status.get("shape_id") + active_states = { + "profiling_current_best", + "agent_running", + "validating_candidate", + "repairing_candidate", + "recording_result", + } + already_recorded = any( + int(item.get("iteration") or 0) == current_iteration + and item.get("shape_id") == current_shape + for item in experiments + ) + active_iteration = None + if ( + current_iteration > 0 + and lane_state in active_states + and not already_recorded + ): + active_iteration = { + "iteration": current_iteration, + "shape_id": current_shape, + "state": lane_state, + "step": step, + "agent_name": current_agent.get("name"), + "agent_status": current_agent.get("status"), + "elapsed_s": current_agent.get("elapsed_s"), + "last_output_age_s": last_output_age, + } + if lane_state == "repairing_candidate": + active_iteration["repair"] = int( + status.get("repair") or 0 + ) + active_iteration["max_repairs"] = int( + status.get("max_repairs") or 4 + ) + lanes.append({ + "worker_id": worker_id, + "gpu": assignment.get("gpu", index), + "assigned": bool(assignment), + "assigned_shapes": assigned_shapes, + "state": lane_state or ( + "not_assigned" if not assignment else "pending" + ), + "current_iteration": current_iteration, + "current_shape": current_shape, + "step": step, + "completed_rounds": completed_rounds, + "target_rounds": target_rounds, + "agent": { + "name": current_agent.get("name"), + "status": current_agent.get("status"), + "elapsed_s": current_agent.get("elapsed_s"), + "last_output_age_s": last_output_age, + "error": current_agent.get("error"), + } if current_agent else None, + "long_running": ( + isinstance(last_output_age, (int, float)) + and last_output_age >= 180 + and agent_status == "running" + ), + "bootstrap_attempts": bootstrap_attempts, + "experiments": experiments, + "active_iteration": active_iteration, + "latest": experiments[-1] if experiments else None, + "guidance": ( + list_guidance(state_dir / "guidance", worker_id) + if state_dir is not None else [] + ), + }) + return {"workers": lanes} + + +def build_router(plugin) -> APIRouter: + router = APIRouter() + + @router.get("/summary") + def summary(task_id: str) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + state_dir = state_dir_for(entry) + workspace_dir = workspace_dir_for(entry) + workers = [] + workers_root = workspace_dir / "workers" + if workers_root.exists(): + for status_path in sorted(workers_root.glob("*/status.json")): + workers.append(_load(status_path, {})) + return { + "run": read_run(state_dir), + "plan": _load(workspace_dir / "plan.json", None), + "workers": workers, + "report": _load(workspace_dir / "final_report.json", None), + } + + @router.get("/iterations") + def iterations(task_id: str) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return read_worker_lanes(workspace_dir_for(entry), state_dir_for(entry)) + + @router.post("/workers/{worker_id}/guidance") + async def submit_guidance( + task_id: str, worker_id: str, request: Request + ) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + plan = _load(workspace_dir_for(entry) / "plan.json", {}) or {} + assigned = { + str(item.get("worker_id")) + for item in (plan.get("assignments") or []) + if isinstance(item, dict) + } + if worker_id not in assigned: + raise HTTPException(status_code=400, detail="worker is not assigned") + try: + body = await request.json() + text = body.get("text", "") if isinstance(body, dict) else "" + guidance = add_guidance( + state_dir_for(entry) / "guidance", worker_id, str(text) + ) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + return {"guidance": guidance} + + @router.post("/workers/{worker_id}/restart") + def restart_worker(task_id: str, worker_id: str) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + if not re.fullmatch(r"worker_[0-3]", worker_id): + raise HTTPException(status_code=400, detail="invalid worker id") + state_dir = state_dir_for(entry) + workspace_dir = workspace_dir_for(entry) + plan = _load(workspace_dir / "plan.json", {}) or {} + assignment = next(( + item for item in (plan.get("assignments") or []) + if isinstance(item, dict) and item.get("worker_id") == worker_id + ), None) + if assignment is None: + raise HTTPException(status_code=400, detail="worker is not assigned") + worker_root = workspace_dir / "workers" / worker_id + status = _load(worker_root / "status.json", {}) or {} + if status.get("state") not in {"failed", "timed_out"}: + raise HTTPException( + status_code=409, + detail="only a failed or timed-out worker can be restarted", + ) + pid_path = worker_root / "restart.pid.json" + active_pid = _running_pid(pid_path) + if active_pid is not None: + raise HTTPException( + status_code=409, + detail=f"worker restart is already running as pid {active_pid}", + ) + requirements = state_dir / "requirements.json" + if not requirements.is_file(): + raise HTTPException(status_code=400, detail="requirements.json is missing") + log_path = state_dir / f"{worker_id}_restart.log" + command, integration_command = _restart_worker_commands( + requirements, state_dir, workspace_dir, worker_id + ) + try: + with log_path.open("ab") as log_file: + process = subprocess.Popen( + command, + cwd=Path(__file__).resolve().parents[4], + stdin=subprocess.DEVNULL, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + close_fds=True, + ) + integration_process = subprocess.Popen( + integration_command, + cwd=Path(__file__).resolve().parents[4], + stdin=subprocess.DEVNULL, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + close_fds=True, + ) + except OSError as exc: + raise HTTPException( + status_code=500, detail=f"failed to start worker: {exc}" + ) from exc + pid_path.write_text(json.dumps({ + "pid": process.pid, + "task_id": task_id, + "worker_id": worker_id, + "physical_gpu": assignment.get("gpu"), + "started_at": time.time(), + "integration_pid": integration_process.pid, + }, indent=2), encoding="utf-8") + return { + "ok": True, + "worker_id": worker_id, + "physical_gpu": assignment.get("gpu"), + "pid": process.pid, + "integration_pid": integration_process.pid, + "isolated": True, + } + + @router.post("/rename-repo") + async def rename_repo(task_id: str, request: Request) -> Dict[str, Any]: + """Rename the task's kernel repository and repair all references. + + Body: ``{"new_name": ""}``. Refuses + while the task's orchestrator is still running. + """ + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + state_dir = state_dir_for(entry) + workspace_dir = workspace_dir_for(entry) + try: + body = await request.json() + except (ValueError, AttributeError) as exc: + raise HTTPException( + status_code=400, detail="invalid JSON body" + ) from exc + new_name = ( + body.get("new_name") if isinstance(body, dict) else None + ) or "" + try: + return rename_kernel_repo( + workspace_dir, new_name, state_dir=state_dir + ) + except (ValueError, RuntimeError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @router.get("/state-graph") + def state_graph(task_id: str) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + state_dir = state_dir_for(entry) + run = read_run(state_dir) + requirements = read_requirements(state_dir) or {} + answers = requirements.get("answers") + source = answers if isinstance(answers, dict) else requirements + return phases.graph_payload( + run.get("current_phase", phases.PREPARE), + run.get("last_outcome"), + run.get("last_transition_label"), + include_baseline=( + source.get("execution_mode") != GEN_AND_OPT_MODE + ), + ) + + @router.get("/skills") + def skills(task_id: str) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + workspace_dir = workspace_dir_for(entry) + library = list_skill_library(workspace_dir) + plan = _load(workspace_dir / "plan.json", {}) or {} + if plan.get("execution_mode") in {LEGACY_SMOKE_MODE, SMOKE_MODE}: + quarantined = library.get("pending") or [] + library["pending"] = [] + library["quarantined_count"] = len(quarantined) + library["publish_disabled_reason"] = ( + "Infrastructure-smoke findings are unrelated to the selected " + "operator and cannot be published as optimization skills." + ) + return library + + @router.post("/skills/{skill_name}/publish") + def publish(task_id: str, skill_name: str) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + workspace_dir = workspace_dir_for(entry) + plan = _load(workspace_dir / "plan.json", {}) or {} + if plan.get("execution_mode") in {LEGACY_SMOKE_MODE, SMOKE_MODE}: + raise HTTPException( + status_code=400, + detail=( + "infrastructure-smoke skills are quarantined and cannot " + "be published" + ), + ) + try: + return { + "skill": publish_skill(workspace_dir, skill_name) + } + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + except FileNotFoundError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + except FileExistsError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + + @router.post("/skills/sync") + def sync_skills(task_id: str) -> Dict[str, Any]: + """Manually mirror the authoritative dsh library into the ccb library.""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + workspace_dir = workspace_dir_for(entry) + return {"sync": sync_skill_libraries(workspace_dir=workspace_dir)} + + @router.post("/skills/fuse") + async def fuse_skill_route(task_id: str, request: Request) -> Dict[str, Any]: + """Trigger the main-agent fusion of one pending skill into the dsh + library (mirrored to ccb afterwards). Runs in the background; poll + ``GET /skills`` ``fuse_status`` for progress.""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + workspace_dir = workspace_dir_for(entry) + state_dir = state_dir_for(entry) + plan = _load(workspace_dir / "plan.json", {}) or {} + if plan.get("execution_mode") in {LEGACY_SMOKE_MODE, SMOKE_MODE}: + raise HTTPException( + status_code=400, + detail=( + "infrastructure-smoke skills are quarantined and cannot " + "be fused" + ), + ) + try: + body = await request.json() + except Exception: # noqa: BLE001 - malformed body + body = {} + skill_name = str(body.get("skill_name") or "").strip() + if not skill_name: + raise HTTPException(status_code=400, detail="missing skill_name") + pending = workspace_dir / "skills" / "pending" / skill_name + if not (pending / "SKILL.md").is_file(): + raise HTTPException( + status_code=404, + detail=f"pending skill not found: {skill_name}", + ) + req = read_requirements(state_dir) + if req is None: + raise HTTPException( + status_code=400, + detail="no requirements.json to read agent framework from", + ) + config = load_config(req) + mark_fuse_running(workspace_dir, skill_name) + + def _run_fusion() -> None: + try: + fuse_skill( + config=config, + workspace_dir=workspace_dir, + state_dir=state_dir, + skill_name=skill_name, + ) + except Exception: # noqa: BLE001 - status file carries the error + pass + + threading.Thread(target=_run_fusion, daemon=True).start() + return { + "ok": True, + "action": "fuse", + "skill_name": skill_name, + "status": "running", + } + + @router.post("/skills/{skill_name}/rollback") + def rollback(task_id: str, skill_name: str) -> Dict[str, Any]: + """Restore the latest SKILL.md backup of a fused skill in the dsh + library and re-mirror to ccb.""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + workspace_dir = workspace_dir_for(entry) + try: + return {"result": rollback_skill(skill_name, workspace_dir=workspace_dir)} + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + except FileNotFoundError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + + @router.get("/variants") + def variants_index(task_id: str) -> Dict[str, Any]: + """Return the fine-grained variant index (which shapes are already + captured in the shared variant library).""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return {"variants": list_variant_index()} + + @router.post("/variants") + async def add_variant_route(task_id: str, request: Request) -> Dict[str, Any]: + """Add one optimized shape's accepted kernel into the shared variant + library, organized by operator type / model / TP / specific operator. + Replaces an existing section for the same shape (with a file backup).""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + state_dir = state_dir_for(entry) + workspace_dir = workspace_dir_for(entry) + try: + body = await request.json() + except Exception: # noqa: BLE001 - malformed body + body = {} + shape_id = str(body.get("shape_id") or "").strip() + if not shape_id: + raise HTTPException(status_code=400, detail="missing shape_id") + req = read_requirements(state_dir) + if req is None: + raise HTTPException( + status_code=400, detail="no requirements.json for this task" + ) + answers = req.get("answers") if isinstance(req.get("answers"), dict) else req + + # Locate the accepted kernel for this shape under any worker lane. + accepted = None + for worker_root in sorted((workspace_dir / "workers").glob("worker_*")): + candidate = worker_root / "accepted" / shape_id / "kernel.hip" + if candidate.is_file(): + accepted = candidate + break + if accepted is None: + raise HTTPException( + status_code=404, + detail=f"no accepted kernel for shape {shape_id}", + ) + manifest_path = accepted.parent / "manifest.json" + manifest = _load(manifest_path, {}) + metrics = dict(manifest.get("metrics") or {}) + # baseline from the task's fixed user-supplied table for the speedup. + initial = _load(workspace_dir / "final_report.json", {}).get("initial_metrics") or {} + baseline = None + b = initial.get(shape_id) + if isinstance(b, dict): + baseline = b.get("median_us") + elif b is not None: + baseline = b + if baseline is not None and metrics.get("median_us"): + metrics["baseline_us"] = baseline + metrics["speedup"] = float(baseline) / float(metrics["median_us"]) + + meta = derive_variant_meta(answers, shape_id) + kernel_source = accepted.read_text(encoding="utf-8", errors="replace") + commit = str(manifest.get("commit") or "") + try: + result = add_variant( + meta=meta, + kernel_source=kernel_source, + commit=commit, + metrics=metrics, + source_task=str(req.get("task_id") or task_id), + backup=True, + # Hard guard: a strictly slower candidate may not replace the + # existing variant for the same shape (equal/faster may). + reject_slower_than_existing=True, + ) + except ValueError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + return {"ok": True, **result} + + return router diff --git a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-agent-fields.js b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-agent-fields.js new file mode 100644 index 00000000..00929b32 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-agent-fields.js @@ -0,0 +1,61 @@ +// Agent framework → model selection for dcu-kernel-auto-opt. +// +// Self-registers one override component ("agent-model") via +// globalThis.__metainferOverrides. The agent_framework field is a plain +// select; this component reads its current value from allValues and shows +// only the models valid for that framework (ccb: Sonnet/Opus; dsh: +// deepseek-v4-flash), auto-correcting the value when the framework changes. +// +// The label→model mapping must stay in sync with +// orchestrator/config.py (AGENT_FRAMEWORKS). + +import { html } from "htm/preact"; +import { useEffect } from "preact/hooks"; + +// Framework label → allowed model labels (mirror of config.py). +const FRAMEWORK_MODELS = { + ccb: ["Opus", "Sonnet"], + dsh: ["deepseek-v4-flash"], +}; +const DEFAULT_MODEL = { + ccb: "Opus", + dsh: "deepseek-v4-flash", +}; + +function AgentModelField({ field, value, onChange, allValues }) { + const framework = String( + (allValues && allValues.agent_framework) || "ccb" + ).toLowerCase(); + const models = FRAMEWORK_MODELS[framework] || FRAMEWORK_MODELS.ccb; + + // Keep the submitted model valid for the active framework: whenever the + // framework changes (or the form first loads) and the current value is not + // an option of that framework, snap it to the framework default. + useEffect(() => { + if (!models.includes(value)) { + onChange(DEFAULT_MODEL[framework] || models[0]); + } + // onChange is stable (setField wrapper); only framework changes matter. + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [framework]); + + return html` + + `; +} + +// ---- register --------------------------------------------------------------- + +var _g = (typeof globalThis !== "undefined" ? globalThis : window); +var _bridge = (_g.__metainferOverrides = _g.__metainferOverrides || {}); +_bridge["agent-model"] = AgentModelField; + +// Named export for the form-renderer to pick up. +export var AgentModelFieldComponent = AgentModelField; diff --git a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-detail.js b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-detail.js new file mode 100644 index 00000000..d2f29238 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-detail.js @@ -0,0 +1,937 @@ +import { html } from "htm/preact"; +import { useCallback, useEffect, useState } from "preact/hooks"; + +async function fetchJson(url) { + const response = await fetch(url, { cache: "no-store" }); + if (!response.ok) throw new Error(`${url} returned HTTP ${response.status}`); + return response.json(); +} + +async function postJson(url, body) { + const response = await fetch(url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }); + const payload = await response.json().catch(() => ({})); + if (!response.ok) { + throw new Error(payload.detail || `${url} returned HTTP ${response.status}`); + } + return payload; +} + +function StateMachine({ graph }) { + const nodes = graph?.nodes || []; + if (!nodes.length) return html`

Waiting for the control plane.

`; + return html` +
+ ${nodes.map((node, index) => html` +
+ ${index + 1} + ${node.label} +
+ ${index < nodes.length - 1 ? html`→` : null} + `)} +
+ `; +} + +function metric(value, digits = 3) { + const number = Number(value); + return Number.isFinite(number) ? number.toFixed(digits) : "—"; +} + +function statusClass(status) { + if (status === "passed" || status === "completed") return "success"; + if (status?.startsWith("bootstrap_")) { + return ["bootstrap_failed", "bootstrap_orphaned"].includes(status) + ? "failed" + : status === "bootstrap_passed" ? "success" : "running"; + } + if ( + status === "running" + || status === "validating" + || status === "retrying" + || status === "agent_running" + || status === "building" + || status === "baseline" + || status === "validating_candidate" + || status === "recording_result" + || status === "optimization_complete" + || status === "skill_writing" + ) { + return "running"; + } + if ( + status === "failed" + || status === "orphaned" + || status === "timed_out" + || status === "skipped" + ) return "failed"; + return "idle"; +} + +function bestRoundPerShape(lanes) { + const best = new Map(); + for (const lane of lanes || []) { + for (const experiment of lane.experiments || []) { + if (experiment.correctness_passed === false) continue; + const median = Number(experiment.metrics?.median_us); + if (!Number.isFinite(median) || !experiment.shape_id) continue; + const accepted = experiment.accepted === true; + let current = best.get(experiment.shape_id); + if (!current) { + current = { accepted: false, iteration: 0, median_us: Infinity }; + } + // Prefer accepted ("best") rounds; among those, the lowest median. + // Before the first acceptance, fall back to the best tested round. + const preferred = accepted && !current.accepted; + const better = accepted === current.accepted && median < current.median_us; + const laterWinner = accepted === current.accepted + && median === current.median_us + && Number(experiment.iteration) > current.iteration; + if (preferred || better || laterWinner) { + best.set(experiment.shape_id, { + accepted, + worker_id: lane.worker_id, + iteration: experiment.iteration, + median_us: median, + metrics: experiment.metrics || {}, + speedup: experiment.speedup, + baseline_us: experiment.baseline_us, + }); + } + } + } + return best; +} + +function ExploreGpuGrid({ lanes }) { + return html` +
+
+ Parallel explore · live GPU lanes + Agent → compile → correctness → benchmark → record → skill +
+
+ ${lanes.map((lane) => { + const completed = Number(lane.completed_rounds || 0); + const target = Number(lane.target_rounds || 0); + const progress = target + ? Math.min(100, Math.round(completed / target * 100)) + : 0; + const lastOutput = lane.agent?.last_output_age_s; + return html` +
+
+
+ GPU ${lane.gpu} + ${lane.worker_id} +
+ ${lane.state} +
+

${lane.step}

+
+
Shape
+
${lane.current_shape || lane.assigned_shapes?.join(", ") || "—"}
+
Round
+
${completed} / ${target || "—"}
+
Agent output
+
${Number.isFinite(Number(lastOutput)) + ? `${metric(lastOutput, 0)} s ago` + : "—"}
+
+
+ +
+ ${lane.long_running ? html` + + No Agent output for over 3 minutes; the lane timeout guard is active. + + ` : null} + ${lane.agent?.error ? html` + ${lane.agent.error} + ` : null} +
+ `; + })} +
+
+ `; +} + +function BootstrapResultCard({ attempt }) { + const files = attempt.generated_files || []; + const shapeEntries = Object.entries(attempt.metrics || {}); + const shapeMetrics = shapeEntries.map(([, item]) => item); + const medians = shapeMetrics + .map((item) => Number(item?.median_us)) + .filter((item) => Number.isFinite(item)); + const medianRange = medians.length + ? `${Math.min(...medians).toFixed(3)}–${Math.max(...medians).toFixed(3)} µs` + : "—"; + return html` +
+
+ Bootstrap ${attempt.attempt} + ${attempt.status} +
+
+
Elapsed
${metric(attempt.elapsed_s, 1)} s
+
Last output
${metric(attempt.last_output_age_s, 1)} s ago
+
Files
${files.length}
+
Verified
${shapeMetrics.length} shapes
+
Median
${medianRange}
+
Code snapshot
+
${attempt.artifact_dir || "—"}
+
+ ${shapeEntries.map(([shapeId, metrics]) => html` +
+ ${shapeId} +
+
Correct
+
+ + ${metrics?.passed ? "PASS" : "FAIL"} + +
+
Median
${metric(metrics?.median_us)} µs
+
P90
${metric(metrics?.p90_us)} µs
+
INT8 TOPS
+
${metric(metrics?.logical_tops ?? metrics?.tflops, 3)}
+
Algorithmic BW
+
${metric(metrics?.algorithmic_bandwidth_gb_s ?? metrics?.bandwidth_gb_s, 1)} GB/s
+
+
+ `)} + ${attempt.error ? html`${attempt.error}` : null} +
+ `; +} + +function IterationCard({ experiment }) { + const metrics = experiment.metrics || {}; + return html` +
+
+ Round ${experiment.iteration} + + ${experiment.accepted ? "best" : "tested"} + +
+
+
Shape
${experiment.shape_id || "—"}
+
Correct
+
+ + ${experiment.correctness_passed ? "PASS" : "FAIL"} + +
+
P90
${metric(metrics.p90_us)} µs
+
INT8 TOPS
+
${metric(metrics.logical_tops ?? metrics.tflops, 3)}
+
Algorithmic BW
+
${metric(metrics.algorithmic_bandwidth_gb_s ?? metrics.bandwidth_gb_s, 1)} GB/s
+
Counter HBM BW
+
${metric(experiment.pmc_evidence?.memory_traffic?.counter_derived_hbm_bandwidth_gb_s, 1)} GB/s
+
Baseline
${metric(experiment.baseline_us)} µs
+
Speedup
${metric(experiment.speedup, 3)}×
+
Test Time
${metric(metrics.median_us)} µs
+
Code snapshot
+
${experiment.artifact_dir || "—"}
+
+ ${experiment.failure_reason + ? html`${experiment.failure_reason}` : null} +
+ `; +} + +function RunningIterationCard({ active }) { + const activity = active.state === "profiling_current_best" + ? "profiling" + : active.state === "repairing_candidate" + ? "repairing" + : active.state === "validating_candidate" + ? "validating" + : "optimizing"; + return html` +
+
+ Round ${active.iteration} + ${activity} +
+
+
Shape
${active.shape_id || "—"}
+
Step
${active.step || "Agent optimizing kernel"}
+ ${active.state === "repairing_candidate" + ? html`
Repair
${active.repair || 0}/${active.max_repairs || 4}
` + : null} +
Agent
${active.agent_status || "running"}
+
Elapsed
${metric(active.elapsed_s, 1)} s
+
Output
+
${Number.isFinite(Number(active.last_output_age_s)) + ? `${metric(active.last_output_age_s, 1)} s ago` + : "waiting"}
+
+
+
+ `; +} + +function PlanCard({ experiment }) { + return html` +
+
+ Round ${experiment.iteration} + ${experiment.manual_guidance + ? html`manual` + : html`agent`} +
+

${experiment.hypothesis || "No optimization plan recorded."}

+ ${(experiment.changes || []).length + ? html`${experiment.changes.join(" · ")}` : null} +
+ `; +} + +function BootstrapPlanCard({ attempt }) { + const files = attempt.generated_files || []; + return html` +
+
+ Bootstrap ${attempt.attempt} + ${attempt.status} +
+

${attempt.hypothesis}

+ ${files.length ? html`${files.join(" · ")}` : null} + ${attempt.error ? html`${attempt.error}` : null} +
+ `; +} + +function GuidanceCard({ item }) { + return html` +
+
+ ${item.status === "pending" ? "Next round" : `Round ${item.consumed_iteration}`} + + ${item.status} + +
+

${item.text}

+ Manual optimization guidance +
+ `; +} + +function WorkerLane({ lane, taskId, onSaved }) { + const [text, setText] = useState(""); + const [saving, setSaving] = useState(false); + const [restarting, setRestarting] = useState(false); + const [message, setMessage] = useState(""); + const submit = async (event) => { + event.preventDefault(); + if (!text.trim() || saving || !lane.assigned) return; + setSaving(true); + setMessage(""); + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + await postJson( + `${base}/workers/${encodeURIComponent(lane.worker_id)}/guidance`, + { text: text.trim() }, + ); + setText(""); + setMessage("Queued for this worker's next round."); + await onSaved(); + } catch (err) { + setMessage(String(err)); + } finally { + setSaving(false); + } + }; + const restart = async () => { + if (restarting || !["failed", "timed_out"].includes(lane.state)) return; + if (!window.confirm( + `Restart ${lane.worker_id} on GPU ${lane.gpu}? Other GPU workers will keep running.`, + )) return; + setRestarting(true); + setMessage(""); + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + await postJson( + `${base}/workers/${encodeURIComponent(lane.worker_id)}/restart`, {}, + ); + setMessage(`Restarted on GPU ${lane.gpu}; sibling workers were not interrupted.`); + await onSaved(); + } catch (err) { + setMessage(String(err)); + } finally { + setRestarting(false); + } + }; + const pending = (lane.guidance || []).filter((item) => item.status === "pending"); + const bootstrap = lane.bootstrap_attempts || []; + const rounds = lane.experiments || []; + const currentIteration = Number(lane.current_iteration || 0); + const currentShape = lane.current_shape; + const activeStates = new Set([ + "profiling_current_best", + "agent_running", + "validating_candidate", + "repairing_candidate", + "recording_result", + ]); + const currentRecorded = rounds.some((item) => + Number(item.iteration || 0) === currentIteration + && item.shape_id === currentShape + ); + const active = lane.active_iteration || ( + currentIteration > 0 + && activeStates.has(lane.state) + && !currentRecorded + ? { + iteration: currentIteration, + shape_id: currentShape, + state: lane.state, + step: lane.step, + agent_name: lane.agent?.name, + agent_status: lane.agent?.status, + elapsed_s: lane.agent?.elapsed_s, + last_output_age_s: lane.agent?.last_output_age_s, + repair: lane.repair, + max_repairs: lane.max_repairs, + } + : null + ); + return html` +
+
+ ${lane.worker_id} + GPU ${lane.gpu} + + ${lane.state} + + ${lane.assigned_shapes?.join(", ") || "No shapes assigned"} + ${["failed", "timed_out"].includes(lane.state) ? html` + + ` : null} +
+
+
+ Iteration results + + ${bootstrap.length} bootstrap · ${rounds.length} completed + ${active ? " · 1 optimizing" : ""} + +
+
+ ${bootstrap.map((item) => html`<${BootstrapResultCard} attempt=${item} />`)} + ${rounds.map((item) => html`<${IterationCard} experiment=${item} />`)} + ${active ? html`<${RunningIterationCard} active=${active} />` : null} + ${!bootstrap.length && !rounds.length && !active + ? html`

${lane.assigned ? "Waiting for iteration 1." : "Worker not assigned."}

` + : null} +
+
+
+
+ Optimization plans + ${bootstrap.length + rounds.length} recorded · ${pending.length} pending +
+
+ ${bootstrap.map((item) => html`<${BootstrapPlanCard} attempt=${item} />`)} + ${rounds.map((item) => html`<${PlanCard} experiment=${item} />`)} + ${(lane.guidance || []) + .filter((item) => item.status === "pending") + .map((item) => html`<${GuidanceCard} item=${item} />`)} + ${!bootstrap.length && !rounds.length && !(lane.guidance || []).length + ? html`

No plan recorded.

` : null} +
+
+ setText(event.currentTarget.value)} + placeholder="Guide this GPU worker in its next round…" + maxlength="4000" + disabled=${!lane.assigned || saving} + /> + +
+ ${message ? html`${message}` : null} +
+
+ `; +} + +function SkillFile({ skill, canPublish, canFuse, onPublish, onFuse, busy }) { + return html` +
+ + + ${skill.name} + ${skill.kind} · ${skill.source} + + ${canFuse ? html` + + ` : null} + ${canPublish ? html` + + ` : html`existing`} + + 
${skill.content}
+
+ `; +} + +function SkillFuseStatus({ status, onRollback }) { + if (!status || status.status === "idle") return null; + if (status.status === "running") { + return html`
融合中…（主 agent 正在扫描 skill 库并决策）
`; + } + if (status.status === "error") { + return html`
融合失败: ${status.error || "unknown error"}
`; + } + const synced = status.synced_to_ccb; + return html` +
+ 融合完成：${status.action === "merge" ? `已融入 "${status.name}"` : `新增 skill "${status.name}"`} + ${synced && (synced.added?.length || synced.updated?.length) + ? ` · 已同步到 ccb 库${synced.added?.length ? `（新增 ${synced.added.length}）` : ""}${synced.updated?.length ? `（更新 ${synced.updated.length}）` : ""}` + : ""} + ${status.action === "merge" ? html` + + ` : null} + ${status.diff ? html` +
查看 diff
${status.diff}
+ ` : null} +
+ `; +} + +function SkillLibrary({ taskId, onClose }) { + const [library, setLibrary] = useState(null); + const [error, setError] = useState(""); + const [busy, setBusy] = useState(""); + const load = useCallback(async () => { + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + setLibrary(await fetchJson(`${base}/skills`)); + setError(""); + } catch (err) { + setError(String(err)); + } + }, [taskId]); + useEffect(() => { load(); }, [load]); + // While a fusion job is running, poll the library so the status flips to + // done/error without a manual refresh. + useEffect(() => { + if (!library || library.fuse_status?.status !== "running") return; + const id = setInterval(load, 4000); + return () => clearInterval(id); + }, [library, load]); + const publish = async (skill) => { + if (!window.confirm(`Add "${skill.name}" to the existing skill library (dsh, mirrored to Claude)?`)) return; + setBusy(`publish-${skill.name}`); + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + await postJson( + `${base}/skills/${encodeURIComponent(skill.name)}/publish`, {}, + ); + await load(); + } catch (err) { + setError(String(err)); + } finally { + setBusy(""); + } + }; + const fuse = async (skill) => { + if (!window.confirm( + `Have the main agent fuse "${skill.name}" into the existing skill library?` + )) return; + setBusy(`fuse-${skill.name}`); + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + await postJson(`${base}/skills/fuse`, { skill_name: skill.name }); + await load(); + } catch (err) { + setError(String(err)); + } finally { + setBusy(""); + } + }; + const syncNow = async () => { + setBusy("sync"); + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + await postJson(`${base}/skills/sync`, {}); + await load(); + } catch (err) { + setError(String(err)); + } finally { + setBusy(""); + } + }; + const rollback = async (name) => { + if (!window.confirm(`Restore the last backup of skill "${name}"?`)) return; + setBusy(`rollback-${name}`); + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + await postJson(`${base}/skills/${encodeURIComponent(name)}/rollback`, {}); + await load(); + } catch (err) { + setError(String(err)); + } finally { + setBusy(""); + } + }; + return html` +
+
+
+

Skill library

+

dsh 库为权威，融合/发布后自动同步到 Claude (ccb) 库。

+
+ + +
+ ${error ? html`
${error}
` : null} + ${library?.publish_disabled_reason ? html` +
+ ${library.publish_disabled_reason} + ${library.quarantined_count + ? ` (${library.quarantined_count} pending files quarantined)` : ""} +
+ ` : null} + <${SkillFuseStatus} status=${library?.fuse_status} onRollback=${rollback} /> + ${!library ? html`

Loading skills…

` : html` +
+
+

Existing skills ${library.existing?.length || 0}

+

${library.existing_root}

+
+ ${(library.existing || []).map((skill) => html` + <${SkillFile} skill=${skill} canPublish=${false} /> + `)} +
+
+
+

Pending skills ${library.pending?.length || 0}

+

Worker skills and the main-agent merged skill.

+
+ ${(library.pending || []).map((skill) => html` + <${SkillFile} + skill=${skill} + canPublish=${true} + canFuse=${true} + busy=${busy} + onPublish=${publish} + onFuse=${fuse} + /> + `)} + ${!(library.pending || []).length + ? html`

No pending skills.

` : null} +
+
+
+ `} +
+ `; +} + +export default function DcuKernelAutoOptDetail({ taskId, data }) { + const [runtime, setRuntime] = useState({ summary: null, graph: null, lanes: [] }); + const [error, setError] = useState(null); + const [showSkills, setShowSkills] = useState(false); + const refresh = useCallback(async () => { + if (!taskId) return; + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const [summary, graph, iterations] = await Promise.all([ + fetchJson(`${base}/summary`), + fetchJson(`${base}/state-graph`), + fetchJson(`${base}/iterations`), + ]); + setRuntime({ summary, graph, lanes: iterations.workers || [] }); + setError(null); + } catch (err) { + setError(String(err)); + } + }, [taskId]); + + useEffect(() => { refresh(); }, [refresh]); + useEffect(() => { + const timer = setInterval(refresh, 3000); + return () => clearInterval(timer); + }, [refresh]); + + // ---- variant library (加入 variant 按钮) ------------------------------ // + const [variantBusy, setVariantBusy] = useState(""); + const [variantMsg, setVariantMsg] = useState(""); + const [variantIndex, setVariantIndex] = useState(null); // null = loading + const loadVariants = useCallback(async () => { + if (!taskId) return; + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const data = await fetchJson(`${base}/variants`); + // Map shape -> variant record (incl. median_us/speedup/source) so the + // update flow can compare the existing variant against this task. + setVariantIndex(new Map((data.variants || []).map((v) => [v.shape, v]))); + } catch (err) { /* non-fatal */ } + }, [taskId]); + useEffect(() => { loadVariants(); }, [loadVariants]); + const addVariant = async (shape) => { + const existing = variantIndex ? variantIndex.get(shape) : null; + if (existing) { + const candidate = bestRoundPerShape(runtime.lanes).get(shape) || {}; + const newUs = candidate.median_us; + const oldUs = existing.median_us; + let verdict = ""; + if (newUs != null && oldUs != null) { + verdict = newUs < oldUs + ? "（本次更快）" + : newUs > oldUs ? "（本次更慢）" : "（持平）"; + } + const ok = window.confirm( + `该算子已有 variant：median ${oldUs ?? "?"} µs` + + `（来源 ${existing.source || "?"}）。\n` + + `本次候选：median ${newUs ?? "?"} µs ${verdict}。\n` + + `用本次候选替换？旧文件会备份（.bak-*）。` + ); + if (!ok) return; + } + setVariantBusy(shape); + setVariantMsg(""); + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const res = await postJson(`${base}/variants`, { shape_id: shape }); + const meta = res.meta || {}; + const tp = meta.tp != null ? `tp${meta.tp}` : "tp?"; + setVariantMsg( + `已加入 variant: ${meta.type} | ${meta.model} | ${tp} | ${meta.operator}` + + `（${res.action === "updated" ? "更新已有分区" : "新增分区"}）` + ); + await loadVariants(); + } catch (err) { + setVariantMsg(`加入 variant 失败: ${String(err)}`); + } finally { + setVariantBusy(""); + } + }; + + const report = runtime.summary?.report; + const run = runtime.summary?.run; + const stopped = run?.final_status === "stopped"; + const stopReason = (run?.notes || []).join(" · ") || "The optimizer stopped before completion."; + const smokeOnly = report?.mode === "real-agent-dcu-smoke" + || ["Real agents + DCU (smoke harness)", "Infrastructure smoke (not operator optimization)"] + .includes(runtime.summary?.plan?.execution_mode); + const repoPath = runtime.summary?.plan?.kernel_repo; + const repoName = repoPath ? String(repoPath).split("/").pop() : null; + // ---- kernel repo rename ------------------------------------------- // + const [renameBusy, setRenameBusy] = useState(false); + const [renameMsg, setRenameMsg] = useState(""); + const renameRepo = async () => { + if (!repoName) return; + const target = window.prompt( + "新的仓库名（kernel-repos/ 下的目录名）:", repoName + ); + if (!target || target.trim() === "" || target.trim() === repoName) return; + setRenameBusy(true); + setRenameMsg(""); + try { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const res = await postJson(`${base}/rename-repo`, { + new_name: target.trim(), + }); + setRenameMsg(`已重命名: ${res.old_name} → ${res.new_name}`); + await refresh(); + } catch (err) { + setRenameMsg(`重命名失败: ${String(err)}`); + } finally { + setRenameBusy(false); + } + }; + const bestRounds = bestRoundPerShape(runtime.lanes); + const shapeOrder = new Map( + (runtime.summary?.plan?.shapes || []).map((item, index) => [item.id, index]), + ); + const bestRoundEntries = Array.from(bestRounds.entries()) + .sort(([a], [b]) => (shapeOrder.get(a) ?? 1e9) - (shapeOrder.get(b) ?? 1e9)); + return html` +
+ ${error ? html`
${error}
` : null} + ${smokeOnly ? html` +
+ Infrastructure smoke only: this task did not execute the selected + operator. Its metrics and generated skills are not valid operator + optimization results. +
+ ` : null} + ${stopped ? html` +
+ Task stopped in ${run?.current_phase || "startup"}: + ${stopReason} +
+ ` : null} + + ${repoName ? html` +
+

+ Kernel repository + — ${repoName} +

+
+ + ${renameMsg ? html`${renameMsg}` : null} +
+
+ ` : null} + +
+

+ State machine + — currently: + + ${(() => { + const nodes = runtime.graph?.nodes || []; + const cur = runtime.graph?.current; + const node = nodes.find(n => n.id === cur); + return node ? node.label : (cur || "starting"); + })()} + +

+ <${StateMachine} graph=${runtime.graph} /> + ${runtime.graph?.current === "parallel_explore" ? html` + <${ExploreGpuGrid} lanes=${runtime.lanes} /> + ` : null} +
+ +
+

Iterations four isolated worker lanes

+
+ ${runtime.lanes.map((lane) => html` + <${WorkerLane} lane=${lane} taskId=${taskId} onSaved=${refresh} /> + `)} +
+
+ +
+

Final serial validation

+ ${bestRoundEntries.length ? html` +

+ Best optimization round per shape + accepted “best” round from each worker lane +

+ + + + + + ${bestRoundEntries.map(([shape, item]) => html` + + + + + + + + + + + `)} + +
ShapeWorkerBest roundPerformanceSpeedupINT8 TOPSAlgorithmic BWVariant
${shape}${item.worker_id}${item.iteration ?? "—"}${metric(item.median_us)} µs${metric(item.speedup, 3)}×${metric(item.metrics.logical_tops ?? item.metrics.tflops, 3)}${metric(item.metrics.algorithmic_bandwidth_gb_s ?? item.metrics.bandwidth_gb_s, 2)} GB/s + +
+ ${variantMsg ? html`

${variantMsg}

` : null} + ` : null} + ${report ? html` + + + + + + ${Object.entries(report.final_validation || {}).map(([shape, item]) => html` + + + + + + + + + `)} + +
ShapeWorkerPerformanceINT8 TOPSAlgorithmic BWResult
${shape}${item.worker_id || bestRounds.get(shape)?.worker_id || "—"}${metric(item.metrics?.median_us ?? item.median_us)} µs${metric(item.metrics?.logical_tops ?? item.metrics?.tflops ?? item.logical_tops, 3)}${metric( + item.metrics?.algorithmic_bandwidth_gb_s + ?? item.metrics?.bandwidth_gb_s + ?? item.algorithmic_bandwidth_gb_s, + 2, + )} GB/s${item.passed ? "PASS" : "FAIL"}
+ ` : html`

Final validation has not started.

`} +
+ +
+

Event timeline

+
+ ${(data?.timeline?.events || []).map((event) => html` +
${event.type}${JSON.stringify(event.payload || {})}
+ `)} +
+
+ +
+
+

Optimization skill library

+

Inspect worker findings and the main-agent synthesis, then publish manually.

+
+ +
+ + ${showSkills ? html` + <${SkillLibrary} taskId=${taskId} onClose=${() => setShowSkills(false)} /> + ` : null} +
+ `; +} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-shape-input.js b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-shape-input.js new file mode 100644 index 00000000..c40b4da7 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-shape-input.js @@ -0,0 +1,1217 @@ +// Operator-specific guided shape input for dcu-kernel-auto-opt. +// Self-registers via globalThis.__metainferOverrides — no imports from +// shared modules, no circular deps, no dynamic import needed. +// +// The form-renderer polls __metainferOverrides and picks us up once this +// file is loaded as a side-effect of the importmap entry. + +import { html } from "htm/preact"; +import { useCallback, useEffect, useMemo, useRef, useState } from "preact/hooks"; + +// ---- helpers --------------------------------------------------------------- + +function parseCsvInts(raw) { + if (!raw || !raw.trim()) return []; + return raw + .split(",") + .map(function (s) { var n = parseInt(s.trim(), 10); return (Number.isFinite(n) && n > 0) ? n : null; }) + .filter(function (n) { return n !== null; }); +} + +// Frontend mirror of the W8A8 GEMM workload catalog per model. The backend +// validates every submitted shape against the frozen operator contract, so a +// stale UI list fails closed instead of starting an invalid optimization. +var W8A8_M_VALUES = [2, 16, 3072]; +// Large-prefill boundary added on 2026-08-06 for TP4; extended to the +// Hy3 / MiniMax M3 / GLM5.2 TP8 catalogs on 2026-08-27 via per-topology +// mValues. DeepSeek TP1/TP8 keep the three original M values. The backend +// validates shapes against the contract. +var W8A8_TP4_M_VALUES = [2, 16, 3072, 4096]; + +// Model label → W8A8 GEMM workload. (M, K) @ (K, N) per weight name and TP +// size, mirroring the model cards in the operator docs. DeepSeek keeps the +// historical ids (tp4_wqkv_a_m4096 …); other models get a model prefix so ids +// never collide across models. +var MODEL_WORKLOADS = { + "DeepSeek V4 Flash": { + idPrefix: "", + topologies: [ + { + tp_size: 1, + operators: [ + { operator: "wqkv_a", K: 4096, N: 1536 }, + { operator: "wq_b", K: 1024, N: 32768 }, + { operator: "indexer.wq_b", K: 1024, N: 8192 }, + { operator: "wo_b", K: 8192, N: 4096 }, + { operator: "shared_gate_up_proj", K: 4096, N: 4096 }, + { operator: "shared_down_proj", K: 2048, N: 4096 }, + ], + }, + { + tp_size: 4, + operators: [ + { operator: "wqkv_a", K: 4096, N: 1536 }, + { operator: "wq_b", K: 1024, N: 8192 }, + { operator: "indexer.wq_b", K: 1024, N: 8192 }, + { operator: "wo_b", K: 2048, N: 4096 }, + { operator: "shared_gate_up_proj", K: 4096, N: 1024 }, + { operator: "shared_down_proj", K: 512, N: 4096 }, + ], + }, + { + tp_size: 8, + operators: [ + { operator: "wqkv_a", K: 4096, N: 1536 }, + { operator: "wq_b", K: 1024, N: 4096 }, + { operator: "indexer.wq_b", K: 1024, N: 8192 }, + { operator: "wo_b", K: 1024, N: 4096 }, + { operator: "shared_gate_up_proj", K: 4096, N: 512 }, + { operator: "shared_down_proj", K: 256, N: 4096 }, + ], + }, + ], + }, + "Hy3 (Hunyuan 3)": { + idPrefix: "hy3_", + topologies: [ + { + tp_size: 1, + operators: [ + { operator: "qkv_proj", K: 4096, N: 10240 }, + { operator: "o_proj", K: 8192, N: 4096 }, + { operator: "shared_gate_up_proj", K: 4096, N: 3072 }, + { operator: "shared_down_proj", K: 1536, N: 4096 }, + ], + }, + { + tp_size: 4, + operators: [ + { operator: "qkv_proj", K: 4096, N: 2560 }, + { operator: "o_proj", K: 2048, N: 4096 }, + { operator: "shared_gate_up_proj", K: 4096, N: 768 }, + { operator: "shared_down_proj", K: 384, N: 4096 }, + ], + }, + { + tp_size: 8, + mValues: [2, 16, 3072, 4096], + operators: [ + { operator: "qkv_proj", K: 4096, N: 1280 }, + { operator: "o_proj", K: 1024, N: 4096 }, + { operator: "shared_gate_up_proj", K: 4096, N: 384 }, + { operator: "shared_down_proj", K: 192, N: 4096 }, + ], + }, + ], + }, + "MiniMax M3": { + idPrefix: "minimax_", + topologies: [ + { + tp_size: 1, + operators: [ + { operator: "qkv_proj", K: 6144, N: 9216 }, + { operator: "qkv_proj_and_indexer_qk", K: 6144, N: 9856 }, + { operator: "o_proj", K: 8192, N: 6144 }, + { operator: "shared_gate_up_proj", K: 6144, N: 6144 }, + { operator: "shared_down_proj", K: 3072, N: 6144 }, + ], + }, + { + tp_size: 4, + operators: [ + { operator: "qkv_proj", K: 6144, N: 2304 }, + { operator: "qkv_proj_and_indexer_qk", K: 6144, N: 2560 }, + { operator: "o_proj", K: 2048, N: 6144 }, + { operator: "shared_gate_up_proj", K: 6144, N: 1536 }, + { operator: "shared_down_proj", K: 768, N: 6144 }, + ], + }, + { + tp_size: 8, + mValues: [2, 16, 3072, 4096], + operators: [ + { operator: "qkv_proj", K: 6144, N: 1280 }, + { operator: "qkv_proj_and_indexer_qk", K: 6144, N: 1536 }, + { operator: "o_proj", K: 1024, N: 6144 }, + { operator: "shared_gate_up_proj", K: 6144, N: 768 }, + { operator: "shared_down_proj", K: 384, N: 6144 }, + ], + }, + ], + }, + "GLM5.2": { + idPrefix: "glm_", + topologies: [ + { + tp_size: 1, + operators: [ + { operator: "fused_qkv_a_proj", K: 6144, N: 2624 }, + { operator: "q_b_proj", K: 2048, N: 16384 }, + { operator: "kv_b_proj", K: 512, N: 28672 }, + { operator: "o_proj", K: 16384, N: 6144 }, + { operator: "shared_gate_up_proj", K: 6144, N: 4096 }, + { operator: "shared_down_proj", K: 2048, N: 6144 }, + ], + }, + { + tp_size: 4, + operators: [ + { operator: "fused_qkv_a_proj", K: 6144, N: 2624 }, + { operator: "q_b_proj", K: 2048, N: 4096 }, + { operator: "kv_b_proj", K: 512, N: 7168 }, + { operator: "o_proj", K: 4096, N: 6144 }, + { operator: "shared_gate_up_proj", K: 6144, N: 1024 }, + { operator: "shared_down_proj", K: 512, N: 6144 }, + ], + }, + { + tp_size: 8, + mValues: [2, 16, 3072, 4096], + operators: [ + { operator: "fused_qkv_a_proj", K: 6144, N: 2624 }, + { operator: "q_b_proj", K: 2048, N: 2048 }, + { operator: "kv_b_proj", K: 512, N: 3584 }, + { operator: "o_proj", K: 2048, N: 6144 }, + { operator: "shared_gate_up_proj", K: 6144, N: 512 }, + { operator: "shared_down_proj", K: 256, N: 6144 }, + ], + }, + ], + }, +}; + +function modelCatalog(modelLabel) { + var model = MODEL_WORKLOADS[modelLabel] || MODEL_WORKLOADS["DeepSeek V4 Flash"]; + var shapes = []; + (model.topologies || []).forEach(function (topology) { + // Per-topology M values override the defaults; TP4 and the Hy3/MiniMax/ + // GLM TP8 catalogs cover the M=4096 large-prefill boundary. + var mValues = topology.mValues || (topology.tp_size === 4 + ? W8A8_TP4_M_VALUES : W8A8_M_VALUES); + topology.operators.forEach(function (item) { + mValues.forEach(function (M) { + shapes.push({ + id: model.idPrefix + "tp" + topology.tp_size + "_" + + item.operator.replace(/\./g, "_") + "_m" + M, + tp_size: topology.tp_size, + operator: item.operator, + M: M, + N: item.N, + K: item.K, + }); + }); + }); + }); + return shapes; +} + +// Backward-compatible DeepSeek-only default catalog. +var DEFAULT_W8A8_SHAPES = modelCatalog("DeepSeek V4 Flash"); + +function currentShapeId(catalog, rawId, modelLabel) { + if (catalog.some(function (shape) { return shape.id === rawId; })) { + return rawId; + } + if (modelLabel === "DeepSeek V4 Flash") { + var legacy = /^m(\d+)_(wqkv_a|wq_b|wo_b|shared_gate_up|shared_down)$/.exec(rawId); + if (!legacy) return null; + var operatorAliases = { + shared_gate_up: "shared_gate_up_proj", + shared_down: "shared_down_proj", + }; + var operator = operatorAliases[legacy[2]] || legacy[2]; + var migrated = "tp4_" + operator + "_m" + legacy[1]; + return catalog.some(function (shape) { + return shape.id === migrated; + }) ? migrated : null; + } + return null; +} + +function selectedShapes(parsedShapes, fallbackToAll, catalog, modelLabel) { + var ids = parsedShapes + .map(function (shape) { return currentShapeId(catalog, shape.id, modelLabel); }) + .filter(Boolean); + if (!ids.length && fallbackToAll) { + ids = catalog.map(function (shape) { return shape.id; }); + } + return selectedMap(ids); +} + +function parseShapeRecords(raw) { + if (!raw) return []; + var shapes = [], re = /\{[^}]+\}/g, match; + while ((match = re.exec(raw)) !== null) { + var text = match[0]; + var idMatch = /\bid:\s*["']?([^,\s}"']+)/.exec(text); + var getInt = function (key) { + var dimMatch = new RegExp("\\b" + key + ":\\s*(\\d+)").exec(text); + return dimMatch ? parseInt(dimMatch[1], 10) : null; + }; + var operatorMatch = /\boperator:\s*["']?([^,\s}"']+)/.exec(text); + var M = getInt("M"), N = getInt("N"), K = getInt("K"); + var tpSize = getInt("tp_size"); + if (idMatch && M != null && N != null && K != null) { + shapes.push({ + id: idMatch[1], + tp_size: tpSize, + operator: operatorMatch ? operatorMatch[1] : null, + M: M, + N: N, + K: K, + }); + } + } + return shapes; +} + +function parseManualOwners(raw) { + var owners = {}; + if (!raw) return owners; + var re = /worker_(\d+):\s*\{[^}]*shapes:\s*\[([^\]]*)\]/g, match; + while ((match = re.exec(raw)) !== null) { + var gpu = parseInt(match[1], 10); + match[2].split(",").forEach(function (item) { + var id = item.trim().replace(/^["']|["']$/g, ""); + if (id) owners[id] = gpu; + }); + } + return owners; +} + +function operatorGroupedOwners(shapes) { + var familyGpu = { + wqkv_a: 0, + fused_qkv_a_proj: 0, + qkv_proj: 0, + wq_b: 1, + "indexer.wq_b": 1, + qkv_proj_and_indexer_qk: 1, + q_b_proj: 1, + kv_b_proj: 1, + wo_b: 2, + o_proj: 2, + shared_gate_up_proj: 3, + shared_down_proj: 3, + }; + var owners = {}; + shapes.forEach(function (shape, index) { + var family = shape.operator || shape.id; + owners[shape.id] = familyGpu[family] != null + ? familyGpu[family] + : index % 4; + }); + return owners; +} + +function mGroupedOwners(shapes) { + var mValues = shapes + .map(function (shape) { return shape.M; }) + .filter(function (value, index, all) { + return all.indexOf(value) === index; + }) + .sort(function (a, b) { return a - b; }); + var owners = {}; + shapes.forEach(function (shape) { + owners[shape.id] = mValues.indexOf(shape.M) % 4; + }); + return owners; +} + +function balancedOwners(shapes) { + var owners = {}; + shapes.forEach(function (shape, index) { + owners[shape.id] = index % 4; + }); + return owners; +} + +function serializeShapesHeader(mode, scope, shapes, modelLabel) { + var lines = []; + if (modelLabel) lines.push("model: " + modelLabel); + lines.push( + "assignment_mode: " + mode, + "shape_scope: " + scope, + "shapes:" + ); + shapes.forEach(function (shape) { + var logical = shape.tp_size != null && shape.operator + ? ", tp_size: " + shape.tp_size + + ", operator: \"" + shape.operator + "\"" + : ""; + lines.push( + " - {id: " + shape.id + ", M: " + shape.M + + ", N: " + shape.N + ", K: " + shape.K + logical + "}" + ); + }); + return lines; +} + +function serializeManualAssignment(shapes, owners, scope, modelLabel) { + var lines = serializeShapesHeader("manual", scope, shapes, modelLabel); + lines.push("assignments:"); + for (var gpu = 0; gpu < 4; gpu++) { + var ids = shapes + .filter(function (shape) { return owners[shape.id] === gpu; }) + .map(function (shape) { return shape.id; }); + if (ids.length) { + lines.push( + " worker_" + gpu + ": {gpu: " + gpu + + ", shapes: [" + ids.join(", ") + "]}" + ); + } + } + return lines.join("\n"); +} + +function selectedMap(ids) { + var selected = {}; + ids.forEach(function (id) { selected[id] = true; }); + return selected; +} + +function TpFilterBar(_a) { + var tpSizes = _a.tpSizes, active = _a.active, onChange = _a.onChange; + return html` +
+ TP + + ${tpSizes.map(function (tp) { return html` + + `; })} +
+ `; +} + +function ShapeCatalog(_a) { + var shapes = _a.shapes, selected = _a.selected, onToggle = _a.onToggle, + readOnly = _a.readOnly; + var tpSizes = shapes + .map(function (shape) { return shape.tp_size; }) + .filter(function (value, index, all) { return all.indexOf(value) === index; }) + .sort(function (a, b) { return a - b; }); + var _b = useState(null), tpFilter = _b[0], setTpFilter = _b[1]; + var visible = tpFilter == null + ? shapes + : shapes.filter(function (shape) { return shape.tp_size === tpFilter; }); + var selectedCount = shapes.filter(function (shape) { + return selected[shape.id]; + }).length; + var visibleSelected = visible.filter(function (shape) { + return selected[shape.id]; + }).length; + return html` +
+
+ Shapes in this task + + ${selectedCount} / ${shapes.length} selected + ${tpFilter != null ? ` · ${visibleSelected} shown` : ""} + +
+ <${TpFilterBar} tpSizes=${tpSizes} active=${tpFilter} onChange=${setTpFilter} /> +
+ ${visible.map(function (shape) { + var checked = Boolean(selected[shape.id]); + return html` + + `; + })} + ${visible.length === 0 ? html`

No shapes in this TP size.

` : null} +
+
+ `; +} + +function ApiDefaultsPreview(_a) { + var catalog = _a.catalog; + var tpSizes = catalog + .map(function (shape) { return shape.tp_size; }) + .filter(function (value, index, all) { return all.indexOf(value) === index; }) + .sort(function (a, b) { return a - b; }); + return html` +
+
+ ${catalog.length} logical shapes. + Workloads are grouped by TP size and kept as separate task identities. + Each operator is optimized at M=2, M=16 and the large-prefill boundary + M=3072; TP4 additionally covers M=4096. +
+
+ ${tpSizes.map(function (tp) { + var topoShapes = catalog.filter(function (shape) { + return shape.tp_size === tp; + }); + var operators = []; + topoShapes.forEach(function (shape) { + if (!operators.some(function (item) { + return item.operator === shape.operator; + })) { + operators.push({ + operator: shape.operator, + K: shape.K, + N: shape.N, + }); + } + }); + var mValues = tp === 4 ? W8A8_TP4_M_VALUES : W8A8_M_VALUES; + return html` +
+
+ TP=${tp} + ${topoShapes.length} shapes +
+
+ ${operators.map(function (item) { + return html` +
+
+ ${item.operator} + (M,${item.K}) @ (${item.K},${item.N}) · M=${mValues.join(",")} +
+
+ `; + })} +
+
+ `; + })} +
+
+ `; +} + +function AiShapeSubset(_a) { + var value = _a.value, onChange = _a.onChange, model = _a.model; + var catalog = modelCatalog(model); + var _b = useState(function () { + var parsed = parseShapeRecords(value || ""); + return selectedShapes(parsed, true, catalog, model); + }), selected = _b[0], setSelected = _b[1]; + + var emit = useCallback(function (next) { + var shapes = catalog.filter(function (shape) { + return next[shape.id]; + }); + setSelected(next); + onChange(serializeShapesHeader("ai", "subset", shapes, model).join("\n")); + }, [onChange, catalog, model]); + + useEffect(function () { emit(selected); }, []); + + return html` +
+
+ Subset task. + The control plane assigns only the checked shapes. Unchecked API + shapes stay on the trusted fallback and are checked again before the + candidate is published. +
+ <${ShapeCatalog} + shapes=${catalog} + selected=${selected} + onToggle=${function (shapeId) { + var next = Object.assign({}, selected); + next[shapeId] = !next[shapeId]; + emit(next); + }} /> +
+ `; +} + +function ModelWorkloadAll(_a) { + var value = _a.value, onChange = _a.onChange, model = _a.model; + var catalog = modelCatalog(model); + var selected = selectedMap(catalog.map(function (shape) { + return shape.id; + })); + var emitted = useRef(false); + useEffect(function () { + if (emitted.current) return; + emitted.current = true; + onChange(serializeShapesHeader("ai", "all", catalog, model).join("\n")); + }, [catalog, onChange, model]); + + return html` +
+
+ Full ${model} workload. + Every shape below is part of this task. The control plane balances + them across the four GPUs automatically; the checkboxes are read-only + because the scope is “All API shapes”. +
+ <${ShapeCatalog} + shapes=${catalog} + selected=${selected} + readOnly=${true} + onToggle=${function () {}} /> +
+ `; +} + +function ManualGpuAssignment(_a) { + var value = _a.value, onChange = _a.onChange, scopeMode = _a.scopeMode, + model = _a.model; + var subset = scopeMode === "Selected shapes only"; + var catalog = modelCatalog(model); + var _b = useState(function () { + var parsedShapes = parseShapeRecords(value || ""); + return subset + ? selectedShapes(parsedShapes, true, catalog, model) + : selectedMap(catalog.map(function (shape) { + return shape.id; + })); + }), selected = _b[0], setSelected = _b[1]; + var _c = useState(function () { + var parsed = parseManualOwners(value || ""); + var migrated = operatorGroupedOwners(catalog); + Object.keys(parsed).forEach(function (shapeId) { + var current = currentShapeId(catalog, shapeId, model); + if (current) migrated[current] = parsed[shapeId]; + }); + return migrated; + }), owners = _c[0], setOwners = _c[1]; + + var emitState = useCallback(function (nextOwners, nextSelected) { + var shapes = catalog.filter(function (shape) { + return !subset || nextSelected[shape.id]; + }); + setOwners(nextOwners); + setSelected(nextSelected); + onChange(serializeManualAssignment( + shapes, nextOwners, subset ? "subset" : "all", model + )); + }, [onChange, subset, catalog, model]); + + useEffect(function () { + emitState(owners, selected); + }, [subset]); + + var moveShape = function (shapeId, gpu) { + var next = Object.assign({}, owners); + next[shapeId] = gpu; + emitState(next, selected); + }; + var shapes = catalog.filter(function (shape) { + return !subset || selected[shape.id]; + }); + + return html` +
+
+ Manual assignment is authoritative. + The coordinator Agent will not repartition these shapes. Every shape + appears exactly once; empty GPU cards are allowed. +
+ ${subset && html` + <${ShapeCatalog} + shapes=${catalog} + selected=${selected} + onToggle=${function (shapeId) { + var nextSelected = Object.assign({}, selected); + nextSelected[shapeId] = !nextSelected[shapeId]; + var nextOwners = Object.assign({}, owners); + if (nextSelected[shapeId] && nextOwners[shapeId] == null) { + nextOwners[shapeId] = operatorGroupedOwners( + catalog + )[shapeId]; + } + emitState(nextOwners, nextSelected); + }} /> + `} +
+ Quick layout: + + + +
+
+ ${[0, 1, 2, 3].map(function (gpu) { + var assigned = shapes.filter(function (shape) { + return owners[shape.id] === gpu; + }); + return html` +
+
+ GPU ${gpu} + ${assigned.length} shapes +
+
+ ${assigned.length + ? assigned.map(function (shape) { + return html` +
+
+ ${shape.id} + TP=${shape.tp_size} · ${shape.operator} · M=${shape.M} · N=${shape.N} · K=${shape.K} +
+ +
+ `; + }) + : html`
No shapes assigned
` + } +
+
+ `; + })} +
+
+ `; +} + +// ---- operator-specific guided forms ---------------------------------------- + +function GemmGuided(_a) { + var value = _a.value, onChange = _a.onChange; + var prev = useMemo(function () { return parseShapeYaml(value || ""); }, [value]); + var _b = useState(function () { return (prev.mVals || ["2", "16", "3072", "4096"]).join(", "); }), + mRaw = _b[0], setMRaw = _b[1]; + var _c = useState(function () { return (prev.nVals || ["1536"]).join(", "); }), + nRaw = _c[0], setNRaw = _c[1]; + var _d = useState(function () { return (prev.kVals || ["4096"]).join(", "); }), + kRaw = _d[0], setKRaw = _d[1]; + + var emit = useCallback(function (m, n, k) { + var mVals = parseCsvInts(m), nVals = parseCsvInts(n), kVals = parseCsvInts(k); + if (!mVals.length || !nVals.length || !kVals.length) return; + var lines = ["shapes:"], dimN = nVals[0], dimK = kVals[0]; + for (var i = 0; i < mVals.length; i++) { + lines.push(" - {id: m" + mVals[i] + ", M: " + mVals[i] + ", N: " + dimN + ", K: " + dimK + "}"); + } + onChange(lines.join("\n")); + }, [onChange]); + + return html` +
+
+ + + 1 ≤ M ≤ 4096; M=3072/4096 covers large prefill +
+
+ + + Output feature dim +
+
+ + + Reduction dim +
+
+ `; +} + +function RmsNormGuided(_a) { + var value = _a.value, onChange = _a.onChange; + var prev = useMemo(function () { return parseShapeYaml(value || ""); }, [value]); + var _b = useState(function () { return (prev.tVals || ["1", "128"]).join(", "); }), + tRaw = _b[0], setTRaw = _b[1]; + var _c = useState(function () { return (prev.hVals || ["7168"]).join(", "); }), + hRaw = _c[0], setHRaw = _c[1]; + + var emit = useCallback(function (t, h) { + var tVals = parseCsvInts(t), hVals = parseCsvInts(h); + if (!tVals.length || !hVals.length) return; + var lines = ["shapes:"], dimH = hVals[0]; + for (var i = 0; i < tVals.length; i++) { + lines.push(" - {id: rmsnorm_t" + tVals[i] + ", T: " + tVals[i] + ", H: " + dimH + "}"); + } + onChange(lines.join("\n")); + }, [onChange]); + + return html` +
+
+ + + T = batch_size × seq_len +
+
+ + + The last dimension being normalized +
+
+ `; +} + +function PrefillAttnGuided(_a) { + var value = _a.value, onChange = _a.onChange; + var prev = useMemo(function () { return parseShapeYaml(value || ""); }, [value]); + var _b = useState(function () { return (prev.tokensVals || ["128", "1024", "4096"]).join(", "); }), + tokensRaw = _b[0], setTokensRaw = _b[1]; + var _c = useState(function () { return (prev.hdVals || ["128"]).join(", "); }), + hdRaw = _c[0], setHdRaw = _c[1]; + var _d = useState(true), kvLenSame = _d[0], setKvLenSame = _d[1]; + var _e = useState(""), kvLenRaw = _e[0], setKvLenRaw = _e[1]; + + var emit = useCallback(function (tokens, hd, same, kv) { + var tVals = parseCsvInts(tokens), hdVals = parseCsvInts(hd); + if (!tVals.length || !hdVals.length) return; + var lines = ["shapes:"], dimHd = hdVals[0]; + for (var i = 0; i < tVals.length; i++) { + var klen = same ? tVals[i] : (parseCsvInts(kv)[0] || tVals[i]); + lines.push(" - {id: prefill_t" + tVals[i] + ", M: " + tVals[i] + ", N: " + klen + ", K: " + dimHd + "}"); + } + onChange(lines.join("\n")); + }, [onChange]); + + return html` +
+
+ + + M = batch × seq_len +
+
+ + + Head dimension, typically 128 +
+
+ +
+ ${!kvLenSame && html` +
+ + + Cross-attention KV length +
+ `} +
+ `; +} + +function DecodeAttnGuided(_a) { + var value = _a.value, onChange = _a.onChange; + var prev = useMemo(function () { return parseShapeYaml(value || ""); }, [value]); + var _b = useState(function () { return (prev.kvVals || ["4096", "32768"]).join(", "); }), + kvRaw = _b[0], setKvRaw = _b[1]; + var _c = useState(function () { return (prev.hdVals || ["128"]).join(", "); }), + hdRaw = _c[0], setHdRaw = _c[1]; + + var emit = useCallback(function (kv, hd) { + var kVals = parseCsvInts(kv), hdVals = parseCsvInts(hd); + if (!kVals.length || !hdVals.length) return; + var lines = ["shapes:"], dimHd = hdVals[0]; + for (var i = 0; i < kVals.length; i++) { + lines.push(" - {id: decode_kv" + kVals[i] + ", M: 1, N: " + kVals[i] + ", K: " + dimHd + "}"); + } + onChange(lines.join("\n")); + }, [onChange]); + + return html` +
+
+ M = 1 (fixed — decode processes 1 token at a time). + Bottleneck is KV cache read; this is memory-bound. +
+
+ + + KV cache length — the decode bottleneck +
+
+ + + Head dimension, typically 128 +
+
+ `; +} + +function RoPeGuided(_a) { + var value = _a.value, onChange = _a.onChange; + var prev = useMemo(function () { return parseShapeYaml(value || ""); }, [value]); + var _b = useState(function () { return (prev.tVals || ["1", "128", "4096"]).join(", "); }), + tRaw = _b[0], setTRaw = _b[1]; + var _c = useState(function () { return (prev.hdVals || ["128"]).join(", "); }), + hdRaw = _c[0], setHdRaw = _c[1]; + + var emit = useCallback(function (t, hd) { + var tVals = parseCsvInts(t), hdVals = parseCsvInts(hd); + if (!tVals.length || !hdVals.length) return; + var lines = ["shapes:"], dimHd = hdVals[0]; + for (var i = 0; i < tVals.length; i++) { + lines.push(" - {id: rope_t" + tVals[i] + ", T: " + tVals[i] + ", H: " + dimHd + "}"); + } + onChange(lines.join("\n")); + }, [onChange]); + + return html` +
+
+ + + T = batch × seq_len +
+
+ + + Per-head rotary dimension +
+
+ `; +} + +// ---- YAML parser (best-effort, for backfilling guided fields on toggle) ----- + +function parseShapeYaml(raw) { + var out = {}; + if (!raw) return out; + var mVals = [], nVals = [], kVals = [], tVals = [], hVals = []; + var tokensVals = [], hdVals = [], kvVals = []; + var re = /\{[^}]+\}/g, m; + while ((m = re.exec(raw)) !== null) { + var inner = m[0]; + var get = function (key) { + var km = new RegExp(key + ":\\s*(\\d+)").exec(inner); + return km ? parseInt(km[1], 10) : null; + }; + var mv = get("M"); if (mv != null) mVals.push(mv); + var nv = get("N"); if (nv != null) nVals.push(nv); + var kv = get("K"); if (kv != null) kVals.push(kv); + var tv = get("T"); if (tv != null) tVals.push(tv); + var hv = get("H"); if (hv != null) hVals.push(hv); + var toks = mv; if (toks != null) tokensVals.push(toks); + var hd = kv; if (hd != null) hdVals.push(hd); + var kvLen = nv; if (kvLen != null) kvVals.push(kvLen); + } + var uniqSort = function (arr) { return arr.filter(function (v, i) { return arr.indexOf(v) === i; }).sort(function (a, b) { return a - b; }); }; + if (mVals.length) out.mVals = uniqSort(mVals); + if (nVals.length) out.nVals = uniqSort(nVals); + if (kVals.length) out.kVals = uniqSort(kVals); + if (tVals.length) out.tVals = uniqSort(tVals); + if (hVals.length) out.hVals = uniqSort(hVals); + if (tokensVals.length) out.tokensVals = uniqSort(tokensVals); + if (hdVals.length) out.hdVals = uniqSort(hdVals); + if (kvVals.length) out.kvVals = uniqSort(kvVals); + return out; +} + +// ---- operator → component map ---------------------------------------------- + +var OPERATOR_GUIDED = { + "Quantized GEMM": GemmGuided, + "Attention": PrefillAttnGuided, + "RMSNorm / LayerNorm": RmsNormGuided, + "RoPE": RoPeGuided, +}; + +// ---- main component --------------------------------------------------------- + +function ShapeInputInner(_a) { + var field = _a.field, value = _a.value, onChange = _a.onChange, allValues = _a.allValues; + var operator = (allValues && allValues.operator) || ""; + var dtype = (allValues && allValues.dtype) || ""; + var model = (allValues && allValues.model) || "DeepSeek V4 Flash"; + var assignmentMode = + (allValues && allValues.shape_assignment_mode) || "AI automatic"; + var manualAssignment = assignmentMode === "Manual by GPU"; + var scopeMode = + (allValues && allValues.shape_scope) || "All API shapes"; + var subsetScope = scopeMode === "Selected shapes only"; + var supportsModelCatalog = + operator === "Quantized GEMM" && dtype === "INT8 W8A8"; + var usesApiDefaults = + supportsModelCatalog && + model === "DeepSeek V4 Flash" && + !(value && value.trim()); + var isCustom = operator === "Custom operator" || !OPERATOR_GUIDED[operator]; + + var _b = useState(function () { + if (usesApiDefaults) return "api"; + if (isCustom) return "raw"; + if (value && /\bassignments\s*:/.test(value)) return "raw"; + if (supportsModelCatalog && model !== "DeepSeek V4 Flash") { + return "catalog"; + } + return "guided"; + }), mode = _b[0], setMode = _b[1]; + + var prevOpRef = useRef(operator); + useEffect(function () { + if (prevOpRef.current !== operator) { + prevOpRef.current = operator; + if (!isCustom) setMode("guided"); + } + }, [operator, isCustom]); + + useEffect(function () { + if (usesApiDefaults) { + setMode("api"); + } else if ( + mode === "api" || + (mode === "catalog" && model === "DeepSeek V4 Flash") + ) { + setMode( + isCustom + ? "raw" + : supportsModelCatalog && model !== "DeepSeek V4 Flash" + ? "catalog" + : "guided" + ); + } + }, [usesApiDefaults, isCustom, supportsModelCatalog, model]); + + var _c = useState(function () { + if (operator === "Attention" && value) { + if (/\bM:\s*1[,\s\}]/.test(value)) return "decode"; + } + return "prefill"; + }), attnSubMode = _c[0], setAttnSubMode = _c[1]; + + var GuidedCmp = OPERATOR_GUIDED[operator]; + + useEffect(function () { + if ( + !manualAssignment && + value && + /\bassignment_mode\s*:\s*manual\b/.test(value) + ) { + onChange(""); + } + }, [manualAssignment]); + + useEffect(function () { + if ( + !manualAssignment && + !subsetScope && + value && + /\bshape_scope\s*:\s*subset\b/.test(value) + ) { + onChange(""); + } + }, [manualAssignment, subsetScope]); + + if (manualAssignment) { + var supportsManualDefaults = supportsModelCatalog; + return html` +
+
+ + + ${supportsManualDefaults + ? "Fixed API workload" + : "Advanced YAML required"} + +
+
+ ${supportsManualDefaults + ? html`<${ManualGpuAssignment} + key=${model} + value=${value} + onChange=${onChange} + scopeMode=${scopeMode} + model=${model} />` + : html` +
+ This operator does not expose a fixed workload for the four-card + editor yet. Define shapes and + assignments explicitly below. +
+ + `} +
+
+ `; + } + + if ( + subsetScope && + supportsModelCatalog + ) { + return html` +
+
+ + Fixed API workload +
+
+ <${AiShapeSubset} + key=${model} + value=${value} + onChange=${onChange} + model=${model} /> +
+
+ `; + } + + return html` +
+
+ ${usesApiDefaults && html` + + `} + ${supportsModelCatalog && model !== "DeepSeek V4 Flash" && html` + + `} + + + ${operator === "Attention" && mode === "guided" && html` + + + + `} +
+
+ ${mode === "api" + ? html` +
+ No manual shape input required. + MetaInfer will read DEFAULT_OPTIMIZATION_SHAPES + from the immutable INT8 W8A8 GEMM API when the task starts. + Choose Guided or Raw only to override that workload. +
+ <${ApiDefaultsPreview} catalog=${modelCatalog(model)} /> + ` + : mode === "catalog" + ? html`<${ModelWorkloadAll} + value=${value} + onChange=${onChange} + model=${model} />` + : mode === "raw" || !GuidedCmp + ? html` + + ` + : operator === "Attention" && attnSubMode === "decode" + ? html`<${DecodeAttnGuided} value=${value} onChange=${onChange} />` + : html`<${GuidedCmp} value=${value} onChange=${onChange} />` + } +
+
+ `; +} + +// ---- register --------------------------------------------------------------- + +var _g = (typeof globalThis !== "undefined" ? globalThis : window); +var _bridge = (_g.__metainferOverrides = _g.__metainferOverrides || {}); +_bridge["shape-input"] = ShapeInputInner; + +// Named export for the form-renderer to pick up. +export var ShapeInput = ShapeInputInner; diff --git a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css new file mode 100644 index 00000000..55f7e00f --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css @@ -0,0 +1,909 @@ +.dkao-detail { + display: grid; + grid-template-columns: minmax(0, 2fr) minmax(320px, 1fr); + gap: var(--sp-3); + align-items: start; +} + +.dkao-state-panel, +.dkao-iterations-panel { + grid-column: 1 / -1; +} + +.dkao-state-machine { + display: flex; + align-items: center; + gap: 10px; + padding: 24px; + overflow-x: auto; +} + +.dkao-state-step { + min-width: 132px; + padding: 13px 16px; + border: 1px solid var(--border); + border-radius: var(--radius); + background: var(--bg); + color: var(--fg-dim); + display: flex; + align-items: center; + gap: 9px; + white-space: nowrap; +} + +.dkao-state-step.active { + color: var(--fg); + border-color: var(--accent); + background: rgba(88, 166, 255, 0.13); + box-shadow: 0 0 0 1px rgba(88, 166, 255, 0.2); +} + +.dkao-state-step.terminal.active { + border-color: var(--ok); + background: rgba(63, 185, 80, 0.12); +} + +.dkao-state-index { + width: 23px; + height: 23px; + border-radius: 50%; + background: var(--panel-2); + color: var(--accent-2); + display: inline-flex; + align-items: center; + justify-content: center; + font-size: var(--fs-xs); + font-weight: 700; +} + +.dkao-state-arrow { + color: var(--muted); + font-size: 20px; +} + +.dkao-explore-summary { + padding: 0 24px 24px; +} + +.dkao-explore-heading { + margin: 0 0 10px; + color: var(--fg-dim); + font-size: var(--fs-xs); + letter-spacing: 0.04em; + text-transform: uppercase; +} + +.dkao-explore-grid { + display: grid; + grid-template-columns: repeat(4, minmax(0, 1fr)); + gap: 10px; +} + +.dkao-explore-gpu { + min-width: 0; + padding: 12px; + border: 1px solid var(--border-soft); + border-radius: var(--radius); + background: var(--bg); +} + +.dkao-explore-gpu.stale { + border-color: var(--warn); +} + +.dkao-explore-gpu > header { + display: flex; + align-items: flex-start; + justify-content: space-between; + gap: 8px; +} + +.dkao-explore-gpu > header > div { + display: flex; + min-width: 0; + flex-direction: column; + gap: 2px; +} + +.dkao-explore-gpu > header strong { + color: var(--accent-2); +} + +.dkao-explore-gpu > header small { + overflow: hidden; + color: var(--muted); + font-family: var(--fs-mono); + text-overflow: ellipsis; + white-space: nowrap; +} + +.dkao-explore-step { + margin: 12px 0 8px; + color: var(--fg); + font-size: var(--fs-sm); +} + +.dkao-explore-gpu dl { + display: grid; + grid-template-columns: 56px minmax(0, 1fr); + gap: 4px 8px; + margin: 0; + font-size: var(--fs-xs); +} + +.dkao-explore-gpu dt { + color: var(--muted); +} + +.dkao-explore-gpu dd { + margin: 0; + overflow: hidden; + color: var(--fg-dim); + font-family: var(--fs-mono); + text-overflow: ellipsis; + white-space: nowrap; +} + +.dkao-explore-progress { + height: 5px; + margin-top: 11px; + overflow: hidden; + border-radius: 3px; + background: var(--panel-2); +} + +.dkao-explore-progress > span { + display: block; + height: 100%; + border-radius: inherit; + background: var(--accent); + transition: width 180ms ease; +} + +.dkao-explore-warning { + margin: 8px 0 0; + color: var(--warn); + font-size: var(--fs-xxs); +} + +.dkao-worker-lanes { + display: flex; + flex-direction: column; +} + +.dkao-worker-lane { + display: grid; + grid-template-columns: 125px minmax(280px, 0.9fr) minmax(340px, 1.1fr); + min-height: 132px; + border-bottom: 1px solid var(--border-soft); +} + +.dkao-worker-lane:last-child { + border-bottom: 0; +} + +.dkao-worker-lane.unassigned { + opacity: 0.65; +} + +.dkao-worker-meta { + padding: 10px; + border-right: 1px solid var(--border-soft); + background: var(--bg-elev); + display: flex; + flex-direction: column; + align-items: flex-start; + gap: 5px; +} + +.dkao-worker-meta strong { + color: var(--accent-2); + font-family: var(--fs-mono); +} + +.dkao-worker-meta small { + color: var(--muted); + word-break: break-word; +} + +.dkao-restart-worker { + width: 100%; + margin-top: auto; + padding: 5px 7px; + border-color: var(--warn); + color: var(--warn); + font-size: var(--fs-xs); +} + +.dkao-iteration-strip { + display: flex; + gap: 7px; + padding: 7px 8px 9px; + overflow-x: auto; + min-height: 89px; + scrollbar-width: thin; +} + +.dkao-iteration-card { + flex: 0 0 148px; + border: 1px solid var(--border-soft); + border-radius: var(--radius-sm); + background: var(--bg); + padding: 7px 8px; +} + +.dkao-iteration-card.accepted { + border-color: rgba(63, 185, 80, 0.55); +} + +.dkao-iteration-card.running { + border-color: rgba(88, 166, 255, 0.7); + background: rgba(88, 166, 255, 0.06); +} + +.dkao-running-pulse { + height: 4px; + margin-top: 7px; + overflow: hidden; + border-radius: 2px; + background: var(--panel-2); +} + +.dkao-running-pulse > span { + display: block; + width: 42%; + height: 100%; + border-radius: inherit; + background: var(--accent); + animation: dkao-iteration-running 1.25s ease-in-out infinite alternate; +} + +@keyframes dkao-iteration-running { + from { transform: translateX(0); } + to { transform: translateX(138%); } +} + +.dkao-iteration-title { + display: flex; + align-items: center; + justify-content: space-between; + margin-bottom: 5px; +} + +.dkao-iteration-card dl { + display: grid; + grid-template-columns: 44px 1fr; + gap: 2px 5px; + margin: 0; + font-size: var(--fs-xs); +} + +.dkao-iteration-card dt { + color: var(--muted); +} + +.dkao-iteration-card dd { + margin: 0; + color: var(--fg-dim); + font-family: var(--fs-mono); +} + +.dkao-lane-section { + min-width: 0; + border-right: 1px solid var(--border-soft); + display: flex; + flex-direction: column; +} + +.dkao-lane-section:last-child { + border-right: 0; +} + +.dkao-lane-heading { + min-height: 27px; + padding: 5px 9px; + border-bottom: 1px solid var(--border-soft); + display: flex; + justify-content: space-between; + color: var(--fg-dim); + font-size: var(--fs-xs); +} + +.dkao-lane-heading span { + color: var(--muted); +} + +.dkao-plan-strip { + display: flex; + gap: 7px; + padding: 7px 8px 9px; + overflow-x: auto; + min-height: 89px; + scrollbar-width: thin; +} + +.dkao-plan-card { + flex: 0 0 190px; + border: 1px solid var(--border-soft); + border-radius: var(--radius-sm); + background: var(--bg); + padding: 7px 8px; +} + +.dkao-plan-card.manual, +.dkao-plan-card.guidance.pending { + border-color: rgba(88, 166, 255, 0.55); +} + +.dkao-plan-title { + display: flex; + align-items: center; + justify-content: space-between; + margin-bottom: 5px; +} + +.dkao-plan-card p { + margin: 0 0 4px; + color: var(--fg-dim); + font-size: var(--fs-xs); + line-height: 1.35; + display: -webkit-box; + -webkit-box-orient: vertical; + -webkit-line-clamp: 2; + overflow: hidden; +} + +.dkao-plan-card small { + color: var(--muted); + display: block; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; +} + +.dkao-guidance-form { + display: flex; + gap: 6px; + padding: 0 8px 7px; +} + +.dkao-guidance-form input { + min-width: 0; + flex: 1; + padding: 6px 8px; + font-size: var(--fs-xs); +} + +.dkao-guidance-form button { + flex: 0 0 auto; + padding: 5px 9px; + font-size: var(--fs-xs); +} + +.dkao-guidance-message { + color: var(--accent-2); + padding: 0 8px 6px; +} + +.dkao-validation-subtitle { + display: flex; + align-items: baseline; + justify-content: space-between; + gap: 12px; + margin: 0 0 10px; + color: var(--fg-dim); + font-size: var(--fs-xs); + letter-spacing: 0.04em; + text-transform: uppercase; +} + +.dkao-validation-subtitle span { + color: var(--muted); + font-weight: 400; + letter-spacing: 0; + text-transform: none; +} + +.dkao-events { + max-height: 320px; + overflow: auto; + padding: 8px 14px; +} + +.dkao-events > div { + display: grid; + grid-template-columns: 150px 1fr; + gap: 12px; + padding: 7px 0; + border-bottom: 1px solid var(--border-soft); + font-size: var(--fs-xs); +} + +.dkao-events code { + color: var(--accent-2); +} + +.dkao-events span { + color: var(--muted); + word-break: break-word; +} + +.dkao-skill-entry, +.dkao-skill-library { + grid-column: 1 / -1; +} + +.dkao-skill-entry, +.dkao-skill-library-title { + display: flex; + align-items: center; + justify-content: space-between; + gap: 18px; + padding: 14px 16px; +} + +.dkao-skill-entry h2, +.dkao-skill-entry p, +.dkao-skill-library-title h2, +.dkao-skill-library-title p { + margin: 0; +} + +.dkao-skill-columns { + display: grid; + grid-template-columns: 1fr 1fr; + min-height: 360px; + border-top: 1px solid var(--border-soft); +} + +.dkao-skill-column { + min-width: 0; + padding: 14px; +} + +.dkao-skill-column + .dkao-skill-column { + border-left: 1px solid var(--border-soft); +} + +.dkao-skill-column h3 { + display: flex; + justify-content: space-between; + margin: 0; +} + +.dkao-skill-column h3 span { + color: var(--muted); + font-family: var(--fs-mono); +} + +.dkao-skill-column > p { + min-height: 20px; + margin: 4px 0 10px; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.dkao-skill-list { + max-height: 560px; + overflow: auto; + display: flex; + flex-direction: column; + gap: 7px; +} + +.dkao-skill-file { + border: 1px solid var(--border-soft); + border-radius: var(--radius-sm); + background: var(--bg); +} + +.dkao-skill-file summary { + cursor: pointer; + padding: 9px 10px; + display: flex; + align-items: center; + justify-content: space-between; + gap: 10px; +} + +.dkao-skill-file summary > span:first-child { + min-width: 0; + display: flex; + flex-direction: column; +} + +.dkao-skill-file summary strong, +.dkao-skill-file summary small { + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.dkao-skill-file summary small { + color: var(--muted); +} + +.dkao-skill-file pre { + max-height: 380px; + overflow: auto; + margin: 0; + padding: 12px; + border-top: 1px solid var(--border-soft); + color: var(--fg-dim); + white-space: pre-wrap; + word-break: break-word; + font-size: var(--fs-xs); +} + +@media (max-width: 1100px) { + .dkao-detail { + grid-template-columns: 1fr; + } + + .dkao-validation-panel, + .dkao-events-panel { + grid-column: 1; + } + + .dkao-worker-lane { + grid-template-columns: 115px minmax(260px, 1fr) minmax(320px, 1fr); + overflow-x: auto; + } + + .dkao-skill-columns { + grid-template-columns: 1fr; + } + + .dkao-skill-column + .dkao-skill-column { + border-left: 0; + border-top: 1px solid var(--border-soft); + } +} + +/* ---- shape input component ------------------------------------------------ */ + +.dkao-shape-input { + border: 1px solid var(--border); + border-radius: var(--radius); + overflow: hidden; +} + +.dkao-shape-tabs { + display: flex; + align-items: center; + padding: 6px; + background: var(--panel-2); + border-bottom: 1px solid var(--border); + gap: 4px; +} + +.dkao-shape-tab { + padding: 5px 11px; + border: 1px solid transparent; + border-radius: var(--radius-sm); + background: transparent; + color: var(--fg-dim); + font-size: var(--fs-xs); + cursor: pointer; + transition: background 0.12s; +} + +.dkao-shape-tab:hover:not(:disabled) { + background: var(--bg); + color: var(--fg); +} + +.dkao-shape-tab.active { + background: var(--bg); + color: var(--fg); + border-color: var(--border); + box-shadow: 0 1px 2px rgba(0,0,0,.06); +} + +.dkao-shape-tab:disabled { + opacity: .45; + cursor: default; +} + +.dkao-shape-spacer { + width: 12px; +} + +.dkao-shape-mode-note { + margin-left: auto; + padding-right: 6px; + color: var(--muted); + font-size: var(--fs-xxs); +} + +.dkao-shape-body { + padding: 10px; +} + +.dkao-shape-body textarea.input { + margin: 0; + border: 0; + border-radius: 0; + resize: vertical; +} + +.dkao-guided-form { + display: flex; + flex-direction: column; + gap: 10px; +} + +.dkao-guided-info { + padding: 9px 11px; + background: var(--bg-elev); + border-radius: var(--radius-sm); + border: 1px solid var(--border-soft); + color: var(--fg-dim); + font-size: var(--fs-xs); + line-height: 1.45; +} + +.dkao-guided-info strong { + color: var(--accent-2); +} + +.dkao-guided-row { + display: flex; + align-items: baseline; + gap: 8px; + flex-wrap: wrap; +} + +.dkao-guided-label { + flex: 0 0 130px; + min-width: 100px; + font-size: var(--fs-xs); + color: var(--fg-dim); + font-family: var(--fs-mono); + text-align: right; + padding-top: 6px; +} + +.dkao-guided-input { + flex: 1 1 160px; + min-width: 120px; + font-family: var(--fs-mono); +} + +.dkao-guided-hint { + flex: 1 1 200px; + min-width: 140px; + font-size: var(--fs-xxs); + color: var(--muted); + padding-top: 6px; +} + +.dkao-guided-check { + display: flex; + align-items: center; + gap: 6px; + font-size: var(--fs-xs); + color: var(--fg-dim); + cursor: pointer; + padding-left: 138px; +} + +.dkao-manual-assignment { + display: flex; + flex-direction: column; + gap: 10px; +} + +.dkao-assignment-actions { + display: flex; + align-items: center; + flex-wrap: wrap; + gap: 5px; + color: var(--muted); + font-size: var(--fs-xxs); +} + +.dkao-shape-catalog { + overflow: hidden; + border: 1px solid var(--border-soft); + border-radius: var(--radius-sm); +} + +.dkao-shape-catalog-header { + display: flex; + align-items: center; + justify-content: space-between; + padding: 8px 10px; + background: var(--panel-2); + font-size: var(--fs-xs); +} + +.dkao-shape-catalog-header span { + color: var(--muted); + font-size: var(--fs-xxs); +} + +.dkao-tp-filter { + display: flex; + align-items: center; + gap: 4px; + padding: 6px 10px; + border-bottom: 1px solid var(--border-soft); + background: var(--panel-2); +} + +.dkao-tp-filter-label { + margin-right: 4px; + color: var(--muted); + font-size: var(--fs-xxs); + letter-spacing: 0.05em; + text-transform: uppercase; +} + +.dkao-shape-catalog-items { + display: grid; + grid-template-columns: repeat(2, minmax(0, 1fr)); + gap: 5px; + padding: 7px; +} + +.dkao-shape-choice { + display: flex; + align-items: center; + min-width: 0; + gap: 7px; + padding: 7px 8px; + border: 1px solid var(--border-soft); + border-radius: var(--radius-sm); + color: var(--fg-dim); + cursor: pointer; +} + +.dkao-shape-choice.selected { + border-color: var(--accent); + background: color-mix(in srgb, var(--accent) 7%, var(--bg)); + color: var(--fg); +} + +.dkao-shape-choice > span { + display: flex; + min-width: 0; + flex-direction: column; + gap: 2px; +} + +.dkao-shape-choice strong, +.dkao-shape-choice small { + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.dkao-shape-choice strong { + font-family: var(--fs-mono); + font-size: var(--fs-xxs); +} + +.dkao-shape-choice small { + color: var(--muted); + font-size: 10px; +} + +.dkao-gpu-grid { + display: grid; + grid-template-columns: repeat(2, minmax(0, 1fr)); + gap: 8px; +} + +.dkao-gpu-card { + min-width: 0; + overflow: hidden; + border: 1px solid var(--border-soft); + border-radius: var(--radius-sm); + background: var(--bg-elev); +} + +.dkao-gpu-card > header { + display: flex; + align-items: center; + justify-content: space-between; + padding: 8px 10px; + border-bottom: 1px solid var(--border-soft); + background: var(--panel-2); + font-size: var(--fs-xs); +} + +.dkao-gpu-card > header span { + color: var(--muted); + font-size: var(--fs-xxs); +} + +.dkao-gpu-shapes { + display: flex; + min-height: 54px; + flex-direction: column; + gap: 6px; + padding: 7px; +} + +.dkao-gpu-shape { + display: flex; + align-items: center; + justify-content: space-between; + gap: 8px; + min-width: 0; + padding: 7px 8px; + border: 1px solid var(--border-soft); + border-radius: var(--radius-sm); + background: var(--bg); +} + +.dkao-gpu-shape > div { + display: flex; + min-width: 0; + flex-direction: column; + gap: 2px; +} + +.dkao-gpu-shape strong, +.dkao-gpu-shape small { + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.dkao-gpu-shape strong { + font-family: var(--fs-mono); + font-size: var(--fs-xxs); +} + +.dkao-gpu-shape small { + color: var(--muted); + font-size: 10px; +} + +.dkao-gpu-shape select { + max-width: 72px; + padding: 3px 4px; + border: 1px solid var(--border); + border-radius: var(--radius-sm); + background: var(--panel-2); + color: var(--fg); + font-size: var(--fs-xxs); +} + +.dkao-gpu-empty { + display: grid; + min-height: 38px; + place-items: center; + color: var(--muted); + font-size: var(--fs-xxs); +} + +@media (max-width: 760px) { + .dkao-shape-catalog-items, + .dkao-gpu-grid, + .dkao-explore-grid { + grid-template-columns: 1fr; + } +} + +@media (min-width: 761px) and (max-width: 1180px) { + .dkao-explore-grid { + grid-template-columns: repeat(2, minmax(0, 1fr)); + } +} + +/* Kernel repository rename panel */ +.dkao-repo-panel .dkao-repo-actions { + display: flex; + align-items: center; + gap: 12px; + margin-top: 6px; +} + +.dkao-repo-panel .dkao-rename-msg { + color: var(--muted); + font-size: var(--fs-xs); +} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/__init__.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/__init__.py new file mode 100644 index 00000000..b0c98baa --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for dcu-kernel-auto-opt.""" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/conftest.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/conftest.py new file mode 100644 index 00000000..31681ab9 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/conftest.py @@ -0,0 +1,39 @@ +"""Shared fixtures for dcu-kernel-auto-opt plugin route tests.""" + +from __future__ import annotations + +import pytest +from fastapi.testclient import TestClient + +from metainfer.testing import isolated_env # noqa: F401 — re-export as fixture +from metainfer.server import app as app_module +from metainfer.server import tasks as _tasks +from metainfer.server.tasks import TaskEntry + + +@pytest.fixture +def app(isolated_env): + return app_module.create_app() + + +@pytest.fixture +def client(app): + return TestClient(app) + + +def register_dkao_task( + state_dir, workspace_dir, task_id: str = "dkao-1" +) -> TaskEntry: + """Register one dcu-kernel-auto-opt task in the WebUI registry.""" + state_dir.mkdir(parents=True, exist_ok=True) + workspace_dir.mkdir(parents=True, exist_ok=True) + entry = TaskEntry( + id=task_id, + type="dcu-kernel-auto-opt", + label="test dkao task", + state_dir=str(state_dir), + workspace_dir=str(workspace_dir), + created_at=0.0, + ) + _tasks.add_task(entry) + return entry diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_config.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_config.py new file mode 100644 index 00000000..c7ccec16 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_config.py @@ -0,0 +1,296 @@ +from __future__ import annotations + +import pytest + +from ..orchestrator.config import ( + DSH_DEFAULT_MODEL_ID, + dsh_model_id, + load_config, + resolve_claude_bin, +) +from ..orchestrator.gpu_binding import ( + bind_worker_gpu, + hide_gpus_from_control_plane, +) + + +def _req(shape_config: str): + return { + "answers": { + "execution_mode": "Mock (no GPU)", + "shape_config": shape_config, + "mock_iterations": "2", + "minimum_improvement_percent": 1.0, + } + } + + +def test_static_assignments_parse(): + cfg = load_config(_req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} + - {id: m16, M: 16, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} + worker_1: {gpu: 1, shapes: [m16]} +""")) + assert set(cfg.shapes) == {"m2", "m16"} + assert [a.gpu for a in cfg.assignments] == [0, 1] + assert cfg.assignment_mode == "manual" + assert cfg.claude_model == "claude-opus-5" + + +def test_new_task_can_select_sonnet(): + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["claude_model"] = "Sonnet" + + assert load_config(req).claude_model == "claude-sonnet-5" + + +def test_unknown_claude_model_is_rejected(): + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["claude_model"] = "default" + + with pytest.raises(ValueError, match="agent_model must be one of"): + load_config(req) + + +def test_dsh_framework_resolves_wrapper_model(monkeypatch): + monkeypatch.setenv("DSH_AGENT_MODEL", "deepseek/deepseek-v4-flash-0731") + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["agent_framework"] = "dsh" + req["answers"]["agent_model"] = "deepseek-v4-flash" + + cfg = load_config(req) + assert cfg.agent_framework == "dsh" + assert cfg.claude_model == "deepseek/deepseek-v4-flash-0731" + + +def test_dsh_default_model_when_agent_model_missing(monkeypatch): + monkeypatch.setenv("DSH_AGENT_MODEL", "deepseek/deepseek-v4-flash-0731") + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["agent_framework"] = "dsh" + + assert load_config(req).claude_model == DSH_DEFAULT_MODEL_ID + + +def test_dsh_rejects_ccb_only_model(monkeypatch): + monkeypatch.setenv("DSH_AGENT_MODEL", "deepseek/deepseek-v4-flash-0731") + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["agent_framework"] = "dsh" + req["answers"]["agent_model"] = "Opus" + + with pytest.raises(ValueError, match="framework 'dsh'"): + load_config(req) + + +def test_unknown_agent_framework_is_rejected(): + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["agent_framework"] = "bogus" + + with pytest.raises(ValueError, match="agent_framework must be one of"): + load_config(req) + + +def test_resolve_claude_bin_per_framework(): + # dsh -> the bundled ccb-compatible DSH wrapper. + dsh_bin = resolve_claude_bin("dsh") + assert dsh_bin.endswith("bridge/dsh/dsh_agent.py") + # Explicit override always wins. + assert resolve_claude_bin("dsh", explicit="/opt/custom/agent") == "/opt/custom/agent" + # ccb defaults to the env / "ccb" (METAINFER_CLAUDE_BIN unset here). + assert resolve_claude_bin("ccb") == "ccb" + + +def test_dsh_model_id_honors_env(monkeypatch): + monkeypatch.setenv("DSH_AGENT_MODEL", "deepseek/deepseek-v4-flash") + assert dsh_model_id() == "deepseek/deepseek-v4-flash" + monkeypatch.delenv("DSH_AGENT_MODEL") + assert dsh_model_id() == DSH_DEFAULT_MODEL_ID + + +def test_explicit_manual_mode_requires_assignments(): + req = _req(""" +assignment_mode: manual +shapes: + - {id: m2, M: 2, N: 16, K: 32} +""") + req["answers"]["execution_mode"] = ( + "Generate & optimize (auto-create kernel repo)" + ) + with pytest.raises(ValueError, match="assignments"): + load_config(req) + + +def test_explicit_ai_mode_can_omit_generate_assignments(): + req = _req(""" +assignment_mode: ai +shapes: + - {id: m2, M: 2, N: 16, K: 32} +""") + req["answers"]["execution_mode"] = ( + "Generate & optimize (auto-create kernel repo)" + ) + cfg = load_config(req) + assert cfg.assignment_mode == "ai" + assert cfg.assignments[0].shape_ids == ["m2"] + + +def test_subset_scope_is_preserved(): + req = _req(""" +shape_scope: subset +assignment_mode: manual +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_2: {gpu: 2, shapes: [m2]} +""") + cfg = load_config(req) + assert cfg.shape_scope == "subset" + assert list(cfg.shapes) == ["m2"] + + +def test_shape_cannot_be_assigned_twice(): + with pytest.raises(ValueError, match="assigned more than once"): + load_config(_req(""" +shapes: + - {id: m2, M: 2} +assignments: + worker_0: {gpu: 0, shapes: [m2]} + worker_1: {gpu: 1, shapes: [m2]} +""")) + + +def test_gpu_cannot_be_shared(): + with pytest.raises(ValueError, match="GPU 0"): + load_config(_req(""" +shapes: + - {id: m2, M: 2} + - {id: m16, M: 16} +assignments: + worker_0: {gpu: 0, shapes: [m2]} + worker_1: {gpu: 0, shapes: [m16]} +""")) + + +def test_worker_gpu_binding_uses_one_filter_only(): + env = {"ROCR_VISIBLE_DEVICES": "3", "UNCHANGED": "yes"} + visible = bind_worker_gpu(env, 2) + assert visible == {"HIP_VISIBLE_DEVICES": "2"} + assert env["HIP_VISIBLE_DEVICES"] == "2" + assert "ROCR_VISIBLE_DEVICES" not in env + assert env["UNCHANGED"] == "yes" + + +def test_control_plane_hides_gpus(): + env = {"ROCR_VISIBLE_DEVICES": "0"} + hide_gpus_from_control_plane(env) + assert env == {"HIP_VISIBLE_DEVICES": ""} + + +def test_real_smoke_mode_is_accepted(): + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["execution_mode"] = "Real agents + DCU (smoke harness)" + assert load_config(req).execution_mode.startswith("Real agents") + + +def test_real_smoke_ignores_legacy_repo_field(): + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"].update({ + "execution_mode": "Real agents + DCU (smoke harness)", + "target_repo_path": ">=1.2x baseline", + }) + assert load_config(req).target_repo_path is None + + +def test_legacy_workers_shape_config_is_normalized(): + cfg = load_config(_req(""" +workers: + worker_0: + gpu: 0 + shapes: + - {id: m2, op: gemm, M: 2, N: 16, K: 32} + worker_1: + gpu: 1 + shapes: + - {id: m16, op: gemm, M: 16, N: 16, K: 32} +""")) + assert set(cfg.shapes) == {"m2", "m16"} + assert cfg.shapes["m2"].params["op"] == "gemm" + assert [(item.worker_id, item.gpu, item.shape_ids) for item in cfg.assignments] == [ + ("worker_0", 0, ["m2"]), + ("worker_1", 1, ["m16"]), + ] + + +def test_real_w8a8_mode_requires_absolute_repo(): + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + # Use an existing directory so it passes the exists() check. + import tempfile + with tempfile.TemporaryDirectory() as td: + req["answers"].update({ + "execution_mode": "Real INT8 W8A8 GEMM", + "target_repo_path": td, + }) + cfg = load_config(req) + assert cfg.execution_mode == "Real INT8 W8A8 GEMM" + assert str(cfg.target_repo_path) == td + + +def test_real_w8a8_mode_accepts_missing_repo_field(): + """W8A8 mode accepts missing target_repo_path (agent auto-generates).""" + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["execution_mode"] = "Real INT8 W8A8 GEMM" + cfg = load_config(req) + assert cfg.execution_mode == "Real INT8 W8A8 GEMM" + assert cfg.target_repo_path is None diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_contract_snapshot.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_contract_snapshot.py new file mode 100644 index 00000000..ca271898 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_contract_snapshot.py @@ -0,0 +1,270 @@ +from __future__ import annotations + +import hashlib +import json + +import pytest + +from ..orchestrator.api_contracts import ( + OperatorAPIContract, + default_optimization_shapes, + resolve_operator_api, + validate_contract_shapes, +) +from ..orchestrator.gen_and_opt_pipeline import _task_local_api_contract + + +def _api_text(shape_id: str, m: int) -> str: + return ( + "DEFAULT_OPTIMIZATION_SHAPES = (" + f"{{'id': '{shape_id}', 'M': {m}, 'N': 16, 'K': 32}}," + ")\n" + "def _check_target_shape(m, n, k):\n" + " if not (1 <= m <= 3072 and n == 16 and k == 32):\n" + " raise ValueError('unsupported')\n" + ) + + +def test_live_contract_accepts_hy3_tp4_shapes(): + contract = resolve_operator_api("Quantized GEMM", "INT8 W8A8") + shapes = { + "qkv": { + "tp_size": 4, + "operator": "qkv_proj", + "M": 4096, + "N": 2560, + "K": 4096, + }, + "o": { + "tp_size": 4, + "operator": "o_proj", + "M": 4096, + "N": 4096, + "K": 2048, + }, + "gate_up": { + "tp_size": 4, + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 768, + "K": 4096, + }, + "down": { + "tp_size": 4, + "operator": "shared_down_proj", + "M": 4096, + "N": 4096, + "K": 384, + }, + } + + validate_contract_shapes(contract, shapes) + + +def test_live_contract_rejects_unregistered_hy3_tp4_variant(): + contract = resolve_operator_api("Quantized GEMM", "INT8 W8A8") + shapes = { + "qkv": { + "tp_size": 4, + "operator": "qkv_proj", + "M": 4096, + "N": 2576, + "K": 4096, + } + } + + with pytest.raises(ValueError, match="requires \\(K, N\\) in"): + validate_contract_shapes(contract, shapes) + + +def test_live_contract_accepts_model_catalog_tp8_m4096_shapes(): + # Model-catalog TP8 large-prefill boundary (2026-08-27): Hy3 / MiniMax M3 + # / GLM5.2 TP8 operators are optimizable at M=4096 ("Selected shapes + # only"); DeepSeek TP8 defaults keep the original three M values. + contract = resolve_operator_api("Quantized GEMM", "INT8 W8A8") + shapes = { + "hy3_tp8_qkv_proj_m4096": { + "tp_size": 8, "operator": "qkv_proj", + "M": 4096, "N": 1280, "K": 4096, + }, + "hy3_tp8_o_proj_m4096": { + "tp_size": 8, "operator": "o_proj", + "M": 4096, "N": 4096, "K": 1024, + }, + "hy3_tp8_shared_gate_up_proj_m4096": { + "tp_size": 8, "operator": "shared_gate_up_proj", + "M": 4096, "N": 384, "K": 4096, + }, + "hy3_tp8_shared_down_proj_m4096": { + "tp_size": 8, "operator": "shared_down_proj", + "M": 4096, "N": 4096, "K": 192, + }, + "minimax_tp8_qkv_proj_m4096": { + "tp_size": 8, "operator": "qkv_proj", + "M": 4096, "N": 1280, "K": 6144, + }, + "minimax_tp8_qkv_proj_and_indexer_qk_m4096": { + "tp_size": 8, "operator": "qkv_proj_and_indexer_qk", + "M": 4096, "N": 1536, "K": 6144, + }, + "minimax_tp8_o_proj_m4096": { + "tp_size": 8, "operator": "o_proj", + "M": 4096, "N": 6144, "K": 1024, + }, + "minimax_tp8_shared_gate_up_proj_m4096": { + "tp_size": 8, "operator": "shared_gate_up_proj", + "M": 4096, "N": 768, "K": 6144, + }, + "minimax_tp8_shared_down_proj_m4096": { + "tp_size": 8, "operator": "shared_down_proj", + "M": 4096, "N": 6144, "K": 384, + }, + "glm52_tp8_fused_qkv_a_proj_m4096": { + "tp_size": 8, "operator": "fused_qkv_a_proj", + "M": 4096, "N": 2624, "K": 6144, + }, + "glm52_tp8_q_b_proj_m4096": { + "tp_size": 8, "operator": "q_b_proj", + "M": 4096, "N": 2048, "K": 2048, + }, + "glm52_tp8_kv_b_proj_m4096": { + "tp_size": 8, "operator": "kv_b_proj", + "M": 4096, "N": 3584, "K": 512, + }, + "glm52_tp8_o_proj_m4096": { + "tp_size": 8, "operator": "o_proj", + "M": 4096, "N": 6144, "K": 2048, + }, + "glm52_tp8_shared_gate_up_proj_m4096": { + "tp_size": 8, "operator": "shared_gate_up_proj", + "M": 4096, "N": 512, "K": 6144, + }, + "glm52_tp8_shared_down_proj_m4096": { + "tp_size": 8, "operator": "shared_down_proj", + "M": 4096, "N": 6144, "K": 256, + }, + } + validate_contract_shapes(contract, shapes) + + +def test_default_optimization_shapes_unchanged_by_model_tp8_boundary(): + # Adding the model-catalog TP8 M=4096 boundary must not change the + # DeepSeek-only default workload or its serial-validation fallback scope. + from ..api.int8w8a8gemm.int8_w8a8_gemm_api import ( + DEFAULT_OPTIMIZATION_SHAPES, + MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES, + ) + + assert MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES == (4096,) + assert len(DEFAULT_OPTIMIZATION_SHAPES) == 42 + assert all( + str(shape["id"]).startswith("tp") for shape in DEFAULT_OPTIMIZATION_SHAPES + ) + + +def test_live_contract_accepts_minimax_m3_tp4_shapes(): + # Regression for minimaxm3-dsh-tp4-m4096-1-0c2f84a9: the MiniMax M3 TP4 + # M=4096 workload stopped in prepare because TP=4 qkv_proj (K, N) + # (6144, 2304) was outside the fixed API contract. The model catalog + # (frontend + baseline table) must be mirrored by the contract. + contract = resolve_operator_api("Quantized GEMM", "INT8 W8A8") + shapes = { + "qkv": { + "tp_size": 4, + "operator": "qkv_proj", + "M": 4096, + "N": 2304, + "K": 6144, + }, + "qkv_indexer": { + "tp_size": 4, + "operator": "qkv_proj_and_indexer_qk", + "M": 4096, + "N": 2560, + "K": 6144, + }, + "o": { + "tp_size": 4, + "operator": "o_proj", + "M": 4096, + "N": 6144, + "K": 2048, + }, + "gate_up": { + "tp_size": 4, + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 1536, + "K": 6144, + }, + "down": { + "tp_size": 4, + "operator": "shared_down_proj", + "M": 4096, + "N": 6144, + "K": 768, + }, + } + + validate_contract_shapes(contract, shapes) + + +def test_live_contract_rejects_unregistered_minimax_tp4_variant(): + contract = resolve_operator_api("Quantized GEMM", "INT8 W8A8") + shapes = { + "qkv": { + "tp_size": 4, + "operator": "qkv_proj", + "M": 4096, + "N": 2320, + "K": 6144, + } + } + + with pytest.raises(ValueError, match="requires \\(K, N\\) in"): + validate_contract_shapes(contract, shapes) + + +def test_task_contract_does_not_follow_live_api_updates(tmp_path): + live = tmp_path / "live.py" + task = tmp_path / "task" + task.mkdir() + snapshot = task / "api.py" + snapshot.write_text(_api_text("old", 16), encoding="utf-8") + digest = hashlib.sha256(snapshot.read_bytes()).hexdigest() + (task / "scaffold_manifest.json").write_text( + json.dumps({"control_plane_files": {"api.py": digest}}), + encoding="utf-8", + ) + live.write_text(_api_text("old", 16), encoding="utf-8") + origin = OperatorAPIContract( + operator="Quantized GEMM", + dtype="INT8 W8A8", + source=live, + destination_name="api.py", + ) + + frozen = _task_local_api_contract(origin, task) + live.write_text(_api_text("new", 3072), encoding="utf-8") + + assert default_optimization_shapes(frozen) == [ + {"id": "old", "M": 16, "N": 16, "K": 32} + ] + + +def test_task_contract_rejects_snapshot_digest_drift(tmp_path): + source = tmp_path / "api.py" + source.write_text(_api_text("old", 16), encoding="utf-8") + (tmp_path / "scaffold_manifest.json").write_text( + json.dumps({"control_plane_files": {"api.py": "wrong"}}), + encoding="utf-8", + ) + origin = OperatorAPIContract( + operator="Quantized GEMM", + dtype="INT8 W8A8", + source=source, + destination_name="api.py", + ) + + with pytest.raises(RuntimeError, match="digest mismatch"): + _task_local_api_contract(origin, tmp_path) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gen_pipeline.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gen_pipeline.py new file mode 100644 index 00000000..3af57879 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gen_pipeline.py @@ -0,0 +1,1799 @@ +from __future__ import annotations + +import importlib.util +import json +import subprocess +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from metainfer.orchestrator.state import StateStore +from metainfer.orchestrator.subagent_manager import SubAgentManager + +from ..orchestrator.config import ( + GEN_AND_OPT_MODE, + load_config, + replace_assignments, + validate_gpu_assignment, + ShapeSpec, + WorkerAssignment, +) +from ..orchestrator.gen_and_opt_pipeline import ( + GenAndOptPipeline, + _COORDINATOR_AGENT_ARGS, + _PERF_GATE_MAX_RETRIES, + _PERF_GATE_RETRY_INTERVAL_S, + _artifact_symbol_prefix, + _final_performance_gate, + _final_synthesis_prompt, + _is_control_plane_artifact, + _render_prebuilt_dispatch, + _require_valid_child_assignments, +) +from ..orchestrator import phases +from ..orchestrator import w8a8_pipeline as pipeline_module +from ..orchestrator.w8a8_pipeline import ( + RealW8A8OptimizationPipeline, + W8A8Runner, + _BENCHMARK_TIMEOUT_S, + _check_required_files, + _REFERENCE_PREPARE_TIMEOUT_S, + archive_iteration_candidate, + candidate_iteration_destination, + publish_iteration_candidate, + snapshot_accepted_kernel_artifact, +) + + +TORCH_AVAILABLE = importlib.util.find_spec("torch") is not None + + +def test_generated_repo_is_valid_existing_repo_seed(tmp_path): + repo = tmp_path / "generated-repo" + for name in ( + "int8_w8a8_gemm_api.py", + "w8a8_backend.py", + "w8a8_bench.py", + "setup.py", + "csrc/bindings.cpp", + "csrc/w8a8_gemm_hip.hip", + ): + path = repo / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("# seed\n", encoding="utf-8") + + assert _check_required_files(repo) + + +def _gen_req(shape_config: str = ""): + if not shape_config: + shape_config = """ +shapes: + - {id: m2, M: 2, N: 1536, K: 4096} + - {id: m16, M: 16, N: 1536, K: 4096} +assignments: + worker_0: {gpu: 0, shapes: [m2]} + worker_1: {gpu: 1, shapes: [m16]} +""" + return { + "task_id": "gen-test", + "task_type": "dcu-kernel-auto-opt", + "answers": { + "execution_mode": GEN_AND_OPT_MODE, + "shape_config": shape_config, + "mock_iterations": "3", + "minimum_improvement_percent": 1.0, + "operator": "Quantized GEMM", + "dtype": "INT8 W8A8", + "target_hardware": "K500SM_AI / gfx928", + "kernel_language": "HIP C++", + }, + } + + +def _shapes_only_req(): + """Request with shapes but no assignments — the generate agent decides.""" + return { + "task_id": "gen-test-auto", + "task_type": "dcu-kernel-auto-opt", + "answers": { + "execution_mode": GEN_AND_OPT_MODE, + "shape_config": """ +shapes: + - {id: m2, M: 2, N: 1536, K: 4096} + - {id: m16, M: 16, N: 1536, K: 4096} + - {id: m64, M: 64, N: 1536, K: 4096} +""", + "mock_iterations": "3", + "minimum_improvement_percent": 1.0, + "operator": "Quantized GEMM", + "dtype": "INT8 W8A8", + "target_hardware": "K500SM_AI / gfx928", + "kernel_language": "HIP C++", + }, + } + + +def _four_worker_req(): + return _gen_req(""" +shapes: + - {id: m2, M: 2, N: 1536, K: 4096} + - {id: m4, M: 4, N: 1536, K: 4096} + - {id: m16, M: 16, N: 1536, K: 4096} + - {id: m4096, M: 4096, N: 1536, K: 4096} +assignments: + worker_0: {gpu: 0, shapes: [m2]} + worker_1: {gpu: 1, shapes: [m4]} + worker_2: {gpu: 2, shapes: [m16]} + worker_3: {gpu: 3, shapes: [m4096]} +""") + + +def _make_pipeline(req, tmp_path): + """Create a minimal GenAndOptPipeline for contract validation testing.""" + state_dir = tmp_path / "state" + workspace_dir = tmp_path / "workspace" + manager = SubAgentManager(claude_bin="ccb") + return GenAndOptPipeline( + req=req, + state_dir=state_dir, + workspace_dir=workspace_dir, + store=StateStore(state_dir), + manager=manager, + ) + + +def test_real_continuation_keeps_fixed_triton_comparison_baseline( + tmp_path, monkeypatch +): + from ..orchestrator import w8a8_pipeline as pipeline_module + + req = _gen_req(""" +shapes: + - {id: m16_wo_b, M: 16, N: 4096, K: 2048} +assignments: + worker_2: {gpu: 2, shapes: [m16_wo_b]} +""") + req["answers"]["execution_mode"] = "Real INT8 W8A8 GEMM" + config = load_config(req) + workspace = tmp_path / "workspace" + (workspace / "workers" / "worker_2").mkdir(parents=True) + (workspace / "shared_baseline").mkdir() + pipeline = RealW8A8OptimizationPipeline( + req=req, + state_dir=tmp_path / "state", + workspace_dir=workspace, + store=StateStore(tmp_path / "state"), + manager=SubAgentManager(claude_bin="ccb"), + ) + + class FakeRunner: + def __init__(self, worker_root, gpu): + pass + + def probe(self): + return {"visible_devices": 1} + + def benchmark(self, shape): + return { + "passed": True, + "median_us": 45.0, + "p90_us": 46.0, + } + + monkeypatch.setattr(pipeline_module, "W8A8Runner", FakeRunner) + + baseline = pipeline._parallel_baseline(config) + + assert baseline["m16_wo_b"]["median_us"] == 54.617 + assert baseline["m16_wo_b"]["baseline_kind"] == "triton_graph" + assert baseline["m16_wo_b"]["bootstrap_metrics"]["median_us"] == 45.0 + + +def test_parallel_explore_continues_with_two_of_four_workers( + tmp_path, monkeypatch +): + pipeline = _make_pipeline(_four_worker_req(), tmp_path) + config = load_config(_four_worker_req()) + for assignment in config.assignments: + ( + pipeline.workspace_dir / "workers" / assignment.worker_id + ).mkdir(parents=True) + + def run_worker(_config, assignment, _baseline): + if assignment.worker_id in {"worker_2", "worker_3"}: + raise RuntimeError("agent stuck timeout") + return {"worker_id": assignment.worker_id, "shapes": {}} + + monkeypatch.setattr(pipeline, "_run_worker", run_worker) + monkeypatch.setattr( + pipeline, + "_author_worker_skill", + lambda _config, assignment: { + "name": f"{assignment.worker_id}-skill" + }, + ) + + workers = pipeline._parallel_agents(config, {}) + + assert set(workers) == {"worker_0", "worker_1"} + assert set(pipeline._worker_failures) == {"worker_2", "worker_3"} + assert all( + item["state"] == "timed_out" + for item in pipeline._worker_failures.values() + ) + + +def test_parallel_explore_continues_with_only_one_of_four_workers( + tmp_path, monkeypatch +): + pipeline = _make_pipeline(_four_worker_req(), tmp_path) + config = load_config(_four_worker_req()) + for assignment in config.assignments: + ( + pipeline.workspace_dir / "workers" / assignment.worker_id + ).mkdir(parents=True) + + def run_worker(_config, assignment, _baseline): + if assignment.worker_id != "worker_0": + raise RuntimeError("worker failed") + return {"worker_id": assignment.worker_id, "shapes": {}} + + monkeypatch.setattr(pipeline, "_run_worker", run_worker) + monkeypatch.setattr( + pipeline, + "_author_worker_skill", + lambda _config, assignment: { + "name": f"{assignment.worker_id}-skill" + }, + ) + + workers = pipeline._parallel_agents(config, {}) + + assert set(workers) == {"worker_0"} + assert set(pipeline._worker_failures) == { + "worker_1", "worker_2", "worker_3", + } + + +def test_one_successful_lane_is_enough_for_main_synthesis( + tmp_path, monkeypatch +): + pipeline = _make_pipeline(_four_worker_req(), tmp_path) + config = load_config(_four_worker_req()) + optimized = [] + + def bootstrap(lane_config): + assignment = lane_config.assignments[0] + if assignment.worker_id != "worker_0": + raise RuntimeError("bootstrap failed") + return {"s0": {"passed": True, "median_us": 1.0}} + + def optimize(lane_config, baseline): + assignment = lane_config.assignments[0] + optimized.append((assignment.worker_id, baseline)) + return { + assignment.worker_id: { + "worker_id": assignment.worker_id, + "shapes": {}, + } + } + + monkeypatch.setattr(pipeline, "_bootstrap_worker_repos", bootstrap) + monkeypatch.setattr(pipeline, "_parallel_agents", optimize) + + baseline, workers = pipeline._parallel_lane_lifecycles(config) + + assert optimized == [ + ("worker_0", {"s0": {"passed": True, "median_us": 1.0}}) + ] + assert baseline == {"s0": {"passed": True, "median_us": 1.0}} + assert set(workers) == {"worker_0"} + + +def test_iteration_agent_timeout_is_recorded_and_next_round_runs( + tmp_path, monkeypatch +): + from ..orchestrator import w8a8_pipeline as pipeline_module + + req = _gen_req(""" +assignment_mode: manual +shapes: + - {id: m2, M: 2, N: 1536, K: 4096} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["mock_iterations"] = "2" + config = load_config(req) + pipeline = _make_pipeline(req, tmp_path) + pipeline.store.init_or_resume("iteration-timeout-test") + root = pipeline.workspace_dir / "workers" / "worker_0" + source = root / "source" + source.joinpath("csrc").mkdir(parents=True) + root.joinpath("logs").mkdir(parents=True) + source.joinpath("int8_w8a8_gemm_api.py").write_text( + "# immutable\n", encoding="utf-8" + ) + source.joinpath("csrc", "w8a8_gemm_hip.hip").write_text( + "// best kernel\n", encoding="utf-8" + ) + subprocess.run( + ["git", "init"], cwd=source, check=True, capture_output=True + ) + subprocess.run( + ["git", "config", "user.name", "test"], cwd=source, check=True + ) + subprocess.run( + ["git", "config", "user.email", "test@example.com"], + cwd=source, + check=True, + ) + subprocess.run(["git", "add", "."], cwd=source, check=True) + subprocess.run( + ["git", "commit", "-m", "best"], + cwd=source, + check=True, + capture_output=True, + ) + + class FakeRunner: + def __init__(self, worker_root, gpu): + self.env = {} + object_path = ( + worker_root / "cache" / "torch" + / "metainfer_w8a8_backend" + / "w8a8_gemm_hip.cuda.o" + ) + object_path.parent.mkdir(parents=True, exist_ok=True) + object_path.write_bytes(b"fake accepted object") + + def probe(self): + return {"visible_devices": 1} + + launched = [] + monkeypatch.setattr(pipeline_module, "W8A8Runner", FakeRunner) + monkeypatch.setattr( + pipeline.manager, "launch", lambda spec: launched.append(spec.name) + ) + monkeypatch.setattr( + pipeline.manager, + "result", + lambda name: SimpleNamespace( + success=False, + error="killed after timeout", + session_id=None, + ), + ) + + baseline = {"m2": {"passed": True, "median_us": 10.0}} + result = pipeline._run_worker( + config, config.assignments[0], baseline + ) + + records = [ + json.loads(line) + for line in ( + root / "runs" / "m2" / "experiments.jsonl" + ).read_text(encoding="utf-8").splitlines() + ] + assert launched == [ + f"worker_0-m2-iter{iteration}" for iteration in range(1, 8) + ] + assert [item["iteration"] for item in records] == list(range(1, 8)) + assert all(not item["accepted"] for item in records) + assert all( + "killed after timeout" in item["failure_reason"] + for item in records + ) + assert result["shapes"]["m2"]["metrics"] == baseline["m2"] + assert source.joinpath("csrc", "w8a8_gemm_hip.hip").read_text( + encoding="utf-8" + ) == "// best kernel\n" + + +def test_control_plane_artifacts_are_not_attributed_to_agent(): + assert _is_control_plane_artifact( + "__pycache__/w8a8_backend.cpython-310.pyc" + ) + assert _is_control_plane_artifact("csrc/bindings_hip.cpp") + assert not _is_control_plane_artifact("csrc/w8a8_gemm_hip.hip") + + +@pytest.mark.skipif(not TORCH_AVAILABLE, reason="PyTorch is not installed") +def test_trusted_harness_pytorch_reference_self_test(): + harness = ( + Path(__file__).resolve().parent.parent + / "assets" / "w8a8_bench.py" + ) + result = subprocess.run( + ["python3", str(harness), "--self-test"], + check=True, + capture_output=True, + text=True, + ) + evidence = json.loads(result.stdout.strip().splitlines()[-1]) + assert evidence["self_test"] == "exact_w8a8_reference" + assert evidence["passed"] is True + assert evidence["actual"] == evidence["expected"] + + +def test_generate_stages_scaffold_and_assignment_without_hip( + tmp_path, monkeypatch +): + from ..orchestrator import gen_and_opt_pipeline as pipeline_module + + kernel_root = tmp_path / "kernel repos" + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(kernel_root)) + req = _gen_req(""" +assignment_mode: manual +shapes: + - {id: m2, M: 2, N: 1536, K: 4096} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["target_repo_path"] = "int8 task with spaces" + pipeline = _make_pipeline(req, tmp_path) + config = load_config(req) + pipeline._validate_contract(config) + monkeypatch.setattr( + pipeline_module, + "_validate_generate_scaffold", + lambda *args, **kwargs: { + "status": "passed", + "implementation_present": False, + }, + ) + launched_specs = [] + + def launch(spec): + launched_specs.append(spec) + spec.workdir.joinpath("proposal.json").write_text( + json.dumps({ + "gpu_assignment": { + "worker_0": {"gpu": 0, "shapes": ["m2"]} + }, + "scaffold_review": { + "preflight_file": "generation_preflight.json", + "preflight_status": "passed", + "harness_reference_self_test_passed": True, + "gpu_probe_passed": True, + "cudagraph_available": True, + "python_graph_wrapper_staged": True, + "pmc_script_checked": True, + "no_hip_implementation": True, + }, + }), + encoding="utf-8", + ) + + monkeypatch.setattr(pipeline.manager, "launch", launch) + monkeypatch.setattr( + pipeline.manager, + "result", + lambda name: SimpleNamespace(success=True, error=None), + ) + + pipeline._prepare_worktrees(config, "gen-test") + assignments = pipeline._generate_kernel_repo(config) + + repo = config.target_repo_path + assert repo is not None + assert assignments == config.assignments + assert launched_specs[0].role == "kernel_coordinator" + assert repo.joinpath("profile_pmc.sh").is_file() + assert "--profile-only" in repo.joinpath( + "profile_pmc.sh" + ).read_text(encoding="utf-8") + assert repo.joinpath("int8_w8a8_gemm_api.py").is_file() + assert repo.joinpath("w8a8_bench.py").is_file() + assert "--reference-cache-dir" in repo.joinpath( + "w8a8_bench.py" + ).read_text(encoding="utf-8") + assert repo.joinpath("w8a8_graph.py").is_file() + assert not repo.joinpath("csrc", "w8a8_gemm_hip.hip").exists() + assert repo.joinpath("generation_preflight.json").is_file() + assert repo.joinpath("generation_review.json").is_file() + manifest = json.loads( + repo.joinpath("scaffold_manifest.json").read_text( + encoding="utf-8" + ) + ) + assert ( + manifest["initial_kernel"] + == "pending_parallel_explore_child_generation" + ) + timeline = pipeline.store.load_timeline() + success = [ + item for item in timeline + if item.get("type") == "generate_success" + ][-1] + assert success["payload"]["role"] == "main_coordinator" + assert success["payload"]["kernel_source_created"] is False + assert ( + success["payload"]["kernel_source"] + == "pending_child_generation" + ) + + +def test_parallel_explore_child_generates_initial_hip( + tmp_path, monkeypatch +): + from ..orchestrator import gen_and_opt_pipeline as pipeline_module + + kernel_root = tmp_path / "kernel repos" + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(kernel_root)) + req = _gen_req(""" +assignment_mode: manual +shapes: + - {id: m2, M: 2, N: 1536, K: 4096} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["target_repo_path"] = "int8 task with spaces" + pipeline = _make_pipeline(req, tmp_path) + config = load_config(req) + pipeline._validate_contract(config) + pipeline._prepare_worktrees(config, "gen-test") + monkeypatch.setattr( + pipeline_module, + "_validate_generate_scaffold", + lambda *args, **kwargs: { + "status": "passed", + "implementation_present": False, + }, + ) + + def coordinate(spec): + spec.workdir.joinpath("proposal.json").write_text( + json.dumps({ + "gpu_assignment": { + "worker_0": {"gpu": 0, "shapes": ["m2"]} + }, + "scaffold_review": { + "preflight_file": "generation_preflight.json", + "preflight_status": "passed", + "harness_reference_self_test_passed": True, + "gpu_probe_passed": True, + "cudagraph_available": True, + "python_graph_wrapper_staged": True, + "pmc_script_checked": True, + "no_hip_implementation": True, + }, + }), + encoding="utf-8", + ) + + monkeypatch.setattr(pipeline.manager, "launch", coordinate) + monkeypatch.setattr( + pipeline.manager, + "result", + lambda name: SimpleNamespace(success=True, error=None), + ) + assignments = pipeline._generate_kernel_repo(config) + config = replace_assignments(config, assignments) + pipeline._create_worker_worktrees(config, "gen-test") + + class FakeRunner: + def __init__(self, root, gpu): + self.env = {} + object_path = ( + root / "cache" / "torch" + / "metainfer_w8a8_backend" + / "w8a8_gemm_hip.cuda.o" + ) + object_path.parent.mkdir(parents=True, exist_ok=True) + object_path.write_bytes(b"fake accepted object") + + def benchmark(self, params): + return { + "passed": True, + "graph_capture_passed": True, + "timing_mode": "cuda_graph_replay", + "median_us": 1.0, + "p90_us": 1.1, + } + + launched_specs = [] + + def launch(spec): + launched_specs.append(spec) + spec.workdir.joinpath("csrc").mkdir(exist_ok=True) + spec.workdir.joinpath( + "csrc", "w8a8_gemm_hip.hip" + ).write_text( + "// generated by child Agent\n", encoding="utf-8" + ) + spec.workdir.joinpath("proposal.json").write_text( + '{"hypothesis":"fresh child kernel"}', + encoding="utf-8", + ) + + monkeypatch.setattr(pipeline_module, "W8A8Runner", FakeRunner) + monkeypatch.setattr(pipeline.manager, "launch", launch) + monkeypatch.setattr( + pipeline.manager, + "result", + lambda name: SimpleNamespace(success=True, error=None), + ) + + metrics = pipeline._bootstrap_worker_repos(config) + + repo = config.target_repo_path + assert repo is not None + worker_source = ( + pipeline.workspace_dir / "workers" / "worker_0" / "source" + ) + assert metrics["m2"]["baseline_kind"] == "triton_graph" + assert metrics["m2"]["median_us"] == 66.924 + assert metrics["m2"]["bootstrap_metrics"]["passed"] is True + assert launched_specs[0].role == "dcu_w8a8_bootstrap_generator" + assert launched_specs[0].workdir == worker_source + assert not any( + char.isspace() for char in str(worker_source.resolve()) + ) + assert not repo.joinpath("csrc", "w8a8_gemm_hip.hip").exists() + assert worker_source.joinpath( + "csrc", "w8a8_gemm_hip.hip" + ).read_text(encoding="utf-8") == "// generated by child Agent\n" + assert b"\r\n" not in worker_source.joinpath( + "profile_pmc.sh" + ).read_bytes() + result = json.loads( + ( + pipeline.workspace_dir / "workers" / "worker_0" + / "bootstrap_result.json" + ).read_text(encoding="utf-8") + ) + assert result["source"] == "child_agent_generated" + assert result["generated_files"] == ["csrc/w8a8_gemm_hip.hip"] + assert ( + result["comparison_baselines"]["m2"]["median_us"] == 66.924 + ) + + +def test_fresh_hip_generation_completes_before_parallel_explore( + tmp_path, monkeypatch +): + from ..orchestrator import gen_and_opt_pipeline as pipeline_module + + req = _gen_req(""" +assignment_mode: manual +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + config = load_config(req) + pipeline = _make_pipeline(req, tmp_path) + observed_phases = [] + synthesis_phases = [] + + monkeypatch.setattr(pipeline_module, "load_config", lambda request: config) + monkeypatch.setattr( + pipeline_module, + "replace_assignments", + lambda current, assignments: current, + ) + monkeypatch.setattr(pipeline, "_validate_contract", lambda current: None) + monkeypatch.setattr( + pipeline, "_prepare_worktrees", lambda current, task_id: None + ) + monkeypatch.setattr(pipeline, "_plan", lambda current: {}) + def generate(current): + observed_phases.append(("generated", pipeline._current_phase)) + return current.assignments + + monkeypatch.setattr(pipeline, "_generate_kernel_repo", generate) + monkeypatch.setattr( + pipeline, + "_create_worker_worktrees", + lambda current, task_id: None, + ) + + def lane_lifecycles(current): + observed_phases.append(("explored", pipeline._current_phase)) + return ( + {"m2": {"passed": True, "median_us": 1.0}}, + {"worker_0": {"worker_id": "worker_0", "shapes": {}}}, + ) + + monkeypatch.setattr( + pipeline, + "_parallel_lane_lifecycles", + lane_lifecycles, + ) + def synthesize_final(current, workers, baseline, task_id): + synthesis_phases.append(("final_candidate", pipeline._current_phase)) + return {"validation": {}} + + def author_skill(current, assignments): + synthesis_phases.append(("merged_skill", pipeline._current_phase)) + return {"name": "merged"} + + monkeypatch.setattr( + pipeline, "_synthesize_final_candidate", synthesize_final + ) + monkeypatch.setattr(pipeline, "_author_merged_skill", author_skill) + + pipeline.run() + + assert observed_phases == [ + ("generated", phases.GENERATE), + ("explored", phases.EXPLORE), + ] + assert synthesis_phases == [ + ("merged_skill", phases.SYNTHESIZE), + ("final_candidate", phases.VALIDATE), + ] + + +# --- Config parsing tests ----------------------------------------------- # + +def test_gen_mode_config_parses(): + """Blank repo name deterministically falls back to the task id.""" + cfg = load_config(_gen_req()) + assert cfg.execution_mode == GEN_AND_OPT_MODE + assert cfg.target_repo_path is not None + assert cfg.target_repo_path.name == "gen-test" + assert cfg.operator == "Quantized GEMM" + assert cfg.dtype == "INT8 W8A8" + assert set(cfg.shapes) == {"m2", "m16"} + assert [a.gpu for a in cfg.assignments] == [0, 1] + + +def test_gen_mode_shapes_only_auto_assigns(): + """When no assignments given in generate mode, auto-assign to worker_0.""" + cfg = load_config(_shapes_only_req()) + assert cfg.execution_mode == GEN_AND_OPT_MODE + assert set(cfg.shapes) == {"m2", "m16", "m64"} + # Auto-assigned: all shapes → worker_0, gpu 0 + assert len(cfg.assignments) == 1 + assert cfg.assignments[0].worker_id == "worker_0" + assert cfg.assignments[0].gpu == 0 + assert set(cfg.assignments[0].shape_ids) == {"m2", "m16", "m64"} + + +def test_gen_mode_requires_quantized_gemm_op(tmp_path): + """Generate mode validates operator=Quantized GEMM in contract check.""" + req = _gen_req() + req["answers"]["operator"] = "Custom operator" + pipeline = _make_pipeline(req, tmp_path) + cfg = load_config(req) + with pytest.raises(ValueError, match="Generate mode requires operator"): + pipeline._validate_contract(cfg) + + +def test_gen_mode_requires_w8a8_dtype(tmp_path): + """Generate mode validates dtype=INT8 W8A8 in contract check.""" + req = _gen_req() + req["answers"]["dtype"] = "FP16 / BF16" + pipeline = _make_pipeline(req, tmp_path) + cfg = load_config(req) + with pytest.raises(ValueError, match="Generate mode requires dtype"): + pipeline._validate_contract(cfg) + + +def test_gen_mode_shapes_need_mnk(tmp_path): + """Every shape must have M, N, K dimensions.""" + req = _gen_req(""" +shapes: + - {id: bad, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [bad]} +""") + pipeline = _make_pipeline(req, tmp_path) + cfg = load_config(req) + with pytest.raises(ValueError, match="is missing"): + pipeline._validate_contract(cfg) + + +def test_gen_mode_resolves_kernel_repo_name_with_spaces( + tmp_path, monkeypatch +): + """The user-facing name resolves below sibling kernel-repos.""" + kernel_root = tmp_path / "kernel-repos" + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(kernel_root)) + req = _gen_req() + req["answers"]["target_repo_path"] = "int8 test2" + cfg = load_config(req) + assert cfg.target_repo_path == kernel_root / "int8 test2" + assert not cfg.target_repo_path.exists() + + +@pytest.mark.parametrize( + "bad_name", ["../escape", "nested/repo", "/tmp/absolute", ".", ".."] +) +def test_gen_mode_rejects_unsafe_kernel_repo_name(bad_name): + req = _gen_req() + req["answers"]["target_repo_path"] = bad_name + with pytest.raises(ValueError, match="single folder name"): + load_config(req) + + +# --- GPU assignment validation ------------------------------------------ # + +def test_validate_gpu_assignment_accepts_valid(): + """Valid assignment dict parses correctly.""" + shapes = { + "m2": ShapeSpec("m2", {"M": 2, "N": 1536, "K": 4096}), + "m16": ShapeSpec("m16", {"M": 16, "N": 1536, "K": 4096}), + "m64": ShapeSpec("m64", {"M": 64, "N": 1536, "K": 4096}), + } + raw = { + "worker_0": {"gpu": 0, "shapes": ["m2", "m16"]}, + "worker_1": {"gpu": 1, "shapes": ["m64"]}, + } + result = validate_gpu_assignment(shapes, raw) + assert len(result) == 2 + assert result[0].worker_id == "worker_0" + assert result[0].gpu == 0 + assert set(result[0].shape_ids) == {"m2", "m16"} + assert result[1].worker_id == "worker_1" + assert result[1].gpu == 1 + assert result[1].shape_ids == ["m64"] + + +def test_validate_gpu_assignment_rejects_missing_shapes(): + """All shapes must be assigned.""" + shapes = {"m2": ShapeSpec("m2", {"M": 2}), "m16": ShapeSpec("m16", {"M": 16})} + raw = {"worker_0": {"gpu": 0, "shapes": ["m2"]}} + with pytest.raises(ValueError, match="unassigned shapes"): + validate_gpu_assignment(shapes, raw) + + +def test_validate_gpu_assignment_rejects_duplicate_gpu(): + """Same GPU can't be used by two workers.""" + shapes = {"m2": ShapeSpec("m2", {"M": 2}), "m16": ShapeSpec("m16", {"M": 16})} + raw = { + "worker_0": {"gpu": 0, "shapes": ["m2"]}, + "worker_1": {"gpu": 0, "shapes": ["m16"]}, + } + with pytest.raises(ValueError, match="assigned to more than one"): + validate_gpu_assignment(shapes, raw) + + +def test_validate_gpu_assignment_rejects_duplicate_shapes(): + """Same shape can't be assigned twice.""" + shapes = {"m2": ShapeSpec("m2", {"M": 2}), "m16": ShapeSpec("m16", {"M": 16})} + raw = { + "worker_0": {"gpu": 0, "shapes": ["m2", "m16"]}, + "worker_1": {"gpu": 1, "shapes": ["m2"]}, + } + with pytest.raises(ValueError, match="assigned more than once"): + validate_gpu_assignment(shapes, raw) + + +def test_validate_gpu_assignment_rejects_empty(): + """Empty assignment dict raises.""" + with pytest.raises(ValueError, match="non-empty"): + validate_gpu_assignment({"m2": ShapeSpec("m2", {"M": 2})}, {}) + + +# --- replace_assignments ----------------------------------------------- # + +def test_replace_assignments_preserves_shapes(): + """replace_assignments only changes the assignment list.""" + cfg = load_config(_gen_req()) + new_assignments = [ + WorkerAssignment("worker_0", 0, ["m2", "m16"]), + ] + new_cfg = replace_assignments(cfg, new_assignments) + assert new_cfg.shapes == cfg.shapes + assert new_cfg.operator == cfg.operator + assert new_cfg.dtype == cfg.dtype + assert new_cfg.assignment_mode == cfg.assignment_mode + assert len(new_cfg.assignments) == 1 + assert set(new_cfg.assignments[0].shape_ids) == {"m2", "m16"} + + +# --- Phase ordering ---------------------------------------------------- # + +def test_generate_phase_is_in_order(): + """Generate mode transitions directly from GENERATE to EXPLORE.""" + graph = phases.graph_payload(phases.PREPARE, include_baseline=False) + node_ids = [n["id"] for n in graph["nodes"]] + gen_idx = node_ids.index(phases.GENERATE) + prep_idx = node_ids.index(phases.PREPARE) + explore_idx = node_ids.index(phases.EXPLORE) + assert phases.BASELINE not in node_ids + assert prep_idx < gen_idx < explore_idx + assert explore_idx == gen_idx + 1 + + +def test_generate_accepts_sparse_one_to_four_worker_gpu_mapping(): + assignments = [ + WorkerAssignment(f"worker_{gpu}", gpu, [f"shape_{gpu}"]) + for gpu in range(4) + ] + _require_valid_child_assignments(assignments) + _require_valid_child_assignments(assignments[:1]) + _require_valid_child_assignments(assignments[:3]) + _require_valid_child_assignments(assignments[1:3]) + with pytest.raises(ValueError, match="map each worker_N"): + _require_valid_child_assignments([ + WorkerAssignment("worker_0", 2, ["shape_2"]) + ]) + with pytest.raises(ValueError, match="between one and four"): + _require_valid_child_assignments([]) + + +def test_generate_accepts_m_variants_on_different_workers(): + shapes = { + "m1_wqkv_a": ShapeSpec( + "m1_wqkv_a", {"M": 1, "N": 1536, "K": 4096} + ), + "m4_wqkv_a": ShapeSpec( + "m4_wqkv_a", {"M": 4, "N": 1536, "K": 4096} + ), + "m16_wqkv_a": ShapeSpec( + "m16_wqkv_a", {"M": 16, "N": 1536, "K": 4096} + ), + } + raw = { + "worker_0": { + "gpu": 0, "shapes": ["m1_wqkv_a", "m4_wqkv_a"] + }, + "worker_2": {"gpu": 2, "shapes": ["m16_wqkv_a"]}, + } + + assignments = validate_gpu_assignment(shapes, raw) + _require_valid_child_assignments(assignments) + + assert assignments[0].shape_ids == ["m1_wqkv_a", "m4_wqkv_a"] + assert assignments[1].shape_ids == ["m16_wqkv_a"] + + +def test_generate_phase_has_label(): + """GENERATE phase has a human-readable label.""" + graph = phases.graph_payload(phases.GENERATE) + labels = {n["id"]: n["label"] for n in graph["nodes"]} + assert labels[phases.GENERATE] == "Generate kernel repo" + + +def test_worker_cache_dirs_exist_before_owner_match(tmp_path, monkeypatch): + """Host-agent cache directories must be included in the ownership pass.""" + from ..orchestrator import gen_and_opt_pipeline as pipeline_module + + pipeline = _make_pipeline(_gen_req(), tmp_path) + (pipeline.workspace_dir / "main").mkdir(parents=True) + config = load_config(_gen_req()) + observed = [] + + def fake_run(command, *, cwd, **kwargs): + if command[:3] == ["git", "worktree", "add"]: + # The source path is the penultimate argument. + Path(command[-2]).mkdir(parents=True) + + class Result: + stdout = "" + + return Result() + + def fake_match_tree_owner(path, owner_source): + observed.append(path) + for name in ("torch", "triton", "xdg", "tmp"): + assert (path / "cache" / name).is_dir() + + monkeypatch.setattr(pipeline_module, "_run", fake_run) + monkeypatch.setattr( + pipeline_module, "_match_tree_owner", fake_match_tree_owner + ) + + pipeline._create_worker_worktrees(config, "gen-test") + + assert len(observed) == len(config.assignments) + for assignment in config.assignments: + candidate = ( + pipeline.workspace_dir / "main" / "candidates" + / assignment.worker_id + ) + assert candidate.is_dir() + assert not candidate.is_symlink() + assert (candidate / "source").is_symlink() + assert (candidate / "source").resolve() == ( + pipeline.workspace_dir / "workers" + / assignment.worker_id / "source" + ).resolve() + assert (candidate / "csrc").is_symlink() + assert "candidates/" in ( + pipeline.workspace_dir / "main" / ".git" / "info" / "exclude" + ).read_text(encoding="utf-8").splitlines() + + +# --- Prompt content ---------------------------------------------------- # + +def test_prompt_renders_shapes_table(): + """generate_kernel_prompt includes shape dims in a readable table.""" + from ..orchestrator.prompts import generate_kernel_prompt + from pathlib import Path + + prompt = generate_kernel_prompt( + operator="Quantized GEMM", + dtype="INT8 W8A8", + shapes={"m2": {"M": 2, "N": 1536, "K": 4096}}, + hardware="K500SM_AI / gfx928", + kernel_language="HIP C++", + source_dir=Path("/tmp/test"), + harness_path=Path("/tmp/harness.py"), + ) + assert "m2" in prompt + assert "1536" in prompt + assert "4096" in prompt + assert "w8a8_gemm_out" in prompt + assert "zth_w8a8.gemm_out" in prompt + assert "immutable" in prompt.lower() + assert "gfx928" in prompt + assert "main coordinator" in prompt.lower() + assert "do not create `.hip`" in prompt.lower() + assert "do not compile" in prompt.lower() + assert "generation_preflight.json" in prompt + assert "w8a8_bench.py" in prompt + assert "pytorch-reference self-test" in prompt.lower() + assert "parallel explore" in prompt.lower() + assert "child implementation agents" in prompt.lower() + assert "initial hip kernels from scratch" in prompt.lower() + assert "simple scalar" not in prompt.lower() + + +def test_prompt_includes_prev_failure(): + """When prev_failure is set, it appears in the prompt.""" + from ..orchestrator.prompts import generate_kernel_prompt + from pathlib import Path + + prompt = generate_kernel_prompt( + operator="Quantized GEMM", + dtype="INT8 W8A8", + shapes={"m2": {"M": 2, "N": 1536, "K": 4096}}, + hardware="K500SM_AI / gfx928", + kernel_language="HIP C++", + source_dir=Path("/tmp/test"), + harness_path=Path("/tmp/harness.py"), + prev_failure="Harness check failed: passed was false", + ) + assert "Previous attempt failed" in prompt + assert "Harness check failed" in prompt + + +def test_prompt_includes_gpu_assignment_instructions(): + """The prompt makes exact-shape assignment authoritative.""" + from ..orchestrator.prompts import generate_kernel_prompt + from pathlib import Path + + prompt = generate_kernel_prompt( + operator="Quantized GEMM", + dtype="INT8 W8A8", + shapes={"m2": {"M": 2, "N": 1536, "K": 4096}}, + hardware="K500SM_AI / gfx928", + kernel_language="HIP C++", + source_dir=Path("/tmp/test"), + harness_path=Path("/tmp/harness.py"), + fixed_assignment={ + "worker_0": {"gpu": 0, "shapes": ["m2"]} + }, + ) + assert "GPU assignment" in prompt + assert "gpu_assignment" in prompt + assert "worker_0" in prompt + assert "preserve it exactly" in prompt.lower() + assert "different m variants" in prompt.lower() + assert "may be assigned to different workers" in prompt.lower() + + +def test_control_plane_assignment_balances_exact_shapes(): + from ..orchestrator.prompts import shape_balanced_assignment + + shapes = { + "m2_wqkv_a": {"M": 2, "N": 1536, "K": 4096}, + "m16_wqkv_a": {"M": 16, "N": 1536, "K": 4096}, + "m2_wq_b": {"M": 2, "N": 8192, "K": 1024}, + "m16_wq_b": {"M": 16, "N": 8192, "K": 1024}, + "m2_wo_b": {"M": 2, "N": 4096, "K": 2048}, + "m16_wo_b": {"M": 16, "N": 4096, "K": 2048}, + "m2_shared_gate_up": {"M": 2, "N": 1024, "K": 4096}, + "m16_shared_gate_up": {"M": 16, "N": 1024, "K": 4096}, + "m2_shared_down": {"M": 2, "N": 4096, "K": 512}, + "m16_shared_down": {"M": 16, "N": 4096, "K": 512}, + } + assignment = shape_balanced_assignment(shapes) + owners = { + shape: worker + for worker, payload in assignment.items() + for shape in payload["shapes"] + } + assert set(owners) == set(shapes) + loads = [] + for payload in assignment.values(): + loads.append(sum( + 2 * shapes[shape]["M"] * shapes[shape]["N"] * shapes[shape]["K"] + for shape in payload["shapes"] + )) + assert max(loads) / min(loads) < 1.5 + + +def test_prompt_assignment_example_omits_empty_subset_workers(): + from ..orchestrator.prompts import shape_balanced_assignment + + assignment = shape_balanced_assignment({ + "m2_wqkv_a": {"M": 2, "N": 1536, "K": 4096}, + }) + + assert assignment == { + "worker_0": {"gpu": 0, "shapes": ["m2_wqkv_a"]} + } + + +def test_parallel_child_bootstrap_prompt_owns_hip_implementation(): + """Only the Parallel explore child owns fresh HIP source.""" + from ..orchestrator.prompts import bootstrap_worker_prompt + from pathlib import Path + + prompt = bootstrap_worker_prompt( + worker_id="worker_2", + gpu=2, + shapes={"m2": {"M": 2, "N": 1536, "K": 4096}}, + hardware="K500SM_AI / gfx928", + kernel_language="HIP C++", + source_dir=Path("/tmp/worker/source"), + harness_path=Path("/tmp/harness.py"), + api_contract_path=Path("/tmp/worker/source/int8_w8a8_gemm_api.py"), + attempt=1, + ) + assert "child kernel implementation agent" in prompt.lower() + assert "parallel explore" in prompt.lower() + assert "assigned shapes" in prompt.lower() + assert "do not copy another task repo" in prompt.lower() + assert "generate kernel repo phase" not in prompt.lower() + assert "csrc/w8a8_gemm_hip.hip" in prompt + assert "w8a8_backend.py" in prompt + assert "w8a8_graph.py" in prompt + assert "torch.cuda.CUDAGraph" in prompt + assert "non-default stream" in prompt + assert "Correctness-first bootstrap strategy" in prompt + assert "Do not run the benchmark or correctness harness" in prompt + assert "trusted control plane" in prompt + assert "blockDim 128 or 256" in prompt + assert "preserve it byte-for-byte" in prompt + assert "load_extension()" in prompt + assert "is_python_module=False" in prompt + assert "You may change only" in prompt + assert "`csrc/w8a8_gemm_hip.hip` plus `proposal.json`" in prompt + assert "launch_pack_w8a8_weight" in prompt + assert "void* workspace" in prompt + assert "main coordinator and\nGenerate phase never write the HIP" in prompt + + +def test_parallel_child_initial_prompt_uses_scalar_correctness_for_m16(): + from ..orchestrator.prompts import bootstrap_worker_prompt + + prompt = bootstrap_worker_prompt( + worker_id="worker_3", + gpu=3, + shapes={"m16": {"M": 16, "N": 8192, "K": 1024}}, + hardware="K500SM_AI / gfx928", + kernel_language="HIP C++", + source_dir=Path("/tmp/worker/source"), + harness_path=Path("/tmp/harness.py"), + api_contract_path=Path("/tmp/worker/source/int8_w8a8_gemm_api.py"), + attempt=1, + ) + + assert "For every assigned shape,\nincluding M=16" in prompt + assert "Do not use DUMMA" in prompt + assert '"path": "scalar_correctness"' in prompt + assert "wavefront size: 64" in prompt + assert "Never use NVIDIA" in prompt + assert "validation_owner" in prompt + assert "scalar generic fallback" in prompt + assert "paired M=2 API shape" in prompt + + +def test_parallel_child_large_prefill_bootstrap_starts_with_dumma(): + from ..orchestrator.prompts import bootstrap_worker_prompt + + prompt = bootstrap_worker_prompt( + worker_id="worker_0", + gpu=0, + shapes={"prefill": {"M": 3072, "N": 1536, "K": 4096}}, + hardware="K500SM_AI / gfx928", + kernel_language="HIP C++", + source_dir=Path("/tmp/worker/source"), + harness_path=Path("/tmp/harness.py"), + api_contract_path=Path( + "/tmp/worker/source/int8_w8a8_gemm_api.py" + ), + attempt=1, + ) + + assert "large-Prefill\nscalar K loop is not a usable" in prompt + assert "INT8 DUMMA m16n16k32" in prompt + assert "references/int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip" in prompt + assert "references/w8a8_gemm_variants.hip" in prompt + assert "neither a whitelist nor a restriction" in prompt + assert '"path": "dumma_prefill_with_scalar_fallback"' in prompt + assert "Keep one simple scalar int8/int32 fallback" in prompt + + +def test_optimization_prompt_enforces_one_evidence_driven_change(): + prompt = RealW8A8OptimizationPipeline._worker_prompt( + WorkerAssignment("worker_0", 0, ["m2"]), + "m2", + {"M": 2, "N": 1536, "K": 4096}, + {"median_us": 12.0, "p90_us": 13.0}, + Path("/tmp/worker"), + 2, + None, + ) + + assert "one falsifiable bottleneck hypothesis" in prompt + assert "already rejected change" in prompt + assert "Do not choose split-K merely because K is large" in prompt + assert "torch.cuda.CUDAGraph.replay()" in prompt + assert "Do not run the harness" in prompt + assert "namespace `du::dumma`" in prompt + assert "filesystem-wide searches" in prompt + assert "hygon-gfx928-memory-isa" not in prompt + assert "hygon-gfx928-compute-isa" not in prompt + assert "Skill tool is disabled" in prompt + assert '"isa_optimization"' not in prompt + assert '"observed_best": {"median_us": 12.0' in prompt + assert "Mandatory decision for this round" in prompt + assert "avoid per-K-tile LDS barriers" in prompt + + +def test_optimization_prompt_repairs_faster_incorrect_candidate(): + prompt = RealW8A8OptimizationPipeline._worker_prompt( + WorkerAssignment("worker_0", 0, ["m2"]), + "m2", + {"M": 2, "N": 1536, "K": 4096}, + {"median_us": 129.0}, + Path("/tmp/worker"), + 2, + None, + [{ + "iteration": 1, + "correctness_passed": False, + "build_success": True, + "speedup": 1.77, + "artifact_dir": "iterations/m2/iteration1", + }], + ) + + assert "repair the faster but incorrect candidate" in prompt + assert "iterations/m2/iteration1" in prompt + assert "Do not replace it with an unrelated architecture" in prompt + + +def test_m16_prompt_uses_exact_installed_dumma_api(): + prompt = RealW8A8OptimizationPipeline._worker_prompt( + WorkerAssignment("worker_0", 0, ["m16"]), + "m16", + {"M": 16, "N": 1536, "K": 4096}, + {"median_us": 200.0}, + Path("/tmp/worker"), + 1, + None, + ) + + assert "du::dumma::du_load_matrix_sync" in prompt + assert "du::dumma::du_mma_sync" in prompt + assert "du::dumma::mem_row_major" in prompt + + +def test_final_synthesis_prompt_requires_one_all_shape_extension(): + prompt = _final_synthesis_prompt( + worker_inputs=[ + { + "worker_id": "worker_0", + "shapes": ["m2_wo_b", "m16_wo_b"], + "source": "/tmp/worker_0/source", + "results": {}, + } + ], + shapes=[ + {"id": "m2_wo_b", "M": 2, "N": 4096, "K": 2048}, + {"id": "m16_wo_b", "M": 16, "N": 4096, "K": 2048}, + ], + source=Path("/tmp/final/source"), + proposal_path=Path("/tmp/final/source/proposal.json"), + previous_failure="correctness failed for m16_wo_b", + fallback_shapes=[ + {"id": "m2_wq_b", "M": 2, "N": 8192, "K": 1024} + ], + ) + assert "one compiled extension" in prompt + assert "generic scalar path as a correctness fallback" in prompt + assert "edit only `csrc/w8a8_gemm_hip.hip`" in prompt + assert "within 5%" in prompt + assert "m2_wo_b" in prompt and "m16_wo_b" in prompt + assert "outside this task's optimization scope" in prompt + assert "m2_wq_b" in prompt + assert "Previous trusted synthesis failure" in prompt + + +def test_prebuilt_dispatch_routes_shape_family_without_hip_code(): + dispatch = _render_prebuilt_dispatch([ + { + "shape_id": "m16_wq_b", + "shape": {"M": 16, "N": 8192, "K": 1024}, + "launch_symbol": "mi_m16_wq_b_launch_w8a8_gemm", + "pack_symbol": "mi_m16_wq_b_launch_pack_w8a8_weight", + }, + { + "shape_id": "m16_wo_b", + "shape": {"M": 16, "N": 4096, "K": 2048}, + "launch_symbol": "mi_m16_wo_b_launch_w8a8_gemm", + "pack_symbol": "mi_m16_wo_b_launch_pack_w8a8_weight", + }, + ]) + + assert "if (n == 8192 && k == 1024)" in dispatch + assert "if (n == 4096 && k == 2048)" in dispatch + assert "mi_m16_wq_b_launch_w8a8_gemm" in dispatch + assert "launch_w8a8_gemm(" in dispatch + assert "__global__" not in dispatch + assert "hipLaunchKernelGGL" not in dispatch + + +def test_artifact_symbol_prefix_is_valid_c_identifier(): + assert _artifact_symbol_prefix("m16-wq/b") == "mi_m16_wq_b_" + + +def test_snapshot_accepted_kernel_artifact_keeps_exact_object(tmp_path): + worker = tmp_path / "worker" + source = worker / "source" / "csrc" + cache = ( + worker / "cache" / "torch" / "metainfer_w8a8_backend" + ) + source.mkdir(parents=True) + cache.mkdir(parents=True) + source.joinpath("w8a8_gemm_hip.hip").write_bytes(b"accepted hip") + cache.joinpath("w8a8_gemm_hip.cuda.o").write_bytes(b"accepted object") + + manifest = snapshot_accepted_kernel_artifact( + worker_root=worker, + shape_id="m16_wq_b", + shape={"M": 16, "N": 8192, "K": 1024}, + metrics={ + "median_us": 30.1, + "p90_us": 30.2, + "graph_capture_passed": True, + }, + commit="abc123", + ) + + assert (worker / manifest["source"]).read_bytes() == b"accepted hip" + assert (worker / manifest["object"]).read_bytes() == b"accepted object" + assert manifest["compile_target"] == "gfx928" + assert manifest["commit"] == "abc123" + + +def test_compile_cache_is_content_addressed(tmp_path): + worker = tmp_path / "worker" + csrc = worker / "source" / "csrc" + csrc.mkdir(parents=True) + csrc.joinpath("bindings.cpp").write_text("// binding\n") + hip = csrc / "w8a8_gemm_hip.hip" + hip.write_text("// candidate one\n") + + runner = W8A8Runner(worker, 0) + first = runner._prepare_compile_cache() + repeated = runner._prepare_compile_cache() + hip.write_text("// candidate two\n") + second = runner._prepare_compile_cache() + + assert first["build_key"] == repeated["build_key"] + assert first["compile_source_dir"] == repeated["compile_source_dir"] + assert first["build_key"] != second["build_key"] + assert Path(first["compile_source_dir"]).joinpath( + "csrc/w8a8_gemm_hip.hip" + ).read_text() == "// candidate one\n" + + +def _runner_with_source(tmp_path: Path) -> tuple[W8A8Runner, Path]: + worker = tmp_path / "worker" + csrc = worker / "source" / "csrc" + csrc.mkdir(parents=True) + csrc.joinpath("bindings.cpp").write_text("// binding\n") + csrc.joinpath("w8a8_gemm_hip.hip").write_text("// kernel\n") + return W8A8Runner(worker, 0), worker + + +def _fake_run_for_reference(records: list): + """Return a _run stand-in that seeds the reference cache on demand.""" + + def fake_run(command, *, cwd, env=None, timeout=None): + records.append((list(command), timeout)) + if "--prepare-reference" in command: + def arg(name): + return command[command.index(name) + 1] + cache_dir = Path(arg("--reference-cache-dir")) + cache_dir.mkdir(parents=True, exist_ok=True) + m, n, k = int(arg("--m")), int(arg("--n")), int(arg("--k")) + (cache_dir / f"exact-int64-v1-m{m}-n{n}-k{k}.pt").write_bytes( + b"ref" + ) + stdout = ( + '{"reference_prepared": true, ' + '"reference_cache_hit": false}\n' + ) + else: + stdout = ( + '{"passed": true, "graph_capture_passed": true, ' + '"median_us": 1.0, "mismatch_count": 0}\n' + ) + return SimpleNamespace(stdout=stdout, returncode=0, stderr="") + + return fake_run + + +def test_benchmark_prepares_reference_cache_when_missing( + tmp_path, monkeypatch +): + runner, worker = _runner_with_source(tmp_path) + records: list = [] + monkeypatch.setattr(pipeline_module, "_run", _fake_run_for_reference(records)) + + metrics = runner.benchmark( + {"M": 4096, "N": 2304, "K": 6144}, check_correctness=True + ) + + # The reference is prepared first, outside the benchmark budget. + assert len(records) == 2 + prepare_command, prepare_timeout = records[0] + assert "--prepare-reference" in prepare_command + assert prepare_timeout == _REFERENCE_PREPARE_TIMEOUT_S + bench_command, bench_timeout = records[1] + assert "--prepare-reference" not in bench_command + assert bench_timeout == _BENCHMARK_TIMEOUT_S + assert runner._reference_cache_path(4096, 2304, 6144).is_file() + assert metrics["median_us"] == 1.0 + + +def test_benchmark_reuses_existing_reference_cache(tmp_path, monkeypatch): + runner, worker = _runner_with_source(tmp_path) + path = runner._reference_cache_path(4096, 2304, 6144) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"ref") + records: list = [] + monkeypatch.setattr(pipeline_module, "_run", _fake_run_for_reference(records)) + + runner.benchmark( + {"M": 4096, "N": 2304, "K": 6144}, check_correctness=True + ) + + assert len(records) == 1 # no prepare-reference call + assert "--prepare-reference" not in records[0][0] + + +def test_benchmark_skips_reference_prep_when_correctness_disabled( + tmp_path, monkeypatch +): + runner, worker = _runner_with_source(tmp_path) + records: list = [] + monkeypatch.setattr(pipeline_module, "_run", _fake_run_for_reference(records)) + + runner.benchmark( + {"M": 4096, "N": 2304, "K": 6144}, check_correctness=False + ) + + assert len(records) == 1 + command, _ = records[0] + assert "--prepare-reference" not in command + assert "--skip-correctness" in command + + +def test_reference_prep_failure_raises(tmp_path, monkeypatch): + runner, worker = _runner_with_source(tmp_path) + records: list = [] + + def failing_run(command, *, cwd, env=None, timeout=None): + records.append((list(command), timeout)) + return SimpleNamespace( + stdout='{"reference_prepared": false}\n', + returncode=0, + stderr="", + ) + + monkeypatch.setattr(pipeline_module, "_run", failing_run) + + with pytest.raises(RuntimeError, match="reference cache preparation failed"): + runner.benchmark( + {"M": 4096, "N": 2304, "K": 6144}, check_correctness=True + ) + + +class _FakeTimelineStore: + def __init__(self): + self.events = [] + + def append_timeline(self, type_: str, payload: dict): + self.events.append((type_, payload)) + + +def _perf_gate(benchmark, median_us, best=100.0, store=None, **kw): + return _final_performance_gate( + shape_id="minimax_tp8_qkv_proj_m16", + best_median=best, + metrics={"passed": True, "median_us": median_us}, + benchmark=benchmark, + max_retries=kw.pop("max_retries", 3), + retry_interval_s=kw.pop("retry_interval_s", 300), + store=store or _FakeTimelineStore(), + ) + + +def test_perf_gate_accepts_without_retry(): + store = _FakeTimelineStore() + result = _perf_gate(lambda: None, median_us=90.0, store=store) + assert result["median_us"] == 90.0 + assert store.events == [] + + +def test_perf_gate_retries_until_best_passes(monkeypatch): + sleeps: list = [] + monkeypatch.setattr("time.sleep", lambda s: sleeps.append(s)) + store = _FakeTimelineStore() + calls = {"n": 0} + + def benchmark(): + calls["n"] += 1 + return {"passed": True, "median_us": 120.0 if calls["n"] == 1 else 95.0} + + result = _perf_gate(benchmark, median_us=130.0, store=store) + # 130 (>105) -> retry1 120 (>105) -> retry2 95 (<=105) -> accept min 95 + assert calls["n"] == 2 + assert result["median_us"] == 95.0 + assert sleeps == [300, 300] + assert [t for t, _ in store.events] == [ + "final_perf_gate_retry", "final_perf_gate_retry", + ] + assert store.events[0][1]["attempt"] == 1 + assert store.events[1][1]["attempt"] == 2 + + +def test_perf_gate_fails_only_after_all_retries(monkeypatch): + sleeps: list = [] + monkeypatch.setattr("time.sleep", lambda s: sleeps.append(s)) + with pytest.raises(RuntimeError, match="after 3 re-measures"): + _perf_gate(lambda: {"passed": True, "median_us": 200.0}, median_us=200.0) + assert len(sleeps) == 3 + + +def test_perf_gate_retry_correctness_failure_raises(monkeypatch): + monkeypatch.setattr("time.sleep", lambda s: None) + with pytest.raises(RuntimeError, match="correctness failed"): + _perf_gate( + lambda: {"passed": False, "median_us": 999.0}, + median_us=200.0, + ) + + +def test_iteration_archive_keeps_kernel_but_shares_api(tmp_path): + source = tmp_path / "source" + source.joinpath("csrc").mkdir(parents=True) + source.joinpath("csrc", "w8a8_gemm_hip.hip").write_text( + "// iteration kernel", encoding="utf-8" + ) + source.joinpath("csrc", "bindings.cpp").write_text( + "// bindings", encoding="utf-8" + ) + source.joinpath("int8_w8a8_gemm_api.py").write_text( + "# immutable shared API", encoding="utf-8" + ) + destination = tmp_path / "iterations" / "m2" / "iteration1" + + archived = archive_iteration_candidate(source, destination) + + assert "csrc/w8a8_gemm_hip.hip" in archived + assert ( + destination / "csrc" / "w8a8_gemm_hip.hip" + ).read_text() == "// iteration kernel" + assert (destination / "csrc" / "bindings.cpp").is_file() + assert not (destination / "int8_w8a8_gemm_api.py").exists() + + +def test_iteration_archive_is_published_to_candidate_repo(tmp_path): + assignment = WorkerAssignment("worker_0", 0, ["m2"]) + candidate = tmp_path / "main" / "candidates" / "worker_0" + candidate.mkdir(parents=True) + archived = tmp_path / "workers" / "worker_0" / "iterations" / ( + "m2/iteration1" + ) + archived.joinpath("csrc").mkdir(parents=True) + archived.joinpath("csrc", "w8a8_gemm_hip.hip").write_text( + "// round one", encoding="utf-8" + ) + + destination = candidate_iteration_destination( + tmp_path, assignment, "m2", 1 + ) + assert destination == candidate / "iteration1" + assert destination is not None + publish_iteration_candidate(archived, destination) + + saved = destination / "csrc" / "w8a8_gemm_hip.hip" + assert saved.read_text(encoding="utf-8") == "// round one" + assert not saved.is_symlink() + + +def test_coordinator_uses_minimal_file_tool_allowlist(): + assert _COORDINATOR_AGENT_ARGS == [ + "--tools", "Read,Glob,Grep,Write", + ] + + +def test_bridge_path_translation_does_not_rewrite_workspaces_segment(): + """The /workspaces directory name must not be translated a second time.""" + from ..bridge.agent_bridge_server import ( + HOST_WORKSPACE_ROOT, + _translate_prompt_paths, + ) + + raw = ( + b"/workspace/MetaInfer/nodes/worker29/workspaces/task/workers/" + b"worker_0/source" + ) + translated = _translate_prompt_paths(raw).decode() + assert translated == ( + f"{HOST_WORKSPACE_ROOT}/MetaInfer/nodes/worker29/workspaces/" + "task/workers/worker_0/source" + ) + + +def test_bridge_accepts_source_only_agent_tool_restrictions(): + from ..bridge.agent_bridge_server import _validated_args + + args = [ + "-p", + "--setting-sources", + "project,local", + "--tools", + "Read,Glob,Grep,Write,Edit", + ] + assert _validated_args(args) == args + + +def test_bridge_enforces_source_only_tools_for_older_orchestrators(): + from ..bridge.agent_bridge_server import _validated_args + + assert _validated_args(["-p"]) == [ + "-p", + "--tools", + "Read,Glob,Grep,Write,Edit", + ] + + +def _write_test_api_contract(api_root): + contract_dir = api_root / "int8w8a8gemm" + contract_dir.mkdir(parents=True) + contract = contract_dir / "int8_w8a8_gemm_api.py" + contract.write_text( + "\n".join([ + "def prepare_weight(*args): pass", + "def allocate_workspace(*args): pass", + "def validate_gemm_out_inputs(*args): pass", + "def w8a8_gemm_out(*args): pass", + "DEFAULT_OPTIMIZATION_SHAPES = (", + " {'id': 'm2', 'M': 2, 'N': 1536, 'K': 4096},", + " {'id': 'm16', 'M': 16, 'N': 1536, 'K': 4096},", + ")", + "def _check_target_shape(m, n, k):", + " if m > 16: raise ValueError('decode M must be <= 16')", + ]), + encoding="utf-8", + ) + return contract + + +def test_fixed_api_rejects_shape_outside_decode_contract( + tmp_path, monkeypatch +): + api_root = tmp_path / "API" + _write_test_api_contract(api_root) + monkeypatch.setenv("METAINFER_OPERATOR_API_ROOT", str(api_root)) + pipeline = _make_pipeline(_shapes_only_req(), tmp_path) + cfg = load_config(_shapes_only_req()) + with pytest.raises(ValueError, match="m64.*outside"): + pipeline._validate_contract(cfg) + + +def test_prepare_uses_named_kernel_repo_and_stages_fixed_api( + tmp_path, monkeypatch +): + api_root = tmp_path / "API" + kernel_root = tmp_path / "kernel-repos" + contract = _write_test_api_contract(api_root) + monkeypatch.setenv("METAINFER_OPERATOR_API_ROOT", str(api_root)) + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(kernel_root)) + req = _gen_req() + req["answers"]["target_repo_path"] = "int8 test2" + pipeline = _make_pipeline(req, tmp_path) + cfg = load_config(req) + pipeline._validate_contract(cfg) + pipeline._prepare_worktrees(cfg, "gen-test") + repo = kernel_root / "int8 test2" + main = tmp_path / "workspace" / "main" + assert cfg.target_repo_path == repo + assert main.is_symlink() + assert not main.readlink().is_absolute() + assert main.resolve() == repo + assert (repo / ".git").is_dir() + staged = ( + main / "int8_w8a8_gemm_api.py" + ) + assert staged.read_bytes() == contract.read_bytes() + assert staged.stat().st_mode & 0o222 == 0 + backend = (repo / "w8a8_backend.py").read_text(encoding="utf-8") + assert "def load_extension()" in backend + assert "is_python_module=False" in backend + assert (repo / "setup.py").is_file() + assert (repo / "csrc" / "bindings.cpp").is_file() + assert (repo / "w8a8_bench.py").is_file() + assert (repo / "w8a8_graph.py").is_file() + assert not (repo / "csrc" / "w8a8_gemm_hip.hip").exists() + manifest = json.loads( + (repo / "scaffold_manifest.json").read_text(encoding="utf-8") + ) + assert manifest["task_id"] == "gen-test" + assert manifest["fresh_repository"] is True + assert manifest["implementation_inherited"] is False + assert ( + manifest["initial_kernel"] + == "pending_parallel_explore_child_generation" + ) + + +def test_new_task_rejects_existing_kernel_repo_without_continuation( + tmp_path, monkeypatch +): + api_root = tmp_path / "API" + kernel_root = tmp_path / "kernel-repos" + _write_test_api_contract(api_root) + monkeypatch.setenv("METAINFER_OPERATOR_API_ROOT", str(api_root)) + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(kernel_root)) + existing = kernel_root / "int8 existing" + existing.mkdir(parents=True) + subprocess.run( + ["git", "init"], cwd=existing, check=True, capture_output=True + ) + (existing / "csrc").mkdir() + (existing / "csrc" / "w8a8_gemm_hip.hip").write_text( + "// existing kernel", encoding="utf-8" + ) + req = _gen_req() + req["answers"]["target_repo_path"] = "int8 existing" + pipeline = _make_pipeline(req, tmp_path) + cfg = load_config(req) + pipeline._validate_contract(cfg) + + with pytest.raises(RuntimeError, match="explicit continuation"): + pipeline._prepare_worktrees(cfg, "gen-test") + + +def test_gen_mode_uses_api_default_shapes_when_omitted( + tmp_path, monkeypatch +): + api_root = tmp_path / "API" + _write_test_api_contract(api_root) + monkeypatch.setenv("METAINFER_OPERATOR_API_ROOT", str(api_root)) + req = _gen_req() + req["answers"].pop("shape_config") + cfg = load_config(req) + assert set(cfg.shapes) == {"m2", "m16"} + assert len(cfg.assignments) == 1 + assert cfg.assignments[0].shape_ids == ["m2", "m16"] + + +def test_mock_mode_can_also_use_api_default_shapes( + tmp_path, monkeypatch +): + api_root = tmp_path / "API" + _write_test_api_contract(api_root) + monkeypatch.setenv("METAINFER_OPERATOR_API_ROOT", str(api_root)) + req = _gen_req() + req["answers"]["execution_mode"] = "Mock (no GPU)" + req["answers"].pop("shape_config") + cfg = load_config(req) + assert set(cfg.shapes) == {"m2", "m16"} + assert cfg.assignments[0].worker_id == "worker_0" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_guidance.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_guidance.py new file mode 100644 index 00000000..2d9d705e --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_guidance.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from ..orchestrator.guidance import ( + add_guidance, + claim_next_guidance, + list_guidance, +) + + +def test_guidance_is_worker_isolated_and_claimed_once(tmp_path): + first = add_guidance(tmp_path, "worker_0", "Use a smaller tile") + add_guidance(tmp_path, "worker_1", "Try split-K") + + claimed = claim_next_guidance(tmp_path, "worker_0", 4) + assert claimed is not None + assert claimed["id"] == first["id"] + assert claimed["consumed_iteration"] == 4 + assert claim_next_guidance(tmp_path, "worker_0", 5) is None + + worker_0 = list_guidance(tmp_path, "worker_0") + worker_1 = list_guidance(tmp_path, "worker_1") + assert worker_0[0]["status"] == "consumed" + assert worker_1[0]["status"] == "pending" + + +def test_guidance_rejects_empty_text(tmp_path): + try: + add_guidance(tmp_path, "worker_0", " ") + except ValueError as exc: + assert "empty" in str(exc) + else: + raise AssertionError("empty guidance should be rejected") diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_isa_analysis.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_isa_analysis.py new file mode 100644 index 00000000..9e51bec8 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_isa_analysis.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +from ..orchestrator.isa_analysis import ( + analyze_inline_asm_source, + evaluate_inline_asm_gate, + resolve_dtk_llvm_bin, + summarize_kernel_symbols, + summarize_isa_text, +) + + +def test_isa_summary_counts_memory_wait_and_mmac(): + summary = summarize_isa_text(""" +global_load_dwordx4 v[0:3], v[4:5], off +s_waitcnt vmcnt(0) +ds_write_b128 v0, v[0:3] +s_barrier +v_mmac_i32_16x16x32_i8 a[0:3], v[0:1], v[2:3], a[0:3] +""") + + assert summary["instruction_counts"]["global_load"] == 1 + assert summary["instruction_counts"]["ds_write"] == 1 + assert summary["instruction_counts"]["waitcnt"] == 1 + assert summary["instruction_counts"]["mmac"] == 1 + assert summary["load_wait0_ds_write_windows"] == 1 + + +def test_isa_summary_attributes_resources_to_profiled_symbols(): + partial = "_Z20w8a8_partial_kernelv" + combine = "_Z20w8a8_combine_kernelv" + disassembly = f""" +0000000000001000 <{partial}>: +global_load_dwordx4 v[0:3], v[4:5], off +v_mmac_i32_16x16x32_i8 a[0:3], v[0:1], v[2:3], a[0:3] +0000000000002000 <{combine}>: +global_load_dword v0, v[1:2], off +global_store_dword v[1:2], v0, off +""" + metadata = f""" +amdhsa.kernels: + - .args: + .group_segment_fixed_size: 10240 + .name: {partial} + .sgpr_count: 27 + .vgpr_count: 40 + .wavefront_size: 64 + - .args: + .group_segment_fixed_size: 0 + .name: {combine} + .sgpr_count: 38 + .vgpr_count: 16 + .wavefront_size: 64 +""" + + summary = summarize_kernel_symbols( + disassembly, + metadata, + kernel_names=[partial, combine], + primary_kernel_name=partial, + ) + + assert summary["kernel_name"] == partial + assert summary["resources"]["vgpr_count"] == 40 + assert summary["resources"]["sgpr_count"] == 27 + assert summary["resources"]["lds_bytes"] == 10240 + assert len(summary["profiled_kernels"]) == 2 + assert summary["profiled_kernels"][1]["resources"]["sgpr_count"] == 38 + + +def test_inline_asm_source_distinguishes_compiler_barrier(): + analysis = analyze_inline_asm_source(r''' +asm volatile("" ::: "memory"); +__asm__ __volatile__("s_waitcnt vmcnt(0)" ::: "memory"); +''') + + assert analysis["asm_block_count"] == 2 + assert analysis["compiler_barrier_count"] == 1 + assert analysis["raw_instruction_asm_count"] == 1 + assert analysis["raw_instruction_fingerprints"] == [ + "s_waitcnt vmcnt(0)" + ] + + +def test_new_inline_asm_requires_plan_and_trusted_isa(): + before = analyze_inline_asm_source("") + after = analyze_inline_asm_source( + 'asm volatile("s_waitcnt vmcnt(0)" ::: "memory");' + ) + + rejected = evaluate_inline_asm_gate( + before=before, + after=after, + proposal={}, + isa_evidence={"available": False}, + ) + accepted = evaluate_inline_asm_gate( + before=before, + after=after, + proposal={ + "isa_optimization": { + "strategy": "inline_asm", + "target_instructions": ["s_waitcnt vmcnt(0)"], + } + }, + isa_evidence={"available": True}, + ) + + assert rejected["required"] is True + assert rejected["passed"] is False + assert accepted["passed"] is True + + +def test_inline_asm_gate_enforces_control_phase_and_verified_targets(): + before = analyze_inline_asm_source("") + after = analyze_inline_asm_source( + 'asm volatile("s_waitcnt vmcnt(0)" ::: "memory");' + ) + proposal = { + "isa_optimization": { + "strategy": "inline_asm", + "target_instructions": ["s_waitcnt vmcnt(0)"], + } + } + + closed = evaluate_inline_asm_gate( + before=before, + after=after, + proposal=proposal, + isa_evidence={"available": True}, + raw_inline_asm_allowed=False, + ) + wrong_target = evaluate_inline_asm_gate( + before=before, + after=after, + proposal=proposal, + isa_evidence={"available": True}, + raw_inline_asm_allowed=True, + verified_target_instructions=["v_pk_add_u16"], + ) + open_gate = evaluate_inline_asm_gate( + before=before, + after=after, + proposal=proposal, + isa_evidence={"available": True}, + raw_inline_asm_allowed=True, + verified_target_instructions=["s_waitcnt vmcnt(0)"], + ) + + assert closed["passed"] is False + assert wrong_target["passed"] is False + assert open_gate["passed"] is True + + +def test_resolve_dtk_llvm_bin_accepts_aillvm_layout(tmp_path): + toolchain = tmp_path / "aillvm" / "bin" + toolchain.mkdir(parents=True) + for name in ( + "llvm-objcopy", + "clang-offload-bundler", + "llvm-readobj", + "llvm-objdump", + ): + (toolchain / name).touch() + + assert resolve_dtk_llvm_bin(toolchain) == toolchain diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_pipeline.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_pipeline.py new file mode 100644 index 00000000..3baf095c --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_pipeline.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import json + +from metainfer.orchestrator.state import StateStore + +from ..orchestrator.guidance import add_guidance +from ..orchestrator.pipeline import MockOptimizationPipeline + + +def _requirements(): + return { + "task_id": "mock-task", + "task_type": "dcu-kernel-auto-opt", + "answers": { + "execution_mode": "Mock (no GPU)", + "mock_iterations": "2", + "minimum_improvement_percent": 1.0, + "shape_config": """ +shapes: + - {id: m2, M: 2, N: 1536, K: 4096} + - {id: m16, M: 16, N: 1536, K: 4096} +assignments: + worker_0: {gpu: 0, shapes: [m2]} + worker_1: {gpu: 1, shapes: [m16]} +""", + }, + } + + +def test_mock_pipeline_is_parallel_and_gpu_free(tmp_path): + state_dir = tmp_path / "state" + workspace_dir = tmp_path / "workspace" + add_guidance( + state_dir / "guidance", "worker_0", "prefer a smaller LDS tile" + ) + report = MockOptimizationPipeline( + req=_requirements(), + state_dir=state_dir, + workspace_dir=workspace_dir, + store=StateStore(state_dir), + ).run() + assert report["status"] == "success" + assert report["real_gpu_used"] is False + assert set(report["workers"]) == {"worker_0", "worker_1"} + assert all(v["passed"] for v in report["final_validation"].values()) + for result in report["final_validation"].values(): + assert result["metrics"]["tflops"] > 0 + assert result["metrics"]["bandwidth_gb_s"] > 0 + for worker in ("worker_0", "worker_1"): + status = json.loads( + (workspace_dir / "workers" / worker / "status.json").read_text() + ) + assert status["state"] == "completed" + assert status["gpu_binding"]["enforced"] is False + worker_0_log = ( + workspace_dir / "workers" / "worker_0" / "runs" / "m2" + / "experiments.jsonl" + ) + first = json.loads(worker_0_log.read_text().splitlines()[0]) + assert first["manual_guidance"] == "prefer a smaller LDS tile" + pending = workspace_dir / "skills" / "pending" + skills = sorted(pending.glob("*/SKILL.md")) + assert len(skills) == 3 + contents = [path.read_text() for path in skills] + manifests = [ + json.loads((path.parent / "manifest.json").read_text()) + for path in skills + ] + assert sum(item["kind"] == "merged" for item in manifests) == 1 + assert sum("## Measured results" in text for text in contents) == 2 + + +def test_dry_run_only_writes_plan(tmp_path): + state_dir = tmp_path / "state" + workspace_dir = tmp_path / "workspace" + report = MockOptimizationPipeline( + req=_requirements(), + state_dir=state_dir, + workspace_dir=workspace_dir, + store=StateStore(state_dir), + ).run(dry_run=True) + assert report["dry_run"] is True + assert not (workspace_dir / "shared_baseline" / "results.json").exists() diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plugin.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plugin.py new file mode 100644 index 00000000..176eb25d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plugin.py @@ -0,0 +1,307 @@ +from __future__ import annotations + +import json + +import metainfer.tasks # noqa: F401 +from metainfer.orchestrator.registry import get_orchestrator +from metainfer.server.forms import load_form_schema +from metainfer.server.registry import get +from ..server.routes import _restart_worker_commands, read_worker_lanes + + +def test_task_and_web_plugins_registered(): + assert get_orchestrator("dcu-kernel-auto-opt").task_type == "dcu-kernel-auto-opt" + plugin = get("dcu-kernel-auto-opt") + assert plugin is not None + assert plugin.detail_view_module == "app/dkao-detail" + + +def test_form_schema_is_available(): + schema = load_form_schema("dcu-kernel-auto-opt") + assert schema is not None + assert schema["type"] == "dcu-kernel-auto-opt" + keys = {field["key"] for field in schema["fields"]} + assert { + "operator", "kernel_language", "target_hardware", "shape_config", + "dtype", "correctness_ref", "perf_target", "max_iterations", + "execution_mode", "target_repo_path", "agent_framework", + "agent_model", + } <= keys + shape = next( + field for field in schema["fields"] + if field["key"] == "shape_config" + ) + assert shape["required"] is False + assert shape["default"] == "" + assert shape["override_component"] == "shape-input" + framework = next( + field for field in schema["fields"] + if field["key"] == "agent_framework" + ) + assert framework["default"] == "ccb" + assert [option["label"] for option in framework["options"]] == [ + "ccb", "dsh" + ] + model = next( + field for field in schema["fields"] + if field["key"] == "agent_model" + ) + assert model["default"] == "Opus" + assert model["override_component"] == "agent-model" + assert [option["label"] for option in model["options"]] == [ + "Opus", "Sonnet", "deepseek-v4-flash" + ] + + +def test_worker_lanes_always_has_four_rows(tmp_path): + lanes = read_worker_lanes(tmp_path) + assert [row["worker_id"] for row in lanes["workers"]] == [ + "worker_0", "worker_1", "worker_2", "worker_3", + ] + + +def test_restart_commands_let_lane_resolve_agent_binary(tmp_path): + # Regression: the restart route used to force METAINFER_CLAUDE_BIN + # (ccb's binary) as an explicit --claude-bin, which breaks dsh-framework + # lanes whose agents must run through bridge/dsh/dsh_agent.py. The lane + # binaries resolve the framework-appropriate binary themselves. + requirements = tmp_path / "requirements.json" + requirements.write_text("{}", encoding="utf-8") + state = tmp_path / "state" + workspace = tmp_path / "workspace" + command, integration = _restart_worker_commands( + requirements, state, workspace, "worker_1" + ) + for cmd in (command, integration): + assert "--claude-bin" not in cmd + assert "--worker-id" in cmd + assert cmd[cmd.index("--worker-id") + 1] == "worker_1" + assert str(requirements) in cmd + assert "restart_worker" in command[2] + assert "integrate_restarted_worker" in integration[2] + + +def test_worker_lanes_surface_live_optimization_step(tmp_path): + workspace = tmp_path / "workspace" + state = tmp_path / "state" + worker = workspace / "workers" / "worker_0" + runs = worker / "runs" / "m2" + runs.mkdir(parents=True) + state.mkdir() + workspace.joinpath("plan.json").write_text(json.dumps({ + "max_iterations": 5, + "assignments": [{ + "worker_id": "worker_0", + "gpu": 0, + "shapes": ["m2"], + }], + }), encoding="utf-8") + worker.joinpath("status.json").write_text(json.dumps({ + "state": "validating_candidate", + "iteration": 2, + "shape_id": "m2", + }), encoding="utf-8") + runs.joinpath("experiments.jsonl").write_text( + json.dumps({"iteration": 1, "shape_id": "m2"}) + "\n", + encoding="utf-8", + ) + state.joinpath("agents.json").write_text(json.dumps({ + "agents": [{ + "name": "worker_0-m2-iter2", + "status": "running", + "started_at": 10, + "last_output_age_s": 181, + }], + }), encoding="utf-8") + + lane = read_worker_lanes(workspace, state)["workers"][0] + + assert lane["step"] == "Compiling and validating candidate" + assert lane["completed_rounds"] == 1 + assert lane["target_rounds"] == 5 + assert lane["long_running"] is True + assert lane["active_iteration"] == { + "iteration": 2, + "shape_id": "m2", + "state": "validating_candidate", + "step": "Compiling and validating candidate", + "agent_name": "worker_0-m2-iter2", + "agent_status": "running", + "elapsed_s": None, + "last_output_age_s": 181, + } + + +def test_worker_lanes_surface_pmc_and_repair_steps(tmp_path): + workspace = tmp_path / "workspace" + state = tmp_path / "state" + worker = workspace / "workers" / "worker_0" + worker.mkdir(parents=True) + state.mkdir() + workspace.joinpath("plan.json").write_text(json.dumps({ + "max_iterations": 5, + "assignments": [{ + "worker_id": "worker_0", + "gpu": 0, + "shapes": ["m2"], + }], + }), encoding="utf-8") + + worker.joinpath("status.json").write_text(json.dumps({ + "state": "profiling_current_best", + "iteration": 2, + "shape_id": "m2", + }), encoding="utf-8") + lane = read_worker_lanes(workspace, state)["workers"][0] + assert lane["active_iteration"]["step"] == ( + "Profiling current best kernel with PMC" + ) + + worker.joinpath("status.json").write_text(json.dumps({ + "state": "repairing_candidate", + "iteration": 2, + "shape_id": "m2", + "repair": 3, + "max_repairs": 4, + }), encoding="utf-8") + lane = read_worker_lanes(workspace, state)["workers"][0] + assert lane["active_iteration"]["step"] == ( + "Repairing compile/correctness failure (3/4)" + ) + assert lane["active_iteration"]["repair"] == 3 + assert lane["active_iteration"]["max_repairs"] == 4 + + +def test_worker_lanes_surface_bootstrap_attempts(tmp_path): + workspace = tmp_path / "workspace" + state = tmp_path / "state" + source = workspace / "workers" / "worker_0" / "source" + source.joinpath("csrc").mkdir(parents=True) + state.mkdir() + workspace.joinpath("plan.json").write_text(json.dumps({ + "assignments": [{ + "worker_id": "worker_0", + "gpu": 0, + "shapes": ["m2_wqkv_a"], + }], + }), encoding="utf-8") + source.joinpath("csrc", "w8a8_gemm_hip.hip").write_text( + "// hip", encoding="utf-8" + ) + state.joinpath("agents.json").write_text(json.dumps({ + "agents": [{ + "name": "worker_0-bootstrap-attempt1", + "status": "running", + "success": None, + "elapsed_s": 12.5, + "last_output_age_s": 0.5, + }], + }), encoding="utf-8") + + lane = read_worker_lanes(workspace, state)["workers"][0] + + assert lane["state"] == "bootstrap_running" + assert lane["bootstrap_attempts"] == [{ + "kind": "bootstrap", + "attempt": 1, + "status": "running", + "hypothesis": ( + "Create and validate the initial HIP implementation for " + "m2_wqkv_a." + ), + "generated_files": ["csrc/w8a8_gemm_hip.hip"], + "metrics": {}, + "artifact_dir": None, + "candidate_files": [], + "error": None, + "elapsed_s": 12.5, + "last_output_age_s": 0.5, + "started_at": None, + }] + + +def test_worker_lanes_surface_live_bootstrap_performance_metrics(tmp_path): + workspace = tmp_path / "workspace" + state = tmp_path / "state" + worker = workspace / "workers" / "worker_0" + source = worker / "source" + source.mkdir(parents=True) + state.mkdir() + workspace.joinpath("plan.json").write_text(json.dumps({ + "assignments": [{ + "worker_id": "worker_0", + "gpu": 0, + "shapes": ["m2_wqkv_a"], + }], + }), encoding="utf-8") + state.joinpath("agents.json").write_text(json.dumps({ + "agents": [{ + "name": "worker_0-bootstrap-attempt1", + "status": "completed", + "success": True, + }], + }), encoding="utf-8") + metrics = { + "passed": True, + "median_us": 8.25, + "p90_us": 8.75, + "tflops": 19.5, + "bandwidth_gb_s": 812.0, + } + worker.joinpath("bootstrap_progress.json").write_text(json.dumps({ + "attempt": 1, + "status": "validating", + "hypothesis": "Use a DUMMA tiled kernel.", + "metrics": {"m2_wqkv_a": metrics}, + }), encoding="utf-8") + + attempt = read_worker_lanes( + workspace, state + )["workers"][0]["bootstrap_attempts"][0] + + assert attempt["status"] == "validating" + assert attempt["hypothesis"] == "Use a DUMMA tiled kernel." + assert attempt["metrics"]["m2_wqkv_a"] == metrics + + +def test_worker_lanes_surface_bootstrap_iteration_snapshot(tmp_path): + workspace = tmp_path / "workspace" + state = tmp_path / "state" + worker = workspace / "workers" / "worker_0" + record = ( + worker / "iterations" / "bootstrap" + / "iteration1" / "iteration.json" + ) + record.parent.mkdir(parents=True) + state.mkdir() + workspace.joinpath("plan.json").write_text(json.dumps({ + "assignments": [{ + "worker_id": "worker_0", + "gpu": 0, + "shapes": ["m2"], + }], + }), encoding="utf-8") + state.joinpath("agents.json").write_text(json.dumps({ + "agents": [{ + "name": "worker_0-bootstrap-attempt1", + "status": "done", + "success": True, + }], + }), encoding="utf-8") + record.write_text(json.dumps({ + "attempt": 1, + "status": "failed", + "error": "trusted compile failed", + "artifact_dir": "iterations/bootstrap/iteration1", + "candidate_files": ["csrc/w8a8_gemm_hip.hip"], + "metrics": {}, + }), encoding="utf-8") + + attempt = read_worker_lanes( + workspace, state + )["workers"][0]["bootstrap_attempts"][0] + + assert attempt["status"] == "failed" + assert attempt["error"] == "trusted compile failed" + assert attempt["artifact_dir"] == "iterations/bootstrap/iteration1" + assert attempt["candidate_files"] == ["csrc/w8a8_gemm_hip.hip"] diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_quality_loop.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_quality_loop.py new file mode 100644 index 00000000..631f4ad2 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_quality_loop.py @@ -0,0 +1,784 @@ +from __future__ import annotations + +import csv +import json +from pathlib import Path + +import pytest + +from ..orchestrator.api_contracts import ( + OperatorAPIContract, + stage_operator_references, +) +from ..orchestrator.config import WorkerAssignment +from ..orchestrator.experience_store import load_verified_experience +from ..orchestrator.pmc_profile import ( + add_unprofiled_bandwidth, + parse_memory_traffic_csv, + parse_pmc_csv, +) +from ..orchestrator.w8a8_pipeline import ( + RealW8A8OptimizationPipeline, + _ISA_AGENT_ARGS, + _SOURCE_ONLY_AGENT_ARGS, + evaluate_candidate_acceptance, + evaluate_final_target, + isa_round_policy, + is_infrastructure_failure, + phase_extension_reason, + pmc_profile_decision, + validate_skill_draft, +) +from ..orchestrator.prompts import ( + split_k_candidate_set, + w8a8_round_strategy, +) + + +def test_pmc_parser_selects_last_w8a8_dispatch(tmp_path): + path = tmp_path / "pmc.csv" + fields = [ + "KernelName", "gpu-id", "grd", "wgr", "lds", "scr", + "arch_vgpr", "accum_vgpr", "sgpr", "wave_size", + "GRBM_COUNT", "GRBM_GUI_ACTIVE", "SQ_ACTIVE_INST_VALU", + "SQ_INSTS_FLAT_LDS_ONLY", "SQ_INSTS_LDS", "SQ_INSTS_VALU", + "SQ_INSTS_VMEM_RD", "SQ_INSTS_VMEM_WR", + "SQ_LDS_BANK_CONFLICT", "SQ_WAIT_INST_LDS", + "TCC_HIT[0]", "TCC_MISS[0]", "BeginNs", "EndNs", + ] + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerow({ + "KernelName": "unrelated", "BeginNs": "1", "EndNs": "2", + }) + writer.writerow({ + "KernelName": "w8a8_gemm_kernel", + "gpu-id": "0", + "grd": "3072", + "wgr": "128", + "lds": "1024", + "scr": "0", + "arch_vgpr": "40", + "accum_vgpr": "0", + "sgpr": "48", + "wave_size": "64", + "GRBM_COUNT": "200", + "GRBM_GUI_ACTIVE": "100", + "SQ_ACTIVE_INST_VALU": "12", + "SQ_INSTS_FLAT_LDS_ONLY": "1", + "SQ_INSTS_LDS": "2", + "SQ_INSTS_VALU": "3", + "SQ_INSTS_VMEM_RD": "4", + "SQ_INSTS_VMEM_WR": "5", + "SQ_LDS_BANK_CONFLICT": "6", + "SQ_WAIT_INST_LDS": "7", + "TCC_HIT[0]": "90", + "TCC_MISS[0]": "10", + "BeginNs": "1000", + "EndNs": "51000", + }) + + evidence = parse_pmc_csv(path) + + assert evidence["kernel_name"] == "w8a8_gemm_kernel" + assert evidence["grid_blocks"] == 24 + assert evidence["profiled_duration_us"] == 50.0 + assert evidence["l2_hit_rate_percent"] == 90.0 + assert evidence["gpu_active_percent"] == 50.0 + assert evidence["counters"]["lds_bank_conflicts"] == 6 + + +def test_pmc_parser_keeps_partial_and_combine_in_one_replay(tmp_path): + path = tmp_path / "pmc.csv" + fields = [ + "KernelName", "grd", "wgr", "lds", "arch_vgpr", "sgpr", + "TCC_HIT[0]", "TCC_MISS[0]", "BeginNs", "EndNs", + ] + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerow({"KernelName": "w8a8_old", "grd": 64, "wgr": 64}) + writer.writerow({"KernelName": "unrelated"}) + writer.writerow({ + "KernelName": "w8a8_splitk_partial", "grd": 30720, + "wgr": 128, "lds": 10240, "arch_vgpr": 40, "sgpr": 27, + "TCC_HIT[0]": 30, "TCC_MISS[0]": 10, + "BeginNs": 1000, "EndNs": 12000, + }) + writer.writerow({ + "KernelName": "w8a8_combine", "grd": 24576, + "wgr": 256, "lds": 0, "arch_vgpr": 16, "sgpr": 38, + "TCC_HIT[0]": 20, "TCC_MISS[0]": 20, + "BeginNs": 12000, "EndNs": 17000, + }) + + evidence = parse_pmc_csv(path) + + assert evidence["primary_kernel_name"] == "w8a8_splitk_partial" + assert evidence["grid_blocks"] == 240 + assert evidence["lds_bytes"] == 10240 + assert evidence["operator_aggregate"]["kernel_count"] == 2 + assert evidence["operator_aggregate"]["profiled_duration_us"] == 16.0 + assert evidence["operator_aggregate"]["l2_hit_rate_percent"] == 62.5 + + +def test_memory_traffic_uses_dtk_request_size_formulas(tmp_path): + fields = [ + "KernelName", + "TCC_EA_RDREQ[0]", "TCC_EA_RDREQ_32B[0]", + "TCC_EA1_RDREQ[0]", "TCC_EA1_RDREQ_32B[0]", + "TCC_EA_WRREQ[0]", "TCC_EA_WRREQ_64B[0]", + "TCC_EA1_WRREQ[0]", "TCC_EA1_WRREQ_64B[0]", + ] + read_path = tmp_path / "pmc-read.csv" + write_path = tmp_path / "pmc-write.csv" + for path, row in ( + (read_path, { + "KernelName": "w8a8_gemm_kernel", + "TCC_EA_RDREQ[0]": "10", + "TCC_EA_RDREQ_32B[0]": "4", + "TCC_EA1_RDREQ[0]": "3", + "TCC_EA1_RDREQ_32B[0]": "1", + }), + (write_path, { + "KernelName": "w8a8_gemm_kernel", + "TCC_EA_WRREQ[0]": "5", + "TCC_EA_WRREQ_64B[0]": "2", + "TCC_EA1_WRREQ[0]": "4", + "TCC_EA1_WRREQ_64B[0]": "1", + }), + ): + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerow(row) + + evidence = { + "memory_traffic": parse_memory_traffic_csv( + read_path, write_path + ) + } + add_unprofiled_bandwidth(evidence, 2.0) + + traffic = evidence["memory_traffic"] + assert traffic["read_bytes_per_dispatch"] == 672 + assert traffic["write_bytes_per_dispatch"] == 384 + assert traffic["total_bytes_per_dispatch"] == 1056 + assert traffic["counter_derived_hbm_bandwidth_gb_s"] == 0.528 + + +def test_memory_traffic_aggregates_multi_kernel_operator(tmp_path): + fields = [ + "KernelName", "TCC_EA_RDREQ[0]", "TCC_EA_RDREQ_32B[0]", + "TCC_EA1_RDREQ[0]", "TCC_EA1_RDREQ_32B[0]", + "TCC_EA_WRREQ[0]", "TCC_EA_WRREQ_64B[0]", + "TCC_EA1_WRREQ[0]", "TCC_EA1_WRREQ_64B[0]", + ] + read_path = tmp_path / "read.csv" + write_path = tmp_path / "write.csv" + for path, counter in ((read_path, "TCC_EA_RDREQ[0]"), + (write_path, "TCC_EA_WRREQ[0]")): + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerow({"KernelName": "unrelated"}) + writer.writerow({"KernelName": "w8a8_partial", counter: 10}) + writer.writerow({"KernelName": "w8a8_combine", counter: 2}) + + traffic = parse_memory_traffic_csv(read_path, write_path) + + assert traffic["kernel_names"] == ["w8a8_partial", "w8a8_combine"] + assert traffic["kernel_count"] == 2 + assert traffic["total_bytes_per_operator_replay"] == 1152 + assert traffic["kernels"][0]["total_bytes"] == 960 + assert traffic["kernels"][1]["total_bytes"] == 192 + + +def test_memory_traffic_marker_excludes_graph_validation_replay(tmp_path): + fields = [ + "KernelName", "TCC_EA_RDREQ[0]", "TCC_EA_RDREQ_32B[0]", + "TCC_EA1_RDREQ[0]", "TCC_EA1_RDREQ_32B[0]", + "TCC_EA_WRREQ[0]", "TCC_EA_WRREQ_64B[0]", + "TCC_EA1_WRREQ[0]", "TCC_EA1_WRREQ_64B[0]", + ] + read_path = tmp_path / "read.csv" + write_path = tmp_path / "write.csv" + for path, counter in ((read_path, "TCC_EA_RDREQ[0]"), + (write_path, "TCC_EA_WRREQ[0]")): + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerow({"KernelName": "w8a8_partial", counter: 10}) + writer.writerow({"KernelName": "w8a8_combine", counter: 2}) + writer.writerow({"KernelName": "vectorized_elementwise_marker"}) + writer.writerow({"KernelName": "w8a8_partial", counter: 10}) + writer.writerow({"KernelName": "w8a8_combine", counter: 2}) + + traffic = parse_memory_traffic_csv(read_path, write_path) + + assert traffic["kernel_names"] == ["w8a8_partial", "w8a8_combine"] + assert traffic["kernel_count"] == 2 + assert traffic["total_bytes_per_operator_replay"] == 1152 + + +def test_memory_traffic_records_tiny_counter_replay_skew(tmp_path): + fields = [ + "KernelName", "TCC_EA_RDREQ[0]", "TCC_EA_RDREQ_32B[0]", + "TCC_EA_WRREQ[0]", "TCC_EA_WRREQ_64B[0]", + ] + read_path = tmp_path / "read.csv" + write_path = tmp_path / "write.csv" + for path, row in ( + (read_path, { + "KernelName": "w8a8_gemm", "TCC_EA_RDREQ[0]": 100, + "TCC_EA_RDREQ_32B[0]": 101, + }), + (write_path, { + "KernelName": "w8a8_gemm", "TCC_EA_WRREQ[0]": 200, + "TCC_EA_WRREQ_64B[0]": 201, + }), + ): + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerow(row) + + traffic = parse_memory_traffic_csv(read_path, write_path) + + assert traffic["read_bytes_per_operator_replay"] == 101 * 32 + assert traffic["write_bytes_per_operator_replay"] == 201 * 64 + assert len(traffic["counter_reconciliation"]) == 2 + + +def test_memory_traffic_rejects_large_counter_mismatch(tmp_path): + fields = [ + "KernelName", "TCC_EA_RDREQ[0]", "TCC_EA_RDREQ_32B[0]", + "TCC_EA_WRREQ[0]", "TCC_EA_WRREQ_64B[0]", + ] + read_path = tmp_path / "read.csv" + write_path = tmp_path / "write.csv" + for path, row in ( + (read_path, { + "KernelName": "w8a8_gemm", "TCC_EA_RDREQ[0]": 100, + "TCC_EA_RDREQ_32B[0]": 120, + }), + (write_path, {"KernelName": "w8a8_gemm"}), + ): + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerow(row) + + with pytest.raises(ValueError, match="exceeds TCC_EA_RDREQ"): + parse_memory_traffic_csv(read_path, write_path) + + +def test_repair_prompt_exposes_four_repairs_and_mismatch(): + prompt = RealW8A8OptimizationPipeline._repair_prompt( + assignment=WorkerAssignment("worker_0", 0, ["m2"]), + shape_id="m2", + shape={"M": 2, "N": 16, "K": 32}, + root=Path("/tmp/worker"), + iteration=1, + repair=4, + error="exact correctness failed", + metrics={ + "mismatch_count": 1, + "first_mismatch": { + "m": 0, "n": 3, "actual": 2.0, "expected": 3.0, + }, + }, + pmc_evidence={"available": True, "arch_vgpr": 40}, + ) + + assert "repair 4/4" in prompt + assert '"mismatch_count": 1' in prompt + assert "Preserve the proposed" in prompt + + +def test_worker_prompt_receives_pmc_evidence(): + prompt = RealW8A8OptimizationPipeline._worker_prompt( + WorkerAssignment("worker_0", 0, ["m2"]), + "m2", + {"M": 2, "N": 16, "K": 32}, + {"median_us": 10.0}, + Path("/tmp/worker"), + 1, + None, + [], + { + "available": True, + "kernel_name": "w8a8_gemm_kernel", + "arch_vgpr": 40, + "isa": { + "available": True, + "instruction_counts": {"mmac": 8, "waitcnt": 4}, + }, + }, + ) + + assert "Trusted PMC evidence" in prompt + assert '"arch_vgpr": 40' in prompt + assert '"mmac": 8' not in prompt + assert "hygon-gfx928-memory-isa" not in prompt + assert "hygon-gfx928-compute-isa" not in prompt + assert "Skill tool is disabled" in prompt + assert "not measured HBM traffic" in prompt + + +def test_pmc_policy_skips_scalar_and_unchanged_best(): + hip_policy = {"skill_allowed": False} + initial = pmc_profile_decision( + iteration=1, + history=[], + source_uses_dumma=False, + isa_policy=hip_policy, + ) + unchanged = pmc_profile_decision( + iteration=3, + history=[{"accepted": False}], + source_uses_dumma=True, + isa_policy=hip_policy, + ) + + assert initial["profile"] is False + assert "scalar bootstrap" in initial["reason"] + assert unchanged["profile"] is False + assert "unchanged" in unchanged["reason"] + + +def test_pmc_policy_profiles_new_best_and_late_isa(): + new_best = pmc_profile_decision( + iteration=2, + history=[{"accepted": True}], + source_uses_dumma=True, + isa_policy={"skill_allowed": False}, + ) + late_isa = pmc_profile_decision( + iteration=9, + history=[{"accepted": False}], + source_uses_dumma=True, + isa_policy={"skill_allowed": True}, + ) + + assert new_best["profile"] is True + assert late_isa["profile"] is True + + +def test_continuation_prompt_is_incremental_and_compact(): + prompt = RealW8A8OptimizationPipeline._worker_prompt( + WorkerAssignment("worker_0", 0, ["m16"]), + "m16", + {"M": 16, "N": 1024, "K": 4096}, + { + "median_us": 10.0, + "p90_us": 10.2, + "latency_samples_us": list(range(30)), + }, + Path("/tmp/worker"), + 4, + None, + [{ + "iteration": 3, + "accepted": False, + "metrics": { + "median_us": 10.1, + "latency_samples_us": list(range(30)), + }, + }], + {"available": False, "skipped": True}, + continuation=True, + ) + + assert "Continue the existing shape-specialized" in prompt + assert "immutable API" in prompt + assert "latency_samples_us" not in prompt + assert "do not\nreread unchanged scaffold files" in prompt.lower() + assert "experiments.jsonl" in prompt + + +def test_m16_rounds_keep_architecture_search_after_iteration_five(): + shape = {"M": 16, "N": 1024, "K": 4096} + + split_k = w8a8_round_strategy(shape, 3, [], { + "grid_blocks": 64, + "device_cu_count": 120, + }) + packed = w8a8_round_strategy(shape, 5, []) + pipeline = w8a8_round_strategy(shape, 6, []) + isa = w8a8_round_strategy(shape, 9, []) + + assert "split-K=2" in split_k + assert "64 blocks for 120 CUs" in split_k + assert "HIP-only packed-weight" in packed + assert "A-only LDS" in pipeline + assert "Late ISA-diagnosis" in isa + + +def test_m16_near_one_block_per_cu_triggers_parallelism_warning(): + strategy = w8a8_round_strategy( + {"M": 16, "N": 4096, "K": 2048}, + 3, + [], + {"grid_blocks": 128, "device_cu_count": 120}, + ) + + assert "two-blocks-per-CU" in strategy + assert "split-K=2" in strategy + + +def test_split_k_candidates_include_non_power_cu_aligned_probe(): + candidates = split_k_candidate_set( + {"M": 16, "N": 1536, "K": 4096}, + cu_count=120, + max_split=16, + ) + + assert 10 in candidates + assert 2 in candidates + assert all(2 <= split <= 16 for split in candidates) + + +def test_optional_variant_is_staged_read_only_under_references(tmp_path): + api = tmp_path / "int8_w8a8_gemm_api.py" + variant = tmp_path / "w8a8_gemm_variants.hip" + api.write_text("# api\n", encoding="utf-8") + variant.write_text("// evidence, not policy\n", encoding="utf-8") + contract = OperatorAPIContract( + operator="Quantized GEMM", + dtype="INT8 W8A8", + source=api, + destination_name=api.name, + reference_sources=(variant,), + ) + + staged = stage_operator_references(contract, tmp_path / "repo") + + assert [path.relative_to(tmp_path / "repo").as_posix() for path in staged] == [ + "references/w8a8_gemm_variants.hip" + ] + assert staged[0].read_text(encoding="utf-8") == "// evidence, not policy\n" + assert staged[0].stat().st_mode & 0o222 == 0 + + +def _valid_history(iteration, improvement, **extra): + record = { + "iteration": iteration, + "build_success": True, + "correctness_passed": True, + "metrics": {"graph_capture_passed": True}, + "acceptance": {"improvement_percent": improvement}, + } + record.update(extra) + return record + + +def test_isa_policy_requires_eight_hip_rounds_and_plateau(): + history = [ + _valid_history(index, improvement, isa_policy={"phase": "hip_only"}) + for index, improvement in enumerate( + [5.0, 3.0, 1.5, 1.0, 0.8, 0.8, -0.2, 1.4], + start=1, + ) + ] + + early = isa_round_policy( + iteration=8, max_iterations=10, history=history[:7] + ) + late = isa_round_policy( + iteration=9, max_iterations=10, history=history + ) + + assert early["skill_allowed"] is False + assert late["phase"] == "isa_guided_hip" + assert late["skill_allowed"] is True + assert late["raw_inline_asm_allowed"] is False + assert "Skill" not in _SOURCE_ONLY_AGENT_ARGS[1].split(",") + assert "Skill" in _ISA_AGENT_ARGS[1].split(",") + + +def test_isa_policy_does_not_count_failed_attempts_as_hip_rounds(): + history = [ + _valid_history(index, 0.5, isa_policy={"phase": "hip_only"}) + for index in range(1, 8) + ] + history.append({ + "iteration": 8, + "build_success": False, + "correctness_passed": False, + "failure_reason": "killed after timeout", + "isa_policy": {"phase": "hip_only"}, + }) + + policy = isa_round_policy( + iteration=9, max_iterations=10, history=history + ) + + assert policy["skill_allowed"] is False + assert policy["valid_hip_rounds"] == 7 + + +def test_timeout_is_an_infrastructure_failure(): + assert is_infrastructure_failure("worker failed: killed") is True + assert is_infrastructure_failure("nonzero exit 143") is True + assert is_infrastructure_failure("exact correctness failed") is False + + +def test_timeout_strategy_does_not_repair_partial_candidate(): + strategy = w8a8_round_strategy( + {"M": 16, "N": 4096, "K": 2048}, + 4, + [{ + "iteration": 3, + "build_success": False, + "failure_reason": "worker failed: killed after timeout", + "artifact_dir": "iterations/m16/iteration3", + }], + {}, + ) + + assert "agent infrastructure" in strategy + assert "do not repair or replay" in strategy + + +def test_late_isa_prompt_exposes_only_selected_isa_skills(): + policy = { + "phase": "isa_guided_hip", + "skill_allowed": True, + "raw_inline_asm_allowed": False, + "plateau": True, + "reason": "test plateau", + } + prompt = RealW8A8OptimizationPipeline._worker_prompt( + WorkerAssignment("worker_0", 0, ["m16"]), + "m16", + {"M": 16, "N": 1024, "K": 4096}, + {"median_us": 10.0}, + Path("/tmp/worker"), + 9, + None, + [], + {"available": True, "isa": {"available": True}}, + isa_policy=policy, + ) + + assert "hygon-gfx928-memory-isa" in prompt + assert "hygon-gfx928-compute-isa" in prompt + assert "Raw inline asm remains forbidden" in prompt + assert '"isa_optimization"' in prompt + + +def test_raw_asm_needs_two_isa_rounds_and_prior_confirmed_limitation(): + history = [ + *[ + _valid_history( + index, + 0.4, + isa_policy={"phase": "hip_only"}, + ) + for index in range(1, 9) + ], + _valid_history( + 9, + 0.2, + isa_policy={"phase": "isa_guided_hip"}, + isa_optimization={ + "compiler_limitation_confirmed": False, + "target_instructions": [], + }, + candidate_isa={"available": True}, + ), + _valid_history( + 10, + 0.1, + isa_policy={"phase": "isa_guided_hip"}, + isa_optimization={ + "compiler_limitation_confirmed": True, + "target_instructions": ["s_waitcnt vmcnt(0)"], + }, + candidate_isa={"available": True}, + ), + ] + + policy = isa_round_policy( + iteration=11, max_iterations=10, history=history + ) + + assert policy["phase"] == "conditional_inline_asm" + assert policy["raw_inline_asm_allowed"] is True + assert policy["verified_target_instructions"] == ["s_waitcnt vmcnt(0)"] + + +def test_one_isa_round_does_not_open_raw_asm(): + history = [ + *[ + _valid_history( + index, 0.4, isa_policy={"phase": "hip_only"} + ) + for index in range(1, 9) + ], + _valid_history( + 9, + 0.2, + isa_policy={"phase": "isa_guided_hip"}, + isa_optimization={ + "compiler_limitation_confirmed": True, + "target_instructions": ["s_waitcnt vmcnt(0)"], + }, + candidate_isa={"available": True}, + ), + ] + + policy = isa_round_policy( + iteration=10, max_iterations=10, history=history + ) + + assert policy["phase"] == "isa_guided_hip" + assert policy["raw_inline_asm_allowed"] is False + assert policy["valid_isa_guided_rounds"] == 1 + + +def test_large_regressions_do_not_prove_plateau(): + history = [ + _valid_history( + index, improvement, isa_policy={"phase": "hip_only"} + ) + for index, improvement in enumerate( + [4.0, 3.0, 1.0, 0.5, 0.2, -21.0, -19.0, -22.0], + start=1, + ) + ] + + policy = isa_round_policy( + iteration=9, max_iterations=10, history=history + ) + + assert policy["phase"] == "hip_only" + assert policy["plateau"] is False + assert "Large regressions" in policy["reason"] + + +def test_phase_extension_reserves_two_valid_isa_rounds(): + hip_history = [ + _valid_history(index, 0.4, isa_policy={"phase": "hip_only"}) + for index in range(1, 9) + ] + one_isa = _valid_history( + 9, + 0.2, + isa_policy={"phase": "isa_guided_hip"}, + isa_optimization={"compiler_limitation_confirmed": False}, + candidate_isa={"available": True}, + ) + + reason = phase_extension_reason( + max_iterations=10, history=[*hip_history, one_isa] + ) + + assert reason == "need 1 more valid ISA-guided HIP experiment(s)" + + +def test_final_target_is_measured_against_fixed_baseline(): + result = evaluate_final_target( + baseline={"m16": {"median_us": 100.0}}, + validation={ + "m16": {"passed": True, "metrics": {"median_us": 80.0}} + }, + target_improvement_percent=20.0, + ) + + assert result["shapes"]["m16"]["improvement_percent"] == 25.0 + assert result["all_shapes_met"] is True + + +def test_acceptance_rejects_median_gain_with_p90_regression(): + result = evaluate_candidate_acceptance( + passed=True, + metrics={"median_us": 90.0, "p90_us": 130.0}, + best_metrics={"median_us": 100.0, "p90_us": 120.0}, + minimum_improvement_percent=2.0, + ) + + assert result["improvement_percent"] > 2.0 + assert result["p90_guard_passed"] is False + assert result["accepted"] is False + + +def test_acceptance_requires_both_median_and_p90(): + result = evaluate_candidate_acceptance( + passed=True, + metrics={"median_us": 90.0, "p90_us": 110.0}, + best_metrics={"median_us": 100.0, "p90_us": 120.0}, + minimum_improvement_percent=2.0, + ) + + assert result["accepted"] is True + + +def test_subthreshold_gain_can_become_shadow_without_updating_best(): + result = evaluate_candidate_acceptance( + passed=True, + metrics={"median_us": 99.4, "p90_us": 118.0}, + best_metrics={"median_us": 100.0, "p90_us": 120.0}, + minimum_improvement_percent=1.0, + ) + + assert result["accepted"] is False + assert result["shadow_eligible"] is True + + +def test_shadow_candidate_must_improve_existing_shadow(): + result = evaluate_candidate_acceptance( + passed=True, + metrics={"median_us": 99.5, "p90_us": 118.0}, + best_metrics={"median_us": 100.0, "p90_us": 120.0}, + shadow_metrics={"median_us": 99.4, "p90_us": 118.0}, + minimum_improvement_percent=1.0, + ) + + assert result["accepted"] is False + assert result["shadow_eligible"] is False + assert result["improves_shadow"] is False + + +def test_verified_experience_ignores_skill_prose_and_filters_shape(tmp_path): + repo = tmp_path / "old-task" + record_path = ( + repo / "candidates" / "worker_0" / "iteration1" + / "iteration.json" + ) + record_path.parent.mkdir(parents=True) + record_path.write_text(json.dumps({ + "shape": {"M": 2, "N": 16, "K": 32}, + "shape_id": "m2", + "iteration": 1, + "hypothesis": "vector load", + "build_success": True, + "correctness_passed": False, + "accepted": False, + "speedup": 1.5, + "metrics": {"median_us": 5.0, "mismatch_count": 1}, + }), encoding="utf-8") + + evidence = load_verified_experience( + tmp_path, {"M": 2, "N": 16, "K": 32} + ) + + assert len(evidence) == 1 + assert evidence[0]["classification"] == "faster_incorrect_repairable" + assert evidence[0]["speedup"] == 1.5 + + +def test_skill_validator_rejects_hardware_claim_from_compile_failure(): + facts = [{ + "accepted": False, + "build_success": False, + "correctness_passed": False, + }] + + with pytest.raises(ValueError, match="unsupported conclusions"): + validate_skill_draft( + "No candidate was accepted. DUMMA is unavailable on this GPU.", + facts, + ) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_rename_kernel_repo.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_rename_kernel_repo.py new file mode 100644 index 00000000..dcb9c05d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_rename_kernel_repo.py @@ -0,0 +1,227 @@ +"""Kernel-repository rename: core logic + Web route tests.""" + +from __future__ import annotations + +import json +import os +import subprocess +from pathlib import Path + +import pytest + +from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.rename_kernel_repo import ( + rename_kernel_repo, +) +from .conftest import register_dkao_task + + +def _git(repo: Path, *args: str) -> subprocess.CompletedProcess: + return subprocess.run( + ["git", "-C", str(repo), *args], + check=True, capture_output=True, text=True, + ) + + +def _make_repo(workspace: Path, repo_root: Path, name: str = "repo-old") -> dict: + """Create a kernel-repos git repo with one linked worktree + main symlink.""" + repo = repo_root / name + repo.mkdir(parents=True) + _git(repo, "init", "-q") + _git(repo, "config", "user.email", "t@t") + _git(repo, "config", "user.name", "t") + (repo / "f.txt").write_text("hi\n", encoding="utf-8") + _git(repo, "add", "f.txt") + _git(repo, "commit", "-qm", "init") + worktree = workspace / "source" + _git(repo, "worktree", "add", "-q", str(worktree)) + main = workspace / "main" + main.symlink_to( + os.path.relpath(repo, start=workspace), target_is_directory=True + ) + return {"repo": repo, "worktree": worktree} + + +def _seed_task_files(workspace: Path, state: Path, repo: Path, task_id: str) -> None: + state.mkdir(parents=True, exist_ok=True) + (state / "requirements.json").write_text(json.dumps({ + "task_id": task_id, + "label": repo.name, + "target_repo_path": repo.name, + }), encoding="utf-8") + (workspace / "plan.json").write_text(json.dumps({ + "kernel_repo": str(repo), + }), encoding="utf-8") + (state / "run.json").write_text( + json.dumps({"finished": True}), encoding="utf-8" + ) + + +# --------------------------------------------------------------------------- # +# Core function +# --------------------------------------------------------------------------- # + +def test_rename_kernel_repo_happy_path(tmp_path, monkeypatch): + repo_root = tmp_path / "kernel-repos" + workspace = tmp_path / "workspaces" / "task-1" + state = tmp_path / "state" / "task-1" + workspace.mkdir(parents=True) + made = _make_repo(workspace, repo_root) + _seed_task_files(workspace, state, made["repo"], "task-1") + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(repo_root)) + + result = rename_kernel_repo(workspace, "repo-new", state_dir=state) + + assert result["renamed"] is True + assert result["old_name"] == "repo-old" + assert result["new_name"] == "repo-new" + new_repo = repo_root / "repo-new" + assert new_repo.is_dir() + assert not (repo_root / "repo-old").exists() + assert result["new_repo"] == str(new_repo) + + # symlink re-pointed + assert workspace.joinpath("main").resolve() == new_repo + + # JSON references updated + req = json.loads((state / "requirements.json").read_text(encoding="utf-8")) + assert req["target_repo_path"] == "repo-new" + assert req["label"] == "repo-new" + plan = json.loads((workspace / "plan.json").read_text(encoding="utf-8")) + assert plan["kernel_repo"] == str(new_repo) + + # linked worktree survives the move and still commits + assert str(made["worktree"]) in _git(new_repo, "worktree", "list").stdout + (made["worktree"] / "f.txt").write_text("hi\nmore\n", encoding="utf-8") + _git(made["worktree"], "add", "f.txt") + _git(made["worktree"], "commit", "-qm", "wt commit after move") + + # timeline event recorded + lines = (state / "timeline.jsonl").read_text(encoding="utf-8").splitlines() + assert any("kernel_repo_renamed" in line for line in lines) + + +@pytest.mark.parametrize("bad_name", [ + "", " ", "a/b", "../escape", "a b", ".hidden", "-lead", +]) +def test_rejects_invalid_names(tmp_path, monkeypatch, bad_name): + repo_root = tmp_path / "kernel-repos" + workspace = tmp_path / "workspaces" / "task-1" + workspace.mkdir(parents=True) + made = _make_repo(workspace, repo_root) + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(repo_root)) + with pytest.raises(ValueError): + rename_kernel_repo(workspace, bad_name) + assert (repo_root / "repo-old").is_dir() # untouched + + +def test_rejects_collision(tmp_path, monkeypatch): + repo_root = tmp_path / "kernel-repos" + workspace = tmp_path / "workspaces" / "task-1" + workspace.mkdir(parents=True) + made = _make_repo(workspace, repo_root) + (repo_root / "repo-new").mkdir() + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(repo_root)) + with pytest.raises(ValueError, match="already exists"): + rename_kernel_repo(workspace, "repo-new") + + +def test_rejects_missing_symlink(tmp_path): + workspace = tmp_path / "workspaces" / "task-1" + workspace.mkdir(parents=True) + with pytest.raises(ValueError, match="not a symlink"): + rename_kernel_repo(workspace, "repo-new") + + +def test_rejects_running_task(tmp_path, monkeypatch): + repo_root = tmp_path / "kernel-repos" + workspace = tmp_path / "workspaces" / "task-1" + state = tmp_path / "state" / "task-1" + workspace.mkdir(parents=True) + made = _make_repo(workspace, repo_root) + state.mkdir(parents=True) + (state / "run.json").write_text( + json.dumps({"finished": False}), encoding="utf-8" + ) + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(repo_root)) + with pytest.raises(RuntimeError, match="still running"): + rename_kernel_repo(workspace, "repo-new", state_dir=state) + assert (repo_root / "repo-old").is_dir() + + +def test_allows_when_no_run_state_and_no_pid(tmp_path, monkeypatch): + repo_root = tmp_path / "kernel-repos" + workspace = tmp_path / "workspaces" / "task-1" + state = tmp_path / "state" / "task-1" + workspace.mkdir(parents=True) + made = _make_repo(workspace, repo_root) + state.mkdir(parents=True) # no run.json, no orchestrator.pid + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(repo_root)) + result = rename_kernel_repo(workspace, "repo-new", state_dir=state) + assert result["renamed"] is True + + +def test_cli_smoke(tmp_path, monkeypatch): + repo_root = tmp_path / "kernel-repos" + workspace = tmp_path / "workspaces" / "task-1" + state = tmp_path / "state" / "task-1" + workspace.mkdir(parents=True) + made = _make_repo(workspace, repo_root) + _seed_task_files(workspace, state, made["repo"], "task-1") + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(repo_root)) + import sys + from metainfer.tasks.dcu_kernel_auto_opt.orchestrator import ( + rename_kernel_repo as module, + ) + proc = subprocess.run( + [sys.executable, "-m", + "metainfer.tasks.dcu_kernel_auto_opt.orchestrator.rename_kernel_repo", + str(state), str(workspace), "repo-cli"], + capture_output=True, text=True, check=True, + ) + assert "repo-cli" in proc.stdout + assert (repo_root / "repo-cli").is_dir() + + +# --------------------------------------------------------------------------- # +# Web route +# --------------------------------------------------------------------------- # + +def test_rename_repo_route(client, isolated_env, monkeypatch, tmp_path): + from metainfer.server import paths as _paths + task_id = "dkao-rename-1" + state_dir = isolated_env["home"] / "tasks" / task_id + workspace_dir = _paths.workspace_dir(task_id) + repo_root = tmp_path / "kernel-repos" + monkeypatch.setenv("METAINFER_KERNEL_REPOS", str(repo_root)) + register_dkao_task(state_dir, workspace_dir, task_id) + made = _make_repo(workspace_dir, repo_root) + _seed_task_files(workspace_dir, state_dir, made["repo"], task_id) + + resp = client.post( + f"/api/dcu-kernel-auto-opt/{task_id}/rename-repo", + json={"new_name": "repo-new"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["new_name"] == "repo-new" + assert (repo_root / "repo-new").is_dir() + assert not (repo_root / "repo-old").exists() + + # Renaming while the orchestrator is running is refused. + (state_dir / "run.json").write_text( + json.dumps({"finished": False}), encoding="utf-8" + ) + resp = client.post( + f"/api/dcu-kernel-auto-opt/{task_id}/rename-repo", + json={"new_name": "repo-while-running"}, + ) + assert resp.status_code == 400 + assert "still running" in resp.json()["detail"] + assert not (repo_root / "repo-while-running").exists() + + # Invalid names are refused with a clear error. + resp = client.post( + f"/api/dcu-kernel-auto-opt/{task_id}/rename-repo", + json={"new_name": "../bad"}, + ) + assert resp.status_code == 400 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_skill_store.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_skill_store.py new file mode 100644 index 00000000..54333e68 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_skill_store.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +import pytest + +from ..orchestrator.config import load_config +from ..orchestrator.skill_store import ( + _apply_fuse_decision, + _parse_fuse_decision, + dsh_skills_root, + generate_merged_skill, + generate_worker_skill, + list_skill_library, + publish_skill, + rollback_skill, + sync_skill_libraries, +) + + +def test_skill_library_lists_and_publishes_without_overwrite( + tmp_path, monkeypatch +): + dsh_root = tmp_path / "dsh-skills" + ccb_root = tmp_path / "claude-skills" + existing = dsh_root / "already-here" + existing.mkdir(parents=True) + (existing / "SKILL.md").write_text( + "---\nname: already-here\ndescription: Existing test skill.\n---\n" + ) + workspace = tmp_path / "workspace" + pending = workspace / "skills" / "pending" / "new-kernel-skill" + pending.mkdir(parents=True) + (pending / "SKILL.md").write_text( + "---\nname: new-kernel-skill\ndescription: New test skill.\n---\n" + ) + (pending / "manifest.json").write_text( + '{"kind":"merged","source":"main_agent"}' + ) + monkeypatch.setenv("DSH_SKILLS_DIR", str(dsh_root)) + monkeypatch.setenv("METAINFER_CLAUDE_SKILLS_DIR", str(ccb_root)) + + library = list_skill_library(workspace) + assert [item["name"] for item in library["existing"]] == ["already-here"] + assert [item["name"] for item in library["pending"]] == ["new-kernel-skill"] + + result = publish_skill(workspace, "new-kernel-skill") + assert result["status"] == "existing" + assert (dsh_root / "new-kernel-skill" / "SKILL.md").is_file() + assert not (dsh_root / "new-kernel-skill" / "manifest.json").exists() + assert not pending.exists() + # publish mirrors the dsh library into the ccb library automatically. + assert (ccb_root / "new-kernel-skill" / "SKILL.md").is_file() + + with pytest.raises(FileNotFoundError): + publish_skill(workspace, "new-kernel-skill") + + +def test_sync_skill_libraries_mirrors_dsh_to_ccb(tmp_path, monkeypatch): + dsh_root = tmp_path / "dsh" + ccb_root = tmp_path / "ccb" + (dsh_root / "alpha").mkdir(parents=True) + (dsh_root / "alpha" / "SKILL.md").write_text("# alpha v2\n", encoding="utf-8") + (dsh_root / "beta").mkdir(parents=True) + (dsh_root / "beta" / "SKILL.md").write_text("# beta\n", encoding="utf-8") + (ccb_root / "alpha").mkdir(parents=True) + (ccb_root / "alpha" / "SKILL.md").write_text("# alpha v1\n", encoding="utf-8") + (ccb_root / "ccb-only").mkdir(parents=True) + (ccb_root / "ccb-only" / "SKILL.md").write_text("# ccb only\n", encoding="utf-8") + monkeypatch.setenv("DSH_SKILLS_DIR", str(dsh_root)) + monkeypatch.setenv("METAINFER_CLAUDE_SKILLS_DIR", str(ccb_root)) + + summary = sync_skill_libraries(workspace_dir=tmp_path / "ws") + + assert summary["added"] == ["beta"] + assert summary["updated"] == ["alpha"] + assert summary["ccb_only"] == ["ccb-only"] + assert (ccb_root / "beta" / "SKILL.md").read_text() == "# beta\n" + assert (ccb_root / "alpha" / "SKILL.md").read_text() == "# alpha v2\n" + # the overwritten ccb skill was backed up + assert len(list((ccb_root / "alpha").glob("SKILL.md.bak-*"))) == 1 + # idempotent + summary2 = sync_skill_libraries(workspace_dir=tmp_path / "ws") + assert summary2["added"] == [] + assert summary2["updated"] == [] + assert summary2["skipped"] == ["alpha", "beta"] + + +def test_rollback_skill_restores_backup(tmp_path, monkeypatch): + dsh_root = tmp_path / "dsh" + ccb_root = tmp_path / "ccb" + (dsh_root / "alpha").mkdir(parents=True) + (dsh_root / "alpha" / "SKILL.md").write_text("# v2\n", encoding="utf-8") + monkeypatch.setenv("DSH_SKILLS_DIR", str(dsh_root)) + monkeypatch.setenv("METAINFER_CLAUDE_SKILLS_DIR", str(ccb_root)) + (dsh_root / "alpha" / "SKILL.md.bak-1").write_text("# v1\n", encoding="utf-8") + sync_skill_libraries(workspace_dir=tmp_path / "ws") + + result = rollback_skill("alpha", workspace_dir=tmp_path / "ws") + assert result["restored_from"].endswith("SKILL.md.bak-1") + assert (dsh_root / "alpha" / "SKILL.md").read_text() == "# v1\n" + assert (ccb_root / "alpha" / "SKILL.md").read_text() == "# v1\n" + + with pytest.raises(FileNotFoundError): + rollback_skill("alpha", workspace_dir=tmp_path / "ws") + + +def test_fuse_decision_parse_and_apply(tmp_path, monkeypatch): + dsh_root = tmp_path / "dsh" + ccb_root = tmp_path / "ccb" + (dsh_root / "existing-skill").mkdir(parents=True) + (dsh_root / "existing-skill" / "SKILL.md").write_text( + "---\nname: existing-skill\ndescription: Old scope.\n---\n# old body\n", + encoding="utf-8", + ) + monkeypatch.setenv("DSH_SKILLS_DIR", str(dsh_root)) + monkeypatch.setenv("METAINFER_CLAUDE_SKILLS_DIR", str(ccb_root)) + + # a "new" decision + decision = _parse_fuse_decision( + '{"action": "new", "name": "my-new-skill", ' + '"description": "New scope", "content": "# new body"}' + ) + assert decision["action"] == "new" + applied = _apply_fuse_decision(decision) + assert applied["name"] == "my-new-skill" + text = (dsh_root / "my-new-skill" / "SKILL.md").read_text(encoding="utf-8") + assert text.startswith("---\nname: my-new-skill\n") + assert "# new body" in text + + # a "merge" decision backs up + replaces the body, keeping frontmatter + merge = _parse_fuse_decision( + '{"action": "merge", "name": "existing-skill", ' + '"content": "# merged body"}' + ) + applied = _apply_fuse_decision(merge) + assert applied["action"] == "merge" + assert applied["backup"] is not None + assert "old body" in applied["diff"] + text = (dsh_root / "existing-skill" / "SKILL.md").read_text(encoding="utf-8") + assert text.startswith("---\nname: existing-skill\n") + assert "# merged body" in text + + # fenced JSON also parses + fenced = _parse_fuse_decision( + '```json\n{"action": "new", "name": "fenced-skill", "content": "x"}\n```' + ) + assert fenced["name"] == "fenced-skill" + + with pytest.raises(ValueError): + _parse_fuse_decision("the agent wrote prose instead of JSON") + + +def test_dsh_skills_root_honors_env(monkeypatch, tmp_path): + monkeypatch.setenv("DSH_SKILLS_DIR", str(tmp_path / "x")) + assert dsh_skills_root() == tmp_path / "x" + + +def test_agent_authored_skill_drafts_keep_managed_frontmatter(tmp_path): + req = { + "task_id": "skill-test", + "answers": { + "execution_mode": "Generate & optimize (auto-create kernel repo)", + "operator": "Quantized GEMM", + "dtype": "INT8 W8A8", + "target_hardware": "gfx928", + "kernel_language": "HIP C++", + "shape_config": """ +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""", + }, + } + config = load_config(req) + assignment = config.assignments[0] + worker = generate_worker_skill( + config=config, + assignment=assignment, + workspace_dir=tmp_path, + agent_draft="# Measured worker evidence\n\nOnly measured facts.", + ) + merged = generate_merged_skill( + config=config, + assignments=[assignment], + workspace_dir=tmp_path, + agent_draft="# Main synthesis\n\nRoute m2 to worker_0.", + ) + + worker_text = ( + tmp_path / "skills" / "pending" / worker["name"] / "SKILL.md" + ).read_text(encoding="utf-8") + merged_text = ( + tmp_path / "skills" / "pending" / merged["name"] / "SKILL.md" + ).read_text(encoding="utf-8") + assert worker_text.startswith("---\nname:") + assert "# Measured worker evidence" in worker_text + assert merged_text.startswith("---\nname:") + assert "# Main synthesis" in merged_text diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_store.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_store.py new file mode 100644 index 00000000..6310e5e2 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_store.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +import pytest + +from ..orchestrator import variant_store as vs +from ..orchestrator.variant_store import ( + derive_variant_meta, + dtype_slug, + model_slug, + operator_type_slug, + parse_shape_meta, +) + + +def test_slugs(): + assert operator_type_slug("Quantized GEMM") == "gemm" + assert operator_type_slug("RMSNorm / LayerNorm") == "rmsnorm" + assert dtype_slug("INT8 W8A8") == "int8w8a8" + assert dtype_slug("FP16 / BF16") == "fp16bf16" + assert model_slug("Hy3 (Hunyuan 3)") == "hy3" + assert model_slug("DeepSeek V4 Flash") == "deepseek-v4" + assert model_slug("MiniMax M3") == "minimax-m3" + assert model_slug("GLM5.2") == "glm52" + + +def test_parse_shape_meta(): + assert parse_shape_meta("hy3_tp4_o_proj_m4096") == { + "tp": 4, "operator": "o_proj", "m": 4096, + } + assert parse_shape_meta("hy3_tp4_shared_gate_up_proj_m4096") == { + "tp": 4, "operator": "shared_gate_up_proj", "m": 4096, + } + assert parse_shape_meta("tp4_wqkv_a_m2") == { + "tp": 4, "operator": "wqkv_a", "m": 2, + } + + +def test_derive_variant_meta(): + answers = {"operator": "Quantized GEMM", "dtype": "INT8 W8A8", "model": "Hy3 (Hunyuan 3)"} + meta = derive_variant_meta(answers, "hy3_tp4_o_proj_m4096") + assert meta["family"] == "int8w8a8-gemm" + assert meta["model"] == "hy3" + assert meta["tp"] == 4 + assert meta["operator_name"] == "o_proj" + assert meta["m"] == 4096 + + +def test_variant_path_matches_requested_layout(tmp_path, monkeypatch): + monkeypatch.setattr(vs, "variant_root", lambda: tmp_path) + answers = {"operator": "Quantized GEMM", "dtype": "INT8 W8A8", "model": "Hy3 (Hunyuan 3)"} + meta = derive_variant_meta(answers, "hy3_tp4_o_proj_m4096") + path = vs.variant_path(meta) + assert path == tmp_path / "int8w8a8-gemm" / "hy3" / "TP4" / "M4096" / "o_proj.hip" + + +def test_add_variant_writes_and_replaces(tmp_path, monkeypatch): + monkeypatch.setattr(vs, "variant_root", lambda: tmp_path) + answers = {"operator": "Quantized GEMM", "dtype": "INT8 W8A8", "model": "Hy3 (Hunyuan 3)"} + meta = derive_variant_meta(answers, "hy3_tp4_o_proj_m4096") + r1 = vs.add_variant( + meta=meta, kernel_source="KERNEL_OLD", commit="abc", + metrics={"median_us": 803.9, "logical_tops": 85.48, "speedup": 24.73}, + source_task="task-1", + ) + assert r1["action"] == "added" + target = vs.variant_path(meta) + assert target.is_file() + assert "KERNEL_OLD" in target.read_text() + assert "@@variant shape=hy3_tp4_o_proj_m4096" in target.read_text() + assert "median_us=803.9" in target.read_text() + + r2 = vs.add_variant( + meta=meta, kernel_source="KERNEL_NEW", commit="def", + metrics={"median_us": 800.0}, + source_task="task-1", + ) + assert r2["action"] == "updated" + assert r2["backup"] is not None + text = target.read_text() + assert "KERNEL_NEW" in text + assert "KERNEL_OLD" not in text + + +def test_list_variant_index_walks_tree(tmp_path, monkeypatch): + monkeypatch.setattr(vs, "variant_root", lambda: tmp_path) + answers = {"operator": "Quantized GEMM", "dtype": "INT8 W8A8", "model": "Hy3 (Hunyuan 3)"} + vs.add_variant( + meta=derive_variant_meta(answers, "hy3_tp4_o_proj_m4096"), + kernel_source="A", commit="x", metrics={}, + ) + vs.add_variant( + meta=derive_variant_meta(answers, "hy3_tp4_qkv_proj_m4096"), + kernel_source="B", commit="y", metrics={}, + ) + index = vs.list_variant_index() + shapes = {item["shape"] for item in index} + assert shapes == {"hy3_tp4_o_proj_m4096", "hy3_tp4_qkv_proj_m4096"} + o_proj = next(item for item in index if item["shape"] == "hy3_tp4_o_proj_m4096") + assert o_proj["family"] == "int8w8a8-gemm" + assert o_proj["model"] == "hy3" + assert o_proj["tp"] == "TP4" + assert o_proj["m"] == "M4096" + assert o_proj["operator"] == "o_proj" + + +def test_parse_variant_header_extracts_metrics(): + text = ( + "// @@variant shape=hy3_tp4_o_proj_m16 commit=abc123 added=2026-08-26\n" + "// median_us=18.22 p90_us=18.24 speedup=3.0\n" + "// source=hy3-dsh-tp4-m16-1-7f1fb1d1\n" + "// body comments like K=4096 must not surface\n" + "// @@end\n" + "kernel source...\n" + ) + parsed = vs._parse_variant_header(text) + assert parsed["shape"] == "hy3_tp4_o_proj_m16" + assert parsed["commit"] == "abc123" + assert parsed["source"] == "hy3-dsh-tp4-m16-1-7f1fb1d1" + assert parsed["median_us"] == 18.22 + assert parsed["p90_us"] == 18.24 + assert parsed["speedup"] == 3.0 + assert "K" not in parsed + assert "shape" not in parsed or parsed["shape"] != "m16" + + +def test_list_variant_index_surfaces_header_metrics(tmp_path, monkeypatch): + monkeypatch.setattr(vs, "variant_root", lambda: tmp_path) + target = ( + tmp_path / "int8w8a8-gemm" / "hy3" / "TP4" / "M16" / "qkv_proj.hip" + ) + target.parent.mkdir(parents=True) + target.write_text( + "// @@variant shape=hy3_tp4_qkv_proj_m16 commit=abc added=2026-08-26\n" + "// median_us=26.05 p90_us=26.09\n" + "// source=task-1\n" + "// @@end\nkernel\n", + encoding="utf-8", + ) + index = vs.list_variant_index() + assert len(index) == 1 + v = index[0] + assert v["shape"] == "hy3_tp4_qkv_proj_m16" + assert v["median_us"] == 26.05 + assert v["p90_us"] == 26.09 + assert v["source"] == "task-1" + assert v["commit"] == "abc" + assert v["operator"] == "qkv_proj" + + +def _add_with_guard(tmp_path, monkeypatch, median_us): + monkeypatch.setattr(vs, "variant_root", lambda: tmp_path) + answers = {"operator": "Quantized GEMM", "dtype": "INT8 W8A8", "model": "Hy3 (Hunyuan 3)"} + meta = derive_variant_meta(answers, "hy3_tp4_qkv_proj_m16") + return vs.add_variant( + meta=meta, + kernel_source="KERNEL", + commit="c", + metrics={"median_us": median_us}, + source_task="task-x", + reject_slower_than_existing=True, + ) + + +def test_reject_slower_replacement_guard(tmp_path, monkeypatch): + assert _add_with_guard(tmp_path, monkeypatch, 26.05)["action"] == "added" + # faster candidate replaces + assert _add_with_guard(tmp_path, monkeypatch, 24.30)["action"] == "updated" + # equal candidate replaces + assert _add_with_guard(tmp_path, monkeypatch, 24.30)["action"] == "updated" + # strictly slower candidate is rejected + with pytest.raises(ValueError, match="rejecting slower variant replacement"): + _add_with_guard(tmp_path, monkeypatch, 30.0) + # the rejected replacement did not overwrite the variant + text = vs.variant_path( + derive_variant_meta( + {"operator": "Quantized GEMM", "dtype": "INT8 W8A8", "model": "Hy3 (Hunyuan 3)"}, + "hy3_tp4_qkv_proj_m16", + ) + ).read_text() + assert "median_us=24.3" in text or "median_us=24.30" in text + + +def test_reject_slower_guard_skips_when_median_unknown(tmp_path, monkeypatch): + monkeypatch.setattr(vs, "variant_root", lambda: tmp_path) + answers = {"operator": "Quantized GEMM", "dtype": "INT8 W8A8", "model": "Hy3 (Hunyuan 3)"} + meta = derive_variant_meta(answers, "hy3_tp4_o_proj_m16") + vs.add_variant( + meta=meta, kernel_source="OLD", commit="a", + metrics={}, source_task="task-1", + ) + # existing has no median in header; candidate median present -> no guard + result = vs.add_variant( + meta=meta, kernel_source="NEW", commit="b", + metrics={"median_us": 999.0}, source_task="task-2", + reject_slower_than_existing=True, + ) + assert result["action"] == "updated" + assert "NEW" in vs.variant_path(meta).read_text() diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_baselines.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_baselines.py new file mode 100644 index 00000000..0eb9d65a --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_baselines.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import pytest + +from ..orchestrator.w8a8_baselines import fixed_triton_graph_baseline + + +@pytest.mark.parametrize( + ("shape", "expected"), + [ + ({"M": 2, "N": 1536, "K": 4096}, 66.924), + ({"M": 16, "N": 1536, "K": 4096}, 66.597), + ({"M": 2, "N": 8192, "K": 1024}, 47.181), + ({"M": 16, "N": 8192, "K": 1024}, 48.049), + ({"M": 2, "N": 4096, "K": 2048}, 54.453), + ({"M": 16, "N": 4096, "K": 2048}, 54.617), + ({"M": 2, "N": 1024, "K": 4096}, 60.389), + ({"M": 16, "N": 1024, "K": 4096}, 60.965), + ({"M": 2, "N": 4096, "K": 512}, 18.975), + ({"M": 16, "N": 4096, "K": 512}, 21.372), + ({"M": 3072, "N": 1536, "K": 4096}, 9488.0), + ({"M": 3072, "N": 8192, "K": 1024}, 15437.0), + ({"M": 3072, "N": 4096, "K": 2048}, 14680.0), + ({"M": 3072, "N": 1024, "K": 4096}, 5985.0), + ({"M": 3072, "N": 4096, "K": 512}, 4199.0), + ({"M": 4096, "N": 1536, "K": 4096}, 13247.253), + ({"M": 4096, "N": 8192, "K": 1024}, 20590.225), + ({"M": 4096, "N": 4096, "K": 2048}, 19881.949), + ({"M": 4096, "N": 1024, "K": 4096}, 8790.192), + ({"M": 4096, "N": 4096, "K": 512}, 5545.821), + ], +) +def test_fixed_triton_graph_baseline(shape, expected): + record = fixed_triton_graph_baseline("shape", shape) + assert record["median_us"] == expected + assert record["baseline_kind"] == "triton_graph" + expected_scope = ( + "prefill_graph_replay" if int(shape["M"]) > 16 + else "decode_graph_replay" + ) + assert record["timing_scope"] == expected_scope + assert record["distribution_stats_available"] is False + + +def test_prefill_baseline_has_prefill_timing_scope(): + record = fixed_triton_graph_baseline( + "tp4_wqkv_a_m3072", + {"M": 3072, "N": 1536, "K": 4096}, + ) + assert record["timing_scope"] == "prefill_graph_replay" + + +def test_fixed_baseline_rejects_unknown_shape(): + with pytest.raises(ValueError, match="no fixed Triton Graph baseline"): + fixed_triton_graph_baseline( + "m4_wqkv_a", {"M": 4, "N": 1536, "K": 4096} + ) + + +@pytest.mark.parametrize( + ("shape", "expected"), + [ + # Hy3 TP8 M=4096 (measured 2026-08-27, CUDA-graph replay). + ({"tp_size": 8, "M": 4096, "N": 1280, "K": 4096}, 19381.383), + ({"tp_size": 8, "M": 4096, "N": 4096, "K": 1024}, 14151.792), + ({"tp_size": 8, "M": 4096, "N": 384, "K": 4096}, 16712.231), + ({"tp_size": 8, "M": 4096, "N": 4096, "K": 192}, 10762.616), + # MiniMax M3 TP8 M=4096. + ({"tp_size": 8, "M": 4096, "N": 1280, "K": 6144}, 41000.618), + ({"tp_size": 8, "M": 4096, "N": 1536, "K": 6144}, 35604.698), + ({"tp_size": 8, "M": 4096, "N": 6144, "K": 1024}, 19897.517), + ({"tp_size": 8, "M": 4096, "N": 768, "K": 6144}, 31882.257), + ({"tp_size": 8, "M": 4096, "N": 6144, "K": 384}, 22389.162), + # GLM5.2 TP8 M=4096. + ({"tp_size": 8, "M": 4096, "N": 2624, "K": 6144}, 68661.417), + ({"tp_size": 8, "M": 4096, "N": 2048, "K": 2048}, 14385.571), + ({"tp_size": 8, "M": 4096, "N": 3584, "K": 512}, 6939.939), + ({"tp_size": 8, "M": 4096, "N": 6144, "K": 2048}, 54775.797), + ({"tp_size": 8, "M": 4096, "N": 512, "K": 6144}, 21254.589), + ({"tp_size": 8, "M": 4096, "N": 6144, "K": 256}, 17700.262), + ], +) +def test_model_catalog_tp8_m4096_baselines(shape, expected): + record = fixed_triton_graph_baseline("shape", shape) + assert record["median_us"] == expected + assert record["timing_scope"] == "prefill_graph_replay" + + +def test_bootstrap_metrics_are_kept_separate(): + bootstrap = {"passed": True, "median_us": 123.0} + record = fixed_triton_graph_baseline( + "m2_wqkv_a", + {"M": 2, "N": 1536, "K": 4096}, + bootstrap_metrics=bootstrap, + ) + assert record["median_us"] == 66.924 + assert record["bootstrap_metrics"] == bootstrap + + +@pytest.mark.parametrize( + ("shape", "eager_median", "eager_p90", "graph_median", "graph_p90"), + [ + ((16, 1536, 4096), 97.660, 98.404, 66.591, 67.347), + ((16, 4096, 1024), 95.284, 96.092, 32.137, 35.969), + ((16, 8192, 1024), 96.204, 96.796, 48.274, 50.794), + ((16, 512, 4096), 96.056, 97.788, 59.210, 59.962), + ((16, 4096, 256), 94.540, 95.740, 14.849, 18.001), + ((3072, 1536, 4096), 9423.243, 9446.877, 9422.539, 9486.942), + ((3072, 4096, 1024), 7637.209, 8033.128, 7617.880, 7686.395), + ((3072, 8192, 1024), 15461.000, 15528.571, 15466.648, 15958.569), + ((3072, 512, 4096), 2775.526, 2780.134, 2740.757, 2751.685), + ((3072, 4096, 256), 2414.377, 2424.569, 2370.583, 2376.343), + ], +) +def test_tp8_hot_cache_baselines( + shape, eager_median, eager_p90, graph_median, graph_p90 +): + m, n, k = shape + record = fixed_triton_graph_baseline( + "tp8_shape", {"tp_size": 8, "M": m, "N": n, "K": k} + ) + assert record["tp_size"] == 8 + assert record["median_us"] == graph_median + assert record["p90_us"] == graph_p90 + assert record["eager_median_us"] == eager_median + assert record["eager_p90_us"] == eager_p90 + assert record["cache_state"] == "hot" + assert record["distribution_stats_available"] is True + expected_protocol = ( + {"warmups": 50, "samples": 50, "launches_per_sample": 20} + if m == 16 + else {"warmups": 10, "samples": 20, "launches_per_sample": 5} + ) + assert record["measurement_protocol"] == expected_protocol + + +def test_same_shape_uses_tp_specific_baseline(): + shape = {"M": 16, "N": 1536, "K": 4096} + assert fixed_triton_graph_baseline("tp4", {**shape, "tp_size": 4})[ + "median_us" + ] == 66.597 + assert fixed_triton_graph_baseline("tp8", {**shape, "tp_size": 8})[ + "median_us" + ] == 66.591 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_bench.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_bench.py new file mode 100644 index 00000000..f19146ec --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_bench.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +try: + import torch +except ModuleNotFoundError: + torch = None + +from ..assets.w8a8_bench import ( + exact_w8a8_reference, + reference_cache_path, + save_reference_cache, + validate_profile_protocol, + w8a8_seed, +) + + +@pytest.mark.skipif(torch is None, reason="PyTorch is not installed") +def test_exact_reference_uses_integer_dot_before_scaling(): + a = torch.tensor([[127, 127, 127], [-127, 1, 2]], dtype=torch.int8) + b = torch.tensor( + [[127, -127], [127, 3], [127, 4]], dtype=torch.int8 + ) + a_scale = torch.tensor([[0.5], [0.25]], dtype=torch.float32) + b_scale = torch.tensor([[0.125], [0.75]], dtype=torch.float32) + + actual = exact_w8a8_reference(a, b, a_scale, b_scale) + dot = torch.tensor( + [ + [3 * 127 * 127, 127 * (-127 + 3 + 4)], + [-127 * 127 + 127 + 2 * 127, 127 * 127 + 3 + 8], + ], + dtype=torch.int64, + ) + expected = ( + dot.float() * a_scale * b_scale.T + ).to(torch.bfloat16) + + assert torch.equal(actual, expected) + + +def test_profile_protocol_requires_one_post_marker_replay(): + validate_profile_protocol(True, 0, 1, 1) + with pytest.raises(ValueError, match="profile-only requires"): + validate_profile_protocol(True, 0, 1, 2) + validate_profile_protocol(False, 100, 30, 100) + + +def test_w8a8_seed_is_deterministic_and_shape_scoped(): + assert w8a8_seed(4096, 2304, 6144) == 20260724 + 4096 + 2304 + 6144 + assert w8a8_seed(4096, 2304, 6144) == w8a8_seed(4096, 2304, 6144) + assert w8a8_seed(4096, 2304, 6144) != w8a8_seed(4096, 2304, 6145) + + +def test_reference_cache_path_matches_benchmark_naming(): + path = reference_cache_path(4096, 2304, 6144, Path("/tmp/cache")) + assert path == Path("/tmp/cache") / "exact-int64-v1-m4096-n2304-k6144.pt" + + +@pytest.mark.skipif(torch is None, reason="PyTorch is not installed") +def test_save_reference_cache_roundtrip(tmp_path): + reference = torch.tensor( + [[1.5, -2.5], [3.0, 4.5]], dtype=torch.bfloat16 + ) + # Parent directory deliberately does not exist yet: serial validation + # uses a fresh cache dir (final/cache/references) and the save must + # create it, or torch.save fails with "Parent directory does not exist". + path = tmp_path / "cache" / "references" / "exact-int64-v1-m2-n2-k2.pt" + save_reference_cache(reference, path) + assert path.is_file() + loaded = torch.load(path, map_location="cpu", weights_only=True) + assert torch.equal(loaded, reference) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_graph.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_graph.py new file mode 100644 index 00000000..2ad88caa --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_graph.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from pathlib import Path + + +def test_python_graph_wrapper_is_staged_control_plane_code(): + wrapper = ( + Path(__file__).resolve().parent.parent + / "assets" / "w8a8_baseline" / "w8a8_graph.py" + ) + text = wrapper.read_text(encoding="utf-8") + assert "def capture_w8a8_graph(" in text + assert "api.w8a8_gemm_out(" in text + assert "torch.cuda.CUDAGraph()" in text + assert "torch.cuda.graph(graph, stream=capture_stream)" in text + assert "def replay(self) -> torch.Tensor:" in text + + +def test_trusted_harness_requires_graph_replay_timing(): + harness = ( + Path(__file__).resolve().parent.parent + / "assets" / "w8a8_bench.py" + ) + text = harness.read_text(encoding="utf-8") + assert "capture_candidate_graph(candidate, out)" in text + assert '"timing_mode": "cuda_graph_replay"' in text + assert "graph_runner.replay()" in text diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/indexer_wq_b.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/indexer_wq_b.hip new file mode 100644 index 00000000..143caa0b --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/indexer_wq_b.hip @@ -0,0 +1,980 @@ +// @@variant shape=tp8_indexer_wq_b_m16 commit=9b5b710d0a2c4cd7ac22e2037dfe8735059eec0b added=2026-08-26 +// median_us=17.85 p90_us=17.88 +// source=hy3-dsh-tp8-m16-1-adf021ab +// W8A8 INT8 GEMM — gfx928 implementation (worker_1, physical GPU 1) +// +// Assigned shapes: +// tp8_wq_b_m16 M=16, N=4096, K=1024 +// tp8_indexer_wq_b_m16 M=16, N=8192, K=1024 +// +// Iteration 0 (correctness-first bootstrap): one simple scalar int8 +// dot-product kernel maps one thread to one output element and computes the +// complete K loop exactly in int32 before applying the two fp32 scales and +// storing bf16. This remains the generic fallback for every shape that does +// not match the exact tp8_wq_b_m16 guard below. +// +// Iteration 1 (DUMMA bootstrap): the exact tp8_wq_b_m16 shape ran a minimal +// 16x16x32 DUMMA kernel, one 64-thread wavefront per 16x16 output tile +// (256 blocks, 1 wave, direct row-major global fragment loads). It built, +// passed exact int32 correctness and Graph capture, and measured 59.9 us +// median (vs the 32.137 us fixed Triton Graph baseline). Its gfx928 ISA +// showed a latency-bound K loop: the library fragment loader expanded into +// 18 global_load_ubyte with 19 serialized s_waitcnt around a single +// v_mmac_i32_16x16x32_i8 per step, and with ~2 blocks resident per CU there +// were too few independent load->MMA streams to hide that latency. +// +// Iteration 2 (architecture round — measure grid parallelism): changed ONLY +// the launch geometry to two wavefronts per block owning two adjacent 16x16 +// N tiles (block tile 16x32, grid N/32 = 128 blocks = 1.07 blocks/CU, two +// co-resident full-K MMA chains per CU). It measured 83.0 us median — FLAT +// OR WORSE than iteration 1 — falsifying grid parallelism as the binding +// lever: the binding constraint is the per-step byte-load fragment path +// itself (18 global_load_ubyte + 19 s_waitcnt per step), and the mandated +// architecture axis is split-K, not more geometry. +// +// Iteration 3 (architecture round — split-K, this file): the unsplit grid +// has 128 blocks / 120 CUs = 1.07 < 2 blocks/CU and K = 1024 >= 1024, so the +// mandate requires split-K=2 plus a CU-aligned candidate, with int32 +// partials written into the caller workspace and a combine+scale kernel in +// the timed Graph. This round implements a ONE-WAVE ZERO-BARRIER workspace +// split-K pair (also the mandate's alternative geometry branch — it reaches +// double the unsplit parallelism with zero barriers): +// (1) w8a8_dumma_m16n16k32_splitk_partial_kernel — one 64-thread +// wavefront per block, grid = (N/16) * kSplitK = 512 blocks for +// split-K=2 (4.27 blocks/CU >= the two-blocks-per-CU target, 512 +// resident wave-tasks vs 256 in the unsplit grid). Block b owns +// (split = b % kSplitK, tile = b / kSplitK): split s sums its +// 32-aligned K-chunk ([0,512) or [512,1024) for split-K=2, 16 +// m16n16k32 steps per wave, same k-ascending operand order as the +// unsplit chain) with the UNCHANGED direct du_load_matrix_sync path, +// and stores the raw int32 accumulator fragment directly into the +// caller workspace plane partials[split][tile][16][16] via the +// verified gfx928 accumulator ownership (row = lane&15, +// col_mod4 = lane>>4, x[i] -> col_mod4 + 4*i). No LDS, no barriers. +// (2) w8a8_dumma_m16n16k32_combine_kernel — one thread per +// output element, grid = 256 blocks x 256 threads, sums the kSplitK +// planes ASCENDING (s = 0..kSplitK-1; int32 addition is exact and +// order-independent below 2^24, so the total is bit-identical to the +// unsplit accumulator), applies (acc * x_scale[row]) * +// weight_scale[col] in the same left-associative float order as the +// reference, and stores round-to-nearest-even bf16. Launched +// immediately after the partial kernel on the same stream, so both +// kernels are captured into the timed Graph. +// The split factor is the single compile-time constant kSplitKMeasured = 2. +// The CU-aligned candidate kSplitKCUAligned = 15 (grid 3840 = exactly +// 32 blocks/CU on the 120-CU device, non-power-of-two, non-uniform 32-aligned +// K chunks) is compiled into the object and selected by flipping the +// constant, so the mandated occupancy sweep is a one-line change. +// +// Iteration 8 (HIP-only resource round — tune ONE occupancy limiter, the +// LDS footprint, from the accepted kernel's PMC): iterations 4-6 proved the +// per-step direct global fragment loads are poison (reg-prefetch 51.0 us, +// B-only LDS 67.0 us) and that only a fully A+B-staged zero-global K loop +// can win (the iteration-6 fused kernel, 25.97 us) — but the fused kernel +// stages the WHOLE K=1024 per block (35,456 B/block: s_a+s_b 16,640 each + +// s_part 2,048 + s_scale 128), pinning it at 1 block/CU co-resident. This +// round moves that staging to the SPLIT GRID (the exact next step prescribed +// by iteration 7): the accepted iteration-3 grid/workspace/combine Graph is +// kept (512 blocks x 64 threads, split-K=2, 2 int32 partial planes, separate +// combine kernel in the timed Graph), but the partial kernel now stages ONLY +// its own 512-K A chunk (16 x 512 B, row-major x_q) and B chunk (packed +// [N][K] P, 512 contiguous B per column) into LDS once: 16,896 B/block +// (s_a/s_b 16 x 528 = 512+16 padded stride) -> 3 blocks/CU co-resident on +// the 64 KiB LDS/CU (vs 1 block/CU for the fused kernel, 4.27 avg for the +// accepted kernel). The K loop then collapses to the fused kernel's proven +// zero-global chain: per step exactly two ds_read_b64 (A and B fragments, +// raw byte order == the du_load_matrix_sync expansion) + one v_mmac, same +// k-ascending chunk [0,512)/[512,1024) as iteration 3. Byte accounting per +// replay is UNCHANGED (no repeated HBM reads: B 4 MiB unique, A 16 KiB +// L2-hot re-read by 256 tiles, 512 KiB partial write+read, 128 KiB out) — +// occupancy is raised by shrinking the per-block LDS footprint, not by +// re-reading global memory. Transport prerequisite (iteration-5/6-validated): +// for (k,n)==(1024,4096) launch_pack_w8a8_weight now writes the [N][K] +// transpose P[n*K+k] = W[k*N+n] once outside the timed region/Graph, and the +// generic scalar fallback decodes it via a b_transposed flag so the paired +// M=2 validation and every other caller of the exact (n,k) stay bit-exact; +// every other shape keeps the identity pack. +// +// Iteration 9 (HIP-only occupancy probe — the mandated one-constant staged +// split sweep): the accepted iteration-8 kernel's PMC (staged partial, +// profiled 14.56 us, 512 blocks x 64 threads, 16,896 B/block LDS -> 3 blocks +// /CU co-resident on the 64 KiB LDS/CU, 0.75 waves/SIMD resident with the +// 4th SIMD of every CU idle, grid 4.27 blocks/CU -> 1.42 serial block rounds) +// shows the staged transport's per-step chain is short (2 ds_read_b64 + +// v_mmac, ~LDS-latency), so the residual limiter is resident-stream count +// over the one-per-block HBM staging wait, not per-step global latency (the +// direct-load split-K=4 regression of iteration 7 does NOT transfer: that +// transport kept a ~500-600 cycle HBM chain inside EVERY K step). This round +// sweeps the dispatched split to S=4 — the exact next step iteration 8 +// prescribed: 1024 blocks x 64 threads, uniform 32-aligned K chunks +// [0,256)/[256,512)/[512,768)/[768,1024), 8 m16n16k32 steps per block, +// staging shrinks to 8,704 B/block (s_a/s_b 16 x 272 = 256+16 padded +// stride, the same 4-dword-mod-32 bank skew as 528) -> 7 blocks/CU +// co-resident (floor(65536/8704) = 7), 1.75 waves/SIMD with ALL FOUR SIMDs +// busy, 1.22 serial block rounds, 4 int32 partial planes = 1 MiB <= the +// API's 16-plane / 4 MiB workspace. Exactness is preserved by construction: +// identical k-ascending m16n16k32 order within each chunk, chunks tile +// [0,1024) exactly once, the UNCHANGED ascending 4-plane combine sums +// (p0+p1+p2)+p3 == the S=2 two-plane total == the unsplit accumulator +// (int32 exact below 2^24), so the output is bit-identical to iteration 8 +// (expected 0 mismatches). Graph contract unchanged: same two kernels, same +// stream order, 4 planes within the same 16-plane allocation, every partial +// overwritten each launch. +// +// Repair 1 (exact-correctness defect of the S=4 sweep, no architecture +// change): the cooperative staging loop decomposed the 16-B chunk index with +// the S=2-fixed mapping r = c>>5, ko = (c&31)*16 (32 chunks per row). That is +// exact at kLocal=512 but at S=4 (kLocal=256, kChunksLocal=256) c only +// reaches 255, so only A rows / B columns 0..7 were staged and the K loop +// read UNINITIALIZED LDS for rows/columns 8..15 (75% of every tile wrong — +// rows 8..15 in all columns plus columns 8..15 of rows 0..7; observed +// mismatch_count 49121/65536, first mismatch m=0 n=8). Fixed in place by +// deriving the per-row chunk count from the actual chunk length, +// kChunksPerRow = kLocal/16 (r = c/kChunksPerRow, ko = (c%kChunksPerRow)*16), +// which is bit-identical to the old mapping at S=2 (512/16 = 32) and stages +// all 16 rows/columns at S=4; kLocal stays a 32-multiple so every 16-B +// global load/store alignment is unchanged. The S=4 mapping/performance idea +// (1024 blocks, 8,704 B/block LDS, 7 blocks/CU) is preserved. +// +// Iteration 11 (HIP-only consolidation — fused combine tail): the accepted +// S=4 tree's operator-aggregate PMC (partial 10.879 us + combine 3.04 us = +// 13.919 us vs the 15.08 timed median) pins the residual wall on the combine +// STRUCTURE: a second kernel launch (~1.16 us inter-kernel/graph gap) plus a +// 1 MiB int32 plane re-read that is HBM-cold (combine L2 hit rate only 19.3%: +// the 4 MiB B stream evicts the partials), while the sibling TP4/M16 lineage +// (qkv_proj.hip rounds 17/23) proves the canonical M<=32 fix is the FUSED +// ATOMIC COMBINE. This round merges the combine into the partial kernel tail: +// every block stores its plane as before, then tid-0 does __threadfence() + +// atomicAdd(&counters[tile], 1), and the block whose atomicAdd returns the +// kSplitK-th arrival of its replay (monotonic mod-kSplitK: r % kSplitK == +// kSplitK-1; int32 wraparound would need 2^31/4 ~ 536M replays) sums the +// tile's kSplitK planes in ascending slice order (identical per-element int32 +// order as the removed combine), applies the identical left-associative +// fp32 scale and RN-even bf16 conversion, and stores byte-identical output +// bytes (0 mismatches expected). The Graph becomes ONE kernel launch per +// replay. The per-tile counters (256 x int32 = 1 KiB) live in the last +// kCountersBytes of the caller workspace (planes 0..3 end at 1 MiB, far below +// the 4 MiB allocation's tail; no overlap), zeroed once per workspace by an +// async hipMemsetAsync on the caller's stream before the first launch (never +// part of steady-state replay; the mod-kSplitK test is also correct if a +// reset were captured, since arrivals 0..kSplitK-1 still identify the last). +// The staged K loop, transport, pack, exact-shape guards, generic scalar +// fallback and the compile-only kSplitKCUAligned=15 two-kernel branch are +// unchanged; the generic fallback never touches the counters. +// +// Indexer iteration 1 (tp8_indexer_wq_b_m16 — first DUMMA fast path for the +// M=16, N=8192, K=1024 assigned shape): until this round the indexer shape +// fell to the generic scalar fallback (bootstrap median ~316 us). The +// mandated minimal 16x16x32 DUMMA geometry — ONE 64-thread wavefront per +// block, ONE 16x16 N tile per block, no split-K, no workspace partials, no +// combine kernel, one launch per replay — is implemented as +// w8a8_dumma_m16n16k32_indexer_kernel with the accepted sibling transport +// (the same chunked LDS staging + zero-global two-ds_read_b64 + v_mmac K +// chain that wins at N=4096): grid = N/16 = 512 blocks (4.27 blocks/CU on +// the 120-CU device), K=1024 staged in 4 chunks of 256 into s_a/s_b 16 x 272 +// padded-stride buffers (8,704 B/block -> 7 blocks/CU co-resident = 840 +// slots >= 512, so EVERY block is resident from t=0, no second round), one +// single-wave __syncthreads() per chunk (no cross-wave coordination), 8 +// zero-global m16n16k32 steps per chunk, and a direct fragment -> scaled +// RN-even bf16 epilogue (validated gfx928 accumulator ownership; no LDS +// round-trip). B is read from the [N][K] transpose pack, so +// launch_pack_w8a8_weight now also transposes (k,n)==(1024,8192), and the +// generic scalar fallback decodes that pack via b_transposed for the paired +// M=2 / prefill shapes with the same (N,K). +// +// Layout contract (logical tensors, contiguous / row-major): +// x_q [M, K] int8 +// packed_weight [K, N] int8 for every shape except the two exact +// (k,n) pairs (1024,4096) and (1024,8192), whose +// pack is the [N, K] transpose P[n*K+k] = W[k*N+n] +// written by launch_pack_w8a8_weight +// x_scale [M] fp32 +// packed_weight_scale [N] fp32 (identity copy of weight_scale) +// out [M, N] bf16 +// workspace caller-owned int32 partial planes (N=4096 split-K +// path only): +// partials[split][tile][16][16], plane stride +// (N/16)*256 int32 elements +// +// Math: +// out[m, n] = bf16( (int32) sum_k x_q[m,k] * W[k,n] * x_scale[m] +// * weight_scale[n] ) +// +// Exactness: the assigned K=1024 keeps every partial int8 dot below 2^24, so +// the int32 sum (in any grouping order) is bit-identical to the fp32 torch +// reference ((A.float() @ B.float()) * x_scale * weight_scale.T) for the +// timed shapes. The generic fallback is exact in int32 for every supported +// shape (max |dot| = K * 127 * 127 < 2^31). + +#include +#include +#include + +#include + +namespace { + +constexpr int kScalarBlock = 256; // 4 wavefronts of 64 (blockDim % 64 == 0) +constexpr int kPackBlock = 256; + +// DUMMA INT8 m16n16k32 constants: one 64-thread wavefront owns one fragment. +constexpr int kDummaM = 16; +constexpr int kDummaN = 16; +constexpr int kDummaK = 32; +constexpr int kDummaWave = 64; + +// Exact-shape constants for tp8_wq_b_m16 (M=16, N=4096, K=1024). +constexpr int kExactK = 1024; +constexpr int kExactN = 4096; +constexpr int kExactNTiles = kExactN / kDummaN; // 256 +// Exact-shape constant for tp8_indexer_wq_b_m16 (M=16, N=8192, K=1024). +constexpr int kExactNIndexer = 8192; + +// Combine kernel block size. +constexpr int kCombineThreads = 256; + +// Iteration 9 staged occupancy sweep — split factor dispatched this round. +// The exact-shape fast path launches the staged partial kernel below, which +// supports every trusted split in {2..15} via runtime 32-aligned chunks, so +// the occupancy sweep stays a one-constant change. S=4: 1024 blocks x 64 +// threads, K-chunk 256 per block, LDS 8,704 B/block (s_a/s_b 16 x 272 = +// 256+16 padded stride) -> 7 blocks/CU co-resident (vs 3 at S=2), grid +// average 8.53 blocks/CU -> 1.22 serial block rounds (vs 1.42), 1.75 +// waves/SIMD resident with all 4 SIMDs busy (vs 0.75 on 3 of 4 SIMDs), 4 +// int32 partial planes = 1 MiB <= the API's 16-plane / 4 MiB workspace. +constexpr int kSplitKMeasured = 4; +// Implemented CU-aligned candidate: grid = 256 * 15 = 3840 blocks = exactly +// 32 blocks/CU on the 120-CU gfx928 (integer blocks per CU, non-power-of-two, +// non-uniform 32-aligned K chunks). Compiled into the object via the +// iteration-3 direct-load partial kernel; never dispatched (the +// workspace_bytes == -1 guard cannot hold and it does not decode the [N][K] +// pack). +constexpr int kSplitKCUAligned = 15; + +// Iteration 11 fused combine tail: per-tile arrival counters (256 tiles x +// int32 = 1 KiB) live in the last kCountersBytes of the caller workspace +// (planes 0..3 end at 1 MiB, far below the counters' offset in the API's +// 16-plane / 4 MiB allocation, so there is no overlap with partial data). +constexpr int64_t kCountersBytes = kExactNTiles * sizeof(int32_t); // 1 KiB + +// Round-to-nearest-even float -> bfloat16 bit pattern (sibling-validated +// helper; identical RN-even conversion to the hip_bfloat16 stores used by +// the old combine kernel, so the packed tail store is byte-identical). +__device__ __forceinline__ uint16_t float_to_bf16_bits(float f) { + uint32_t u = 0; + __builtin_memcpy(&u, &f, sizeof(u)); + const uint32_t bias = 0x7FFFu + ((u >> 16) & 1u); + u += bias; + return static_cast(u >> 16); +} + +// One-time async zero of the fused-combine arrival counters, guarded by the +// workspace pointer: issued once per workspace before the first launch on +// the caller's stream, so the captured Graph contains only the single fused +// kernel launch and every steady-state replay is memset-free. +static void* g_fused_counters_zeroed = nullptr; + +// --------------------------------------------------------------------------- +// Iteration 3 split-K partial GEMM for the exact tp8_wq_b_m16 shape. +// One 64-thread wavefront per block, grid = (N/16) * kSplitK blocks. Block b +// owns (split = b % kSplitK, tile = b / kSplitK): split s computes the +// 16x16 output tile over its 32-aligned K-chunk [kStart(s), kEnd(s)) in +// k-ascending m16n16k32 steps (16 steps for split-K=2), with the unchanged +// direct row-major global fragment loads, and stores the raw int32 +// accumulator fragment into workspace plane s (tile-major, row stride 16) +// via the verified gfx928 accumulator ownership. Zero barriers, zero LDS. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kDummaWave) +void w8a8_dumma_m16n16k32_splitk_partial_kernel( + const int8_t* __restrict__ x_q, // [16, K] row-major + const int8_t* __restrict__ weight, // [K, N] row-major (identity pack) + int32_t* __restrict__ partials, // [kSplitK][N/16][16][16] int32 + const int n) { + const int split = static_cast(blockIdx.x) % kSplitK; + const int tile = static_cast(blockIdx.x) / kSplitK; + const int lane = static_cast(threadIdx.x); + const int n0 = tile * kDummaN; + + // 32-aligned non-uniform K chunks: kStart(s) = floor(s*K/S) rounded down + // to a DUMMA-K multiple; the last chunk absorbs the remainder. For + // split-K=2 this is exactly [0,512) and [512,1024). + const int kStart = ((split * kExactK) / kSplitK) & ~(kDummaK - 1); + const int kEnd = (split == kSplitK - 1) + ? kExactK + : ((((split + 1) * kExactK) / kSplitK) & + ~(kDummaK - 1)); + + du::dumma::DUFragment a_frag; + du::dumma::DUFragment b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Same per-step instruction stream as iterations 1/2 (direct global + // fragment loads, exact int32 k-ascending accumulation); only the K range + // per wave shrinks by the split factor. Fully unrolled (16 steps for + // split-K=2) like the validated split-K lineage kernels. +#pragma unroll + for (int k0 = kStart; k0 < kEnd; k0 += kDummaK) { + du::dumma::du_load_matrix_sync(a_frag, x_q + k0, kExactK); + du::dumma::du_load_matrix_sync( + b_frag, weight + static_cast(k0) * n + n0, n); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // gfx928 int8 m16n16k32 accumulator ownership, established against + // du_store_matrix_sync: lane%16 selects the row, lane/16 selects col%4, + // and x[i] holds the columns col%4 + 4*i. Store the raw int32 partial + // (no scaling) into the caller workspace plane. + const int row = lane & 15; + const int col_mod4 = lane >> 4; + const int n_tiles = n / kDummaN; + int32_t* __restrict__ p = partials + + (static_cast(split) * n_tiles + tile) * + (kDummaM * kDummaN); +#pragma unroll + for (int i = 0; i < 4; ++i) { + p[row * kDummaN + col_mod4 + 4 * i] = acc_frag.x[i]; + } +} + +// --------------------------------------------------------------------------- +// Iteration 8 split-grid staged partial GEMM for the exact tp8_wq_b_m16 +// shape — the LDS-footprint occupancy tune. Same grid as iteration 3 (512 +// blocks x 64 threads, split = b % kSplitK, tile = b / kSplitK, 32-aligned +// non-uniform chunks [kStart, kEnd)), same workspace partial-plane contract, +// same combine kernel in the timed Graph. The per-step direct global +// fragment loads are replaced by ONE cooperative staging pass per block: +// each lane moves 16-B chunks of the block's own A chunk (x_q rows, 512 B +// each at stride K) and B chunk (packed [N][K] P, 512 contiguous B per +// column) with coalesced dwordx4 global reads and 16-B aligned b128 LDS +// stores — 16 global loads + 16 LDS stores per lane at split-K=2 — followed +// by ONE __syncthreads() and a zero-global K loop whose every step is +// exactly two ds_read_b64 (A and B fragments; raw byte order == the +// du_load_matrix_sync expansion) + one v_mmac (the iteration-6-validated +// fused chain, minus the barrier/s_part handoff). LDS = 16,896 B/block at +// the dispatched S=2 (s_a/s_b 16 x 528 = 512+16 padded stride) -> 3 blocks +// /CU co-resident on the 64 KiB LDS/CU. Template supports every trusted +// split in {2..15}: the buffer covers ceil(K/S) rounded up to a DUMMA-K +// multiple and both loops use runtime chunk bounds, so the mandated sweep +// stays a one-constant change (kSplitKMeasured). +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kDummaWave) +void w8a8_dumma_m16n16k32_splitk_staged_partial_kernel( + const int8_t* __restrict__ x_q, // [16, K] row-major + const int8_t* __restrict__ weight, // packed [N, K] n-major (exact) + int32_t* __restrict__ partials, // [kSplitK][N/16][16][16] int32 + const float* __restrict__ x_scale, // [16] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [16, N] bf16 + int32_t* __restrict__ counters, // [N/16] per-tile arrival counts + const int n) { + constexpr int kChunkMax = + (((kExactK + kSplitK - 1) / kSplitK + kDummaK - 1) / kDummaK) * + kDummaK; // 512 at split-K=2 + constexpr int kStageStride = kChunkMax + 16; // 528 = 512+16 at split-K=2 + + __shared__ __align__(16) int8_t s_a[kDummaM * kStageStride]; + __shared__ __align__(16) int8_t s_b[kDummaN * kStageStride]; + + const int split = static_cast(blockIdx.x) % kSplitK; + const int tile = static_cast(blockIdx.x) / kSplitK; + const int lane = static_cast(threadIdx.x); + const int n0 = tile * kDummaN; + + // Same 32-aligned non-uniform K chunks as the iteration-3 kernel; for + // split-K=2 this is exactly [0,512) and [512,1024). + const int kStart = ((split * kExactK) / kSplitK) & ~(kDummaK - 1); + const int kEnd = (split == kSplitK - 1) + ? kExactK + : ((((split + 1) * kExactK) / kSplitK) & + ~(kDummaK - 1)); + const int kLocal = kEnd - kStart; // 32-multiple, <= kChunkMax + + // Cooperative staging of this block's A and B chunks: every 16-B chunk of + // each matrix is moved exactly once. Per pass each lane issues one + // contiguous dwordx4 global load per matrix (coalesced 16-B sectors: lanes + // 0..31 cover one 512-B row/column segment, lanes 32..63 the next) and one + // 16-B aligned b128 LDS store (stride 528 = 16*33 keeps every store + // 16-B-aligned). c enumerates the kChunksLocal 16-B chunks of ONE matrix + // (A rows == B columns, both 16 x kLocal) and the chunk -> (row, k-offset) + // decomposition uses kChunksPerRow = kLocal/16 so ALL 16 rows/columns are + // staged at every split; the fixed 32-chunks-per-row shift mapping only + // covered rows/columns 0..7 once kLocal dropped to 256 at split-K=4 + // (repair 1: uninitialized-LDS reads for rows/columns 8..15). + const int kChunksLocal = (kLocal * kDummaM) / 16; // 16-B chunks per matrix + const int kChunksPerRow = kLocal / 16; // 16-B chunks per A row / B column + const int64_t b_base = + static_cast(n0) * kExactK + kStart; // packed P column base +#pragma unroll + for (int c = lane; c < kChunksLocal; c += kDummaWave) { + const int r = c / kChunksPerRow; // A row == B column in the 16 + const int ko = (c % kChunksPerRow) * 16; // k offset in the row/column + const int4 va = *reinterpret_cast( + x_q + r * kExactK + kStart + ko); + const int4 vb = *reinterpret_cast( + weight + b_base + static_cast(r) * kExactK + ko); + *reinterpret_cast(s_a + r * kStageStride + ko) = va; + *reinterpret_cast(s_b + r * kStageStride + ko) = vb; + } + __syncthreads(); + + du::dumma::DUFragment a_frag; + du::dumma::DUFragment b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Fragment ownership == the du_load_matrix_sync expansion (validated in + // iterations 3/5/6): A: row = lane&15, kg = lane>>4, x[i] = A[row][kg*8+i]; + // B: col = lane&15, kg = lane>>4, x[i] = P[n0+col][kg*8+i] -- 8 contiguous + // staged bytes per fragment, one ds_read_b64 each, raw byte order (the + // loader's byte-reassembly is a no-op on the 8-byte path, so the manual + // fill is bit-identical). Zero-global, zero-barrier k-ascending chain over + // the block's own chunk. + const int frag_row = lane & 15; // A row == B column + const int kg = lane >> 4; + const int lds_off = frag_row * kStageStride + kg * 8; + const int kSteps = kLocal / kDummaK; +#pragma unroll + for (int s = 0; s < kSteps; ++s) { + *reinterpret_cast(&a_frag.x[0]) = + *reinterpret_cast(s_a + lds_off + s * kDummaK); + *reinterpret_cast(&b_frag.x[0]) = + *reinterpret_cast(s_b + lds_off + s * kDummaK); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Same raw int32 partial-plane store as iteration 3 (verified gfx928 + // accumulator ownership: row = lane&15, col_mod4 = lane>>4, x[i] -> + // col_mod4 + 4*i). + const int row = lane & 15; + const int col_mod4 = lane >> 4; + const int n_tiles = n / kDummaN; + int32_t* __restrict__ p = partials + + (static_cast(split) * n_tiles + tile) * + (kDummaM * kDummaN); +#pragma unroll + for (int i = 0; i < 4; ++i) { + p[row * kDummaN + col_mod4 + 4 * i] = acc_frag.x[i]; + } + + // Iteration 11 fused combine tail (replaces the separate combine kernel in + // the timed Graph; sibling-validated arrival protocol): every block's tid-0 + // releases with __threadfence() + atomicAdd(&counters[tile], 1); the LAST + // arrival for its tile (atomicAdd return value r with r % kSplitK == + // kSplitK-1 — hardware-ordered after all kSplitK sibling stores+fences, no + // spin loop, no residency assumption) sums the tile's planes in ascending + // slice order — the identical per-element int32 accumulation order as the + // old combine — applies the identical (acc * x_scale[row]) * + // weight_scale[col] fp32 scale order, and stores RN-even bf16 at the same + // addresses. The counters are MONOTONIC across replays: replay r's arrivals + // return (r-1)*kSplitK + 0..kSplitK-1, so the kSplitK-th arrival of EVERY + // replay is exactly the r % kSplitK == kSplitK-1 one (int32 wraparound + // would need 2^31/4 ~ 536M replays); launch_w8a8_gemm zeroes them once per + // workspace before the first launch and the generic fallback never touches + // them. Element mapping: lin0 = lane*4 covers the 256-element tile exactly + // once (row = lin0>>4, col0 = lin0&15, 4 consecutive columns per lane — 4 | + // 16 so a run never crosses the 16-column row boundary), each plane is read + // as one perfectly coalesced dwordx4 per lane, and the four 2-B bf16 + // results pack into one 8-B aligned uint64 store — bit-identical output + // bytes to the removed combine kernel (0 mismatches expected). + __syncthreads(); + __shared__ int s_is_last; + if (threadIdx.x == 0) { + __threadfence(); // release: this block's plane stores are visible to + // the observer of its arrival atomic + const int arrived = atomicAdd(&counters[tile], 1); + __threadfence(); // acquire: the reads below (after the barrier) see + // every sibling store released before prior arrivals + s_is_last = ((arrived % kSplitK) == kSplitK - 1); + } + __syncthreads(); + if (s_is_last) { + const int lin0 = lane * 4; // 4 consecutive columns of one row + const int crow = lin0 >> 4; // 0..15 + const int col0 = lin0 & 15; // 0..12 (multiple of 4) + int32_t sums[4] = {0, 0, 0, 0}; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + const int4 plane4 = *reinterpret_cast( + partials + (static_cast(s) * n_tiles + tile) * + (kDummaM * kDummaN) + + lin0); + sums[0] += plane4.x; + sums[1] += plane4.y; + sums[2] += plane4.z; + sums[3] += plane4.w; + } + const int out_col0 = tile * kDummaN + col0; + uint64_t packed = 0; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float scaled = static_cast(sums[e]) * x_scale[crow] * + weight_scale[out_col0 + e]; + packed |= static_cast(float_to_bf16_bits(scaled)) << (16 * e); + } + *reinterpret_cast(out + static_cast(crow) * n + + out_col0) = packed; + } +} + +// --------------------------------------------------------------------------- +// Indexer iteration 1 minimal DUMMA fast path for the exact +// tp8_indexer_wq_b_m16 shape (M=16, N=8192, K=1024): one 64-thread wavefront +// per block, one 16x16 N tile per block, grid = N/16 = 512 blocks. No +// split-K, no workspace partial planes, no combine kernel — the block writes +// its scaled bf16 tile directly to out, so the timed Graph contains exactly +// ONE kernel launch per replay. +// Transport = the accepted sibling staged chain: K is staged in 4 chunks of +// 256 into s_a/s_b (16 x 272 padded stride, 8,704 B/block -> 7 blocks/CU +// co-resident on the 64 KiB LDS/CU = 840 slots >= the 512-block grid, so +// every block is resident from t=0 with no second round). Per chunk: one +// cooperative staging pass (per lane two coalesced 16-B dwordx4 global loads +// — A from the x_q rows, B from the packed [N][K] P columns — and two 16-B +// aligned b128 LDS stores, chunk -> (row, k-offset) decomposition with +// kChunksPerRow = 256/16 = 16 so all 16 rows/columns are staged) + ONE +// single-wave __syncthreads() + 8 zero-global m16n16k32 steps (per step +// exactly two ds_read_b64 + one v_mmac, raw byte order == the +// du_load_matrix_sync expansion). Epilogue: direct fragment -> out store via +// the verified gfx928 accumulator ownership (row = lane&15, col_mod4 = +// lane>>4, x[i] -> col_mod4 + 4*i), same left-associative fp32 scale order +// and RN-even bf16 conversion as the reference. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDummaWave) +void w8a8_dumma_m16n16k32_indexer_kernel( + const int8_t* __restrict__ x_q, // [16, K] row-major + const int8_t* __restrict__ weight, // packed [N, K] n-major + const float* __restrict__ x_scale, // [16] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [16, N] bf16 + const int n) { + constexpr int kChunk = 256; // staged K per chunk (8 DUMMA steps) + constexpr int kStageStride = kChunk + 16; // 272 (same bank skew as 528) + constexpr int kChunksPerRow = kChunk / 16; // 16-B chunks per A row / B col + constexpr int kChunksLocal = (kChunk * kDummaM) / 16; // 256 per matrix + + __shared__ __align__(16) int8_t s_a[kDummaM * kStageStride]; + __shared__ __align__(16) int8_t s_b[kDummaN * kStageStride]; + + const int tile = static_cast(blockIdx.x); + const int lane = static_cast(threadIdx.x); + const int n0 = tile * kDummaN; + + du::dumma::DUFragment a_frag; + du::dumma::DUFragment b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Fragment ownership == the du_load_matrix_sync expansion (validated in the + // sibling lineage): A: row = lane&15, kg = lane>>4, x[i] = A[row][kg*8+i]; + // B: col = lane&15, kg = lane>>4, x[i] = P[n0+col][kg*8+i] -- 8 contiguous + // staged bytes per fragment, one ds_read_b64 each. + const int frag_row = lane & 15; // A row == B column + const int kg = lane >> 4; + const int lds_off = frag_row * kStageStride + kg * 8; + +#pragma unroll + for (int c0 = 0; c0 < kExactK; c0 += kChunk) { + // Cooperative staging of this chunk's A rows and B columns: every 16-B + // chunk of each matrix moved exactly once (4 iterations for 64 lanes). +#pragma unroll + for (int c = lane; c < kChunksLocal; c += kDummaWave) { + const int r = c / kChunksPerRow; // A row == B column + const int ko = (c % kChunksPerRow) * 16; // k offset in the row/column + const int4 va = *reinterpret_cast( + x_q + r * kExactK + c0 + ko); + const int4 vb = *reinterpret_cast( + weight + (static_cast(n0) + r) * kExactK + c0 + ko); + *reinterpret_cast(s_a + r * kStageStride + ko) = va; + *reinterpret_cast(s_b + r * kStageStride + ko) = vb; + } + __syncthreads(); + + // Zero-global k-ascending m16n16k32 chain over this chunk (8 steps). +#pragma unroll + for (int s = 0; s < kChunk / kDummaK; ++s) { + *reinterpret_cast(&a_frag.x[0]) = + *reinterpret_cast(s_a + lds_off + s * kDummaK); + *reinterpret_cast(&b_frag.x[0]) = + *reinterpret_cast(s_b + lds_off + s * kDummaK); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + __syncthreads(); // protect the next chunk's overwrite of s_a/s_b + } + + // Direct scaled epilogue: same gfx928 accumulator ownership and same + // left-associative (acc * x_scale[row]) * weight_scale[col] fp32 order and + // RN-even bf16 conversion as the scalar/combine reference paths. + const int row = lane & 15; + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = n0 + col_mod4 + 4 * i; + const float scaled = static_cast(acc_frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[static_cast(row) * n + col] = + hip_bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// Iteration 3 split-K combine + scale for the exact tp8_wq_b_m16 shape: +// reduce the kSplitK int32 partial planes ascending (each plane is +// [n_tiles][16][16], tile-major with row stride 16), apply +// (acc * x_scale[row]) * weight_scale[col] in the same left-associative +// float order as the reference, and store round-to-nearest-even bf16. +// One thread per output element, grid = ceil(16*n / 256) = 256 blocks. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kCombineThreads) +void w8a8_dumma_m16n16k32_combine_kernel( + const int32_t* __restrict__ partials, + const float* __restrict__ x_scale, // [16] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [16, N] bf16 + const int n) { + const int idx = static_cast(blockIdx.x) * kCombineThreads + + static_cast(threadIdx.x); + const int total = kDummaM * n; + if (idx >= total) { + return; + } + + const int row = idx / n; + const int col = idx - row * n; + const int tile = col / kDummaN; + const int tile_base = tile * (kDummaM * kDummaN) + + row * kDummaN + (col & (kDummaN - 1)); + const int n_tiles = n / kDummaN; + int32_t sum = 0; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + sum += partials[static_cast(s) * n_tiles * + (kDummaM * kDummaN) + + tile_base]; + } + + // Same left-to-right fp32 evaluation order as the reference. + const float scaled = + static_cast(sum) * x_scale[row] * weight_scale[col]; + out[idx] = hip_bfloat16(scaled); +} + +// One thread computes exactly one out[row, col] element. Adjacent lanes map +// to adjacent addresses in the fastest-changing N dimension. +// Iteration 8: b_transposed selects the exact-shape packed [N][K] layout +// (P[n*K + k] = W[k*N + n], produced by launch_pack_w8a8_weight), so the +// logical column col starts at P[col*K] with unit stride; the identity pack +// keeps the legacy row-major [K][N] layout with stride n. The branch is +// hoisted out of the K loop by the compiler. +__global__ __launch_bounds__(kScalarBlock) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ x_q, // [M, K] row-major + const int8_t* __restrict__ weight, // [K, N] row-major (identity pack) + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [M, N] bf16 + const int m, + const int n, + const int k, + const int b_transposed) { + const int64_t tid = + static_cast(blockIdx.x) * kScalarBlock + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (tid >= total) { + return; + } + const int row = static_cast(tid / n); + const int col = static_cast(tid - static_cast(row) * n); + const int8_t* a_row = x_q + static_cast(row) * k; + const int8_t* b_col = + weight + (b_transposed ? static_cast(col) * k : col); + const int64_t b_stride = b_transposed ? 1 : static_cast(n); + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * b_stride]); + } + // Same left-to-right fp32 evaluation order as the torch reference. + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[tid] = hip_bfloat16(scaled); +} + +// Generic element-wise device-to-device copy used by the identity pack. +template +__global__ __launch_bounds__(kPackBlock) void w8a8_identity_copy_kernel( + const T* __restrict__ src, + T* __restrict__ dst, + const int64_t count) { + const int64_t tid = + static_cast(blockIdx.x) * kPackBlock + threadIdx.x; + if (tid < count) { + dst[tid] = src[tid]; + } +} + +// Iteration 8 exact-shape pack: writes the [N][K] transpose +// dst[n*K + k] = src[k*N + n] (used for K=1024, N=4096), so the staged +// partial kernel reads each block's 16 B columns as 16 contiguous 512-B +// segments and every B fragment as one 8-byte ds_read_b64. One thread per +// 16-byte k-chunk of one column; runs once per weight outside the timed +// region, so the strided byte loads are acceptable. All other shapes keep +// the identity copy kernel above. +__global__ __launch_bounds__(kPackBlock) void w8a8_pack_transpose_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + const int k, + const int n) { + const int kchunks = (k + 15) >> 4; + const int64_t total = static_cast(n) * kchunks; + const int64_t stride = + static_cast(gridDim.x) * blockDim.x; + for (int64_t t = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + t < total; t += stride) { + const int col = static_cast(t / kchunks); + const int k0 = static_cast(t % kchunks) * 16; + alignas(16) int8_t chunk[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) { + chunk[j] = src[(static_cast(k0) + j) * n + col]; + } + *reinterpret_cast(dst + static_cast(col) * k + k0) = + *reinterpret_cast(chunk); + } +} + +inline int blocks_for(const int64_t count, const int block) { + const int64_t b = (count + block - 1) / block; + return b < 1 ? 1 : static_cast(b); +} + +} // namespace + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // Exact-shape fast path: tp8_wq_b_m16 (M=16, N=4096, K=1024). Iteration 9: + // the split-grid STAGED partial kernel at the dispatched split S=4 — 1024 + // blocks x 64 threads, each block stages its own 256-K A and B chunks into + // LDS once (8,704 B/block -> 7 blocks/CU co-resident, all 4 SIMDs busy) + // and then runs a zero-global two-ds_read_b64 + v_mmac K loop — writes the + // 4 int32 partial planes into the caller workspace (4 * 16 * 4096 * 4 = + // 1 MiB; the API allocates 16 planes / 4 MiB for this shape, so this + // always qualifies). Iteration 11: the combine is FUSED into the partial + // kernel tail (per-tile arrival counters in the last 1 KiB of the caller + // workspace; the last arrival per tile sums the planes and stores the + // scaled bf16 output), so the timed Graph contains exactly ONE kernel + // launch per replay on the same stream, every partial element overwritten + // each launch (no clear), and the counters are zeroed once per workspace + // by an async hipMemsetAsync before the first launch (never captured). + // B is read from the exact-shape [N][K] pack P[n*K+k] = W[k*N+n] produced + // once by launch_pack_w8a8_weight outside the timed region/Graph. + if (m == kDummaM && n == kExactN && k == kExactK) { + constexpr int kSplitK = kSplitKMeasured; + constexpr int64_t kPartialsBytes = + static_cast(kSplitK) * kDummaM * kExactN * sizeof(int32_t); + // Iteration 11: the fused path needs the 4 int32 partial planes PLUS the + // per-tile arrival counters (last kCountersBytes of the caller + // workspace); the API's 16-plane / 4 MiB allocation always qualifies. + if (workspace != nullptr && + workspace_bytes >= kPartialsBytes + kCountersBytes) { + int32_t* counters = reinterpret_cast( + static_cast(workspace) + workspace_bytes - kCountersBytes); + // One-time async zero of the arrival counters on the caller's stream, + // guarded by the workspace pointer: issued before the first launch with + // a given workspace, so steady-state replays (and the captured Graph) + // contain only the single fused kernel launch. + if (g_fused_counters_zeroed != workspace) { + hipMemsetAsync(counters, 0, kCountersBytes, stream); + g_fused_counters_zeroed = workspace; + } + // ONE kernel per replay: the staged partial kernel now fuses the + // combine into its tail (last arrival per tile sums the planes and + // stores the scaled bf16 output). Same stream, same workspace + // addresses, every partial overwritten each launch. + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_m16n16k32_splitk_staged_partial_kernel), + dim3(static_cast(kExactNTiles * kSplitK)), + dim3(kDummaWave), + 0, + stream, + a, + b, + static_cast(workspace), + x_scale, + weight_scale, + static_cast(out), + counters, + n); + return; + } + // Workspace smaller than the fused budget (never true for the API's + // 16-plane allocation): fall through to the generic scalar fallback. + } + + // Exact-shape fast path: tp8_indexer_wq_b_m16 (M=16, N=8192, K=1024). + // Indexer iteration 1: the minimal one-wave-per-block, one-N-tile DUMMA + // geometry — 512 blocks x 64 threads, K staged in 4 chunks of 256 into + // 8,704 B/block LDS (7 blocks/CU co-resident, 840 slots >= 512, every + // block resident from t=0), zero-global two-ds_read_b64 + v_mmac K chain, + // direct fragment -> scaled bf16 epilogue. No split-K, no workspace, no + // combine — ONE kernel launch per replay. B is read from the exact-shape + // [N][K] pack P[n*K+k] = W[k*N+n] produced once by launch_pack_w8a8_weight + // outside the timed region/Graph. + if (m == kDummaM && n == kExactNIndexer && k == kExactK) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16n16k32_indexer_kernel), + dim3(static_cast(kExactNIndexer / kDummaN)), + dim3(kDummaWave), + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast(out), + n); + return; + } + + // Compile-only reference for the mandated CU-aligned candidate + // (kSplitK=15 -> grid 3840 = exactly 32 blocks/CU, non-power-of-two, + // non-uniform 32-aligned K chunks). Never dispatched: the workspace_bytes + // == -1 guard cannot hold for any caller-provided byte count, and this + // iteration-3 direct-load instantiation does NOT decode the exact-shape + // [N][K] pack (it predates it), so it must never be launched against the + // packed buffer. It keeps the iteration-3 kernel compiled (and correct) + // while kSplitKMeasured = 2 dispatches the staged partial kernel; the + // staged kernel supports every trusted split in {2..15} with runtime + // 32-aligned chunks, so the occupancy sweep stays a one-constant change. + if (workspace_bytes == -1 && workspace != nullptr) { + constexpr int kSplitK = kSplitKCUAligned; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16n16k32_splitk_partial_kernel), + dim3(static_cast(kExactNTiles * kSplitK)), + dim3(kDummaWave), + 0, + stream, + a, + b, + static_cast(workspace), + n); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16n16k32_combine_kernel), + dim3(static_cast( + (kDummaM * n + kCombineThreads - 1) / kCombineThreads)), + dim3(kCombineThreads), + 0, + stream, + static_cast(workspace), + x_scale, + weight_scale, + static_cast(out), + n); + } + + // Every other (m, n, k) — including the paired M=2 API shape with the same + // (N, K) and the M=3072 prefill shape — reaches this single generic scalar + // launch. Exact-shape guards sit BEFORE this fallback and never remove it. + // Iteration 8 + indexer iteration 1: the exact (k == 1024 && n == 4096) and + // (k == 1024 && n == 8192) weights are packed as the [N][K] transpose by + // launch_pack_w8a8_weight, so every non-M16 caller of these pairs — + // including the paired M=2 validations — reads them through the scalar + // fallback with the transposed flag. All other shapes use the identity pack + // and the legacy row-major read. + const int b_transposed = + (k == kExactK && (n == kExactN || n == kExactNIndexer)) ? 1 : 0; + const int64_t total = static_cast(m) * n; + const int blocks = blocks_for(total, kScalarBlock); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_kernel), + dim3(blocks), + dim3(kScalarBlock), + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast(out), + m, + n, + k, + b_transposed); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Iteration 8 + indexer iteration 1 exact-shape pack: for + // (k,n)==(1024,4096) and (k,n)==(1024,8192) write the [N][K] transpose + // P[n*K + k] = W[k*N + n], so the staged kernels read each block's 16 B + // columns as 16 contiguous K segments and every B fragment as one 8-byte + // ds_read_b64. Runs once per weight outside the timed region/Graph. All + // other shapes keep the identity layout (row-major [K][N]) consumed by the + // DUMMA and scalar kernels. + const int64_t w_count = static_cast(k) * n; + const int64_t s_count = n; + if (w_count > 0) { + if (k == kExactK && (n == kExactN || n == kExactNIndexer)) { + const int kchunks = (k + 15) >> 4; + const int64_t chunks = static_cast(n) * kchunks; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_pack_transpose_i8_kernel), + dim3(blocks_for(chunks, kPackBlock)), + dim3(kPackBlock), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + dim3(blocks_for(w_count, kPackBlock)), + dim3(kPackBlock), + 0, + stream, + raw_weight, + packed_weight, + w_count); + } + } + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + dim3(blocks_for(s_count, kPackBlock)), + dim3(kPackBlock), + 0, + stream, + weight_scale, + packed_weight_scale, + s_count); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_down_proj.hip new file mode 100644 index 00000000..487c6b3a --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_down_proj.hip @@ -0,0 +1,997 @@ +// @@variant shape=tp8_shared_down_proj_m16 commit=119c281feefbb681f88dba89897484424bb07937 added=2026-08-26 +// median_us=8.453 p90_us=9.13 +// source=hy3-dsh-tp8-m16-1-adf021ab +// W8A8 INT8 GEMM bootstrap kernel for Hygon K500SM_AI / gfx928. +// +// Worker_3 (physical GPU 3) assigned shapes: +// tp8_shared_gate_up_proj_m16 : M=16, N=512, K=4096 +// tp8_shared_down_proj_m16 : M=16, N=4096, K=256 +// +// TP8 shared_down_proj (M=16, N=4096, K=256): iteration 1 was the minimal +// gfx928 INT8 DUMMA m16n16k32 bootstrap (w8a8_dumma_m16_n16_k256_kernel, +// one 64-thread wavefront per block, 8 k-ascending steps, zero barriers). +// Iteration 2 (architecture round) replaces it with the launch-geometry +// probe w8a8_dumma_m16_n16_k256_sk2_kernel: 2 wavefronts per block (128 +// threads), one 16x16 output tile per block, grid = N/16 = 256 blocks, +// in-block split-K = 2 along K (K-half = 128 per wave, 4 m16n16k32 steps), +// one END-of-K barrier, LDS-plane ascending split combine. The generic +// scalar kernel remains the fallback for every other (m,n,k). +// +// Iteration 1 strategy: scalar correctness bootstrap + minimal DUMMA path. +// * w8a8_dumma_m16_n16_kernel: exact M=16,N=512,K=4096 gate_up shape, +// one 64-thread wavefront per block, one 16x16 output tile per block, +// LDS-staged A/B, m16n16k32 int8 DUMMA with explicit int32 accumulation +// and a direct register epilogue (scale + bf16 store). +// * w8a8_scalar_gemm_kernel: generic fallback for every unmatched (m,n,k). +// * No split-K, double buffering, or inline asm in this round. +// +// Iteration 2 strategy (architecture round): launch-geometry exploration. +// * w8a8_dumma_m16_n16_k4_kernel: same exact M=16,N=512,K=4096 gate_up +// shape, now 4 wavefronts per block (256 threads) with in-block +// split-K = 4: each wave accumulates one 1024-row K quarter in +// k-ascending m16n16k32 steps from its own private 8 KiB LDS staging +// region; the four 16x16 int32 partials are combined in ascending split +// order through a 4 KiB LDS plane (bit-exact int32). Grid stays +// N/16 = 32 blocks (the N axis caps independent blocks at 32 of 120 CUs; +// grid split-K is the next lever if per-CU chains alone do not close the +// gap vs the 59.21 us Triton baseline). Accepted at 30.87 us median. +// +// Iteration 5 strategy (packed-weight/staging round): replace the B path of +// the accepted kernel with the packed [N,K] n-major layout +// (w8a8_dumma_m16_n16_k4_packedb_kernel): the iteration-2 B fragment costs 8 +// ds_read_u8 per v_mmac (row-major [k][16] stage samples 8 k-rows at 16-byte +// stride per lane; 7.7 LDS bank conflicts per LDS instruction). The packed +// layout + matrix_b col_major fragments turn each B fragment into one 8-byte +// LDS read. One-time pack (w8a8_pack_gateup_nmajor_kernel) runs +// out-of-timed-region/out-of-Graph for (k,n)==(4096,512); the generic scalar +// fallback decodes the same pack for paired M=2 API shapes via b_packed. +// +// The host contract is fixed by csrc/bindings.cpp. This file owns the two +// stable C symbols: +// launch_w8a8_gemm(...) -- timed GEMM (all shapes, generic fallback) +// launch_pack_w8a8_weight(...) -- optional out-of-timed-region packing +// (bootstrap: identity device-to-device +// copy, valid for every (K,N)) +// +// Header order is load-bearing for this DTK: hip_runtime first, then +// hip_bfloat16, then du_mma (du_mma.h is not self-contained when included +// before the HIP runtime headers). du_mma.h is included up front so later +// DUMMA optimization rounds only add kernels, not header reordering. + +#include +#include +#include + +#include + +namespace { + +// gfx928 wavefront size is 64 lanes; every blockDim below is a multiple of +// it (128 / 256 / 512 threads). + +// --------------------------------------------------------------------------- +// Scalar correctness kernel: one thread per output element. +// +// idx -> (row = idx / n, col = idx % n). Consecutive threads therefore walk +// consecutive columns of the same row: B reads b[k*n + col] are byte-adjacent +// across lanes and out writes are adjacent bf16, while the A row is broadcast +// within a warp (every lane of one row reads the same a[row*k + kk]). +// +// This is also the generic scalar fallback for every unmatched (m,n,k): +// later rounds must keep any exact-shape specialization guarded AFTER this +// path in the dispatch so paired M=2 API shapes with the same (N,K) still +// reach the scalar kernel. b_packed selects the packed [N,K] n-major layout +// produced by launch_pack_w8a8_weight for (k,n)==(4096,512): +// b[col*k + kk] == W[kk*n + col] (P[n][k] = W[k][n]). Every other (k,n) +// keeps the byte-identical original [K,N] row-major decode. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(128) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k, + int b_packed) { + const int idx = static_cast(blockIdx.x) * static_cast(blockDim.x) + + static_cast(threadIdx.x); + const int total = m * n; + if (idx >= total) { + return; + } + + const int row = idx / n; + const int col = idx - row * n; + + // Exact int32 dot product over the full K dimension. + const int8_t* __restrict__ a_row = a + static_cast(row) * k; + const int8_t* __restrict__ b_col = + b_packed ? (b + static_cast(col) * k) : (b + col); + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_packed + ? b_col[kk] + : b_col[static_cast(kk) * n]); + } + + // out = bf16(dot * x_scale[m] * weight_scale[n]); scales applied + // left-to-right exactly like the Python reference expression. + const float scaled = static_cast(acc) * x_scale[row] * weight_scale[col]; + out[static_cast(row) * n + col] = + static_cast(scaled); +} + +// --------------------------------------------------------------------------- +// TP8 shared_down_proj (M=16, N=4096, K=256) iteration 2: architecture +// round -- launch-geometry probe (2 waves per block, 1 adjacent N tile). +// Two 64-thread wavefronts per block, one 16x16 output tile per block, +// grid = N/16 = 256 blocks (2.13 blocks/CU on 120 CUs; 256 independent +// blocks keep all CUs active), in-block split-K = 2 along K: wave w owns K +// rows [w*128, (w+1)*128) and runs a 4-step k-ascending m16n16k32 chain -- +// half the iteration-1 serial 8-step chain -- doubling wavefronts per CU +// from 2.13 to 4.27 and giving the scheduler two independent MMAC chains per +// block to hide DUMMA/LDS latency behind. The two 16x16 int32 partials are +// published to two 1 KiB LDS planes; one END-of-K __syncthreads() orders +// wave0's plane before wave1's ascending split sum (s0 + s1), which +// reproduces the scalar reference's k-ascending int32 accumulation +// bit-exactly (max |acc| = 256*128*128 << 2^31). Wave1 only runs the fused +// scale+bf16 epilogue and one merged 8-byte store per lane; wave0 exits +// after the barrier. +// +// Data path (byte-identical to iteration 1): the whole K=256 slice is +// staged ONCE into LDS with cooperative 16-byte int4 loads -- each wave +// stages only its own K half (2 x int4 loads per lane per matrix), so no +// cross-wave staging barrier is needed. A stays logical row-major [16][256] +// (4 KiB, L2-hot across all 256 blocks); B is read from the packed [N,K] +// n-major layout (P[n*256+k] = W[k*4096+n], produced one-time out-of-timed- +// region for (k,n)==(256,4096)), which makes every matrix_b col_major +// fragment one contiguous 8-byte LDS read. Both stages use the 272 = 256+16 +// bank-skewed stride (same skew class as the accepted gate_up 400-stride +// stages). Zero barriers per k-step; exactly one barrier per block. +// +// Epilogue: wave1 reads both planes with the iteration-1 lane mapping +// (row = lane>>2, col0 = (lane&3)*4) and finalizes 4 consecutive columns of +// one row: (s_part0 + s_part1) * x_scale[row] * weight_scale[n0+col0+i] -> +// bf16, four adjacent 2-byte stores the compiler merges into one 8-byte +// store (16 rows x 32 B per wavefront instruction). No workspace use (the +// split is in-block). LDS/block = 2 x 4,352 + 2 x 1,024 = 10,752 B -> ~6 +// blocks/CU by LDS, ~48 VGPR x 128 thr -> no residency cliff. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(128) void w8a8_dumma_m16_n16_k256_sk2_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, // packed [N,K] n-major for this shape + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n, + int k) { + constexpr int kTileM = 16; + constexpr int kTileN = 16; + constexpr int kTileK = 32; + constexpr int kK = 256; + constexpr int kSplit = 2; + constexpr int kHalfK = kK / kSplit; // 128 K rows per wave + constexpr int kLdsStride = kK + 16; // 272 = 256 + 16 bank skew + const int tid = static_cast(threadIdx.x); // 0..127 + const int wave = tid >> 6; // 0..1 (split index) + const int lane = tid & 63; // 0..63 + const int n0 = static_cast(blockIdx.x) * kTileN; + + __shared__ __align__(16) int8_t lds_a[kTileM * kLdsStride]; // 4,352 B + __shared__ __align__(16) int8_t lds_b[kTileN * kLdsStride]; // 4,352 B + __shared__ __align__(16) int32_t s_part0[kTileM * kTileN]; // 1,024 B + __shared__ __align__(16) int32_t s_part1[kTileM * kTileN]; // 1,024 B + + // Stage the whole K slice once, wave w covering its own K half: + // lane -> (row = lane>>2, kk16 = (lane&3)<<4); i covers the half's two + // 64-k quarters, so 16 rows x 128 k = 2,048 B per wave per matrix, each + // (row, chunk) written exactly once (4,096 B per matrix total, same global + // traffic as iteration 1). Wave-private regions: no staging barrier. + const int row = lane >> 2; + const int kk16 = (lane & 3) << 4; + const int kBase = wave * kHalfK; + int4 va[2], vb[2]; +#pragma unroll + for (int i = 0; i < 2; ++i) { + va[i] = *reinterpret_cast(a + row * k + kBase + i * 64 + kk16); + vb[i] = *reinterpret_cast( + b + (n0 + row) * k + kBase + i * 64 + kk16); + } +#pragma unroll + for (int i = 0; i < 2; ++i) { + *reinterpret_cast(lds_a + row * kLdsStride + kBase + i * 64 + kk16) = + va[i]; + *reinterpret_cast(lds_b + row * kLdsStride + kBase + i * 64 + kk16) = + vb[i]; + } + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Zero-in-loop-barrier LDS-only K-half loop: 4 k-ascending m16n16k32 steps + // per wave (each wave reads only its own staged region). +#pragma unroll + for (int kk = kBase; kk < kBase + kHalfK; kk += kTileK) { + du::dumma::du_load_matrix_sync(a_frag, lds_a + kk, kLdsStride); + du::dumma::du_load_matrix_sync(b_frag, lds_b + kk, kLdsStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Publish the wave's 16x16 int32 partial to its private LDS plane, then + // one END-of-K barrier: orders wave0's plane before wave1's read (wave1 + // reads its own plane in same-wave program order; the compiler emits the + // lgkmcnt wait). Wave0 exits after the barrier. + if (wave == 0) { + du::dumma::du_store_matrix_sync(s_part0, acc_frag, kTileN, + du::dumma::mem_row_major); + } else { + du::dumma::du_store_matrix_sync(s_part1, acc_frag, kTileN, + du::dumma::mem_row_major); + } + __syncthreads(); + + if (wave == 1) { + // Fused epilogue: ascending split sum (s0 + s1) * x_scale * weight_scale + // -> bf16; each lane finalizes 4 consecutive columns of one row. + const int orow = lane >> 2; // 0..15 + const int ocol0 = (lane & 3) * 4; // 0,4,8,12 + const float xs = x_scale[orow]; + const float4 ws = + *reinterpret_cast(weight_scale + n0 + ocol0); + const float wsf[4] = {ws.x, ws.y, ws.z, ws.w}; + hip_bfloat16* op = out + orow * n + n0 + ocol0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int32_t s = + s_part0[orow * kTileN + ocol0 + i] + + s_part1[orow * kTileN + ocol0 + i]; + const float scaled = static_cast(s) * xs * wsf[i]; + op[i] = static_cast(scaled); + } + } +} + +// --------------------------------------------------------------------------- +// Iteration 2 (architecture round): launch-geometry exploration for the exact +// M=16, N=512, K=4096 gate_up shape (TP8). Geometry: 4 wavefronts per block +// (256 threads), one independent 16x16 output tile per block (grid = N/16 = +// 32 blocks), in-block split-K = 4: wave w accumulates K rows +// [w*1024, (w+1)*1024) in k-ascending m16n16k32 steps. Every wave owns a +// private 8 KiB LDS staging region (A 16x256 + B 256x16) reloaded in four +// 256-row sub-stages with the same cooperative 16-byte int4 global loads as +// iteration 1; one __syncthreads() per sub-stage orders each wave's staging +// before its fragment reads (regions are wave-private, so the reuse barrier +// only guards the wave's own write-after-read). The four 16x16 int32 +// accumulators are published to a 4 KiB LDS partial plane, then summed in +// ascending split order by all four waves (4 rows each), which reproduces the +// scalar reference's k-ascending int32 dot exactly. LDS total = 36 KiB -> 1 +// block per CU. The grid is still capped at 32 independent blocks (N axis), +// i.e. 32 of 120 CUs; this round measures whether 4 concurrent MMAC chains +// per CU (128 total wavefronts, 9 barriers/block vs 32 in iteration 1) close +// the latency gap vs the 59.21 us Triton baseline. If not, the wall is +// memory-side and the next lever is grid split-K (fill all 120 CUs) plus a +// packed-B layout. +// +// Iteration 5 (packed-weight/staging round): the accepted iteration-2 kernel +// costs ~9.25 LDS reads per v_mmac in the steady state, of which 8 are +// ds_read_u8 for the matrix_b fragment: B is staged [k][16] row-major and the +// row_major m16n16k32 B fragment samples 8 k-rows at 16-byte LDS stride per +// lane (k-step stride 512 = 32 k-rows x 16 B), so every B fragment is eight +// single-byte reads with 4-way+ bank conflicts (7.7 conflicts per LDS +// instruction overall). This round replaces ONLY the B path with the packed +// [N,K] n-major layout (P[n*K+k] = W[k*N+n], one-time out-of-timed-region +// transpose by w8a8_pack_gateup_nmajor_kernel for (k,n)==(4096,512)): B is +// staged per wave as [16 n][272 k] (272 = 256 + 16 bank skew) and loaded with +// matrix_b col_major fragments, turning each B fragment into ONE 8-byte +// ds_read2_b32/b64 (lane's 8 k bytes are contiguous; 64 lanes x 8 B = 512 B +// spread over 64 banks, max 2-way conflict). Geometry, A path (row-major +// [16][256], already one ds_read2_b32 per fragment), barrier structure, +// in-block split-K=4 exact int32 combine, and the direct register epilogue +// are byte-identical to iteration 2. Expected LDS reads per MMAC: ~9.25 -> ~2 +// (1 A + 1 B). Falsifiable: if the median does not drop below 30.87 us, the +// 30.87 us wall is not the B-fragment LDS read path / bank conflicts, and the +// next lever is occupancy (grid split-K with a fused or measured combine) or +// A-side staging. +// +// Iteration 7 (pipeline round): single-buffering vs double-buffering +// comparison. All three gates hold (K=4096>=1024; fresh PMC L2 hit 48.716% < +// 70%; doubled LDS budget 40,960 B < 48 KiB after the per-stage chunk is +// shrunk from 256 to 128 K rows). The exact accepted-code-object ISA (source +// digest bb625614..., 32 blocks x 256 thr, 40 arch VGPR / 37,888 B LDS) shows +// the K loop is a serialized latency chain: each of the 8 staging vectors per +// 256-k sub-stage issues as global_load_dwordx4 immediately followed by +// s_waitcnt vmcnt(0) + ds_write_b128 (loop body at +0x148/+0x1c4), so every +// global load's full L2/DRAM latency is exposed with only 128 wavefronts on +// 120 CUs (1.07 waves/CU) to hide behind -- the dominant cost at 25.82 us. +// This round restructures staging into a legal double-buffered pipeline on +// the accepted geometry (32 blocks x 4 waves, in-block split-K=4, packed B): +// per-wave A/B stages become two 16-row x 144-stride (128 + 16 skew) buffers +// of 128-K-row chunks; the loop issues chunk c+1's 4 int4 loads per lane +// BEFORE chunk c's 4-MMAC burst and writes them to the other buffer after the +// burst, so the vmcnt wait overlaps the MMACs and one barrier per 128-k chunk +// suffices (1 prologue + 7 chunk + 1 combine = 9 barriers/block -- the same +// total as the single-buffered 9, i.e. 0.00220 per k-step both ways, but each +// new barrier gates a chunk whose loads were already in flight). A gains the +// same 144 bank skew (single-buffered A was [16][256] = exactly 64 LDS banks, +// all 16 rows aliased onto the same banks, ~16-way fragment conflicts): A +// fragment conflicts drop to ~4-way. Exact int32 k-ascending order per wave, +// split order, in-block combine, epilogue, exact-shape guard, packed layout, +// and the generic scalar fallback (incl. the b_packed M=2 decode) are +// byte-identical. No async-copy claim: the overlap is the buffer rotation +// plus compiler-scheduled early load issue (waits land at the ds_write). +// +// Iteration 8 (occupancy/resource round, mandated: tune one occupancy +// limiter -- waves per block / VGPR live range / LDS footprint / spill +// removal -- using the fresh PMC evidence, without trading repeated HBM +// reads for occupancy). Fresh PMC of the accepted iteration-7 kernel +// (source digest 7cbc3b41..., 32 blocks x 256 thr, 56 arch VGPR / 32 SGPR / +// 0 scratch, 40,960 B static LDS, profiled 18.72 us) shows the only +// occupancy limiter with headroom is waves per block: registers fit ~18 +// waves (56 VGPR x 64 lanes = 3,584 VGPRs/wave vs 65,536/CU), scratch is 0, +// and the 40,960 B LDS caps residency at 1 block/CU -- but the grid is only +// 32 blocks, so the block-level LDS limit is not binding and the per-CU +// wave count is the lever. 8 is in the trusted split-K probe set for the +// 120-CU part, and the control-plane warning (32 blocks < 2 blocks/CU) is +// answered in wavefront terms: this probe doubles waves per block 4 -> 8, +// giving 32 blocks x 8 waves = 256 wavefronts = 2.13 waves/CU (the +// two-blocks-per-CU latency-hiding target) with NO grid split-K, NO combine +// kernel, and NO extra HBM reads (B stays the packed [N,K] once-read 2 MiB +// DRAM stream; A stays the L2-hot 64 KiB logical re-read; per-replay global +// reads unchanged at 2 MiB A (L2) + 2 MiB B (DRAM)). Per-wave K slices +// halve (1024 -> 512), so the double-buffer chunk shrinks 128 -> 64 K rows +// to keep LDS under 64 KiB/CU: per-wave A/B stages become 2 x [16][80] +// (64 + 16 bank skew), 8 chunks of 64 k per wave, ONE A + ONE B int4 load +// per lane per chunk (64 lanes x 16 B = 16 rows x 64 k per matrix), 2 MMACs +// per chunk, 1 prologue + 7 chunk + 1 combine = 9 barriers/block (same +// total as iteration 7, each gating a 64-k chunk with 8 independent waves +// behind it). The partial plane grows to 8 x 16 x 16 int32 = 8,192 B; the +// eight planes are summed in ascending split order (bit-exact int32: +// max |partial| = 512*128*128 = 8,388,608; full-K total 67,108,864 < 2^31, +// so no overflow and any grouping is exact) and each wave finalizes 2 rows +// (row = wave*2 + lane/32; each output written by 2 lanes with the same +// value). LDS/block 40,960 -> 49,152 B (still 1 block/CU +// < 64 KiB). Exact-shape guard, packed layout, generic scalar fallback +// (incl. the b_packed M=2 decode), and the Graph/current-stream contract +// are untouched. Falsifiable: if the median does not drop below 20.0658 us +// (or p90 fails <= 20.0786), doubling waves per block at 64-k chunks does +// not beat the 128-k double-buffered 4-wave chain, and the next lever is +// grid split-K with a measured combine (trusted sweep [2,3,4,7,8,10,11,15]) +// or a reverted geometry. Raw inline asm: none (policy forbids it this +// round). +// +// Iteration 12 (register depth-2 prefetch round, final conditional inline-asm +// round: raw asm stays forbidden -- plateau=false, recent_valid_improvements +// [-4.52, -1.47, +0.28] has a regression, not three within [-2%, +2%); one +// HIP-only consolidation change). The control-plane pre-micro-optimization +// mandate was settled in iteration 9 (finer one-wave zero-barrier grid +// split-K=8 + separate combine: 12.5149/12.7389, rejected -4.52%) and the +// barrier/conflict levers were measured in iterations 10 (barrier removal +// 9->1: 12.1278/12.3229, median regressed) and 11 (per-wave LDS bank-phase +// rotation: 11.9168/12.5016, +0.28% -- noise, not accepted). The accepted +// kernel's exact-source ISA (digest 12fc6260..., 32 blocks x 512 thr, 40 arch +// VGPR / 32 SGPR / 0 scratch, 49,152 B LDS, profiled 14.4 us vs unprofiled +// 11.9497 us) pins the remaining wall: every 64-k chunk iteration is a +// serialized convoy -- chunk c+1's two global_load_dwordx4 issue at +// 0x36D0/0x36DC, the 2 v_mmac run at 0x3790/0x3800, then s_waitcnt vmcnt(1) + +// ds_write_b128 (A) at 0x3820/0x3824 and s_waitcnt vmcnt(0) + ds_write_b128 +// (B) at 0x3830/0x3834 expose the full L2/DRAM latency at the staging store +// with only 2 x 16 B in flight per lane (16 KiB in flight per block vs ~76 +// KiB needed to sustain the 175 GB/s request rate across ~700 cycles of DRAM +// latency -- latency-bound by ~5x). Iteration 9 proved the batch-load +// mechanism works mechanically (its loss was the 256-block geometry/combine +// cost, not the staging); this round consolidates that finding onto the +// accepted geometry. Change (single variable: register staging depth 1 -> 2 +// chunks in flight per lane): the loop now issues chunk c+2's loads at the +// top of iteration c (a_pre2/b_pre2) and the prologue preloads chunk 1 +// (a_pre/b_pre) before the prologue barrier, so each chunk's data lives in +// registers for one full extra chunk period + a barrier before its ds_write; +// the vmcnt wait at the staging store is then covered by the previous chunk's +// MMAC burst and the barrier slack instead of stalling the wave. 4 x int4 in +// flight per lane (64 B), 16 loads in flight per wave, 32 KiB per block; VGPR +// 40 -> ~48-56, scratch stays 0, and residency is unchanged (49,152 B LDS +// still caps at 1 block/CU; 8 waves x 64 lanes x ~56 VGPR = 28,672 << 65,536 +// VGPRs/CU). Everything else is byte-identical: grid 32 x 512 thr, in-block +// split-K=8, 64-k double-buffer chunks, 80-stride stages, 9 barriers/block, +// A 2 MiB (L2-hot 64 KiB re-read 32x) + B 2 MiB (packed [N,K] once-read +// DRAM), exact int32 k-ascending per-wave order + ascending split combine, +// direct register epilogue, exact-shape guard (m==16 && n==512 && k==4096), +// one-time out-of-timed-region pack, generic scalar fallback (incl. the +// b_packed M=2 decode), and the Graph/current-stream contract. No +// asynchronous-copy claim: the overlap is compiler-scheduled load issue at +// depth 2 only. Falsifiable gate: normal-benchmark median_us < +// 11.949679851531982 AND p90_us <= 12.632080316543579 (accepted best). +// Predicted PMC/ISA signature if the mechanism is real: the vmcnt waits at +// the staging ds_write (0x3820/0x3830) move off the critical path -- lds_wait +// dropping from 5,212 and the profiled-vs-unprofiled gap shrinking -- with +// vmem_read_instructions unchanged at 4,608 (same loads, same bytes), LDS +// 49,152 B unchanged, grid 32, wg 512, 9 barriers, 40 -> ~48-56 VGPR, 0 +// scratch. If flat or slower, the wall is the DRAM request service rate / +// barrier-locked slowest-wave tail at depth 2, and the next lever is a deeper +// occupancy probe (split-K=16, 512 blocks, planes exactly filling the 512 KiB +// workspace) or a reverted kernel. Raw inline asm: none (policy forbids it +// this round). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(512) void w8a8_dumma_m16_n16_k8_packedb_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, // packed [N,K] n-major for this shape + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n, + int k) { + constexpr int kTileM = 16; + constexpr int kTileN = 16; + constexpr int kTileK = 32; + constexpr int kSplitK = 8; // wavefronts per block + constexpr int kStageK = 64; // K rows per double-buffer chunk + constexpr int kLdsStride = kStageK + 16; // 80: 64 + 16 bank skew + const int tid = static_cast(threadIdx.x); + const int wave = tid >> 6; // 0..7 + const int lane = tid & 63; + const int n0 = static_cast(blockIdx.x) * kTileN; + + // Per-wave private double-buffered A/B stages: 8 waves x 2 buffers x + // 16 x 80 = 20,480 B each; partial plane 8 x 16 x 16 int32 = 8,192 B; + // total 49,152 B/block (under the 64 KiB LDS/CU limit, 1 block/CU). + __shared__ __align__(16) int8_t lds_a[kSplitK * 2 * kTileM * kLdsStride]; + __shared__ __align__(16) int8_t lds_b[kSplitK * 2 * kTileN * kLdsStride]; + __shared__ __align__(16) int32_t s_part[kSplitK * kTileM * kTileN]; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Each wave stages only its own K quarter (no cross-wave LDS sharing), so + // the block-wide barriers act as ordering/pacing fences exactly as in the + // single-buffered iteration-5 kernel; per-chunk hazards are covered by the + // rotation: buffer (c+1)&1 is last read by chunk c-1, which every wave + // finished before the previous barrier. + int8_t* a_buf = lds_a + wave * (2 * kTileM * kLdsStride); + int8_t* b_buf = lds_b + wave * (2 * kTileN * kLdsStride); + + const int k_lo = wave * (k / kSplitK); + const int kChunks = (k / kSplitK) / kStageK; // 8 for k=4096, split 8 + + // Prologue: issue chunk 0's 1 A + 1 B int4 load per lane (64 lanes x + // 16 B = 16 rows x 64 k per matrix; lane -> (row = lane>>2, kk16 = + // (lane&3)<<4)) and stage them into buffer 0. Iteration 12 (register + // depth-2 prefetch): chunk 1's loads (a_pre/b_pre) are also issued here, + // BEFORE the prologue barrier, so they are in flight across the barrier and + // land while chunk 0 computes -- the vmcnt wait at their staging store in + // iteration 0 is covered instead of exposed. The rotation (a_pre2/b_pre2 + // loaded at the top of each loop iteration) keeps two chunks in flight per + // lane for the whole K loop. + int4 a_pre, b_pre, a_pre2, b_pre2; + { + const int row = lane >> 2; + const int kk16 = (lane & 3) << 4; + *reinterpret_cast(a_buf + row * kLdsStride + kk16) = + *reinterpret_cast(a + row * k + k_lo + kk16); + *reinterpret_cast(b_buf + row * kLdsStride + kk16) = + *reinterpret_cast(b + (n0 + row) * k + k_lo + kk16); + if (kChunks > 1) { + const int k1 = k_lo + kStageK; + a_pre = *reinterpret_cast(a + row * k + k1 + kk16); + b_pre = + *reinterpret_cast(b + (n0 + row) * k + k1 + kk16); + } + } + __syncthreads(); // chunk 0 staged before any wave computes it + + for (int c = 0; c < kChunks; ++c) { + const bool has_next = (c + 1 < kChunks); + const bool has_next2 = (c + 2 < kChunks); + + // Iteration 12 (register depth-2 prefetch): issue chunk c+2's global + // loads NOW, before this chunk's MMAC burst, so two chunks (c+1 in + // a_pre/b_pre, loaded one full iteration ago; c+2 in a_pre2/b_pre2, in + // flight now) are outstanding per lane -- 4 x 16 B vs 2 today. The + // compiler's vmcnt waits land at the ds_write below, but by then the + // waited-on loads have had a full extra chunk period + the previous + // barrier to complete, so the staging store no longer exposes the raw + // L2/DRAM latency. + if (has_next2) { + const int k_abs = k_lo + (c + 2) * kStageK; + const int row = lane >> 2; + const int kk16 = (lane & 3) << 4; + a_pre2 = *reinterpret_cast(a + row * k + k_abs + kk16); + b_pre2 = + *reinterpret_cast(b + (n0 + row) * k + k_abs + kk16); + } + + // Compute this chunk from buffer (c & 1): 2 k-ascending m16n16k32 steps. + int8_t* a_cur = a_buf + (c & 1) * (kTileM * kLdsStride); + int8_t* b_cur = b_buf + (c & 1) * (kTileN * kLdsStride); +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du::dumma::du_load_matrix_sync(a_frag, a_cur + kk, kLdsStride); + du::dumma::du_load_matrix_sync(b_frag, b_cur + kk, kLdsStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Stage chunk c+1 into the other buffer from the registers loaded one + // full iteration ago (uniform branch: safe to fence). + if (has_next) { + int8_t* a_dst = a_buf + ((c + 1) & 1) * (kTileM * kLdsStride); + int8_t* b_dst = b_buf + ((c + 1) & 1) * (kTileN * kLdsStride); + const int row = lane >> 2; + const int kk16 = (lane & 3) << 4; + *reinterpret_cast(a_dst + row * kLdsStride + kk16) = a_pre; + *reinterpret_cast(b_dst + row * kLdsStride + kk16) = b_pre; + __syncthreads(); // chunk c+1 visible; buffer (c&1) free for c+2 staging + } + + // Rotate the register pipeline: chunk c+2 becomes the next one-ahead. + a_pre = a_pre2; + b_pre = b_pre2; + } + + // Publish the eight 16x16 int32 partials to LDS, then sum them in + // ascending split order (bit-exact int32; max |partial| = 512*128*128 = + // 8,388,608 and the full-K total 67,108,864 stays below 2^31, so no + // overflow) and store the scaled bf16 output directly. Each wave + // finalizes 2 rows of the tile: row = wave*2 + lane/32, col = lane%16 + // (each of the 256 outputs is written by 2 lanes with the identical + // value, a benign duplicate store), and the partial plane layout is + // row-major [16][16] int32, so the read address is linear in the lane. + du::dumma::du_store_matrix_sync(s_part + wave * (kTileM * kTileN), acc_frag, + kTileN, du::dumma::mem_row_major); + __syncthreads(); + + const int row = wave * 2 + (lane >> 5); // 2 rows per wave, 16 total + const int col = lane & 15; + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + acc += s_part[s * (kTileM * kTileN) + row * kTileN + col]; + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[n0 + col]; + out[row * n + n0 + col] = static_cast(scaled); +} + +// --------------------------------------------------------------------------- +// Iteration 14 (final conditional inline-asm round: raw asm stays forbidden +// because the control plane has NOT confirmed a HIP plateau -- plateau=false, +// recent_valid_improvements_percent [+0.28, +2.93, +9.03] is not three values +// within [-2%, +2%) -- so this is one HIP-only consolidation change). The +// accepted iteration-13 kernel (10.6473/10.7777 us; fresh/exact-source PMC: +// grid 480 x 64 thr, 64 arch VGPR/32 SGPR/0 scratch, 12,800 B LDS, profiled +// 9.76 us, lds_bank_conflicts 65,536 on 12,288 LDS instructions, vmem_read +// 4,096, L2 42.74%) established the one-wave zero-barrier grid family's +// measured occupancy trend: split-K=8 (256 blocks, 2.13 blocks/CU) measured +// 12.5149/12.7389 us (rejected, iter 9) while split-K=15 (480 blocks, exactly +// 4.0 blocks/CU) measured 10.6473/10.7777 us (+9.03% vs the 32-block in-block +// kernel) -- more blocks/CU within the same mechanism is the measured +// direction, and iteration 12's accepted hypothesis explicitly named the next +// lever as "split-K=16, 512 blocks = 4.27 blocks/CU, planes exactly filling +// the 512 KiB workspace capacity". This round probes exactly that point, +// outside the trusted probe set {2,3,4,7,8,10,11,15} because evidence +// supports it: (a) the 2.13 -> 4.0 blocks/CU sweep is monotonic in the +// measured direction; (b) S=15's non-uniform slicing (13 x 256 + 2 x 384) +// leaves two 1.5x-work 384-k straggler blocks (s=13,14) whose host CUs define +// the grid-kernel tail, and S=16 makes ALL 512 blocks uniform 256-k work (8 +// MMAC steps, 4 int4 loads per lane per matrix) -- the stragglers disappear; +// (c) the shape contract allocates workspace_split_k_capacity = min(16, +// budget, k/32) = 16 planes of m*n int32 = exactly 524,288 B, so S=16 fits +// with zero slack (the "exact-fill, zero-margin" concern that selected 15 +// last round is now falsifiable; if the guard fails, the accepted in-block +// kernel runs and correctness is preserved). Change (single variable: split-K +// 15 -> 16 on the accepted one-wave grid mechanism; everything else +// byte-identical): grid 480 -> 512 blocks x 64 thr (4.27 blocks/CU; 5 +// co-resident blocks on 32 CUs: LDS 5 x 12,800 = 64,000 B <= 64 KiB/CU, VGPR +// 64 x 64 x 5 = 20,480 <= 65,536/CU, so no residency cliff); uniform +// ascending 64-aligned slices k0 = s*256, sliceK = 256 (max |partial| = +// 256*128*128 = 4,194,304; full-K total 67,108,864 < 2^31; ascending split +// sum order preserved, bit-exact int32 identical to the accepted combine); +// the staging code collapses to the single 4-load path (no 6-load branch; +// one static unrolled 8-step zero-barrier K loop, one staging wait per +// block); LDS stride 400 unchanged (12,800 B/block); each block publishes +// its 16x16 int32 partial to workspace plane [tile][s] (32 x 16 x 1,024 B = +// 524,288 B = exactly the 512 KiB this shape allocates); the stateless +// combine kernel (grid 32 x 256 thr, same stream, inside the Graph -- the +// timed region includes the combine per the operator contract) sums the 16 +// planes per tile in ascending split order and applies x_scale/weight_scale/ +// bf16. Byte accounting: global reads unchanged (A 2 MiB = L2-hot 64 KiB +// re-read 32x, B 2 MiB once from DRAM); workspace plane writes 512 KiB + +// combine plane reads 512 KiB (L2-hot) + bf16 out 16 KiB per replay. The +// accepted iteration-12 in-block kernel stays in the file byte-identical as +// the workspace-insufficient fallback (defensive). Exact-shape guard +// (m==16 && n==512 && k==4096), packed [N,K] layout + one-time pack, generic +// scalar fallback (incl. the b_packed M=2 decode), and the Graph/current- +// stream contract are untouched. No asynchronous-copy claim: the batch load +// is compiler-scheduled memory-level parallelism only. Falsifiable gate: +// normal-benchmark median_us < 10.647284984588623 AND p90_us <= +// 10.777679681777954 (accepted best). Predicted PMC signature if the +// mechanism is real: primary kernel grid_blocks 480 -> 512 with kernel name +// w8a8_dumma_m16_n16_sk16_grid_kernel, workgroup_size 64, lds_bytes 12,800, +// barriers/block 0, vmem_read 4,096, lds_bank_conflicts unchanged (~65k: the +// 400-stride residual is not the lever -- iter 11 measured conflict removal +// as noise), and the per-block uniformity (no 384-k blocks) shortening the +// grid-kernel tail -- the normal-benchmark median/P90 are the score. If flat +// or slower, the wall is the grid-level DRAM request service rate / 5-block- +// CU imbalance at 4.27 blocks/CU and the occupancy family is exhausted at +// this shape: the next lever is the combine launch (a fused last-arrival +// combine needs an out-of-Graph counter pre-zero in the Python layer) or a +// revert to the accepted S=15 kernel. Raw inline asm: none (policy forbids +// it this round). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(64) void w8a8_dumma_m16_n16_sk16_grid_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, // packed [N,K] n-major for this shape + int32_t* __restrict__ planes, // [32 tiles][16 splits][16][16] int32 + int n, + int k) { + constexpr int kTileM = 16; + constexpr int kTileN = 16; + constexpr int kTileK = 32; + constexpr int kSplitK = 16; + constexpr int kStride = 384 + 16; // 400: 256-k slice + bank skew + const int lane = static_cast(threadIdx.x); // 64 lanes, one wavefront + const int tile = static_cast(blockIdx.x) / kSplitK; // 0..31 + const int s = static_cast(blockIdx.x) % kSplitK; // 0..15 + const int n0 = tile * kTileN; + // Uniform ascending 64-aligned slicing: 16 x 256 = 4,096. + const int sliceK = 256; + const int k0 = s * 256; + + __shared__ __align__(16) int8_t lds_a[kTileM * kStride]; + __shared__ __align__(16) int8_t lds_b[kTileN * kStride]; + + // Stage the whole slice once: all int4 global loads issue back-to-back + // before any ds_write (no __syncthreads -- single wavefront; the compiler + // orders the LDS write-before-read with s_waitcnt lgkmcnt). Uniform 256-k + // slice -> single static unrolled 4-load path for every block. + const int row = lane >> 2; // 0..15 + const int kk16 = (lane & 3) << 4; // 0,16,32,48 + int4 va[4], vb[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + va[i] = *reinterpret_cast(a + row * k + k0 + i * 64 + kk16); + vb[i] = *reinterpret_cast( + b + (n0 + row) * k + k0 + i * 64 + kk16); + } +#pragma unroll + for (int i = 0; i < 4; ++i) { + *reinterpret_cast(lds_a + row * kStride + i * 64 + kk16) = va[i]; + *reinterpret_cast(lds_b + row * kStride + i * 64 + kk16) = vb[i]; + } + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Zero-barrier LDS-only K loop: 8 k-ascending m16n16k32 steps. + for (int kk = 0; kk < sliceK; kk += kTileK) { + du::dumma::du_load_matrix_sync(a_frag, lds_a + kk, kStride); + du::dumma::du_load_matrix_sync(b_frag, lds_b + kk, kStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Publish the 16x16 int32 partial to workspace plane [tile][s]; the + // separate combine kernel sums the 16 planes in ascending split order. + du::dumma::du_store_matrix_sync( + planes + (tile * kSplitK + s) * (kTileM * kTileN), acc_frag, kTileN, + du::dumma::mem_row_major); +} + +// --------------------------------------------------------------------------- +// Iteration 14: stateless combine for the sk16 one-wave grid. One thread per +// output element (grid 32 tiles x 256 threads = 8,192 outputs); each thread +// reads its 16 int32 partials from the [tile][s] planes in ascending split +// order (bit-exact int32, same order as the accepted combine), applies +// x_scale/weight_scale, and stores bf16. Launched on the same stream inside +// the Graph immediately after the grid kernel, so the timed region includes +// the split-K combine exactly as the operator contract requires. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_gateup_m16_sk16_combine_kernel( + const int32_t* __restrict__ planes, // [32][16][16][16] int32 + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n) { + constexpr int kTileM = 16; + constexpr int kTileN = 16; + constexpr int kSplitK = 16; + const int tile = static_cast(blockIdx.x); + const int row = static_cast(threadIdx.x) >> 4; + const int col = static_cast(threadIdx.x) & 15; + const int32_t* p = + planes + tile * (kSplitK * kTileM * kTileN) + row * kTileN + col; + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + acc += p[s * (kTileM * kTileN)]; + } + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[tile * kTileN + col]; + out[row * n + tile * kTileN + col] = static_cast(scaled); +} + +// --------------------------------------------------------------------------- +// Iteration 5: one-time out-of-timed-region packing for (k,n)==(4096,512). +// Transposes the logical [K,N] weight into the packed [N,K] n-major layout +// packed[n*K + k] = W[k*N + n] so the exact-shape kernel's matrix_b col_major +// fragment loads read 8 contiguous k bytes (one 8-byte LDS read per fragment) +// and the generic scalar fallback decodes b[col*K + kk] for the same (k,n). +// 64(k) x 64(n) tile per block via an LDS transpose (grid = (K/64, N/64)). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_gateup_nmajor_kernel( + const int8_t* __restrict__ src, // W[K][N] row-major + int8_t* __restrict__ dst, // P[N][K] n-major + int k, + int n) { + constexpr int kTile = 64; + constexpr int kStride = 80; // 64 + 16 bank skew + __shared__ __align__(16) int8_t lds[kTile * kStride]; + const int k0 = static_cast(blockIdx.x) * kTile; + const int n0 = static_cast(blockIdx.y) * kTile; + const int tid = static_cast(threadIdx.x); + + // Load: W row k0+kk, 16 consecutive n columns -> LDS row kk. + const int kk = tid >> 2; // 0..63 + const int nn16 = (tid & 3) * 16; // 0,16,32,48 + *reinterpret_cast(&lds[kk * kStride + nn16]) = + *reinterpret_cast(src + (k0 + kk) * n + n0 + nn16); + __syncthreads(); + + // Store: P[n0+nn][k0 + kk16*16 .. +15] is 16 contiguous k bytes. + const int nn = tid & 63; // 0..63 + const int kk16 = tid >> 6; // 0..3 + int4 outv; + int8_t* op = reinterpret_cast(&outv); +#pragma unroll + for (int i = 0; i < 16; ++i) { + op[i] = lds[(kk16 * 16 + i) * kStride + nn]; + } + *reinterpret_cast(dst + (n0 + nn) * k + k0 + kk16 * 16) = outv; +} + +// --------------------------------------------------------------------------- +// Identity device-to-device packing helpers (bootstrap pack_weight). +// packed_weight[i] = raw_weight[i] (K*N int8), packed_scale[i] = scale[i] +// (N fp32). Later Parallel explore rounds may replace these kernels with a +// real packed layout plus the matching GEMM interpretation; the generic +// identity copy stays the fallback for unmatched (K,N). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_identity_pack_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t idx = + static_cast(blockIdx.x) * static_cast(blockDim.x) + + static_cast(threadIdx.x); + if (idx < count) { + dst[idx] = src[idx]; + } +} + +__global__ __launch_bounds__(256) void w8a8_identity_pack_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int count) { + const int idx = static_cast(blockIdx.x) * static_cast(blockDim.x) + + static_cast(threadIdx.x); + if (idx < count) { + dst[idx] = src[idx]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Timed GEMM entry point. Launches only on the caller-provided stream; no +// allocation, no synchronization inside the timed region. The iteration-13 +// exact-shape path uses the caller-provided workspace (preallocated before +// Graph capture) as the split-K partial-plane buffer; the scalar path needs +// no split-K partials and leaves workspace unused. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + + // Exact-shape specialization: TP8 shared_gate_up_proj decode (M=16). + // Iteration 14: one-wave zero-barrier grid, split-K=16 (512 blocks x 64 + // thr = 4.27 blocks/CU, uniform ascending 256-k slices, whole-slice + // 400-stride LDS staging, zero barriers) writing 512 KiB of int32 partial + // planes to the shape's workspace (capacity = exactly 16 planes), plus a + // stateless combine kernel on the same stream inside the Graph (the timed + // region includes the combine exactly per the operator contract). The + // accepted iteration-12 in-block split-K=8 depth-2-prefetch kernel (32 + // blocks x 512 thr, workspace-free) remains the fallback if the workspace + // is smaller than the plane contract (defensive: the Python shape contract + // allocates 512 KiB, so this is defensive only). b is the packed [N,K] + // n-major layout produced one-time out-of-timed-region by + // launch_pack_w8a8_weight for this (k,n). The generic scalar path below + // remains the fallback for every other shape, including the paired M=2 + // API shapes with the same (N,K), which decode the packed [N,K] B layout + // through the b_packed flag. + if (m == 16 && n == 512 && k == 4096) { + constexpr int64_t kPlanesBytes = 32 * 16 * 16 * 16 * 4; // 524,288 + if (workspace != nullptr && workspace_bytes >= kPlanesBytes) { + hipLaunchKernelGGL( + w8a8_dumma_m16_n16_sk16_grid_kernel, + dim3(512), + dim3(64), + 0, + stream, + a, + b, + reinterpret_cast(workspace), + n, + k); + hipLaunchKernelGGL( + w8a8_gateup_m16_sk16_combine_kernel, + dim3(32), + dim3(256), + 0, + stream, + reinterpret_cast(workspace), + x_scale, + weight_scale, + reinterpret_cast(out), + n); + } else { + hipLaunchKernelGGL( + w8a8_dumma_m16_n16_k8_packedb_kernel, + dim3(static_cast(n / 16)), + dim3(512), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast(out), + n, + k); + } + return; + } + + // Exact-shape specialization: TP8 shared_down_proj decode (M=16, N=4096, + // K=256). Iteration 2 (architecture round): launch-geometry probe -- 2 + // wavefronts per block (128 thr), one 16x16 output tile per block (grid = + // N/16 = 256 blocks), in-block split-K = 2 along K (K-half = 128 per wave, + // 4 m16n16k32 steps each), whole K=256 slice staged once into 272-stride + // LDS from the packed [N,K] n-major B layout, one END-of-K barrier, + // LDS-plane ascending split combine + fused scale/bf16 epilogue. b is the + // packed [N,K] layout produced one-time out-of-timed-region by + // launch_pack_w8a8_weight for (k,n)==(256,4096); the generic scalar path + // below stays the fallback for every other shape, including the paired + // M=2 API shapes with the same (N,K), which decode the packed [N,K] B + // layout through the b_packed flag. + if (m == 16 && n == 4096 && k == 256) { + hipLaunchKernelGGL( + w8a8_dumma_m16_n16_k256_sk2_kernel, + dim3(static_cast(n / 16)), + dim3(128), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast(out), + n, + k); + return; + } + + constexpr int kBlockThreads = 128; // multiple of wavefront size 64 + const int64_t total = static_cast(m) * n; + const int grid = + static_cast((total + kBlockThreads - 1) / kBlockThreads); + const int b_packed = + (k == 4096 && n == 512) || (k == 256 && n == 4096) ? 1 : 0; + + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + dim3(static_cast(grid)), + dim3(static_cast(kBlockThreads)), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast(out), + m, + n, + k, + b_packed); +} + +// --------------------------------------------------------------------------- +// Optional out-of-timed-region weight packing. For the TP8 gate_up shape +// (K=4096, N=512) and the TP8 shared_down_proj shape (K=256, N=4096) this +// writes the packed [N,K] n-major transpose (P[n*K+k] = W[k*N+n]) consumed +// by the exact-shape DUMMA kernels (col_major B fragments) and decoded by +// the generic scalar fallback (b_packed flag, incl. paired M=2 API shapes). +// Every other (K,N) keeps the identity device-to-device copy, valid for all +// shapes. +// --------------------------------------------------------------------------- +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kBlockThreads = 256; + const int64_t count = static_cast(k) * n; + + if ((k == 4096 && n == 512) || (k == 256 && n == 4096)) { + hipLaunchKernelGGL( + w8a8_pack_gateup_nmajor_kernel, + dim3(static_cast(k / 64), static_cast(n / 64)), + dim3(static_cast(kBlockThreads)), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else if (count > 0) { + const int grid = + static_cast((count + kBlockThreads - 1) / kBlockThreads); + hipLaunchKernelGGL( + w8a8_identity_pack_i8_kernel, + dim3(static_cast(grid)), + dim3(static_cast(kBlockThreads)), + 0, + stream, + raw_weight, + packed_weight, + count); + } + + if (n > 0) { + const int grid = (n + kBlockThreads - 1) / kBlockThreads; + hipLaunchKernelGGL( + w8a8_identity_pack_f32_kernel, + dim3(static_cast(grid)), + dim3(static_cast(kBlockThreads)), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_gate_up_proj.hip new file mode 100644 index 00000000..15e74f81 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_gate_up_proj.hip @@ -0,0 +1,819 @@ +// @@variant shape=tp8_shared_gate_up_proj_m16 commit=ea0f87b2e586e93141907abbdfc541b1585f4c50 added=2026-08-26 +// median_us=10.28 p90_us=10.33 +// source=hy3-dsh-tp8-m16-1-adf021ab +// W8A8 INT8 GEMM bootstrap kernel for Hygon K500SM_AI / gfx928. +// +// Worker_3 (physical GPU 3) assigned shapes: +// tp8_shared_gate_up_proj_m16 : M=16, N=512, K=4096 +// tp8_shared_down_proj_m16 : M=16, N=4096, K=256 +// +// Iteration 1 strategy: scalar correctness bootstrap + minimal DUMMA path. +// * w8a8_dumma_m16_n16_kernel: exact M=16,N=512,K=4096 gate_up shape, +// one 64-thread wavefront per block, one 16x16 output tile per block, +// LDS-staged A/B, m16n16k32 int8 DUMMA with explicit int32 accumulation +// and a direct register epilogue (scale + bf16 store). +// * w8a8_scalar_gemm_kernel: generic fallback for every unmatched (m,n,k). +// * No split-K, double buffering, or inline asm in this round. +// +// Iteration 2 strategy (architecture round): launch-geometry exploration. +// * w8a8_dumma_m16_n16_k4_kernel: same exact M=16,N=512,K=4096 gate_up +// shape, now 4 wavefronts per block (256 threads) with in-block +// split-K = 4: each wave accumulates one 1024-row K quarter in +// k-ascending m16n16k32 steps from its own private 8 KiB LDS staging +// region; the four 16x16 int32 partials are combined in ascending split +// order through a 4 KiB LDS plane (bit-exact int32). Grid stays +// N/16 = 32 blocks (the N axis caps independent blocks at 32 of 120 CUs; +// grid split-K is the next lever if per-CU chains alone do not close the +// gap vs the 59.21 us Triton baseline). Accepted at 30.87 us median. +// +// Iteration 5 strategy (packed-weight/staging round): replace the B path of +// the accepted kernel with the packed [N,K] n-major layout +// (w8a8_dumma_m16_n16_k4_packedb_kernel): the iteration-2 B fragment costs 8 +// ds_read_u8 per v_mmac (row-major [k][16] stage samples 8 k-rows at 16-byte +// stride per lane; 7.7 LDS bank conflicts per LDS instruction). The packed +// layout + matrix_b col_major fragments turn each B fragment into one 8-byte +// LDS read. One-time pack (w8a8_pack_gateup_nmajor_kernel) runs +// out-of-timed-region/out-of-Graph for (k,n)==(4096,512); the generic scalar +// fallback decodes the same pack for paired M=2 API shapes via b_packed. +// +// The host contract is fixed by csrc/bindings.cpp. This file owns the two +// stable C symbols: +// launch_w8a8_gemm(...) -- timed GEMM (all shapes, generic fallback) +// launch_pack_w8a8_weight(...) -- optional out-of-timed-region packing +// (bootstrap: identity device-to-device +// copy, valid for every (K,N)) +// +// Header order is load-bearing for this DTK: hip_runtime first, then +// hip_bfloat16, then du_mma (du_mma.h is not self-contained when included +// before the HIP runtime headers). du_mma.h is included up front so later +// DUMMA optimization rounds only add kernels, not header reordering. + +#include +#include +#include + +#include + +namespace { + +// gfx928 wavefront size is 64 lanes; every blockDim below is a multiple of +// it (128 / 256 / 512 threads). + +// --------------------------------------------------------------------------- +// Scalar correctness kernel: one thread per output element. +// +// idx -> (row = idx / n, col = idx % n). Consecutive threads therefore walk +// consecutive columns of the same row: B reads b[k*n + col] are byte-adjacent +// across lanes and out writes are adjacent bf16, while the A row is broadcast +// within a warp (every lane of one row reads the same a[row*k + kk]). +// +// This is also the generic scalar fallback for every unmatched (m,n,k): +// later rounds must keep any exact-shape specialization guarded AFTER this +// path in the dispatch so paired M=2 API shapes with the same (N,K) still +// reach the scalar kernel. b_packed selects the packed [N,K] n-major layout +// produced by launch_pack_w8a8_weight for (k,n)==(4096,512): +// b[col*k + kk] == W[kk*n + col] (P[n][k] = W[k][n]). Every other (k,n) +// keeps the byte-identical original [K,N] row-major decode. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(128) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k, + int b_packed) { + const int idx = static_cast(blockIdx.x) * static_cast(blockDim.x) + + static_cast(threadIdx.x); + const int total = m * n; + if (idx >= total) { + return; + } + + const int row = idx / n; + const int col = idx - row * n; + + // Exact int32 dot product over the full K dimension. + const int8_t* __restrict__ a_row = a + static_cast(row) * k; + const int8_t* __restrict__ b_col = + b_packed ? (b + static_cast(col) * k) : (b + col); + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_packed + ? b_col[kk] + : b_col[static_cast(kk) * n]); + } + + // out = bf16(dot * x_scale[m] * weight_scale[n]); scales applied + // left-to-right exactly like the Python reference expression. + const float scaled = static_cast(acc) * x_scale[row] * weight_scale[col]; + out[static_cast(row) * n + col] = + static_cast(scaled); +} + +// --------------------------------------------------------------------------- +// Iteration 2 (architecture round): launch-geometry exploration for the exact +// M=16, N=512, K=4096 gate_up shape (TP8). Geometry: 4 wavefronts per block +// (256 threads), one independent 16x16 output tile per block (grid = N/16 = +// 32 blocks), in-block split-K = 4: wave w accumulates K rows +// [w*1024, (w+1)*1024) in k-ascending m16n16k32 steps. Every wave owns a +// private 8 KiB LDS staging region (A 16x256 + B 256x16) reloaded in four +// 256-row sub-stages with the same cooperative 16-byte int4 global loads as +// iteration 1; one __syncthreads() per sub-stage orders each wave's staging +// before its fragment reads (regions are wave-private, so the reuse barrier +// only guards the wave's own write-after-read). The four 16x16 int32 +// accumulators are published to a 4 KiB LDS partial plane, then summed in +// ascending split order by all four waves (4 rows each), which reproduces the +// scalar reference's k-ascending int32 dot exactly. LDS total = 36 KiB -> 1 +// block per CU. The grid is still capped at 32 independent blocks (N axis), +// i.e. 32 of 120 CUs; this round measures whether 4 concurrent MMAC chains +// per CU (128 total wavefronts, 9 barriers/block vs 32 in iteration 1) close +// the latency gap vs the 59.21 us Triton baseline. If not, the wall is +// memory-side and the next lever is grid split-K (fill all 120 CUs) plus a +// packed-B layout. +// +// Iteration 5 (packed-weight/staging round): the accepted iteration-2 kernel +// costs ~9.25 LDS reads per v_mmac in the steady state, of which 8 are +// ds_read_u8 for the matrix_b fragment: B is staged [k][16] row-major and the +// row_major m16n16k32 B fragment samples 8 k-rows at 16-byte LDS stride per +// lane (k-step stride 512 = 32 k-rows x 16 B), so every B fragment is eight +// single-byte reads with 4-way+ bank conflicts (7.7 conflicts per LDS +// instruction overall). This round replaces ONLY the B path with the packed +// [N,K] n-major layout (P[n*K+k] = W[k*N+n], one-time out-of-timed-region +// transpose by w8a8_pack_gateup_nmajor_kernel for (k,n)==(4096,512)): B is +// staged per wave as [16 n][272 k] (272 = 256 + 16 bank skew) and loaded with +// matrix_b col_major fragments, turning each B fragment into ONE 8-byte +// ds_read2_b32/b64 (lane's 8 k bytes are contiguous; 64 lanes x 8 B = 512 B +// spread over 64 banks, max 2-way conflict). Geometry, A path (row-major +// [16][256], already one ds_read2_b32 per fragment), barrier structure, +// in-block split-K=4 exact int32 combine, and the direct register epilogue +// are byte-identical to iteration 2. Expected LDS reads per MMAC: ~9.25 -> ~2 +// (1 A + 1 B). Falsifiable: if the median does not drop below 30.87 us, the +// 30.87 us wall is not the B-fragment LDS read path / bank conflicts, and the +// next lever is occupancy (grid split-K with a fused or measured combine) or +// A-side staging. +// +// Iteration 7 (pipeline round): single-buffering vs double-buffering +// comparison. All three gates hold (K=4096>=1024; fresh PMC L2 hit 48.716% < +// 70%; doubled LDS budget 40,960 B < 48 KiB after the per-stage chunk is +// shrunk from 256 to 128 K rows). The exact accepted-code-object ISA (source +// digest bb625614..., 32 blocks x 256 thr, 40 arch VGPR / 37,888 B LDS) shows +// the K loop is a serialized latency chain: each of the 8 staging vectors per +// 256-k sub-stage issues as global_load_dwordx4 immediately followed by +// s_waitcnt vmcnt(0) + ds_write_b128 (loop body at +0x148/+0x1c4), so every +// global load's full L2/DRAM latency is exposed with only 128 wavefronts on +// 120 CUs (1.07 waves/CU) to hide behind -- the dominant cost at 25.82 us. +// This round restructures staging into a legal double-buffered pipeline on +// the accepted geometry (32 blocks x 4 waves, in-block split-K=4, packed B): +// per-wave A/B stages become two 16-row x 144-stride (128 + 16 skew) buffers +// of 128-K-row chunks; the loop issues chunk c+1's 4 int4 loads per lane +// BEFORE chunk c's 4-MMAC burst and writes them to the other buffer after the +// burst, so the vmcnt wait overlaps the MMACs and one barrier per 128-k chunk +// suffices (1 prologue + 7 chunk + 1 combine = 9 barriers/block -- the same +// total as the single-buffered 9, i.e. 0.00220 per k-step both ways, but each +// new barrier gates a chunk whose loads were already in flight). A gains the +// same 144 bank skew (single-buffered A was [16][256] = exactly 64 LDS banks, +// all 16 rows aliased onto the same banks, ~16-way fragment conflicts): A +// fragment conflicts drop to ~4-way. Exact int32 k-ascending order per wave, +// split order, in-block combine, epilogue, exact-shape guard, packed layout, +// and the generic scalar fallback (incl. the b_packed M=2 decode) are +// byte-identical. No async-copy claim: the overlap is the buffer rotation +// plus compiler-scheduled early load issue (waits land at the ds_write). +// +// Iteration 8 (occupancy/resource round, mandated: tune one occupancy +// limiter -- waves per block / VGPR live range / LDS footprint / spill +// removal -- using the fresh PMC evidence, without trading repeated HBM +// reads for occupancy). Fresh PMC of the accepted iteration-7 kernel +// (source digest 7cbc3b41..., 32 blocks x 256 thr, 56 arch VGPR / 32 SGPR / +// 0 scratch, 40,960 B static LDS, profiled 18.72 us) shows the only +// occupancy limiter with headroom is waves per block: registers fit ~18 +// waves (56 VGPR x 64 lanes = 3,584 VGPRs/wave vs 65,536/CU), scratch is 0, +// and the 40,960 B LDS caps residency at 1 block/CU -- but the grid is only +// 32 blocks, so the block-level LDS limit is not binding and the per-CU +// wave count is the lever. 8 is in the trusted split-K probe set for the +// 120-CU part, and the control-plane warning (32 blocks < 2 blocks/CU) is +// answered in wavefront terms: this probe doubles waves per block 4 -> 8, +// giving 32 blocks x 8 waves = 256 wavefronts = 2.13 waves/CU (the +// two-blocks-per-CU latency-hiding target) with NO grid split-K, NO combine +// kernel, and NO extra HBM reads (B stays the packed [N,K] once-read 2 MiB +// DRAM stream; A stays the L2-hot 64 KiB logical re-read; per-replay global +// reads unchanged at 2 MiB A (L2) + 2 MiB B (DRAM)). Per-wave K slices +// halve (1024 -> 512), so the double-buffer chunk shrinks 128 -> 64 K rows +// to keep LDS under 64 KiB/CU: per-wave A/B stages become 2 x [16][80] +// (64 + 16 bank skew), 8 chunks of 64 k per wave, ONE A + ONE B int4 load +// per lane per chunk (64 lanes x 16 B = 16 rows x 64 k per matrix), 2 MMACs +// per chunk, 1 prologue + 7 chunk + 1 combine = 9 barriers/block (same +// total as iteration 7, each gating a 64-k chunk with 8 independent waves +// behind it). The partial plane grows to 8 x 16 x 16 int32 = 8,192 B; the +// eight planes are summed in ascending split order (bit-exact int32: +// max |partial| = 512*128*128 = 8,388,608; full-K total 67,108,864 < 2^31, +// so no overflow and any grouping is exact) and each wave finalizes 2 rows +// (row = wave*2 + lane/32; each output written by 2 lanes with the same +// value). LDS/block 40,960 -> 49,152 B (still 1 block/CU +// < 64 KiB). Exact-shape guard, packed layout, generic scalar fallback +// (incl. the b_packed M=2 decode), and the Graph/current-stream contract +// are untouched. Falsifiable: if the median does not drop below 20.0658 us +// (or p90 fails <= 20.0786), doubling waves per block at 64-k chunks does +// not beat the 128-k double-buffered 4-wave chain, and the next lever is +// grid split-K with a measured combine (trusted sweep [2,3,4,7,8,10,11,15]) +// or a reverted geometry. Raw inline asm: none (policy forbids it this +// round). +// +// Iteration 12 (register depth-2 prefetch round, final conditional inline-asm +// round: raw asm stays forbidden -- plateau=false, recent_valid_improvements +// [-4.52, -1.47, +0.28] has a regression, not three within [-2%, +2%); one +// HIP-only consolidation change). The control-plane pre-micro-optimization +// mandate was settled in iteration 9 (finer one-wave zero-barrier grid +// split-K=8 + separate combine: 12.5149/12.7389, rejected -4.52%) and the +// barrier/conflict levers were measured in iterations 10 (barrier removal +// 9->1: 12.1278/12.3229, median regressed) and 11 (per-wave LDS bank-phase +// rotation: 11.9168/12.5016, +0.28% -- noise, not accepted). The accepted +// kernel's exact-source ISA (digest 12fc6260..., 32 blocks x 512 thr, 40 arch +// VGPR / 32 SGPR / 0 scratch, 49,152 B LDS, profiled 14.4 us vs unprofiled +// 11.9497 us) pins the remaining wall: every 64-k chunk iteration is a +// serialized convoy -- chunk c+1's two global_load_dwordx4 issue at +// 0x36D0/0x36DC, the 2 v_mmac run at 0x3790/0x3800, then s_waitcnt vmcnt(1) + +// ds_write_b128 (A) at 0x3820/0x3824 and s_waitcnt vmcnt(0) + ds_write_b128 +// (B) at 0x3830/0x3834 expose the full L2/DRAM latency at the staging store +// with only 2 x 16 B in flight per lane (16 KiB in flight per block vs ~76 +// KiB needed to sustain the 175 GB/s request rate across ~700 cycles of DRAM +// latency -- latency-bound by ~5x). Iteration 9 proved the batch-load +// mechanism works mechanically (its loss was the 256-block geometry/combine +// cost, not the staging); this round consolidates that finding onto the +// accepted geometry. Change (single variable: register staging depth 1 -> 2 +// chunks in flight per lane): the loop now issues chunk c+2's loads at the +// top of iteration c (a_pre2/b_pre2) and the prologue preloads chunk 1 +// (a_pre/b_pre) before the prologue barrier, so each chunk's data lives in +// registers for one full extra chunk period + a barrier before its ds_write; +// the vmcnt wait at the staging store is then covered by the previous chunk's +// MMAC burst and the barrier slack instead of stalling the wave. 4 x int4 in +// flight per lane (64 B), 16 loads in flight per wave, 32 KiB per block; VGPR +// 40 -> ~48-56, scratch stays 0, and residency is unchanged (49,152 B LDS +// still caps at 1 block/CU; 8 waves x 64 lanes x ~56 VGPR = 28,672 << 65,536 +// VGPRs/CU). Everything else is byte-identical: grid 32 x 512 thr, in-block +// split-K=8, 64-k double-buffer chunks, 80-stride stages, 9 barriers/block, +// A 2 MiB (L2-hot 64 KiB re-read 32x) + B 2 MiB (packed [N,K] once-read +// DRAM), exact int32 k-ascending per-wave order + ascending split combine, +// direct register epilogue, exact-shape guard (m==16 && n==512 && k==4096), +// one-time out-of-timed-region pack, generic scalar fallback (incl. the +// b_packed M=2 decode), and the Graph/current-stream contract. No +// asynchronous-copy claim: the overlap is compiler-scheduled load issue at +// depth 2 only. Falsifiable gate: normal-benchmark median_us < +// 11.949679851531982 AND p90_us <= 12.632080316543579 (accepted best). +// Predicted PMC/ISA signature if the mechanism is real: the vmcnt waits at +// the staging ds_write (0x3820/0x3830) move off the critical path -- lds_wait +// dropping from 5,212 and the profiled-vs-unprofiled gap shrinking -- with +// vmem_read_instructions unchanged at 4,608 (same loads, same bytes), LDS +// 49,152 B unchanged, grid 32, wg 512, 9 barriers, 40 -> ~48-56 VGPR, 0 +// scratch. If flat or slower, the wall is the DRAM request service rate / +// barrier-locked slowest-wave tail at depth 2, and the next lever is a deeper +// occupancy probe (split-K=16, 512 blocks, planes exactly filling the 512 KiB +// workspace) or a reverted kernel. Raw inline asm: none (policy forbids it +// this round). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(512) void w8a8_dumma_m16_n16_k8_packedb_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, // packed [N,K] n-major for this shape + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n, + int k) { + constexpr int kTileM = 16; + constexpr int kTileN = 16; + constexpr int kTileK = 32; + constexpr int kSplitK = 8; // wavefronts per block + constexpr int kStageK = 64; // K rows per double-buffer chunk + constexpr int kLdsStride = kStageK + 16; // 80: 64 + 16 bank skew + const int tid = static_cast(threadIdx.x); + const int wave = tid >> 6; // 0..7 + const int lane = tid & 63; + const int n0 = static_cast(blockIdx.x) * kTileN; + + // Per-wave private double-buffered A/B stages: 8 waves x 2 buffers x + // 16 x 80 = 20,480 B each; partial plane 8 x 16 x 16 int32 = 8,192 B; + // total 49,152 B/block (under the 64 KiB LDS/CU limit, 1 block/CU). + __shared__ __align__(16) int8_t lds_a[kSplitK * 2 * kTileM * kLdsStride]; + __shared__ __align__(16) int8_t lds_b[kSplitK * 2 * kTileN * kLdsStride]; + __shared__ __align__(16) int32_t s_part[kSplitK * kTileM * kTileN]; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Each wave stages only its own K quarter (no cross-wave LDS sharing), so + // the block-wide barriers act as ordering/pacing fences exactly as in the + // single-buffered iteration-5 kernel; per-chunk hazards are covered by the + // rotation: buffer (c+1)&1 is last read by chunk c-1, which every wave + // finished before the previous barrier. + int8_t* a_buf = lds_a + wave * (2 * kTileM * kLdsStride); + int8_t* b_buf = lds_b + wave * (2 * kTileN * kLdsStride); + + const int k_lo = wave * (k / kSplitK); + const int kChunks = (k / kSplitK) / kStageK; // 8 for k=4096, split 8 + + // Prologue: issue chunk 0's 1 A + 1 B int4 load per lane (64 lanes x + // 16 B = 16 rows x 64 k per matrix; lane -> (row = lane>>2, kk16 = + // (lane&3)<<4)) and stage them into buffer 0. Iteration 12 (register + // depth-2 prefetch): chunk 1's loads (a_pre/b_pre) are also issued here, + // BEFORE the prologue barrier, so they are in flight across the barrier and + // land while chunk 0 computes -- the vmcnt wait at their staging store in + // iteration 0 is covered instead of exposed. The rotation (a_pre2/b_pre2 + // loaded at the top of each loop iteration) keeps two chunks in flight per + // lane for the whole K loop. + int4 a_pre, b_pre, a_pre2, b_pre2; + { + const int row = lane >> 2; + const int kk16 = (lane & 3) << 4; + *reinterpret_cast(a_buf + row * kLdsStride + kk16) = + *reinterpret_cast(a + row * k + k_lo + kk16); + *reinterpret_cast(b_buf + row * kLdsStride + kk16) = + *reinterpret_cast(b + (n0 + row) * k + k_lo + kk16); + if (kChunks > 1) { + const int k1 = k_lo + kStageK; + a_pre = *reinterpret_cast(a + row * k + k1 + kk16); + b_pre = + *reinterpret_cast(b + (n0 + row) * k + k1 + kk16); + } + } + __syncthreads(); // chunk 0 staged before any wave computes it + + for (int c = 0; c < kChunks; ++c) { + const bool has_next = (c + 1 < kChunks); + const bool has_next2 = (c + 2 < kChunks); + + // Iteration 12 (register depth-2 prefetch): issue chunk c+2's global + // loads NOW, before this chunk's MMAC burst, so two chunks (c+1 in + // a_pre/b_pre, loaded one full iteration ago; c+2 in a_pre2/b_pre2, in + // flight now) are outstanding per lane -- 4 x 16 B vs 2 today. The + // compiler's vmcnt waits land at the ds_write below, but by then the + // waited-on loads have had a full extra chunk period + the previous + // barrier to complete, so the staging store no longer exposes the raw + // L2/DRAM latency. + if (has_next2) { + const int k_abs = k_lo + (c + 2) * kStageK; + const int row = lane >> 2; + const int kk16 = (lane & 3) << 4; + a_pre2 = *reinterpret_cast(a + row * k + k_abs + kk16); + b_pre2 = + *reinterpret_cast(b + (n0 + row) * k + k_abs + kk16); + } + + // Compute this chunk from buffer (c & 1): 2 k-ascending m16n16k32 steps. + int8_t* a_cur = a_buf + (c & 1) * (kTileM * kLdsStride); + int8_t* b_cur = b_buf + (c & 1) * (kTileN * kLdsStride); +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du::dumma::du_load_matrix_sync(a_frag, a_cur + kk, kLdsStride); + du::dumma::du_load_matrix_sync(b_frag, b_cur + kk, kLdsStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Stage chunk c+1 into the other buffer from the registers loaded one + // full iteration ago (uniform branch: safe to fence). + if (has_next) { + int8_t* a_dst = a_buf + ((c + 1) & 1) * (kTileM * kLdsStride); + int8_t* b_dst = b_buf + ((c + 1) & 1) * (kTileN * kLdsStride); + const int row = lane >> 2; + const int kk16 = (lane & 3) << 4; + *reinterpret_cast(a_dst + row * kLdsStride + kk16) = a_pre; + *reinterpret_cast(b_dst + row * kLdsStride + kk16) = b_pre; + __syncthreads(); // chunk c+1 visible; buffer (c&1) free for c+2 staging + } + + // Rotate the register pipeline: chunk c+2 becomes the next one-ahead. + a_pre = a_pre2; + b_pre = b_pre2; + } + + // Publish the eight 16x16 int32 partials to LDS, then sum them in + // ascending split order (bit-exact int32; max |partial| = 512*128*128 = + // 8,388,608 and the full-K total 67,108,864 stays below 2^31, so no + // overflow) and store the scaled bf16 output directly. Each wave + // finalizes 2 rows of the tile: row = wave*2 + lane/32, col = lane%16 + // (each of the 256 outputs is written by 2 lanes with the identical + // value, a benign duplicate store), and the partial plane layout is + // row-major [16][16] int32, so the read address is linear in the lane. + du::dumma::du_store_matrix_sync(s_part + wave * (kTileM * kTileN), acc_frag, + kTileN, du::dumma::mem_row_major); + __syncthreads(); + + const int row = wave * 2 + (lane >> 5); // 2 rows per wave, 16 total + const int col = lane & 15; + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + acc += s_part[s * (kTileM * kTileN) + row * kTileN + col]; + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[n0 + col]; + out[row * n + n0 + col] = static_cast(scaled); +} + +// --------------------------------------------------------------------------- +// Iteration 14 (final conditional inline-asm round: raw asm stays forbidden +// because the control plane has NOT confirmed a HIP plateau -- plateau=false, +// recent_valid_improvements_percent [+0.28, +2.93, +9.03] is not three values +// within [-2%, +2%) -- so this is one HIP-only consolidation change). The +// accepted iteration-13 kernel (10.6473/10.7777 us; fresh/exact-source PMC: +// grid 480 x 64 thr, 64 arch VGPR/32 SGPR/0 scratch, 12,800 B LDS, profiled +// 9.76 us, lds_bank_conflicts 65,536 on 12,288 LDS instructions, vmem_read +// 4,096, L2 42.74%) established the one-wave zero-barrier grid family's +// measured occupancy trend: split-K=8 (256 blocks, 2.13 blocks/CU) measured +// 12.5149/12.7389 us (rejected, iter 9) while split-K=15 (480 blocks, exactly +// 4.0 blocks/CU) measured 10.6473/10.7777 us (+9.03% vs the 32-block in-block +// kernel) -- more blocks/CU within the same mechanism is the measured +// direction, and iteration 12's accepted hypothesis explicitly named the next +// lever as "split-K=16, 512 blocks = 4.27 blocks/CU, planes exactly filling +// the 512 KiB workspace capacity". This round probes exactly that point, +// outside the trusted probe set {2,3,4,7,8,10,11,15} because evidence +// supports it: (a) the 2.13 -> 4.0 blocks/CU sweep is monotonic in the +// measured direction; (b) S=15's non-uniform slicing (13 x 256 + 2 x 384) +// leaves two 1.5x-work 384-k straggler blocks (s=13,14) whose host CUs define +// the grid-kernel tail, and S=16 makes ALL 512 blocks uniform 256-k work (8 +// MMAC steps, 4 int4 loads per lane per matrix) -- the stragglers disappear; +// (c) the shape contract allocates workspace_split_k_capacity = min(16, +// budget, k/32) = 16 planes of m*n int32 = exactly 524,288 B, so S=16 fits +// with zero slack (the "exact-fill, zero-margin" concern that selected 15 +// last round is now falsifiable; if the guard fails, the accepted in-block +// kernel runs and correctness is preserved). Change (single variable: split-K +// 15 -> 16 on the accepted one-wave grid mechanism; everything else +// byte-identical): grid 480 -> 512 blocks x 64 thr (4.27 blocks/CU; 5 +// co-resident blocks on 32 CUs: LDS 5 x 12,800 = 64,000 B <= 64 KiB/CU, VGPR +// 64 x 64 x 5 = 20,480 <= 65,536/CU, so no residency cliff); uniform +// ascending 64-aligned slices k0 = s*256, sliceK = 256 (max |partial| = +// 256*128*128 = 4,194,304; full-K total 67,108,864 < 2^31; ascending split +// sum order preserved, bit-exact int32 identical to the accepted combine); +// the staging code collapses to the single 4-load path (no 6-load branch; +// one static unrolled 8-step zero-barrier K loop, one staging wait per +// block); LDS stride 400 unchanged (12,800 B/block); each block publishes +// its 16x16 int32 partial to workspace plane [tile][s] (32 x 16 x 1,024 B = +// 524,288 B = exactly the 512 KiB this shape allocates); the stateless +// combine kernel (grid 32 x 256 thr, same stream, inside the Graph -- the +// timed region includes the combine per the operator contract) sums the 16 +// planes per tile in ascending split order and applies x_scale/weight_scale/ +// bf16. Byte accounting: global reads unchanged (A 2 MiB = L2-hot 64 KiB +// re-read 32x, B 2 MiB once from DRAM); workspace plane writes 512 KiB + +// combine plane reads 512 KiB (L2-hot) + bf16 out 16 KiB per replay. The +// accepted iteration-12 in-block kernel stays in the file byte-identical as +// the workspace-insufficient fallback (defensive). Exact-shape guard +// (m==16 && n==512 && k==4096), packed [N,K] layout + one-time pack, generic +// scalar fallback (incl. the b_packed M=2 decode), and the Graph/current- +// stream contract are untouched. No asynchronous-copy claim: the batch load +// is compiler-scheduled memory-level parallelism only. Falsifiable gate: +// normal-benchmark median_us < 10.647284984588623 AND p90_us <= +// 10.777679681777954 (accepted best). Predicted PMC signature if the +// mechanism is real: primary kernel grid_blocks 480 -> 512 with kernel name +// w8a8_dumma_m16_n16_sk16_grid_kernel, workgroup_size 64, lds_bytes 12,800, +// barriers/block 0, vmem_read 4,096, lds_bank_conflicts unchanged (~65k: the +// 400-stride residual is not the lever -- iter 11 measured conflict removal +// as noise), and the per-block uniformity (no 384-k blocks) shortening the +// grid-kernel tail -- the normal-benchmark median/P90 are the score. If flat +// or slower, the wall is the grid-level DRAM request service rate / 5-block- +// CU imbalance at 4.27 blocks/CU and the occupancy family is exhausted at +// this shape: the next lever is the combine launch (a fused last-arrival +// combine needs an out-of-Graph counter pre-zero in the Python layer) or a +// revert to the accepted S=15 kernel. Raw inline asm: none (policy forbids +// it this round). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(64) void w8a8_dumma_m16_n16_sk16_grid_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, // packed [N,K] n-major for this shape + int32_t* __restrict__ planes, // [32 tiles][16 splits][16][16] int32 + int n, + int k) { + constexpr int kTileM = 16; + constexpr int kTileN = 16; + constexpr int kTileK = 32; + constexpr int kSplitK = 16; + constexpr int kStride = 384 + 16; // 400: 256-k slice + bank skew + const int lane = static_cast(threadIdx.x); // 64 lanes, one wavefront + const int tile = static_cast(blockIdx.x) / kSplitK; // 0..31 + const int s = static_cast(blockIdx.x) % kSplitK; // 0..15 + const int n0 = tile * kTileN; + // Uniform ascending 64-aligned slicing: 16 x 256 = 4,096. + const int sliceK = 256; + const int k0 = s * 256; + + __shared__ __align__(16) int8_t lds_a[kTileM * kStride]; + __shared__ __align__(16) int8_t lds_b[kTileN * kStride]; + + // Stage the whole slice once: all int4 global loads issue back-to-back + // before any ds_write (no __syncthreads -- single wavefront; the compiler + // orders the LDS write-before-read with s_waitcnt lgkmcnt). Uniform 256-k + // slice -> single static unrolled 4-load path for every block. + const int row = lane >> 2; // 0..15 + const int kk16 = (lane & 3) << 4; // 0,16,32,48 + int4 va[4], vb[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + va[i] = *reinterpret_cast(a + row * k + k0 + i * 64 + kk16); + vb[i] = *reinterpret_cast( + b + (n0 + row) * k + k0 + i * 64 + kk16); + } +#pragma unroll + for (int i = 0; i < 4; ++i) { + *reinterpret_cast(lds_a + row * kStride + i * 64 + kk16) = va[i]; + *reinterpret_cast(lds_b + row * kStride + i * 64 + kk16) = vb[i]; + } + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Zero-barrier LDS-only K loop: 8 k-ascending m16n16k32 steps. + for (int kk = 0; kk < sliceK; kk += kTileK) { + du::dumma::du_load_matrix_sync(a_frag, lds_a + kk, kStride); + du::dumma::du_load_matrix_sync(b_frag, lds_b + kk, kStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Publish the 16x16 int32 partial to workspace plane [tile][s]; the + // separate combine kernel sums the 16 planes in ascending split order. + du::dumma::du_store_matrix_sync( + planes + (tile * kSplitK + s) * (kTileM * kTileN), acc_frag, kTileN, + du::dumma::mem_row_major); +} + +// --------------------------------------------------------------------------- +// Iteration 14: stateless combine for the sk16 one-wave grid. One thread per +// output element (grid 32 tiles x 256 threads = 8,192 outputs); each thread +// reads its 16 int32 partials from the [tile][s] planes in ascending split +// order (bit-exact int32, same order as the accepted combine), applies +// x_scale/weight_scale, and stores bf16. Launched on the same stream inside +// the Graph immediately after the grid kernel, so the timed region includes +// the split-K combine exactly as the operator contract requires. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_gateup_m16_sk16_combine_kernel( + const int32_t* __restrict__ planes, // [32][16][16][16] int32 + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n) { + constexpr int kTileM = 16; + constexpr int kTileN = 16; + constexpr int kSplitK = 16; + const int tile = static_cast(blockIdx.x); + const int row = static_cast(threadIdx.x) >> 4; + const int col = static_cast(threadIdx.x) & 15; + const int32_t* p = + planes + tile * (kSplitK * kTileM * kTileN) + row * kTileN + col; + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + acc += p[s * (kTileM * kTileN)]; + } + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[tile * kTileN + col]; + out[row * n + tile * kTileN + col] = static_cast(scaled); +} + +// --------------------------------------------------------------------------- +// Iteration 5: one-time out-of-timed-region packing for (k,n)==(4096,512). +// Transposes the logical [K,N] weight into the packed [N,K] n-major layout +// packed[n*K + k] = W[k*N + n] so the exact-shape kernel's matrix_b col_major +// fragment loads read 8 contiguous k bytes (one 8-byte LDS read per fragment) +// and the generic scalar fallback decodes b[col*K + kk] for the same (k,n). +// 64(k) x 64(n) tile per block via an LDS transpose (grid = (K/64, N/64)). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_gateup_nmajor_kernel( + const int8_t* __restrict__ src, // W[K][N] row-major + int8_t* __restrict__ dst, // P[N][K] n-major + int k, + int n) { + constexpr int kTile = 64; + constexpr int kStride = 80; // 64 + 16 bank skew + __shared__ __align__(16) int8_t lds[kTile * kStride]; + const int k0 = static_cast(blockIdx.x) * kTile; + const int n0 = static_cast(blockIdx.y) * kTile; + const int tid = static_cast(threadIdx.x); + + // Load: W row k0+kk, 16 consecutive n columns -> LDS row kk. + const int kk = tid >> 2; // 0..63 + const int nn16 = (tid & 3) * 16; // 0,16,32,48 + *reinterpret_cast(&lds[kk * kStride + nn16]) = + *reinterpret_cast(src + (k0 + kk) * n + n0 + nn16); + __syncthreads(); + + // Store: P[n0+nn][k0 + kk16*16 .. +15] is 16 contiguous k bytes. + const int nn = tid & 63; // 0..63 + const int kk16 = tid >> 6; // 0..3 + int4 outv; + int8_t* op = reinterpret_cast(&outv); +#pragma unroll + for (int i = 0; i < 16; ++i) { + op[i] = lds[(kk16 * 16 + i) * kStride + nn]; + } + *reinterpret_cast(dst + (n0 + nn) * k + k0 + kk16 * 16) = outv; +} + +// --------------------------------------------------------------------------- +// Identity device-to-device packing helpers (bootstrap pack_weight). +// packed_weight[i] = raw_weight[i] (K*N int8), packed_scale[i] = scale[i] +// (N fp32). Later Parallel explore rounds may replace these kernels with a +// real packed layout plus the matching GEMM interpretation; the generic +// identity copy stays the fallback for unmatched (K,N). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_identity_pack_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t idx = + static_cast(blockIdx.x) * static_cast(blockDim.x) + + static_cast(threadIdx.x); + if (idx < count) { + dst[idx] = src[idx]; + } +} + +__global__ __launch_bounds__(256) void w8a8_identity_pack_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int count) { + const int idx = static_cast(blockIdx.x) * static_cast(blockDim.x) + + static_cast(threadIdx.x); + if (idx < count) { + dst[idx] = src[idx]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Timed GEMM entry point. Launches only on the caller-provided stream; no +// allocation, no synchronization inside the timed region. The iteration-13 +// exact-shape path uses the caller-provided workspace (preallocated before +// Graph capture) as the split-K partial-plane buffer; the scalar path needs +// no split-K partials and leaves workspace unused. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + + // Exact-shape specialization: TP8 shared_gate_up_proj decode (M=16). + // Iteration 14: one-wave zero-barrier grid, split-K=16 (512 blocks x 64 + // thr = 4.27 blocks/CU, uniform ascending 256-k slices, whole-slice + // 400-stride LDS staging, zero barriers) writing 512 KiB of int32 partial + // planes to the shape's workspace (capacity = exactly 16 planes), plus a + // stateless combine kernel on the same stream inside the Graph (the timed + // region includes the combine exactly per the operator contract). The + // accepted iteration-12 in-block split-K=8 depth-2-prefetch kernel (32 + // blocks x 512 thr, workspace-free) remains the fallback if the workspace + // is smaller than the plane contract (defensive: the Python shape contract + // allocates 512 KiB, so this is defensive only). b is the packed [N,K] + // n-major layout produced one-time out-of-timed-region by + // launch_pack_w8a8_weight for this (k,n). The generic scalar path below + // remains the fallback for every other shape, including the paired M=2 + // API shapes with the same (N,K), which decode the packed [N,K] B layout + // through the b_packed flag. + if (m == 16 && n == 512 && k == 4096) { + constexpr int64_t kPlanesBytes = 32 * 16 * 16 * 16 * 4; // 524,288 + if (workspace != nullptr && workspace_bytes >= kPlanesBytes) { + hipLaunchKernelGGL( + w8a8_dumma_m16_n16_sk16_grid_kernel, + dim3(512), + dim3(64), + 0, + stream, + a, + b, + reinterpret_cast(workspace), + n, + k); + hipLaunchKernelGGL( + w8a8_gateup_m16_sk16_combine_kernel, + dim3(32), + dim3(256), + 0, + stream, + reinterpret_cast(workspace), + x_scale, + weight_scale, + reinterpret_cast(out), + n); + } else { + hipLaunchKernelGGL( + w8a8_dumma_m16_n16_k8_packedb_kernel, + dim3(static_cast(n / 16)), + dim3(512), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast(out), + n, + k); + } + return; + } + + constexpr int kBlockThreads = 128; // multiple of wavefront size 64 + const int64_t total = static_cast(m) * n; + const int grid = + static_cast((total + kBlockThreads - 1) / kBlockThreads); + const int b_packed = (k == 4096 && n == 512) ? 1 : 0; + + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + dim3(static_cast(grid)), + dim3(static_cast(kBlockThreads)), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast(out), + m, + n, + k, + b_packed); +} + +// --------------------------------------------------------------------------- +// Optional out-of-timed-region weight packing. For the TP8 gate_up shape +// (K=4096, N=512) this writes the packed [N,K] n-major transpose consumed by +// the exact-shape DUMMA kernel (col_major B fragments) and decoded by the +// generic scalar fallback (b_packed flag, incl. paired M=2 API shapes). +// Every other (K,N) keeps the identity device-to-device copy, valid for all +// shapes. +// --------------------------------------------------------------------------- +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kBlockThreads = 256; + const int64_t count = static_cast(k) * n; + + if (k == 4096 && n == 512) { + hipLaunchKernelGGL( + w8a8_pack_gateup_nmajor_kernel, + dim3(static_cast(k / 64), static_cast(n / 64)), + dim3(static_cast(kBlockThreads)), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else if (count > 0) { + const int grid = + static_cast((count + kBlockThreads - 1) / kBlockThreads); + hipLaunchKernelGGL( + w8a8_identity_pack_i8_kernel, + dim3(static_cast(grid)), + dim3(static_cast(kBlockThreads)), + 0, + stream, + raw_weight, + packed_weight, + count); + } + + if (n > 0) { + const int grid = (n + kBlockThreads - 1) / kBlockThreads; + hipLaunchKernelGGL( + w8a8_identity_pack_f32_kernel, + dim3(static_cast(grid)), + dim3(static_cast(kBlockThreads)), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wo_b.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wo_b.hip new file mode 100644 index 00000000..f420f456 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wo_b.hip @@ -0,0 +1,683 @@ +// @@variant shape=tp8_wo_b_m16 commit=eb0fbbad60c58b1cf09e6ceff6cb9aaa68a2c1cb added=2026-08-26 +// median_us=9.81 p90_us=9.84 +// source=hy3-dsh-tp8-m16-1-adf021ab +// MetaInfer INT8 W8A8 GEMM for gfx928 (worker_2 / tp8_wo_b_m16). +// +// Iteration 9: exact-shape DUMMA specialization for (m, n, k) == (16, 4096, +// 1024), plus the scalar generic fallback for every other shape. +// +// Iteration 9 change (prefetch round, stage-K=64 double-buffered A+B LDS +// prefetch with packed-B fragment slots): the accepted iteration-2 kernel +// (256 blocks x 128 threads, in-block split-K=2, per-wave stage-K=128, +// 17.073 us median / 17.096 us p90) is serialized on the per-stage global +// load latency. The exact iteration-2 ISA (source digest a59a4a28...) shows +// the stage loop as: 4 global_load_dwordx4 -> s_waitcnt vmcnt(3..0) (one +// wait in front of EACH staging ds_write_b128) -> 32 ds_read_u8 + 4 +// ds_read2_b32 -> 4 v_mmac, with the next stage's loads issued only at the +// top of the next iteration: every one of the 4 stages per wave exposes the +// full L2 latency of its global loads with nothing to hide it. Iteration 7 +// tried a depth-1 register pipeline (hold arrays + guarded rotation) and +// collapsed to 48.03 us; its exact ISA shows why: the compiler refused to +// keep the cross-iteration hold registers and materialized the rotation as +// a global-memory round trip (4 buffer_store_dwordx4 + 4 +// buffer_load_dwordx4 per stage in the steady loop), i.e. MORE traffic with +// the same serialization. This round adopts the prefetch pattern validated +// on the same DTK/gfx928 by the worker_0 wqkv_a lineage (accepted 18.55 us, +// ISA-verified): per K=64 chunk, the NEXT chunk's global loads (1 x 16-B A +// dwordx4 + 2 x 8-B B dwordx2 per lane, scalar variables only - no arrays, +// no cross-iteration registers) are issued at the TOP of the iteration and +// committed to the ALTERNATE LDS buffer AFTER the current chunk's MMAC +// burst, so the compiler's vmcnt wait lands at the alternate-buffer ds_write +// and the load latency overlaps the current chunk's LDS/MMAC work instead +// of stalling in front of it. B is consumed from the one-time packed +// [n_tile][k_step][lane][8] fragment-slot layout (w8a8_pack_b_fragslot_ +// kernel, produced by launch_pack_w8a8_weight outside the timed region and +// Graph capture, same byte count and graph-stable buffer), which replaces +// B's current 8x-over-fetching strided 16-B global reads with one aligned, +// coalesced dwordx2 per lane per k32 step; the scalar fallback decodes the +// pack for the paired M=2 shape. Everything else is byte-identical to the +// accepted winner: grid = N/16 = 256 blocks x 128 threads, in-block +// split-K=2 (wave w accumulates k in [w*K/2, (w+1)*K/2) k-ascending in 32-B +// DUMMA steps), LDS 10240 B/block (2 waves x 2 buffers x 2048 B of A/B +// chunks + 2 int32 partial planes -> 6 blocks/CU, residency unchanged), +// zero in-K-loop barriers, one END-of-K __syncthreads combine on wave 0 in +// ascending split order (mod-2^32 int32 addition groups exactly -> +// bit-identical, 0 mismatches), direct transposed-lane epilogue with scales +// before the bf16 store, and the scalar fallback for the paired M=2 shape +// and every unmatched (m,n,k). Graph safety is unchanged (single kernel +// launch on the caller stream inside capture, static LDS, no +// allocation/synchronization; every LDS byte is overwritten before every +// read per replay). Predicted signature if the mechanism is real: the +// steady K loop shows next-chunk global_load_dwordx4/dwordx2 BEFORE the +// current chunk's ds_reads/v_mmac with the vmcnt wait at the +// alternate-buffer ds_write_b128/b64 (vs today's vmcnt(3..0)-before- +// ds_write at stage top); vmem_read_instructions rises ~9472 -> ~13.5K (B +// loads split into narrower but coalesced dwordx2); lds_bank_conflicts +// drops (A row stride 64 B -> 8-way instead of 8-16-way; B slots +// contiguous -> conflict-free). Falsifiable gate: normal-benchmark +// median_us < 17.073 AND p90_us < 17.096 (accepted best). If flat or +// slower, the ~17 us floor is not the per-stage global-load exposure at +// 4.27 waves/CU (the co-resident wave already hides it) and the next round +// targets the launch/dispatch floor or the epilogue; the round still counts +// as a valid HIP experiment (build + correctness + graph capture). +// +// Iteration 15 change (split-epilogue round): the accepted iteration-9 +// kernel (256 blocks x 128 threads, in-block split-K=2, per-wave stage-K=64 +// double-buffer prefetch + packed-B, 10.695 us median / 11.145 us p90, 40 +// arch VGPR / 32 SGPR / 10240 B LDS, 512 wavefronts) ends with a +// wave-0-ONLY epilogue: after the END-of-K barrier, wave 1 exits while wave +// 0 serially runs 4 combine iterations (8 ds_read2st64_b32 + 4 stores) plus +// its 5 on-demand global scale loads (1 x_scale[r] + 4 weight_scale[n0+c], +// ISA 0x4434/0x445C/0x4510/0x45E4/0x46B0) with s_waitcnt vmcnt(1)/vmcnt(0) +// between the two scale multiplies of each of the 4 unrolled output +// iterations - the block's completion time is wave 0's full epilogue path. +// Iteration 14 (scales staged into LDS at kernel top, rejected at +// 10.822/10.866) proved the epilogue tail is live: removing the 5 scale +// loads cut p90 by 0.28 us (11.145 -> 10.866) while the prologue staging +// cost pushed the median +0.127 us. This round splits the SAME combine +// across both waves instead: wave w owns output words [w*128, w*128+128) +// (i = 2*w, 2*w+1), so each lane runs 2 combine iterations instead of 4, +// each wave needs only 3 scale loads (1 x_scale + 2 weight_scale), and wave +// 1's post-barrier idle becomes useful work. Per-element math is unchanged: +// the two-term int32 sum s_partial[0][idx] + s_partial[1][idx] and the +// float chain float(acc) * x_scale[r] * weight_scale[n0+c] are computed +// exactly once, by one wave, in the same expression order -> output +// bit-identical (0 mismatches). Everything else is byte-identical to the +// accepted iteration-9 winner: grid 256 x 128 threads, split-K=2 wave +// ownership and k-ascending DUMMA order, stage-K=64 prefetch, packed-B, +// LDS 10240 B/block unchanged (6 blocks/CU capacity >= 2.13 actual, all +// 256 blocks co-resident), the single END-of-K barrier (still orders both +// planes before either wave's combine reads), and the scalar fallback for +// the paired M=2 shape and every unmatched (m,n,k). Graph safety is +// unchanged (single kernel launch on the caller stream inside capture, +// static LDS 10240 B <= 64 KiB, no allocation/synchronization; every LDS +// byte is overwritten before every read per replay). Predicted signature if +// the mechanism is real: profiled kernel duration ~8.64 -> ~8.3-8.5 us; +// epilogue per wave shows 2 unrolled iterations of (2 x ds_read2st64_b32 + +// float mul + d16 store) with only 3 global_load_dword, no s_waitcnt +// vmcnt(1) between the scale multiplies; vmem_read_instructions ~13.5K -> +// ~14.3K (wave 0's 5 epilogue loads + wave 1's new 3 = 8 per block vs 5 +// today), lds_instructions ~23.3K unchanged (same reads, spread over 2 +// waves), VGPR ~40 -> ~44 (wave 1 holds 3 scale floats), LDS 10240 +// B/block unchanged. Falsifiable gate: normal-benchmark median_us < 10.695 +// AND p90_us < 11.145 (accepted best). If flat or slower, the epilogue is +// NOT on the block's critical path (the residual tail is the launch/ +// dispatch floor ~2 us above the 8.64 us profiled kernel and/or the HBM +// streaming rate at 40.6% L2 hit with B read exactly once), the next round +// targets the B-side slot swizzle or a register-only scale transport; the +// round still counts as a valid HIP experiment (build + correctness + +// graph capture). +// +// Iteration 18 repair 1 (compile fix): the paired 4-B bf16 epilogue store +// read the bf16 raw bits via __float2bfloat16(s).data, which does not exist +// on this DTK's __hip_bfloat16 (amd_hip_bf16.h stores the bits in a +// protected __x; the only raw-bit accessor is operator unsigned short()). +// Fixed by casting through static_cast so overload +// resolution selects the raw-bit conversion operator and NOT the value-based +// operator unsigned int() (an exact-match candidate that would silently +// convert the scaled float value instead of reinterpreting the bf16 bits). +// Semantics unchanged: the two 16-bit bf16 patterns still land in one 32-B +// word (lo = s0, hi = s1), output bit-identical to the previous two +// 2-B scalar stores; no kernel logic, mapping, or launch geometry changed. +// +// Iteration 19 change (B chunk-slot transport consolidation): the accepted +// iteration-18 kernel (256 blocks x 128 threads, in-block split-K=2, +// per-wave stage-K=64 double-buffer prefetch + packed-B, 9.949 us median / +// 9.962 us p90, 40 arch VGPR / 32 SGPR / 10240 B LDS, kernel +// w8a8_dumma_m16_n16_2wave_dbuf_kernelILi64EEE, 512 wavefronts) has B as the +// ONLY remaining 8-B-width global stream: the exact iteration-18 ISA shows +// per chunk the next-chunk B fetched as 2 x global_load_dwordx2 (0x42C8 / +// 0x42D0, 8 B per lane, 512 B apart -> not fuseable in the current +// [n_tile][k_step][lane][8] layout) while A is already ONE +// global_load_dwordx4 (0x42C0) per lane per chunk, and B's LDS write/read +// are already compiler-fused (ds_write2st64_b64 0x43C0 / ds_read2st64_b64 +// 0x42F8). Iteration 17 (register-only B transport, 11.529 us / 12.477 p90) +// proved the B LDS staging round trip is load-bearing and iteration 13 +// (A-side 8-way -> 2-way bank-conflict pad, 12.30 us) proved LDS conflict +// counts are NOT the limiter, so the steady loop's remaining HIP-visible +// inefficiency is the per-chunk 2 x 512-B B request stream itself. This +// round re-packs B from [n_tile][k_step][lane][8] to +// [n_tile][k_chunk][lane][16] (k_chunk = K/64 = 16 chunks; each 16-B lane +// slot holds BOTH k32 steps of one K=64 wavefront chunk, lo = step 2ch, hi = +// step 2ch+1, identical byte count K*N = 4,194,304), so each lane fetches +// the whole chunk with ONE 16-B aligned load (BChunkSlot alignas(16) -> +// global_load_dwordx4, matching the validated worker_0 wqkv recipe "one +// 16-B cooperative load/lane/stage" and the kernel's own A int4 path), and +// the LDS side becomes one ds_write_b128 / one ds_read_b128 per chunk with +// contiguous 16-B lane slots (clean banks vs the current 4-way 8-B slot +// stream). Per-element math is unchanged: identical global bytes in +// identical k-ascending 2 x du_mma_sync order per chunk (the two k32 steps' +// fragment bytes are the same 8-B values, just regrouped into one 16-B +// slot), identical split ownership, identical ascending split sum -> int32 +// bit-identical (0 mismatches). Everything else is byte-identical to the +// accepted iteration-18 winner: grid 256 x 128 threads, split-K=2 wave +// ownership and k-ascending DUMMA order, stage-K=64 A prefetch structure +// (A path untouched, ldm = kStageK = 64 unchanged), LDS 10240 B/block +// unchanged (B chunk buffer stays 1024 B per wave per buffer: 64 lanes x 16 +// B), the single END-of-K __syncthreads, the iteration-18 split two-wave +// combine with one 4-B bf16 store and one 8-B float2 weight_scale load per +// lane, and the scalar fallback for the paired M=2 shape and every unmatched +// (m,n,k) (packed_b_element updated to decode the new 16-B slot layout). The +// one-time pack kernel (launch_pack_w8a8_weight, out of the timed region and +// Graph capture, same byte count, graph-stable buffer) writes the new +// layout; the paired M=2 validation decodes it via packed_b_element, so the +// packed-shape scalar path stays exact. Graph safety is unchanged (single +// kernel launch on the caller stream inside capture, static LDS 10240 B <= +// 64 KiB, no allocation/synchronization; every LDS byte is overwritten +// before every read per replay; B bytes move global -> LDS -> registers +// ordered by the existing vmcnt/lgkmcnt waits). Predicted signature if the +// mechanism is real: the steady chunk loop shows ONE next-chunk +// global_load_dwordx4 for B instead of two global_load_dwordx2 (vmem_read_ +// instructions 13,312 -> ~11,264: B loads 8192 -> 4096), one +// ds_write_b128 + one ds_read_b128 per chunk replacing the st64-fused +// 8-B-slot pair (lds_instructions ~22.5K unchanged - same op count, cleaner +// banks: lds_bank_conflicts 233,472 -> ~180-200K as the B 4-way slot stream +// disappears; A conflicts remain and are known non-binding per iteration +// 13), VGPR ~40 -> ~40-42 (BChunkSlot = 4 VGPRs, same live set as the two +// uint64 today), LDS 10240 B/block unchanged, grid 256 / workgroup 128 / +// 512 waves unchanged, profiled kernel duration ~8.0 us -> ~7.8-8.0 us if +// the 1-KB per-chunk request stream improves L2/DRAM efficiency vs today's +// 2 x 512-B requests. Falsifiable gate: normal-benchmark median_us < 9.949 +// AND p90_us < 9.962 (accepted best). If flat or slower, the per-chunk +// request width is NOT the residual (the kernel is already streaming B at +// ~530-550 GB/s real HBM, the highest measured rate in this lineage and +// consistent with the o_proj 8.2-MB shape's ~450-500 GB/s cap; the residual +// is the HBM wall + the ~1.9 us launch/dispatch floor - not addressable +// from HIP at this shape), and the next round documents the HIP floor; the +// round still counts as a valid HIP experiment (build + correctness + graph +// capture). +// +// DUMMA path (w8a8_dumma_m16_n16_2wave_dbuf_kernel): +// * native gfx928 DUMMA INT8 m16n16k32 tiles, one 16x16 output tile per +// block, two 64-thread wavefronts per block (128 threads, 2 x 64); +// * in-block split-K=2 along K: wave 0 accumulates k in [0, K/2), wave 1 +// accumulates k in [K/2, K). Both publish their int32 partial plane to +// LDS; one END-of-K block barrier; both waves split the plane-sum +// combine (wave w owns output words [w*128, w*128+128)), so each +// element's two-term int32 sum is computed exactly once by one wave +// (int32 addition groups exactly -> bit-identical to the k-ascending +// order) and stored scaled bf16; +// * grid = N/16 = 256 blocks (>= 120 CUs), 512 wavefronts total +// (~4.3 waves/CU); +// * per-wave private double-buffered A/B stage chunks (K=64 per chunk, +// 2 chunks per buffer, 2 buffers per wave) so the K loop needs no +// barrier at all (within-wave LDS write->read ordering is emitted by +// the compiler as lgkmcnt waits); only the END-of-K combine barrier +// remains (1 barrier per block); +// * each chunk: 2 x m16n16k32 du_mma_sync in k-ascending order, with the +// next chunk's A/B global loads prefetched into scalar registers before +// the current chunk's MMAC burst and committed to the alternate buffer +// after it (vmcnt wait lands at the commit); +// * direct epilogue: transposed lane-contiguous int32 partial planes +// (element (row, col) -> word row + 16*col), combine reads are +// conflict-free, scales applied before the bf16 store; the 256-word +// combine is split across both waves (2 iterations per lane each) so +// the kernel-end critical path halves and wave 1 does not idle; +// * B is consumed in the packed [n_tile][k_chunk][lane][16] fragment-slot +// layout for (K, N) == (1024, 4096) (iteration-19 chunk-slot repack: +// each 16-B lane slot holds the two k32 steps of one K=64 chunk, one +// one-time pack in launch_pack_w8a8_weight, out of the timed +// region/Graph); every other (K, N) keeps the identity pack. +// +// Scalar path: unchanged correctness-first fallback used by the paired M=2 +// shape and every other unmatched (m, n, k); for (K, N) == (1024, 4096) it +// decodes the packed-B layout elementwise. +// +// Header order is fixed by the toolchain: HIP runtime, bf16, then du_mma +// (du_mma.h is not self-contained before the HIP runtime headers). + +#include +#include +#include + +#include + +namespace { + +constexpr int kScalarThreads = 128; // multiple of the gfx928 wavefront (64) +constexpr int kPackThreads = 256; + +// Fragment-slot B pack constants for the exact (K, N) == (1024, 4096) shape. +constexpr int kDummaTileK = 32; +constexpr int kPackedSlotBytes = 8; // one lane's matrix_b fragment bytes + // (one k32 step) +constexpr int kChunkSlotBytes = 16; // one lane's full K=64 chunk bytes + // (2 k32 steps, iteration-19 layout) +constexpr int kChunkBytes = 64 * kChunkSlotBytes; // 1024 B per wavefront chunk + +// One lane's 16-B B chunk slot in the [n_tile][k_chunk][lane][16] layout: +// lo = k-step 2ch fragment bytes, hi = k-step 2ch+1. alignas(16) makes the +// compiler emit 16-B vector loads/stores (global_load_dwordx4 / +// ds_read_b128 / ds_write_b128) exactly like the A int4 path; slot addresses +// are 16-B aligned by construction (b_base multiples of 1024 + lane*16). +struct alignas(16) BChunkSlot { + uint64_t lo; + uint64_t hi; +}; + +// Elementwise decode of the packed [n_tile][k_chunk][lane][16] fragment-slot +// layout for (K, N) == (1024, 4096) (iteration-19 chunk-slot repack: each +// 16-B lane slot holds one wavefront's two k32 steps of a K=64 chunk, lo = +// step 2ch, hi = step 2ch+1): returns logical weight[kk, n_col] with the +// same (tile, chunk, lane, byte) -> logical (k, n) mapping as the pack +// kernel and the DUMMA kernel. Used by the generic scalar fallback (paired +// M=2 validation) for the packed shape only. +__device__ __forceinline__ int8_t packed_b_element( + const int8_t* __restrict__ packed, int n_col, int kk, int k) { + const int k_chunks = k / (kDummaTileK * 2); // 16 for K=1024 + const int nt = n_col >> 4; + const int nn = n_col & 15; + const int ch = kk >> 6; + const int kk64 = kk & 63; + const int kk32 = kk64 & 31; + const int lane = (kk32 >> 3) * 16 + nn; + const int i = (kk32 & 7) + ((kk64 >> 5) << 3); + return packed[(static_cast(nt) * k_chunks + ch) * kChunkBytes + + lane * kChunkSlotBytes + i]; +} + +// Exact-shape DUMMA kernel for (M=16, N=4096, K=1024), tp8 wo_b. +// +// Two wavefronts per block, one 16x16 N tile per block, in-block split-K=2: +// wave 0 accumulates k in [0, K/2), wave 1 accumulates k in [K/2, K). Each +// wave owns private double-buffered A/B stage chunks (stage-K=64), so the K +// loop runs without any cross-wave barrier; both waves publish their int32 +// partial plane to LDS, one END-of-K block barrier orders publication, and +// both waves split the plane-sum combine (wave w owns output words +// [w*128, w*128+128)); each element's two-term int32 sum is computed exactly +// once by one wave (mod-2^32 int32 addition groups exactly, so the grouped +// sum is bit-identical to the k-ascending scalar order), scaled by +// x_scale/weight_scale, and stored bf16. +// +// Steady chunk loop (validated worker_0 prefetch pattern): the NEXT chunk's +// A/B global loads are issued into scalar registers at the top of the +// iteration (their latency is covered by the current chunk's ds_read/MMAC +// burst), and are committed to the ALTERNATE LDS buffer after the burst, so +// the compiler's vmcnt wait lands at the alternate-buffer ds_write instead +// of stalling before the current stage's LDS writes (the iteration-2 ISA +// serialization) and without the iteration-7 cross-iteration register +// rotation that the compiler spilled through global memory. +template +__global__ __launch_bounds__(128) void w8a8_dumma_m16_n16_2wave_dbuf_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ packed_b, // [n_tile][k_chunk][lane][16] + const float* __restrict__ x_scale, // [16, 1] fp32 + const float* __restrict__ weight_scale, // [N, 1] fp32 + hip_bfloat16* __restrict__ out, // [16, N] bf16 + int k, + int n) { + constexpr int kTileM = 16; + constexpr int kTileN = 16; + constexpr int kTileK = 32; + constexpr int kWaveSize = 64; + constexpr int kSplits = 2; + constexpr int kStageSteps = kStageK / kTileK; // 2 k32 steps per chunk + const int lane = threadIdx.x & (kWaveSize - 1); + const int wave = threadIdx.x >> 6; // 0 or 1 + const int n0 = static_cast(blockIdx.x) * kTileN; + const int ntile = static_cast(blockIdx.x); + + // Per-wave private double-buffered staging: A chunk 16 x 64 (1024 B) and B + // chunk kStageSteps x 64 x 8 (1024 B) per buffer, 2 buffers per wave + // (2 x 2 x 2048 = 8192 B), plus 2 x 16 x 16 x 4 = 2048 B of int32 partial + // planes -> 10240 B LDS per block (identical to the accepted iteration-2 + // kernel; 6 blocks/CU, all 256 blocks resident). + __shared__ __align__(16) int8_t a_tile[kSplits][2][kTileM * kStageK]; + __shared__ __align__(16) int8_t b_tile[kSplits][2][kChunkBytes]; + __shared__ __align__(16) int32_t s_partial[kSplits][kTileM * kTileN]; + + du::dumma::DUFragment< + du::dumma::matrix_a, kTileM, kTileN, kTileK, signed char, + du::dumma::row_major> + a_frag; + du::dumma::DUFragment< + du::dumma::matrix_b, kTileM, kTileN, kTileK, signed char, + du::dumma::row_major> + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // A chunk roles: 4 consecutive lanes own one 64-B row (16 B each), so the + // staging loads are 16-B vectorized and coalesced. + const int a_row = lane >> 2; + const int a_col16 = (lane & 3) * 16; + + const int kPerWave = k / kSplits; // 512 for K=1024 + const int kBase = wave * kPerWave; // 0 or 512 + const int kChunksTotal = k / kStageK; // 16 for K=1024 + const int64_t b_base = + static_cast(ntile) * kChunksTotal * kChunkBytes + + (kBase / kStageK) * kChunkBytes + lane * kChunkSlotBytes; + + // Prologue: stage chunk 0 into buffer 0 (single wait before the loop). + *reinterpret_cast(a_tile[wave][0] + a_row * kStageK + a_col16) = + *reinterpret_cast(a + a_row * k + kBase + a_col16); + *reinterpret_cast(b_tile[wave][0] + lane * kChunkSlotBytes) = + *reinterpret_cast(packed_b + b_base); + + const int nchunks = kPerWave / kStageK; // 8 for K=1024, K=64 chunks + for (int c = 0; c < nchunks; ++c) { + const int buf = c & 1; + + // Prefetch chunk c+1 from global into registers (no wait yet). The + // values are consumed only by the alternate-buffer stores below, so the + // compiler issues these loads before (or at worst alongside) the MMAC + // burst and lands the vmcnt wait at the ds_write - the overlap that + // removes the iteration-2 per-stage vmcnt(3..0)-before-ds_write + // exposure. Scalar variables only (no arrays, no cross-iteration + // registers), so the iteration-7 global-memory spill cannot recur. + int4 a_v; + BChunkSlot b_v; + const bool has_next = (c + 1 < nchunks); + if (has_next) { + const int k1 = kBase + (c + 1) * kStageK; + a_v = *reinterpret_cast(a + a_row * k + k1 + a_col16); + const int64_t boff = b_base + static_cast(c + 1) * kChunkBytes; + b_v = *reinterpret_cast(packed_b + boff); + } + + // Compute the current chunk from LDS: kStageSteps k-ascending k32 + // steps (unchanged DUMMA order per wave). No barrier: within-wave LDS + // write->read ordering is guaranteed by the compiler's lgkmcnt waits, + // and each wave only touches its own stage buffers. + const int8_t* a_buf = a_tile[wave][buf]; + const int8_t* b_buf = b_tile[wave][buf]; + // One 16-B LDS read per lane per chunk covers BOTH k32 steps of the + // chunk (lo = step 0, hi = step 1 of the chunk's 2-step K=64 span). + const BChunkSlot bv = + *reinterpret_cast(b_buf + lane * kChunkSlotBytes); +#pragma unroll + for (int j = 0; j < kStageSteps; ++j) { + du::dumma::du_load_matrix_sync(a_frag, a_buf + j * kTileK, kStageK); + const uint64_t b64 = j ? bv.hi : bv.lo; + __builtin_memcpy(b_frag.x, &b64, sizeof(b64)); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Commit the prefetched chunk into the alternate buffer (the compiler + // inserts the vmcnt wait here, after the MMAC burst above). + if (has_next) { + const int nb = (c + 1) & 1; + *reinterpret_cast(a_tile[wave][nb] + a_row * kStageK + + a_col16) = a_v; + *reinterpret_cast(b_tile[wave][nb] + + lane * kChunkSlotBytes) = b_v; + } + } + + // Publish the int32 partial plane in a transposed layout (element (row,col) + // -> word row + 16*col). Lane L's x[i] is (row = L&15, col = (L>>4) + 4*i), + // so it lands at word L + 64*i: consecutive words per step, conflict-free. + const int row = lane & 15; + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s_partial[wave][row + kTileN * (col_mod4 + 4 * i)] = acc_frag.x[i]; + } + __syncthreads(); // END-of-K: both partial planes visible before combine + + // Combine on BOTH waves (split in iteration 15 so the kernel-end critical + // path halves and wave 1's post-barrier idle becomes useful work). + // Wave w owns output words [w*128, w*128+128) = the full 16 rows x the + // 8-column band [8*w, 8*w+8) (word idx = row + 16*col in the transposed + // layout). Lane L owns the SAME-ROW ADJACENT-COLUMN pair + // (r = L&15, c = 8*wave + 2*(L>>4)) and (r, c+1): words + // A = w*128 + 32*(L>>4) + r and A+16 (an exact cover of the 128-word wave + // span, each word read by exactly one lane). The two bf16 results land + // side by side in out[r][n0+c .. n0+c+2), so both are emitted as ONE 4-B + // store (was two 2-B scalar stores) and the two adjacent fp32 weight + // scales come from ONE 8-B float2 load (was two 4-B loads); x_scale[r] is + // hoisted (loop-invariant). Per-plane reads stay at the same 64-B word + // separation (A, A+16) the compiler already fuses to ds_read2st64_b32, + // with lane-contiguous 16-word spans per 16-lane group (bank spread 2-way + // -> 4-way on this tail-only path, ~4 ns, not in the steady loop). Each + // element's two-term int32 sum s_partial[0][idx] + s_partial[1][idx] and + // the left-associative float chain float(acc) * x_scale[r] * + // weight_scale[n0+c] are computed exactly once, by one lane, in the same + // expression order as before -> output bit-identical (0 mismatches). + const int r = lane & 15; + const int c = 8 * wave + 2 * (lane >> 4); + const float xs = x_scale[r]; + const float2 ws = *reinterpret_cast(&weight_scale[n0 + c]); + const int idx0 = 2 * wave * kWaveSize + 32 * (lane >> 4) + r; + const int idx1 = idx0 + 16; + const int32_t acc0 = s_partial[0][idx0] + s_partial[1][idx0]; + const int32_t acc1 = s_partial[0][idx1] + s_partial[1][idx1]; + const float s0 = static_cast(acc0) * xs * ws.x; + const float s1 = static_cast(acc1) * xs * ws.y; + *reinterpret_cast(&out[r * n + n0 + c]) = + (static_cast( + static_cast(__float2bfloat16(s0)))) | + (static_cast( + static_cast(__float2bfloat16(s1))) + << 16); +} + +// out[m, n] = bf16( int32(sum_k a[m,k] * b[k,n]) * x_scale[m] * weight_scale[n] ) +// +// One thread per output element. Grid covers M*N; every supported (m, n, k) +// from the fixed contract reaches this kernel (generic fallback). For +// (K, N) == (1024, 4096) the weight buffer holds the packed +// [n_tile][k_step][lane][8] fragment-slot layout, which is decoded +// elementwise here so the paired M=2 validation stays exact. +__global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (idx >= total) { + return; + } + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + + // (K, N) == (1024, 4096) weights are packed into the fragment-slot layout + // by launch_pack_w8a8_weight; every other shape keeps the logical layout. + const bool b_packed = (k == 1024 && n == 4096); + + const int8_t* __restrict__ a_row = a + static_cast(row) * k; + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + const int8_t b_val = + b_packed ? packed_b_element(b, col, kk, k) + : b[static_cast(kk) * n + col]; + acc += static_cast(a_row[kk]) * + static_cast(b_val); + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); +} + +// Identity device-to-device copy (bootstrap packing layout). +__global__ void w8a8_pack_identity_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t numel) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < numel) { + dst[idx] = src[idx]; + } +} + +// One-time device permutation of the raw logical [K, N] row-major weight +// into the packed [n_tile][k_chunk][lane][16] B fragment-slot layout for the +// exact (K, N) == (1024, 4096) shape (iteration-19 chunk-slot repack: each +// 16-B lane slot holds the two k32 steps of one K=64 wavefront chunk, so the +// DUMMA kernel fetches the whole chunk with ONE 16-B global load per lane +// instead of two 8-B loads). One thread owns one 16-byte lane slot; slot +// ((ntile * (k/64) + chunk) * 64 + lane) * 16 + i gets logical +// weight[chunk*64 + (i>>3)*32 + (lane>>4)*8 + (i&7)][ntile*16 + (lane&15)]. +// Byte count is unchanged (K*N), so packed_weight keeps the same numel and +// graph-stable address; runs outside the timed region and Graph capture. +__global__ __launch_bounds__(kPackThreads) void +w8a8_pack_b_fragslot_kernel(const int8_t* __restrict__ raw, // [K, N] + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t tid = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int k_chunks = k / (kDummaTileK * 2); + const int64_t total = static_cast(n / 16) * k_chunks * 64; + if (tid >= total) { + return; + } + const int lane = static_cast(tid & 63); + const int64_t slot = tid >> 6; // ntile * k_chunks + chunk + const int chunk = static_cast(slot % k_chunks); + const int ntile = static_cast(slot / k_chunks); + const int nn = lane & 15; + const int k8 = (lane >> 4) << 3; + const int8_t* src0 = + raw + (static_cast(chunk) * (2 * kDummaTileK) + k8) * n + + ntile * 16 + nn; + const int8_t* src1 = src0 + static_cast(kDummaTileK) * n; + int8_t* dst = packed + tid * kChunkSlotBytes; +#pragma unroll + for (int i = 0; i < kPackedSlotBytes; ++i) { + dst[i] = src0[static_cast(i) * n]; + dst[kPackedSlotBytes + i] = src1[static_cast(i) * n]; + } +} + +// Identity device-to-device copy of the per-column fp32 scales. +__global__ void w8a8_pack_scale_identity_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int numel) { + const int idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < numel) { + dst[idx] = src[idx]; + } +} + +} // namespace + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + // Exact-shape DUMMA specialization: tp8 wo_b (M=16, N=4096, K=1024). + // Two wavefronts per block (in-block split-K=2 along K), one 16x16 N tile + // per block; grid = N/16 = 256 blocks >= 120 CUs -> 512 wavefronts total + // (~4.3 waves/CU). Stage-K=64 double-buffered A+B prefetch with packed-B + // fragment-slot loads (b is the iteration-19 packed [n_tile][k_chunk] + // [lane][16] layout for this shape: one 16-B load per lane per K=64 + // chunk). + if (m == 16 && n == 4096 && k == 1024) { + constexpr int kStageK = 64; + const dim3 grid(static_cast(n / 16)); + const dim3 block(128); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_n16_2wave_dbuf_kernel), + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast(out), + k, + n); + return; + } + // Generic scalar fallback for every other (m, n, k), including the paired + // M=2 shape with the same (N, K); the scalar kernel decodes the packed-B + // layout for (k, n) == (1024, 4096). + const int64_t total = static_cast(m) * n; + const int64_t blocks = (total + kScalarThreads - 1) / kScalarThreads; + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + dim3(static_cast(blocks)), + dim3(kScalarThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast(out), + m, + n, + k); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_elems = static_cast(k) * n; + if (weight_elems > 0) { + if (k == 1024 && n == 4096) { + // Exact (K, N): pack B into the [n_tile][k_step][lane][8] fragment- + // slot layout consumed by the iteration-9 DUMMA kernel (one-time, + // out of the timed region and Graph capture, same byte count). + const int64_t slots = static_cast(n / 16) * (k / 32) * 64; + const unsigned grid = static_cast( + (slots + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_pack_b_fragslot_kernel), + dim3(grid), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + // Bootstrap: identity device-to-device packing for every other (K, N). + const int64_t blocks = (weight_elems + kPackThreads - 1) / kPackThreads; + hipLaunchKernelGGL( + w8a8_pack_identity_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_elems); + } + } + if (n > 0) { + const int blocks = (n + kPackThreads - 1) / kPackThreads; + hipLaunchKernelGGL( + w8a8_pack_scale_identity_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wq_b.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wq_b.hip new file mode 100644 index 00000000..ff0598e7 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wq_b.hip @@ -0,0 +1,825 @@ +// @@variant shape=tp8_wq_b_m16 commit=bfaef2fc76d0e4a095b2ffae2b466f2863511d43 added=2026-08-26 +// median_us=14.67 p90_us=14.7 +// source=hy3-dsh-tp8-m16-1-adf021ab +// W8A8 INT8 GEMM — gfx928 implementation (worker_1, physical GPU 1) +// +// Assigned shapes: +// tp8_wq_b_m16 M=16, N=4096, K=1024 +// tp8_indexer_wq_b_m16 M=16, N=8192, K=1024 +// +// Iteration 0 (correctness-first bootstrap): one simple scalar int8 +// dot-product kernel maps one thread to one output element and computes the +// complete K loop exactly in int32 before applying the two fp32 scales and +// storing bf16. This remains the generic fallback for every shape that does +// not match the exact tp8_wq_b_m16 guard below. +// +// Iteration 1 (DUMMA bootstrap): the exact tp8_wq_b_m16 shape ran a minimal +// 16x16x32 DUMMA kernel, one 64-thread wavefront per 16x16 output tile +// (256 blocks, 1 wave, direct row-major global fragment loads). It built, +// passed exact int32 correctness and Graph capture, and measured 59.9 us +// median (vs the 32.137 us fixed Triton Graph baseline). Its gfx928 ISA +// showed a latency-bound K loop: the library fragment loader expanded into +// 18 global_load_ubyte with 19 serialized s_waitcnt around a single +// v_mmac_i32_16x16x32_i8 per step, and with ~2 blocks resident per CU there +// were too few independent load->MMA streams to hide that latency. +// +// Iteration 2 (architecture round — measure grid parallelism): changed ONLY +// the launch geometry to two wavefronts per block owning two adjacent 16x16 +// N tiles (block tile 16x32, grid N/32 = 128 blocks = 1.07 blocks/CU, two +// co-resident full-K MMA chains per CU). It measured 83.0 us median — FLAT +// OR WORSE than iteration 1 — falsifying grid parallelism as the binding +// lever: the binding constraint is the per-step byte-load fragment path +// itself (18 global_load_ubyte + 19 s_waitcnt per step), and the mandated +// architecture axis is split-K, not more geometry. +// +// Iteration 3 (architecture round — split-K, this file): the unsplit grid +// has 128 blocks / 120 CUs = 1.07 < 2 blocks/CU and K = 1024 >= 1024, so the +// mandate requires split-K=2 plus a CU-aligned candidate, with int32 +// partials written into the caller workspace and a combine+scale kernel in +// the timed Graph. This round implements a ONE-WAVE ZERO-BARRIER workspace +// split-K pair (also the mandate's alternative geometry branch — it reaches +// double the unsplit parallelism with zero barriers): +// (1) w8a8_dumma_m16n16k32_splitk_partial_kernel — one 64-thread +// wavefront per block, grid = (N/16) * kSplitK = 512 blocks for +// split-K=2 (4.27 blocks/CU >= the two-blocks-per-CU target, 512 +// resident wave-tasks vs 256 in the unsplit grid). Block b owns +// (split = b % kSplitK, tile = b / kSplitK): split s sums its +// 32-aligned K-chunk ([0,512) or [512,1024) for split-K=2, 16 +// m16n16k32 steps per wave, same k-ascending operand order as the +// unsplit chain) with the UNCHANGED direct du_load_matrix_sync path, +// and stores the raw int32 accumulator fragment directly into the +// caller workspace plane partials[split][tile][16][16] via the +// verified gfx928 accumulator ownership (row = lane&15, +// col_mod4 = lane>>4, x[i] -> col_mod4 + 4*i). No LDS, no barriers. +// (2) w8a8_dumma_m16n16k32_combine_kernel — one thread per +// output element, grid = 256 blocks x 256 threads, sums the kSplitK +// planes ASCENDING (s = 0..kSplitK-1; int32 addition is exact and +// order-independent below 2^24, so the total is bit-identical to the +// unsplit accumulator), applies (acc * x_scale[row]) * +// weight_scale[col] in the same left-associative float order as the +// reference, and stores round-to-nearest-even bf16. Launched +// immediately after the partial kernel on the same stream, so both +// kernels are captured into the timed Graph. +// The split factor is the single compile-time constant kSplitKMeasured = 2. +// The CU-aligned candidate kSplitKCUAligned = 15 (grid 3840 = exactly +// 32 blocks/CU on the 120-CU device, non-power-of-two, non-uniform 32-aligned +// K chunks) is compiled into the object and selected by flipping the +// constant, so the mandated occupancy sweep is a one-line change. +// +// Iteration 8 (HIP-only resource round — tune ONE occupancy limiter, the +// LDS footprint, from the accepted kernel's PMC): iterations 4-6 proved the +// per-step direct global fragment loads are poison (reg-prefetch 51.0 us, +// B-only LDS 67.0 us) and that only a fully A+B-staged zero-global K loop +// can win (the iteration-6 fused kernel, 25.97 us) — but the fused kernel +// stages the WHOLE K=1024 per block (35,456 B/block: s_a+s_b 16,640 each + +// s_part 2,048 + s_scale 128), pinning it at 1 block/CU co-resident. This +// round moves that staging to the SPLIT GRID (the exact next step prescribed +// by iteration 7): the accepted iteration-3 grid/workspace/combine Graph is +// kept (512 blocks x 64 threads, split-K=2, 2 int32 partial planes, separate +// combine kernel in the timed Graph), but the partial kernel now stages ONLY +// its own 512-K A chunk (16 x 512 B, row-major x_q) and B chunk (packed +// [N][K] P, 512 contiguous B per column) into LDS once: 16,896 B/block +// (s_a/s_b 16 x 528 = 512+16 padded stride) -> 3 blocks/CU co-resident on +// the 64 KiB LDS/CU (vs 1 block/CU for the fused kernel, 4.27 avg for the +// accepted kernel). The K loop then collapses to the fused kernel's proven +// zero-global chain: per step exactly two ds_read_b64 (A and B fragments, +// raw byte order == the du_load_matrix_sync expansion) + one v_mmac, same +// k-ascending chunk [0,512)/[512,1024) as iteration 3. Byte accounting per +// replay is UNCHANGED (no repeated HBM reads: B 4 MiB unique, A 16 KiB +// L2-hot re-read by 256 tiles, 512 KiB partial write+read, 128 KiB out) — +// occupancy is raised by shrinking the per-block LDS footprint, not by +// re-reading global memory. Transport prerequisite (iteration-5/6-validated): +// for (k,n)==(1024,4096) launch_pack_w8a8_weight now writes the [N][K] +// transpose P[n*K+k] = W[k*N+n] once outside the timed region/Graph, and the +// generic scalar fallback decodes it via a b_transposed flag so the paired +// M=2 validation and every other caller of the exact (n,k) stay bit-exact; +// every other shape keeps the identity pack. +// +// Iteration 9 (HIP-only occupancy probe — the mandated one-constant staged +// split sweep): the accepted iteration-8 kernel's PMC (staged partial, +// profiled 14.56 us, 512 blocks x 64 threads, 16,896 B/block LDS -> 3 blocks +// /CU co-resident on the 64 KiB LDS/CU, 0.75 waves/SIMD resident with the +// 4th SIMD of every CU idle, grid 4.27 blocks/CU -> 1.42 serial block rounds) +// shows the staged transport's per-step chain is short (2 ds_read_b64 + +// v_mmac, ~LDS-latency), so the residual limiter is resident-stream count +// over the one-per-block HBM staging wait, not per-step global latency (the +// direct-load split-K=4 regression of iteration 7 does NOT transfer: that +// transport kept a ~500-600 cycle HBM chain inside EVERY K step). This round +// sweeps the dispatched split to S=4 — the exact next step iteration 8 +// prescribed: 1024 blocks x 64 threads, uniform 32-aligned K chunks +// [0,256)/[256,512)/[512,768)/[768,1024), 8 m16n16k32 steps per block, +// staging shrinks to 8,704 B/block (s_a/s_b 16 x 272 = 256+16 padded +// stride, the same 4-dword-mod-32 bank skew as 528) -> 7 blocks/CU +// co-resident (floor(65536/8704) = 7), 1.75 waves/SIMD with ALL FOUR SIMDs +// busy, 1.22 serial block rounds, 4 int32 partial planes = 1 MiB <= the +// API's 16-plane / 4 MiB workspace. Exactness is preserved by construction: +// identical k-ascending m16n16k32 order within each chunk, chunks tile +// [0,1024) exactly once, the UNCHANGED ascending 4-plane combine sums +// (p0+p1+p2)+p3 == the S=2 two-plane total == the unsplit accumulator +// (int32 exact below 2^24), so the output is bit-identical to iteration 8 +// (expected 0 mismatches). Graph contract unchanged: same two kernels, same +// stream order, 4 planes within the same 16-plane allocation, every partial +// overwritten each launch. +// +// Repair 1 (exact-correctness defect of the S=4 sweep, no architecture +// change): the cooperative staging loop decomposed the 16-B chunk index with +// the S=2-fixed mapping r = c>>5, ko = (c&31)*16 (32 chunks per row). That is +// exact at kLocal=512 but at S=4 (kLocal=256, kChunksLocal=256) c only +// reaches 255, so only A rows / B columns 0..7 were staged and the K loop +// read UNINITIALIZED LDS for rows/columns 8..15 (75% of every tile wrong — +// rows 8..15 in all columns plus columns 8..15 of rows 0..7; observed +// mismatch_count 49121/65536, first mismatch m=0 n=8). Fixed in place by +// deriving the per-row chunk count from the actual chunk length, +// kChunksPerRow = kLocal/16 (r = c/kChunksPerRow, ko = (c%kChunksPerRow)*16), +// which is bit-identical to the old mapping at S=2 (512/16 = 32) and stages +// all 16 rows/columns at S=4; kLocal stays a 32-multiple so every 16-B +// global load/store alignment is unchanged. The S=4 mapping/performance idea +// (1024 blocks, 8,704 B/block LDS, 7 blocks/CU) is preserved. +// +// Iteration 11 (HIP-only consolidation — fused combine tail): the accepted +// S=4 tree's operator-aggregate PMC (partial 10.879 us + combine 3.04 us = +// 13.919 us vs the 15.08 timed median) pins the residual wall on the combine +// STRUCTURE: a second kernel launch (~1.16 us inter-kernel/graph gap) plus a +// 1 MiB int32 plane re-read that is HBM-cold (combine L2 hit rate only 19.3%: +// the 4 MiB B stream evicts the partials), while the sibling TP4/M16 lineage +// (qkv_proj.hip rounds 17/23) proves the canonical M<=32 fix is the FUSED +// ATOMIC COMBINE. This round merges the combine into the partial kernel tail: +// every block stores its plane as before, then tid-0 does __threadfence() + +// atomicAdd(&counters[tile], 1), and the block whose atomicAdd returns the +// kSplitK-th arrival of its replay (monotonic mod-kSplitK: r % kSplitK == +// kSplitK-1; int32 wraparound would need 2^31/4 ~ 536M replays) sums the +// tile's kSplitK planes in ascending slice order (identical per-element int32 +// order as the removed combine), applies the identical left-associative +// fp32 scale and RN-even bf16 conversion, and stores byte-identical output +// bytes (0 mismatches expected). The Graph becomes ONE kernel launch per +// replay. The per-tile counters (256 x int32 = 1 KiB) live in the last +// kCountersBytes of the caller workspace (planes 0..3 end at 1 MiB, far below +// the 4 MiB allocation's tail; no overlap), zeroed once per workspace by an +// async hipMemsetAsync on the caller's stream before the first launch (never +// part of steady-state replay; the mod-kSplitK test is also correct if a +// reset were captured, since arrivals 0..kSplitK-1 still identify the last). +// The staged K loop, transport, pack, exact-shape guards, generic scalar +// fallback and the compile-only kSplitKCUAligned=15 two-kernel branch are +// unchanged; the generic fallback never touches the counters. +// +// Layout contract (logical tensors, contiguous / row-major): +// x_q [M, K] int8 +// packed_weight [K, N] int8 (identity D2D copy of raw_weight made +// by launch_pack_w8a8_weight) +// x_scale [M] fp32 +// packed_weight_scale [N] fp32 (identity copy of weight_scale) +// out [M, N] bf16 +// workspace caller-owned int32 partial planes: +// partials[split][tile][16][16], plane stride +// (N/16)*256 int32 elements +// +// Math: +// out[m, n] = bf16( (int32) sum_k x_q[m,k] * W[k,n] * x_scale[m] +// * weight_scale[n] ) +// +// Exactness: the assigned K=1024 keeps every partial int8 dot below 2^24, so +// the int32 sum (in any grouping order) is bit-identical to the fp32 torch +// reference ((A.float() @ B.float()) * x_scale * weight_scale.T) for the +// timed shapes. The generic fallback is exact in int32 for every supported +// shape (max |dot| = K * 127 * 127 < 2^31). + +#include +#include +#include + +#include + +namespace { + +constexpr int kScalarBlock = 256; // 4 wavefronts of 64 (blockDim % 64 == 0) +constexpr int kPackBlock = 256; + +// DUMMA INT8 m16n16k32 constants: one 64-thread wavefront owns one fragment. +constexpr int kDummaM = 16; +constexpr int kDummaN = 16; +constexpr int kDummaK = 32; +constexpr int kDummaWave = 64; + +// Exact-shape constants for tp8_wq_b_m16 (M=16, N=4096, K=1024). +constexpr int kExactK = 1024; +constexpr int kExactN = 4096; +constexpr int kExactNTiles = kExactN / kDummaN; // 256 + +// Combine kernel block size. +constexpr int kCombineThreads = 256; + +// Iteration 9 staged occupancy sweep — split factor dispatched this round. +// The exact-shape fast path launches the staged partial kernel below, which +// supports every trusted split in {2..15} via runtime 32-aligned chunks, so +// the occupancy sweep stays a one-constant change. S=4: 1024 blocks x 64 +// threads, K-chunk 256 per block, LDS 8,704 B/block (s_a/s_b 16 x 272 = +// 256+16 padded stride) -> 7 blocks/CU co-resident (vs 3 at S=2), grid +// average 8.53 blocks/CU -> 1.22 serial block rounds (vs 1.42), 1.75 +// waves/SIMD resident with all 4 SIMDs busy (vs 0.75 on 3 of 4 SIMDs), 4 +// int32 partial planes = 1 MiB <= the API's 16-plane / 4 MiB workspace. +constexpr int kSplitKMeasured = 4; +// Implemented CU-aligned candidate: grid = 256 * 15 = 3840 blocks = exactly +// 32 blocks/CU on the 120-CU gfx928 (integer blocks per CU, non-power-of-two, +// non-uniform 32-aligned K chunks). Compiled into the object via the +// iteration-3 direct-load partial kernel; never dispatched (the +// workspace_bytes == -1 guard cannot hold and it does not decode the [N][K] +// pack). +constexpr int kSplitKCUAligned = 15; + +// Iteration 11 fused combine tail: per-tile arrival counters (256 tiles x +// int32 = 1 KiB) live in the last kCountersBytes of the caller workspace +// (planes 0..3 end at 1 MiB, far below the counters' offset in the API's +// 16-plane / 4 MiB allocation, so there is no overlap with partial data). +constexpr int64_t kCountersBytes = kExactNTiles * sizeof(int32_t); // 1 KiB + +// Round-to-nearest-even float -> bfloat16 bit pattern (sibling-validated +// helper; identical RN-even conversion to the hip_bfloat16 stores used by +// the old combine kernel, so the packed tail store is byte-identical). +__device__ __forceinline__ uint16_t float_to_bf16_bits(float f) { + uint32_t u = 0; + __builtin_memcpy(&u, &f, sizeof(u)); + const uint32_t bias = 0x7FFFu + ((u >> 16) & 1u); + u += bias; + return static_cast(u >> 16); +} + +// One-time async zero of the fused-combine arrival counters, guarded by the +// workspace pointer: issued once per workspace before the first launch on +// the caller's stream, so the captured Graph contains only the single fused +// kernel launch and every steady-state replay is memset-free. +static void* g_fused_counters_zeroed = nullptr; + +// --------------------------------------------------------------------------- +// Iteration 3 split-K partial GEMM for the exact tp8_wq_b_m16 shape. +// One 64-thread wavefront per block, grid = (N/16) * kSplitK blocks. Block b +// owns (split = b % kSplitK, tile = b / kSplitK): split s computes the +// 16x16 output tile over its 32-aligned K-chunk [kStart(s), kEnd(s)) in +// k-ascending m16n16k32 steps (16 steps for split-K=2), with the unchanged +// direct row-major global fragment loads, and stores the raw int32 +// accumulator fragment into workspace plane s (tile-major, row stride 16) +// via the verified gfx928 accumulator ownership. Zero barriers, zero LDS. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kDummaWave) +void w8a8_dumma_m16n16k32_splitk_partial_kernel( + const int8_t* __restrict__ x_q, // [16, K] row-major + const int8_t* __restrict__ weight, // [K, N] row-major (identity pack) + int32_t* __restrict__ partials, // [kSplitK][N/16][16][16] int32 + const int n) { + const int split = static_cast(blockIdx.x) % kSplitK; + const int tile = static_cast(blockIdx.x) / kSplitK; + const int lane = static_cast(threadIdx.x); + const int n0 = tile * kDummaN; + + // 32-aligned non-uniform K chunks: kStart(s) = floor(s*K/S) rounded down + // to a DUMMA-K multiple; the last chunk absorbs the remainder. For + // split-K=2 this is exactly [0,512) and [512,1024). + const int kStart = ((split * kExactK) / kSplitK) & ~(kDummaK - 1); + const int kEnd = (split == kSplitK - 1) + ? kExactK + : ((((split + 1) * kExactK) / kSplitK) & + ~(kDummaK - 1)); + + du::dumma::DUFragment a_frag; + du::dumma::DUFragment b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Same per-step instruction stream as iterations 1/2 (direct global + // fragment loads, exact int32 k-ascending accumulation); only the K range + // per wave shrinks by the split factor. Fully unrolled (16 steps for + // split-K=2) like the validated split-K lineage kernels. +#pragma unroll + for (int k0 = kStart; k0 < kEnd; k0 += kDummaK) { + du::dumma::du_load_matrix_sync(a_frag, x_q + k0, kExactK); + du::dumma::du_load_matrix_sync( + b_frag, weight + static_cast(k0) * n + n0, n); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // gfx928 int8 m16n16k32 accumulator ownership, established against + // du_store_matrix_sync: lane%16 selects the row, lane/16 selects col%4, + // and x[i] holds the columns col%4 + 4*i. Store the raw int32 partial + // (no scaling) into the caller workspace plane. + const int row = lane & 15; + const int col_mod4 = lane >> 4; + const int n_tiles = n / kDummaN; + int32_t* __restrict__ p = partials + + (static_cast(split) * n_tiles + tile) * + (kDummaM * kDummaN); +#pragma unroll + for (int i = 0; i < 4; ++i) { + p[row * kDummaN + col_mod4 + 4 * i] = acc_frag.x[i]; + } +} + +// --------------------------------------------------------------------------- +// Iteration 8 split-grid staged partial GEMM for the exact tp8_wq_b_m16 +// shape — the LDS-footprint occupancy tune. Same grid as iteration 3 (512 +// blocks x 64 threads, split = b % kSplitK, tile = b / kSplitK, 32-aligned +// non-uniform chunks [kStart, kEnd)), same workspace partial-plane contract, +// same combine kernel in the timed Graph. The per-step direct global +// fragment loads are replaced by ONE cooperative staging pass per block: +// each lane moves 16-B chunks of the block's own A chunk (x_q rows, 512 B +// each at stride K) and B chunk (packed [N][K] P, 512 contiguous B per +// column) with coalesced dwordx4 global reads and 16-B aligned b128 LDS +// stores — 16 global loads + 16 LDS stores per lane at split-K=2 — followed +// by ONE __syncthreads() and a zero-global K loop whose every step is +// exactly two ds_read_b64 (A and B fragments; raw byte order == the +// du_load_matrix_sync expansion) + one v_mmac (the iteration-6-validated +// fused chain, minus the barrier/s_part handoff). LDS = 16,896 B/block at +// the dispatched S=2 (s_a/s_b 16 x 528 = 512+16 padded stride) -> 3 blocks +// /CU co-resident on the 64 KiB LDS/CU. Template supports every trusted +// split in {2..15}: the buffer covers ceil(K/S) rounded up to a DUMMA-K +// multiple and both loops use runtime chunk bounds, so the mandated sweep +// stays a one-constant change (kSplitKMeasured). +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kDummaWave) +void w8a8_dumma_m16n16k32_splitk_staged_partial_kernel( + const int8_t* __restrict__ x_q, // [16, K] row-major + const int8_t* __restrict__ weight, // packed [N, K] n-major (exact) + int32_t* __restrict__ partials, // [kSplitK][N/16][16][16] int32 + const float* __restrict__ x_scale, // [16] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [16, N] bf16 + int32_t* __restrict__ counters, // [N/16] per-tile arrival counts + const int n) { + constexpr int kChunkMax = + (((kExactK + kSplitK - 1) / kSplitK + kDummaK - 1) / kDummaK) * + kDummaK; // 512 at split-K=2 + constexpr int kStageStride = kChunkMax + 16; // 528 = 512+16 at split-K=2 + + __shared__ __align__(16) int8_t s_a[kDummaM * kStageStride]; + __shared__ __align__(16) int8_t s_b[kDummaN * kStageStride]; + + const int split = static_cast(blockIdx.x) % kSplitK; + const int tile = static_cast(blockIdx.x) / kSplitK; + const int lane = static_cast(threadIdx.x); + const int n0 = tile * kDummaN; + + // Same 32-aligned non-uniform K chunks as the iteration-3 kernel; for + // split-K=2 this is exactly [0,512) and [512,1024). + const int kStart = ((split * kExactK) / kSplitK) & ~(kDummaK - 1); + const int kEnd = (split == kSplitK - 1) + ? kExactK + : ((((split + 1) * kExactK) / kSplitK) & + ~(kDummaK - 1)); + const int kLocal = kEnd - kStart; // 32-multiple, <= kChunkMax + + // Cooperative staging of this block's A and B chunks: every 16-B chunk of + // each matrix is moved exactly once. Per pass each lane issues one + // contiguous dwordx4 global load per matrix (coalesced 16-B sectors: lanes + // 0..31 cover one 512-B row/column segment, lanes 32..63 the next) and one + // 16-B aligned b128 LDS store (stride 528 = 16*33 keeps every store + // 16-B-aligned). c enumerates the kChunksLocal 16-B chunks of ONE matrix + // (A rows == B columns, both 16 x kLocal) and the chunk -> (row, k-offset) + // decomposition uses kChunksPerRow = kLocal/16 so ALL 16 rows/columns are + // staged at every split; the fixed 32-chunks-per-row shift mapping only + // covered rows/columns 0..7 once kLocal dropped to 256 at split-K=4 + // (repair 1: uninitialized-LDS reads for rows/columns 8..15). + const int kChunksLocal = (kLocal * kDummaM) / 16; // 16-B chunks per matrix + const int kChunksPerRow = kLocal / 16; // 16-B chunks per A row / B column + const int64_t b_base = + static_cast(n0) * kExactK + kStart; // packed P column base +#pragma unroll + for (int c = lane; c < kChunksLocal; c += kDummaWave) { + const int r = c / kChunksPerRow; // A row == B column in the 16 + const int ko = (c % kChunksPerRow) * 16; // k offset in the row/column + const int4 va = *reinterpret_cast( + x_q + r * kExactK + kStart + ko); + const int4 vb = *reinterpret_cast( + weight + b_base + static_cast(r) * kExactK + ko); + *reinterpret_cast(s_a + r * kStageStride + ko) = va; + *reinterpret_cast(s_b + r * kStageStride + ko) = vb; + } + __syncthreads(); + + du::dumma::DUFragment a_frag; + du::dumma::DUFragment b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Fragment ownership == the du_load_matrix_sync expansion (validated in + // iterations 3/5/6): A: row = lane&15, kg = lane>>4, x[i] = A[row][kg*8+i]; + // B: col = lane&15, kg = lane>>4, x[i] = P[n0+col][kg*8+i] -- 8 contiguous + // staged bytes per fragment, one ds_read_b64 each, raw byte order (the + // loader's byte-reassembly is a no-op on the 8-byte path, so the manual + // fill is bit-identical). Zero-global, zero-barrier k-ascending chain over + // the block's own chunk. + const int frag_row = lane & 15; // A row == B column + const int kg = lane >> 4; + const int lds_off = frag_row * kStageStride + kg * 8; + const int kSteps = kLocal / kDummaK; +#pragma unroll + for (int s = 0; s < kSteps; ++s) { + *reinterpret_cast(&a_frag.x[0]) = + *reinterpret_cast(s_a + lds_off + s * kDummaK); + *reinterpret_cast(&b_frag.x[0]) = + *reinterpret_cast(s_b + lds_off + s * kDummaK); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Same raw int32 partial-plane store as iteration 3 (verified gfx928 + // accumulator ownership: row = lane&15, col_mod4 = lane>>4, x[i] -> + // col_mod4 + 4*i). + const int row = lane & 15; + const int col_mod4 = lane >> 4; + const int n_tiles = n / kDummaN; + int32_t* __restrict__ p = partials + + (static_cast(split) * n_tiles + tile) * + (kDummaM * kDummaN); +#pragma unroll + for (int i = 0; i < 4; ++i) { + p[row * kDummaN + col_mod4 + 4 * i] = acc_frag.x[i]; + } + + // Iteration 11 fused combine tail (replaces the separate combine kernel in + // the timed Graph; sibling-validated arrival protocol): every block's tid-0 + // releases with __threadfence() + atomicAdd(&counters[tile], 1); the LAST + // arrival for its tile (atomicAdd return value r with r % kSplitK == + // kSplitK-1 — hardware-ordered after all kSplitK sibling stores+fences, no + // spin loop, no residency assumption) sums the tile's planes in ascending + // slice order — the identical per-element int32 accumulation order as the + // old combine — applies the identical (acc * x_scale[row]) * + // weight_scale[col] fp32 scale order, and stores RN-even bf16 at the same + // addresses. The counters are MONOTONIC across replays: replay r's arrivals + // return (r-1)*kSplitK + 0..kSplitK-1, so the kSplitK-th arrival of EVERY + // replay is exactly the r % kSplitK == kSplitK-1 one (int32 wraparound + // would need 2^31/4 ~ 536M replays); launch_w8a8_gemm zeroes them once per + // workspace before the first launch and the generic fallback never touches + // them. Element mapping: lin0 = lane*4 covers the 256-element tile exactly + // once (row = lin0>>4, col0 = lin0&15, 4 consecutive columns per lane — 4 | + // 16 so a run never crosses the 16-column row boundary), each plane is read + // as one perfectly coalesced dwordx4 per lane, and the four 2-B bf16 + // results pack into one 8-B aligned uint64 store — bit-identical output + // bytes to the removed combine kernel (0 mismatches expected). + __syncthreads(); + __shared__ int s_is_last; + if (threadIdx.x == 0) { + __threadfence(); // release: this block's plane stores are visible to + // the observer of its arrival atomic + const int arrived = atomicAdd(&counters[tile], 1); + __threadfence(); // acquire: the reads below (after the barrier) see + // every sibling store released before prior arrivals + s_is_last = ((arrived % kSplitK) == kSplitK - 1); + } + __syncthreads(); + if (s_is_last) { + const int lin0 = lane * 4; // 4 consecutive columns of one row + const int crow = lin0 >> 4; // 0..15 + const int col0 = lin0 & 15; // 0..12 (multiple of 4) + int32_t sums[4] = {0, 0, 0, 0}; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + const int4 plane4 = *reinterpret_cast( + partials + (static_cast(s) * n_tiles + tile) * + (kDummaM * kDummaN) + + lin0); + sums[0] += plane4.x; + sums[1] += plane4.y; + sums[2] += plane4.z; + sums[3] += plane4.w; + } + const int out_col0 = tile * kDummaN + col0; + uint64_t packed = 0; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float scaled = static_cast(sums[e]) * x_scale[crow] * + weight_scale[out_col0 + e]; + packed |= static_cast(float_to_bf16_bits(scaled)) << (16 * e); + } + *reinterpret_cast(out + static_cast(crow) * n + + out_col0) = packed; + } +} + +// --------------------------------------------------------------------------- +// Iteration 3 split-K combine + scale for the exact tp8_wq_b_m16 shape: +// reduce the kSplitK int32 partial planes ascending (each plane is +// [n_tiles][16][16], tile-major with row stride 16), apply +// (acc * x_scale[row]) * weight_scale[col] in the same left-associative +// float order as the reference, and store round-to-nearest-even bf16. +// One thread per output element, grid = ceil(16*n / 256) = 256 blocks. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kCombineThreads) +void w8a8_dumma_m16n16k32_combine_kernel( + const int32_t* __restrict__ partials, + const float* __restrict__ x_scale, // [16] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [16, N] bf16 + const int n) { + const int idx = static_cast(blockIdx.x) * kCombineThreads + + static_cast(threadIdx.x); + const int total = kDummaM * n; + if (idx >= total) { + return; + } + + const int row = idx / n; + const int col = idx - row * n; + const int tile = col / kDummaN; + const int tile_base = tile * (kDummaM * kDummaN) + + row * kDummaN + (col & (kDummaN - 1)); + const int n_tiles = n / kDummaN; + int32_t sum = 0; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + sum += partials[static_cast(s) * n_tiles * + (kDummaM * kDummaN) + + tile_base]; + } + + // Same left-to-right fp32 evaluation order as the reference. + const float scaled = + static_cast(sum) * x_scale[row] * weight_scale[col]; + out[idx] = hip_bfloat16(scaled); +} + +// One thread computes exactly one out[row, col] element. Adjacent lanes map +// to adjacent addresses in the fastest-changing N dimension. +// Iteration 8: b_transposed selects the exact-shape packed [N][K] layout +// (P[n*K + k] = W[k*N + n], produced by launch_pack_w8a8_weight), so the +// logical column col starts at P[col*K] with unit stride; the identity pack +// keeps the legacy row-major [K][N] layout with stride n. The branch is +// hoisted out of the K loop by the compiler. +__global__ __launch_bounds__(kScalarBlock) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ x_q, // [M, K] row-major + const int8_t* __restrict__ weight, // [K, N] row-major (identity pack) + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [M, N] bf16 + const int m, + const int n, + const int k, + const int b_transposed) { + const int64_t tid = + static_cast(blockIdx.x) * kScalarBlock + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (tid >= total) { + return; + } + const int row = static_cast(tid / n); + const int col = static_cast(tid - static_cast(row) * n); + const int8_t* a_row = x_q + static_cast(row) * k; + const int8_t* b_col = + weight + (b_transposed ? static_cast(col) * k : col); + const int64_t b_stride = b_transposed ? 1 : static_cast(n); + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * b_stride]); + } + // Same left-to-right fp32 evaluation order as the torch reference. + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[tid] = hip_bfloat16(scaled); +} + +// Generic element-wise device-to-device copy used by the identity pack. +template +__global__ __launch_bounds__(kPackBlock) void w8a8_identity_copy_kernel( + const T* __restrict__ src, + T* __restrict__ dst, + const int64_t count) { + const int64_t tid = + static_cast(blockIdx.x) * kPackBlock + threadIdx.x; + if (tid < count) { + dst[tid] = src[tid]; + } +} + +// Iteration 8 exact-shape pack: writes the [N][K] transpose +// dst[n*K + k] = src[k*N + n] (used for K=1024, N=4096), so the staged +// partial kernel reads each block's 16 B columns as 16 contiguous 512-B +// segments and every B fragment as one 8-byte ds_read_b64. One thread per +// 16-byte k-chunk of one column; runs once per weight outside the timed +// region, so the strided byte loads are acceptable. All other shapes keep +// the identity copy kernel above. +__global__ __launch_bounds__(kPackBlock) void w8a8_pack_transpose_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + const int k, + const int n) { + const int kchunks = (k + 15) >> 4; + const int64_t total = static_cast(n) * kchunks; + const int64_t stride = + static_cast(gridDim.x) * blockDim.x; + for (int64_t t = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + t < total; t += stride) { + const int col = static_cast(t / kchunks); + const int k0 = static_cast(t % kchunks) * 16; + alignas(16) int8_t chunk[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) { + chunk[j] = src[(static_cast(k0) + j) * n + col]; + } + *reinterpret_cast(dst + static_cast(col) * k + k0) = + *reinterpret_cast(chunk); + } +} + +inline int blocks_for(const int64_t count, const int block) { + const int64_t b = (count + block - 1) / block; + return b < 1 ? 1 : static_cast(b); +} + +} // namespace + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // Exact-shape fast path: tp8_wq_b_m16 (M=16, N=4096, K=1024). Iteration 9: + // the split-grid STAGED partial kernel at the dispatched split S=4 — 1024 + // blocks x 64 threads, each block stages its own 256-K A and B chunks into + // LDS once (8,704 B/block -> 7 blocks/CU co-resident, all 4 SIMDs busy) + // and then runs a zero-global two-ds_read_b64 + v_mmac K loop — writes the + // 4 int32 partial planes into the caller workspace (4 * 16 * 4096 * 4 = + // 1 MiB; the API allocates 16 planes / 4 MiB for this shape, so this + // always qualifies). Iteration 11: the combine is FUSED into the partial + // kernel tail (per-tile arrival counters in the last 1 KiB of the caller + // workspace; the last arrival per tile sums the planes and stores the + // scaled bf16 output), so the timed Graph contains exactly ONE kernel + // launch per replay on the same stream, every partial element overwritten + // each launch (no clear), and the counters are zeroed once per workspace + // by an async hipMemsetAsync before the first launch (never captured). + // B is read from the exact-shape [N][K] pack P[n*K+k] = W[k*N+n] produced + // once by launch_pack_w8a8_weight outside the timed region/Graph. + if (m == kDummaM && n == kExactN && k == kExactK) { + constexpr int kSplitK = kSplitKMeasured; + constexpr int64_t kPartialsBytes = + static_cast(kSplitK) * kDummaM * kExactN * sizeof(int32_t); + // Iteration 11: the fused path needs the 4 int32 partial planes PLUS the + // per-tile arrival counters (last kCountersBytes of the caller + // workspace); the API's 16-plane / 4 MiB allocation always qualifies. + if (workspace != nullptr && + workspace_bytes >= kPartialsBytes + kCountersBytes) { + int32_t* counters = reinterpret_cast( + static_cast(workspace) + workspace_bytes - kCountersBytes); + // One-time async zero of the arrival counters on the caller's stream, + // guarded by the workspace pointer: issued before the first launch with + // a given workspace, so steady-state replays (and the captured Graph) + // contain only the single fused kernel launch. + if (g_fused_counters_zeroed != workspace) { + hipMemsetAsync(counters, 0, kCountersBytes, stream); + g_fused_counters_zeroed = workspace; + } + // ONE kernel per replay: the staged partial kernel now fuses the + // combine into its tail (last arrival per tile sums the planes and + // stores the scaled bf16 output). Same stream, same workspace + // addresses, every partial overwritten each launch. + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_m16n16k32_splitk_staged_partial_kernel), + dim3(static_cast(kExactNTiles * kSplitK)), + dim3(kDummaWave), + 0, + stream, + a, + b, + static_cast(workspace), + x_scale, + weight_scale, + static_cast(out), + counters, + n); + return; + } + // Workspace smaller than the fused budget (never true for the API's + // 16-plane allocation): fall through to the generic scalar fallback. + } + + // Compile-only reference for the mandated CU-aligned candidate + // (kSplitK=15 -> grid 3840 = exactly 32 blocks/CU, non-power-of-two, + // non-uniform 32-aligned K chunks). Never dispatched: the workspace_bytes + // == -1 guard cannot hold for any caller-provided byte count, and this + // iteration-3 direct-load instantiation does NOT decode the exact-shape + // [N][K] pack (it predates it), so it must never be launched against the + // packed buffer. It keeps the iteration-3 kernel compiled (and correct) + // while kSplitKMeasured = 2 dispatches the staged partial kernel; the + // staged kernel supports every trusted split in {2..15} with runtime + // 32-aligned chunks, so the occupancy sweep stays a one-constant change. + if (workspace_bytes == -1 && workspace != nullptr) { + constexpr int kSplitK = kSplitKCUAligned; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16n16k32_splitk_partial_kernel), + dim3(static_cast(kExactNTiles * kSplitK)), + dim3(kDummaWave), + 0, + stream, + a, + b, + static_cast(workspace), + n); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16n16k32_combine_kernel), + dim3(static_cast( + (kDummaM * n + kCombineThreads - 1) / kCombineThreads)), + dim3(kCombineThreads), + 0, + stream, + static_cast(workspace), + x_scale, + weight_scale, + static_cast(out), + n); + } + + // Every other (m, n, k) — including the paired M=2 API shape with the same + // (N, K) and the indexer N=8192 shape — reaches this single generic scalar + // launch. Exact-shape guards sit BEFORE this fallback and never remove it. + // Iteration 8: the exact (n == 4096 && k == 1024) weight is packed as the + // [N][K] transpose by launch_pack_w8a8_weight, so every non-M16 caller of + // this pair — including the paired M=2 validation — reads it through the + // scalar fallback with the transposed flag. All other shapes use the + // identity pack and the legacy row-major read. + const int b_transposed = (n == kExactN && k == kExactK) ? 1 : 0; + const int64_t total = static_cast(m) * n; + const int blocks = blocks_for(total, kScalarBlock); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_kernel), + dim3(blocks), + dim3(kScalarBlock), + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast(out), + m, + n, + k, + b_transposed); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Iteration 8 exact-shape pack: for (k,n)==(1024,4096) write the [N][K] + // transpose P[n*K + k] = W[k*N + n], so the staged partial kernel reads + // each block's 16 B columns as 16 contiguous 512-B segments and every B + // fragment as one 8-byte ds_read_b64. Runs once per weight outside the + // timed region/Graph. All other shapes keep the identity layout + // (row-major [K][N]) consumed by the DUMMA and scalar kernels. + const int64_t w_count = static_cast(k) * n; + const int64_t s_count = n; + if (w_count > 0) { + if (k == kExactK && n == kExactN) { + const int kchunks = (k + 15) >> 4; + const int64_t chunks = static_cast(n) * kchunks; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_pack_transpose_i8_kernel), + dim3(blocks_for(chunks, kPackBlock)), + dim3(kPackBlock), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + dim3(blocks_for(w_count, kPackBlock)), + dim3(kPackBlock), + 0, + stream, + raw_weight, + packed_weight, + w_count); + } + } + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + dim3(blocks_for(s_count, kPackBlock)), + dim3(kPackBlock), + 0, + stream, + weight_scale, + packed_weight_scale, + s_count); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wqkv_a.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wqkv_a.hip new file mode 100644 index 00000000..bfa665ac --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wqkv_a.hip @@ -0,0 +1,1966 @@ +// @@variant shape=tp8_wqkv_a_m16 commit=b67e227c20b3349d46020f9f52cd2484b2142ff3 added=2026-08-26 +// median_us=17.5 p90_us=17.52 +// source=hy3-dsh-tp8-m16-1-adf021ab +// MetaInfer W8A8 GEMM - worker_0 HIP implementation (gfx928). +// +// Task: hy3-dsh-tp8-m16-1-adf021ab, worker_0 owns shape +// tp8_wqkv_a_m16: M=16, N=1536, K=4096 (physical GPU 0) +// +// Iteration 1 established the minimal gfx928 DUMMA INT8 tile for the +// assigned shape (one 64-thread wavefront per 16x16 N tile, grid=96, +// direct row-major fragment loads, direct register epilogue). It was +// accepted at 194.6 us median (vs 588.9 us scalar bootstrap), but its +// exact-source ISA is load-wait bound: per 32-wide K step it emits ~17 +// global_load_ubyte + 6 global_load_dword with per-load vmcnt waits in +// front of a single v_mmac_i32_16x16x32_i8. +// +// Iteration 2 (accepted, 142.46 us median / 143.91 us p90) was the first +// mandated ARCHITECTURE round: WAVES=2 (blockDim=128), in-block split-K=2 +// along K (wave 0: [0,2048), wave 1: [2048,4096)), one 16-wide N tile per +// block, grid = N/16 = 96 blocks (0.8 of 120 CUs), LDS int32 plane combine +// (2 x 16 x 20), one END-of-K barrier, direct register epilogue. Its +// exact-source ISA still shows the iteration-1 per-load vmcnt cascade +// (~17 global_load_ubyte + 6 global_load_dword per K step) in front of one +// v_mmac_i32_16x16x32_i8 per step; the second wavefront per CU recovered +// 194.6 -> 142.5 us but the kernel remains load-latency bound. +// +// Iteration 3 is the second mandated ARCHITECTURE round: the unsplit grid +// (96 blocks) has fewer than two blocks per device CU and K=4096 >= 1024, +// so the K loop is now partitioned across CROSS-BLOCK splits: +// * partial kernel: one 64-thread wavefront per block, ZERO barriers, +// ZERO LDS. grid = (N/16) * split_k blocks; blockIdx.x encodes +// split-major / ntile-minor (ntile = x % ntiles, split = x / ntiles). +// split_k is runtime-selectable in [1, min(16, K/32)] (default 2; env +// override METAINFER_W8A8_SPLIT_K selects any trusted candidate). K is +// tiled into 32-aligned step-based slices +// [(s*steps)/split_k, ((s+1)*steps)/split_k) that cover [0, K) exactly +// once ascending (non-uniform for non-power-of-two split_k, e.g. +// split_k=5 -> 25/26/25/26/26 steps), so every split boundary stays +// aligned to the DUMMA k32 unit. Each block keeps the unchanged +// k-ascending du_load_matrix_sync + du_mma_sync loop over its slice +// (per-slice int32 partials bit-identical to iterations 1-2) and +// publishes the int32 tile with du_store_matrix_sync (mem_row_major, +// ldm=16) into caller workspace plane (split, ntile) at +// partials[(split*ntiles + ntile) * 16 * 16]. +// * combine+scale kernel: grid = N/16 blocks x 256 threads (4 +// wavefronts), 1 element per thread; sums the split_k planes in +// ascending split order (bit-identical int32 to the serial k-ascending +// accumulation - integer addition is exact and +// |full dot| <= 4096*127*127 ~= 6.6e7 << 2^31), applies +// x_scale[m] * weight_scale[n], stores bf16. Both kernels launch on +// the caller stream inside the timed Graph, so the combine cost is +// included in the measured median/P90. +// * split_k=2 -> grid 192 blocks (1.6 blocks/CU); the CU-aligned +// non-power-of-two candidate split_k=5 -> grid 480 = 4 x 120 CUs +// exactly is the same code path with a different split count (trusted +// occupancy-probe candidates {2,3,5,6,8,10,16} are all env-selectable). +// * expected mechanism: cross-block splits multiply the resident +// one-wave blocks per CU (each block is an independent load/MMAC +// stream that covers sibling blocks' vmcnt stalls), the partial +// kernel has no barrier or LDS combine on its critical path, and the +// combine kernel is a tiny 96-block latency-tolerant tail. Exact bf16 +// equality is expected (identical int32 accumulation order: +// k-ascending per split, ascending split sum). +// +// Iteration 4 is the mandated A-only-staging pipeline round. The +// iteration-3 partial kernel's exact-source PMC (192 blocks, split_k=2) +// shows 16 vmem-read instructions and ~63 issued VALU per 32-wide K step +// with per-load vmcnt waits in front of the single v_mmac - the same +// byte-load fragment reassembly poison as iterations 1-2 (A and B +// fragments each expand to ~8-12 scalar global loads). The chosen change +// is A-ONLY staging (B stays direct; B vectorization/packing is the later +// round): each one-wave block stages its own A slice [16, slice_k] once at +// block start into LDS with vectorized 16-byte int4 global loads and +// ds_write_b128 stores (row stride = slice_k + 16 bytes: 16-B-aligned, +// non-power-of-two bank skew, down_proj-validated rule), and the K loop +// then reads the A fragment from LDS via du_load_matrix_sync(smem, ldm = +// local_stride) - A's ~8-12 global loads + vmcnt waits per step leave the +// loop, replaced by low-latency ds_read. Default split_k is raised 2 -> 5: +// grid = 96 x 5 = 480 one-wave zero-barrier blocks = exactly 4 x 120 CUs, +// each block's stage is 16 x (832 + 16) = 13,568 B dynamic LDS so four +// blocks fit per CU (54,272 B <= 64 KiB), directly answering the +// control-plane warning that 192 blocks is below the two-blocks-per-CU +// latency-hiding target (480 = 4 blocks/CU). A bytes (64 KiB total, +// L2-hot) are still re-read once per N-tile block (96x per split pass) but +// as one vectorized burst per block instead of per-step scalar loads; B +// bytes (6.29 MiB cold) are still read exactly once per owning +// (split, ntile) block, unchanged. Exact bf16 equality is preserved: the +// staged data is the same A bytes in the same row-major order and the +// int32 accumulation order is unchanged (k-ascending per split, ascending +// split sum). +// +// Iteration 5 is the mandated HIP-only packed-weight/staging comparison +// round (chosen design: ONE PACKED LAYOUT - B fragment slots; the A-only +// LDS staging of iteration 4 stays as-is; B-only staging was not chosen +// because it would add a second ~13.3 KiB LDS slice per block and drop +// co-residency from 4 to 2 blocks/CU). The identity pack_weight bootstrap +// becomes a real one-time device permutation for the exact +// (K, N) == (4096, 1536) weight: +// packed[n_tile][k_step][lane][8]: slot +// ((ntile * (k/32) + step) * 64 + lane) * 8 + i +// holds logical weight[step*32 + (lane>>4)*8 + i][ntile*16 + (lane&15)]. +// Byte count is unchanged (K*N = 6,291,456 B), so the packed buffer keeps +// the caller's numel and graph-stable address; the pack runs outside the +// timed region and Graph capture (prepare_weight path). +// The pack makes each lane's matrix_b m16n16k32 row_major fragment for one +// k32 step ONE contiguous aligned 8-byte (dwordx2) global load in the exact +// du_mma.hpp byte order (x[i] = p[(col + i)*ldm + row], row = lane&15, +// col = (lane>>4) << 3; validated o_proj lineage pattern): B's per-step +// scalar byte loads + fragment reassembly + per-load vmcnt cascade (8 vmem +// instructions and ~50 VALU per step in the iteration-4 exact ISA) leave +// the K loop, which becomes ds_read (A, from LDS) -> 1 dwordx2 (B, packed) +// -> v_mmac. The DUMMA kernels and the generic scalar fallback decode the +// SAME (tile, k_step, lane, byte) -> logical (k, n) mapping (coherent +// repack, o_proj-validated rule), so the paired M=2 fallback validation +// stays exact. Geometry is unchanged from iteration 4: split_k=5, grid 480 +// = 4 x 120 CUs exactly, one-wave zero-in-loop-barrier blocks, 13,568 B +// dynamic LDS per block, separate combine kernel in the timed Graph. Exact +// bf16 equality is preserved: every fragment feeds the same bytes to the +// same v_mmac in the same order, and the int32 accumulation order is +// unchanged (k-ascending per split, ascending split sum). +// +// Iteration 6 is the mandated pipeline/staging-family round: the +// iteration-5 kernel's exact-source ISA (primary kernel +// w8a8_dumma_m16n16k32_sk_astage_packedb_partial_kernel, 480 blocks, 16 +// VGPR/32 SGPR) shows the steady loop as +// global_load_dwordx2 (B, packed) -> ds_read2_b32 (A, from LDS) -> +// s_waitcnt lgkmcnt(0) -> s_waitcnt vmcnt(0) -> v_mmac - one exposed +// vmcnt(0) wait in front of every v_mmac for the cold once-read B stream +// (B = 6.29 MB, read exactly once; L2 hit rate 53.3%, l2_misses 107,607), +// so the per-step critical path is B's global latency, not issue rate +// (18,528 vmem reads = exactly 1 B dwordx2 per k32 step + the one-time A +// stage; 12,288 loop steps). The chosen staging family is A+B LDS with a +// stage-K=64 DOUBLE BUFFER (the validated winner mechanism of the TP4 +// reference for this exact (M=16, K=4096, N=1536) shape: 17.344 us Graph): +// * the K slice is processed in 2-step (K=64) chunks; dynamic LDS holds +// two ping-pong buffers of 2,560 B each (A 16x64 with a padded 96-B row +// stride - 8-B A-fragment reads at the 2-way LDS floor vs the ~4.4 +// conflicts of the stride-848 layout - plus B 1 KiB of packed fragment +// slots), 5,120 B per block total (constant, independent of split_k), +// so the default split_k=5 grid 480 = 4 x 120 CUs still fits 4 +// blocks/CU (20,480 B <= 64 KiB); +// * per chunk each lane issues 3 cooperative global loads (1 x 16-B A +// dwordx4 + 2 x 8-B B dwordx2) into registers BEFORE the current +// chunk's 1-2 v_mmac_i32_16x16x32_i8 and commits them to the alternate +// LDS buffer AFTER the MMAC burst, so the vmcnt wait lands at the +// ds_write and the cold B latency overlaps the current chunk's compute +// (the harness ISA extraction must show next-chunk global loads before +// the current v_mmac and the wait near the alternate-buffer ds_write; +// the software pipeline is NOT claimed from HIP source alone); +// * the steady K loop then has ZERO global loads: per k32 step ds_read +// (A) -> ds_read (B) -> v_mmac. The zero-LDS direct-A packed-B kernel +// (split_k=1 64-KiB overflow fallback) and the A-only astage packed-B +// kernel are preserved as symbols but no longer dispatched. Exact bf16 +// equality is preserved: staged bytes are the same data fed to the same +// v_mmac in the same order; int32 accumulation order unchanged +// (k-ascending per split, ascending split sum). +// * iteration-6 repair 1: the per-chunk step count was clamped to +// kStageSteps. The original `nsteps = slice_steps - s0` is the remaining +// count and is 3-4 for the second-to-last chunk (25/26 - 2*11 at +// split_k=5), so the j-loop read A/B beyond the staged 64-column chunk +// (stale LDS) and polluted every accumulator. Fix: nsteps = +// min(kStageSteps, slice_steps - s0) -> 1 or 2, matching the staged +// chunk width; int32 accumulation order unchanged. +// * iteration-6 repair 2: the B slot base must index the slice's first +// GLOBAL k32 step (k_lo/kDummaTileK) because the packed +// [n_tile][k_step][lane][8] layout is indexed by global step (pack +// kernel slot = ntile*k_steps + step, step in [0, K/32)). The original +// b_slot_base pointed at global step 0 and added only the slice-relative +// chunk/step offsets, so every split with k_lo > 0 (splits 1-4 of the +// default split_k=5) fed B bytes of global steps [0, slice_steps) into +// the accumulator while A was staged correctly - the dominant +// near-total mismatch (24,538/24,576) that survived repair 1. Fix: +// b_slot_base += (k_lo / kDummaTileK) * kPackedTileBytes; A staging +// (a_base = k_lo) and the int32 accumulation order are unchanged. +// +// Iteration 7 is the mandated pipeline round: COMPARE SINGLE buffering with +// DOUBLE buffering (all three gates hold for this shape: K=4096 >= 1024, +// L2 hit rate 54.2% < 70%, and the doubled LDS budget 2 x 2,560 = 5,120 B +// < 48 KiB). Iteration 6 already measured and accepted the DOUBLE-buffered +// stage-K=64 arm (18.551 us median / 18.583 us p90; 480 one-wave +// zero-barrier blocks = 4 x 120 CUs; 5,120 B dynamic LDS/block; the +// per-chunk register prefetch issues chunk c+1's 3 cooperative global loads +// before chunk c's MMAC burst and the vmcnt wait lands at the post-MMAC +// alternate-buffer ds_write). This round measures the SINGLE-buffered arm: +// * the two ping-pong 2,560-B buffers collapse into ONE 2,560-B buffer - +// the chunk c+1 prefetch still goes through registers and commits AFTER +// the current chunk's MMAC burst, but now to the same LDS addresses the +// current chunk was just read from (write-after-read ordering is program +// order in the single wave; the compiler keeps the same +// global_load_dwordx4/dwordx2-before-v_mmac and vmcnt-wait-at-ds_write +// placement, with the ds_write at a FIXED base instead of the +// (c&1)*0xa00 toggle); +// * dynamic LDS drops 5,120 -> 2,560 B per block (per CU 20,480 -> +// 10,240 B; occupancy is unchanged because the grid is 480 blocks = +// 4 blocks/CU by design and neither footprint binds the 64 KiB LDS); +// * barriers per K step: 0 in BOTH variants (one 64-lane wave per block, +// no cross-wave LDS producer/consumer - the pipeline synchronization is +// s_waitcnt vmcnt/lgkmcnt only; the TP4 4-wave reference needs ~2 +// barriers per stage-K=64 chunk, i.e. ~1 per K step, for cross-wave LDS +// visibility, which does not apply to the one-wave blocks here); +// * launch geometry (split_k=5, grid 480, one wave/block, block-N=16), +// K slicing, packed-B fragment-slot layout, int32 partial publish, +// separate combine kernel in the timed Graph, and the exact bf16 +// accumulation order are all unchanged. The iteration-6 double-buffered +// kernel stays in the source as a preserved symbol but is no longer +// dispatched. +// +// Iteration 14 (shadow, not accepted) fused the iteration-3 combine/scale +// kernel into the iteration-7 single-buffered partial kernel as a +// qkv-proven last-arrival atomic tail (monotonic counters in the last +// counter_bytes of the caller workspace, one launch per Graph replay). It +// measured 18.382 us median / 18.409 us p90 - only +0.55% over the 18.484 +// us iteration-7 shadow - proving the ~2 us partial-to-Graph gap was the +// FUSED TAIL's own fence/atomic/tail cost, not the combine kernel's +// launch+L2 round trip (consistent with the TP4 lineage where the atomic +// last-arriver fused finalize on the 4-wave design measured +1.52 us). +// +// Iteration 15 (this round, HIP-only; plateau not proven - the last three +// valid improvements were -50.78% / -6.65% / +0.92%) adopts the +// TP4-VALIDATED 4-wave design for this exact (M=16, K=4096, N=1536) local +// shape, measured at 17.344 us median / 17.616 us p90 on the same gfx928: +// * block-N=64 (4 x wave64, 256 threads), one 16x16 quadrant per wave, +// all four waves SHARE one staged A chunk (A logical rereads drop +// 6.4 MB -> 1.57 MB; each block stages A once per K=64 chunk instead of +// once per 16-wide tile); +// * stage-K=64 DOUBLE buffer, 10,240 B/block; stage i+1's global int4 +// loads issue before stage i's 2 x v_mmac per wave and commit to the +// alternate buffer with the vmcnt wait at the ds_write; two +// __syncthreads() per stage (pure-HIP form of the TP4 raw-first- +// barrier sync mode - NO raw asm this round); +// * B is staged from THIS session's packed [n_tile][k_step][lane][8] +// fragment-slot layout (one 16-B int4 per lane = two adjacent 8-B +// slots, one 16-B ds_write_b128 into an LDS mirror of the slot layout) +// instead of the TP4 raw [K,N] row loads - same staged bytes, same +// fragment reads (one ds_read_b64 per lane per k32 step); +// * non-uniform stage-based K slicing, default split_k=10 -> grid = +// 24*10 = 240 blocks = exactly 2 x 120 CUs; +// * the separate 96-block x 256-thread combine kernel returns to the +// timed Graph (the TP4 lineage measured the fused finalize at +1.52 us +// on this design, so the iteration-14 fused 1-wave kernel is now a +// preserved symbol); +// * exact bf16: same staged bytes, same k-ascending per-split / +// ascending-split-sum int32 order. +// +// Iteration 16 (this round, HIP-only repair of the iteration-15 candidate): +// the iteration-15 4-wave kernel measured 60.158 us (single sample) with +// 24,555/24,576 mismatches. The SOLE correctness defect was the B-staging +// global address in w8a8_dumma_m16n16k32_4wave_sk_astage_packedb_dbuf_ +// partial_kernel: each stage chunk was read as one contiguous 4,096-B +// window of the packed [n_tile][k_step][lane][8] layout, but that layout is +// indexed by the GLOBAL step (each ntile's 128 steps form a contiguous +// 64 KiB run), so a 4-ntile x 2-step chunk is NOT contiguous - only wave 0 +// of chunk 0 read correct bytes. The repair changes ONLY the B global +// offsets: each thread's 16-B slot pair is addressed with the exact +// per-(ntile, step) packed offset ((4*tile64 + w)*k_steps + k_lo/32 + +// 2c + s)*512 + (tid&31)*16, and the chunk stride becomes +// kStageSteps*kPackedTileBytes = 1,024 B. The fast mapping is preserved +// byte-for-byte: same 4-wave block-N=64 stage-K=64 double-buffer pipeline, +// same one 16-B global load + one 16-B ds_write_b128 per thread per chunk +// into the unchanged slot-mirror LDS layout, same two __syncthreads() per +// stage, same fragment reads (one ds_read_b64 per lane per k32 step), same +// plane publish, same separate 96-block x 256-thread combine kernel in the +// timed Graph, same shape guards and generic fallbacks. Exact bf16: +// identical staged bytes fed to the same v_mmac in the same order; int32 +// accumulation order unchanged (k-ascending per split, ascending split +// sum). Falsifiable gate vs the 18.382 us iteration-14 shadow: if the +// TP4-validated occupancy transfers once correctness is restored, the +// median lands ~17.3-17.7 us; if the 2-blocks/CU co-residency or the +// 2-barrier-per-stage sync costs more on this packed-B path, it lands at or +// above 18.382. +// +// Dispatch policy: +// * The DUMMA fast path is reachable only for the exact assigned shape; +// every other (m, n, k) falls back to the generic scalar kernel. +// * launch_pack_w8a8_weight performs a real fragment-slot pack for the +// exact (K, N) == (4096, 1536) weight and stays an identity +// device-to-device copy for every other (K, N). +// * The timed operator uses only the caller-provided out/workspace (the +// split-K partial + combine path stores int32 partials in the caller +// workspace), launches only on the caller's hipStream_t, and performs +// no allocation, compilation, autotuning, packing, or host/device +// synchronization. If the caller provides fewer than one int32 plane +// (16*16*4 bytes), the previous workspace-free in-block split-K=2 +// kernel (iteration 2) keeps the fast path correct. + +// Known-good include order for this DTK: HIP runtime first, then +// hip_bfloat16.h, then du_mma.h (du_mma.h is not self-contained if it is +// included before the HIP runtime headers). +#include +#include +#include + +#include +#include + +namespace { + +// One wavefront on gfx928 is 64 lanes; blockDim must be a multiple of 64. +constexpr int kScalarBlockThreads = 128; +constexpr int kCopyBlockThreads = 256; + +// Minimal DUMMA INT8 tile: m16n16k32 with int32 accumulation. +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; +constexpr int kDummaWaveSize = 64; + +// Exact assigned shape (M=16, N=1536, K=4096) used by the packed-B +// interpretation and dispatch guards. +constexpr int kTargetK = 4096; +constexpr int kTargetN = 1536; + +// Iteration 5 packed-B fragment-slot layout (exact shape only): +// packed[n_tile][k_step][lane][8]: slot +// ((ntile * (k/32) + step) * 64 + lane) * 8 + i +// holds logical weight[step*32 + (lane>>4)*8 + i][ntile*16 + (lane&15)] - +// the 8 bytes of one lane's matrix_b m16n16k32 row_major fragment in the +// exact du_mma.hpp byte order (x[i] = p[(col + i) * ldm + row], row = +// lane&15, col = (lane>>4) << 3). Byte count is unchanged (K*N), so the +// packed buffer keeps the caller's numel and graph-stable address. +constexpr int kPackedSlotBytes = 8; // bytes of one lane's B fragment +constexpr int kPackedTileBytes = kDummaWaveSize * kPackedSlotBytes; // 512 + +// Iteration 6 stage-K=64 double-buffer constants (exact shape only). One +// stage chunk covers 2 k32 steps (K=64): A 16 rows x 64 B row-major with a +// padded 96-B row stride (16-B-aligned non-power-of-two bank skew; 96/4 = +// 24 dwords, 24*row mod 64 spreads the 16 rows onto 16 bank phases with the +// 8-B A-fragment reads at the 2-way LDS floor, vs the ~4.4 conflicts of the +// iteration-4/5 stride-848 layout) and B 1 KiB of packed fragment slots +// (8 B per lane per step). Two ping-pong buffers = 5,120 B per block, so +// four blocks per CU (20,480 B <= 64 KiB) at the default split_k=5 grid. +constexpr int kStageSteps = 2; // k32 steps per stage chunk +constexpr int kStageK = kDummaTileK * kStageSteps; // 64 K-cols per chunk +constexpr int kStageAStride = 96; // padded A row stride (bytes) +constexpr int kStageABufBytes = + kDummaTileM * kStageAStride; // 1536 +constexpr int kStageBBufBytes = + kDummaWaveSize * kPackedSlotBytes * kStageSteps; // 1024 +constexpr int kStageBufBytes = kStageABufBytes + kStageBBufBytes; // 2560 +// Iteration 6 double buffer: two ping-pong buffers = 5,120 B per block. +constexpr int kStageTotalBytes = 2 * kStageBufBytes; // 5120 +// Iteration 7 single-buffer arm: ONE 2,560-B buffer per block (the sbuf +// kernel below passes kStageBufBytes as its dynamic LDS size). + +// Iteration 15 4-wave block-N=64 stage-K=64 DOUBLE-buffer constants (the +// TP4-validated design for this EXACT local shape (M=16, K=4096, N=1536): +// 17.344 us median / 17.616 us p90 on the same gfx928 at split_k=10, grid +// 240 = exactly 2 x 120 CUs, 4 wave64 per block, one 16x16 quadrant per +// wave sharing one staged A chunk; see +// references/w8a8_gemm_variants.hip and the split10 reference). Per stage +// chunk (K=64): A raw [16][64] staged by wave 0 (1,024 B, stride 64) + B +// staged as a 4,096-B mirror of the packed fragment-slot layout (4 ntiles +// x 2 steps x 512 B); two ping-pong buffers = 10,240 B per block (2,560 B +// per wave-pair-of-CUs at 2 blocks/CU = 20,480 B/CU <= 64 KiB). +constexpr int k4wBlockNTiles = 4; // block-N = 64 +constexpr int k4wBlockThreads = k4wBlockNTiles * kDummaWaveSize; // 256 +constexpr int k4wStageABytes = kDummaTileM * kStageK; // 1024 +constexpr int k4wStageBBytes = + k4wBlockNTiles * kStageSteps * kPackedTileBytes; // 4096 +constexpr int k4wStageBufBytes = k4wStageABytes + k4wStageBBytes; // 5120 +constexpr int k4wStageTotalBytes = 2 * k4wStageBufBytes; // 10240 + +// --------------------------------------------------------------------------- +// Iteration 5 packed-B helpers (exact shape (K, N) == (4096, 1536) only). +// --------------------------------------------------------------------------- + +// Elementwise decode of the packed [n_tile][k_step][lane][8] fragment-slot +// layout: returns logical weight[kk, n_col] with the same (tile, k_step, +// lane, byte) -> logical (k, n) mapping as the pack kernel and the DUMMA +// kernels. Used by the generic scalar fallback (M=2 paired validation). +__device__ __forceinline__ int8_t +packed_b_element(const int8_t* __restrict__ packed, int n_col, int kk) { + const int nt = n_col >> 4; + const int nn = n_col & 15; + const int step = kk >> 5; + const int kk8 = kk & 31; + const int lane = (kk8 >> 3) * 16 + nn; + const int i = kk8 & 7; + return packed[(static_cast(nt) * (kTargetK / kDummaTileK) + step) * + kPackedTileBytes + + lane * kPackedSlotBytes + i]; +} + +// One aligned 8-byte (dwordx2) global load per lane per k32 step from the +// packed fragment-slot layout; the loaded uint64 feeds b_frag.x directly in +// the exact du_mma.hpp matrix_b row_major byte order (validated o_proj +// lineage pattern: __builtin_memcpy of the 8 fragment bytes, same DTK). +__device__ __forceinline__ void +load_packed_b_fragment( + du::dumma::DUFragment& + b_frag, + const int8_t* __restrict__ packed_b, int ntile, int k_steps, int step, + int lane) { + const int64_t off = + (static_cast(ntile) * k_steps + step) * kPackedTileBytes + + lane * kPackedSlotBytes; + const uint64_t bv = *reinterpret_cast(packed_b + off); + __builtin_memcpy(b_frag.x, &bv, sizeof(bv)); +} + +// Generic scalar INT8 W8A8 GEMM. +// +// out[m, n] = bf16(int32_dot(a[m, :], b[:, n]) * x_scale[m] * weight_scale[n]) +// +// One thread owns one output element; adjacent lanes map to adjacent n +// columns (fastest-changing N dimension), so B reads and bf16 stores are +// contiguous across lanes within a row. +__global__ __launch_bounds__(kScalarBlockThreads) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, // [M, K] row-major int8 + const int8_t* __restrict__ b, // [K, N] row-major int8 + const float* __restrict__ x_scale, // [M, 1] fp32 + const float* __restrict__ weight_scale,// [N, 1] fp32 + hip_bfloat16* __restrict__ out, // [M, N] bf16 + int m, + int n, + int k) { + const int64_t tid = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (tid >= total) { + return; + } + const int row = static_cast(tid / n); + const int col = static_cast(tid % n); + + const int8_t* a_row = a + static_cast(row) * k; + // Iteration 5: for the exact (K=4096, N=1536) shape the weight buffer is + // the packed [n_tile][k_step][lane][8] fragment-slot layout (see + // w8a8_pack_b_fragslot_kernel), so the generic fallback decodes it + // elementwise (keeps the paired M=2 validation exact with the new pack). + const bool packed_b = (k == kTargetK && n == kTargetN); + + // Exact integer dot product, k-ascending (bit-identical to the reference + // int64 accumulation order; the int32 accumulator cannot overflow for any + // supported K). + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + const int8_t b_val = + packed_b ? packed_b_element(b, col, kk) + : b[static_cast(kk) * n + col]; + acc += static_cast(a_row[kk]) * static_cast(b_val); + } + + // Scale exactly like the reference: (dot * x_scale[m]) * weight_scale[n]. + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[tid] = hip_bfloat16(scaled); +} + +// gfx928 DUMMA INT8 kernel for the assigned shape tp8_wqkv_a_m16 +// (M=16, N=1536, K=4096) - iteration 2 launch geometry: +// * blockDim = 128 = 2 wavefronts; WAVES_PER_BLOCK = 2 with in-block +// split-K = 2: each wavefront accumulates one contiguous K half +// (wave 0: [0, 2048), wave 1: [2048, 4096)) k-ascending in 32-wide +// DUMMA steps, preserving the exact int32 accumulation order per split; +// * one 16-wide N tile per block (N-TILES = 1), grid = N/16 = 96 blocks +// (the maximum grid reachable on the N-tile axis; reported coverage is +// 96 of 120 CUs); +// * partial int32 planes combined in LDS (2 x 16 x 16 int32 = 2 KiB) in +// ascending split order p0 + p1 after one END-of-K barrier - bit- +// identical to the serial k-ascending accumulation (integer addition +// is associative), matching the validated "k-ascending per split, +// ascending split sum" pattern; +// * direct register epilogue on wave 0 with the verified gfx928 +// accumulator ownership (row = lane&15, col_mod4 = lane>>4, +// acc_frag.x[i] owns column col_mod4 + 4*i): apply x_scale[m] * +// weight_scale[n] in float and store bf16; +// * no A/B LDS staging, no vectorized fragment loads, no inline asm - +// the load strategy is unchanged from iteration 1 so this round +// isolates the launch-geometry axis (expected: the second wavefront +// per CU overlaps the per-load vmcnt waits that dominate the +// iteration-1 ISA). +__global__ __launch_bounds__(2 * kDummaWaveSize) void +w8a8_dumma_m16n16k32_2wave_sk2_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ b, // [K, N] row-major int8 + const float* __restrict__ x_scale, // [16, 1] fp32 + const float* __restrict__ weight_scale, // [N, 1] fp32 + hip_bfloat16* __restrict__ out, // [16, N] bf16 + int n, + int k) { + // One int32 partial plane per wave; combined in ascending split order. + // Row stride is padded 16 -> 20 words (16-B-aligned, non-power-of-two): + // with the natural stride the plane word index (row*16 + col) mod 64 + // aliases rows differing by 4 onto the same LDS bank phase (4-way + // conflict on every plane access); stride 20 spreads the 16 rows of a + // column group onto 16 distinct banks (20*row mod 64 is a permutation + // for row in [0,16)). + __shared__ int32_t s_part[2][kDummaTileM][kDummaTileN + 4]; + + const int lane = static_cast(threadIdx.x) & (kDummaWaveSize - 1); + const int wave = static_cast(threadIdx.x) >> 6; // 0 or 1 + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Each wave accumulates one contiguous K half, k-ascending in 32-wide + // DUMMA steps. The split boundary (k/2 = 2048) is a multiple of 32, so + // every step stays aligned to the DUMMA/staging unit. + const int k_half = k / 2; + const int k_lo = wave * k_half; + for (int k0 = k_lo; k0 < k_lo + k_half; k0 += kDummaTileK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); + du::dumma::du_load_matrix_sync(b_frag, b + k0 * n + n0, n); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Publish the int32 partial (k-ascending within the split). + const int row = lane & 15; + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s_part[wave][row][col_mod4 + 4 * i] = acc_frag.x[i]; + } + // One END-of-K barrier: both partial planes are complete before the sum. + __syncthreads(); + + // Wave 0 finalizes: ascending split sum p0 + p1 (bit-identical int32 + // order), scale, bf16 store. + if (wave == 0) { + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = col_mod4 + 4 * i; + const int32_t total = s_part[0][row][col] + s_part[1][row][col]; + const float scaled = static_cast(total) * xs * + weight_scale[n0 + col]; + out[row * n + n0 + col] = hip_bfloat16(scaled); + } + } +} + +// Iteration 5: workspace-free in-block split-K=2 kernel with the packed-B +// fragment-slot load (same geometry and epilogue as +// w8a8_dumma_m16n16k32_2wave_sk2_kernel; B per k32 step becomes one aligned +// dwordx2 per lane instead of the direct row-major byte loads). Keeps the +// fast path correct for callers without a usable workspace once the weight +// is packed. +__global__ __launch_bounds__(2 * kDummaWaveSize) void +w8a8_dumma_m16n16k32_2wave_sk2_packedb_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ packed_b, // [n_tile][k_step][lane][8] + const float* __restrict__ x_scale, // [16, 1] fp32 + const float* __restrict__ weight_scale, // [N, 1] fp32 + hip_bfloat16* __restrict__ out, // [16, N] bf16 + int n, + int k) { + __shared__ int32_t s_part[2][kDummaTileM][kDummaTileN + 4]; + + const int lane = static_cast(threadIdx.x) & (kDummaWaveSize - 1); + const int wave = static_cast(threadIdx.x) >> 6; // 0 or 1 + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + const int k_half = k / 2; + const int k_lo = wave * k_half; + const int k_steps = k / kDummaTileK; + const int ntile = static_cast(blockIdx.x); + for (int k0 = k_lo; k0 < k_lo + k_half; k0 += kDummaTileK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); + load_packed_b_fragment(b_frag, packed_b, ntile, k_steps, k0 / kDummaTileK, + lane); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + const int row = lane & 15; + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s_part[wave][row][col_mod4 + 4 * i] = acc_frag.x[i]; + } + __syncthreads(); + + if (wave == 0) { + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = col_mod4 + 4 * i; + const int32_t total = s_part[0][row][col] + s_part[1][row][col]; + const float scaled = static_cast(total) * xs * + weight_scale[n0 + col]; + out[row * n + n0 + col] = hip_bfloat16(scaled); + } + } +} + +// gfx928 DUMMA INT8 cross-block split-K partial kernel for the assigned +// shape tp8_wqkv_a_m16 (M=16, N=1536, K=4096) - iteration 3 architecture: +// * one 64-thread wavefront per block (blockDim = 64), ZERO barriers, +// ZERO LDS: each block is a self-contained (k-slice, 16x16 N-tile) +// partial; +// * grid = (N/16) * split_k blocks; blockIdx.x encodes split-major / +// ntile-minor (ntile = x % ntiles, split = x / ntiles), so the first +// wave of blocks streams each split's B slice once; +// * split_k is a runtime parameter in [1, min(16, K/32)]; K is tiled +// into 32-aligned step-based slices +// [(s*steps)/split_k, ((s+1)*steps)/split_k) that cover [0, K) exactly +// once ascending (non-uniform for non-power-of-two split_k), so every +// split boundary stays aligned to the DUMMA k32 unit; +// * per-split accumulation is k-ascending with the unchanged +// du_load_matrix_sync + du_mma_sync loop (per-slice int32 partials +// bit-identical to iterations 1-2; the hardware m16n16k32 i8 +// accumulation order is unchanged); +// * the int32 partial tile is published with du_store_matrix_sync +// (mem_row_major, ldm = 16) into caller workspace plane (split, ntile) +// at partials[(split*ntiles + ntile) * 16 * 16] - no LDS round trip, +// no barrier. +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16n16k32_sk_partial_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ b, // [K, N] row-major int8 + int32_t* __restrict__ partials, // [split_k * ntiles][16][16] + int n, + int k, + int split_k) { + const int ntiles = n / kDummaTileN; + const int ntile = static_cast(blockIdx.x) % ntiles; + const int split = static_cast(blockIdx.x) / ntiles; + const int n0 = ntile * kDummaTileN; + const int total_steps = k / kDummaTileK; + // Step-based 32-aligned slice boundaries (validated gate_up pattern): + // slice s owns DUMMA steps [s*total_steps/split_k, (s+1)*total_steps/ + // split_k), tiling [0, K) exactly once ascending. + const int k_lo = + (static_cast(split) * total_steps / split_k) * kDummaTileK; + const int k_hi = + (static_cast(split + 1) * total_steps / split_k) * + kDummaTileK; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + for (int k0 = k_lo; k0 < k_hi; k0 += kDummaTileK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); + du::dumma::du_load_matrix_sync( + b_frag, b + static_cast(k0) * n + n0, n); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + du::dumma::du_store_matrix_sync( + partials + + (static_cast(split) * ntiles + ntile) * + (kDummaTileM * kDummaTileN), + acc_frag, kDummaTileN, du::dumma::mem_row_major); +} + +// Iteration 5: zero-LDS split-K partial kernel with the packed-B +// fragment-slot load (same geometry, K slicing and A direct-load path as +// w8a8_dumma_m16n16k32_sk_partial_kernel; B per k32 step becomes one +// aligned dwordx2 per lane). Used when the A-slice LDS stage would exceed +// the 64 KiB limit (split_k=1) now that the weight is packed. +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16n16k32_sk_directa_packedb_partial_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ packed_b, // [n_tile][k_step][lane][8] + int32_t* __restrict__ partials, // [split_k * ntiles][16][16] + int n, + int k, + int split_k) { + const int ntiles = n / kDummaTileN; + const int ntile = static_cast(blockIdx.x) % ntiles; + const int split = static_cast(blockIdx.x) / ntiles; + const int total_steps = k / kDummaTileK; + const int k_lo = + (static_cast(split) * total_steps / split_k) * kDummaTileK; + const int k_hi = + (static_cast(split + 1) * total_steps / split_k) * + kDummaTileK; + + const int lane = static_cast(threadIdx.x) & (kDummaWaveSize - 1); + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + for (int k0 = k_lo; k0 < k_hi; k0 += kDummaTileK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); + load_packed_b_fragment(b_frag, packed_b, ntile, total_steps, + k0 / kDummaTileK, lane); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + du::dumma::du_store_matrix_sync( + partials + + (static_cast(split) * ntiles + ntile) * + (kDummaTileM * kDummaTileN), + acc_frag, kDummaTileN, du::dumma::mem_row_major); +} + +// gfx928 DUMMA INT8 cross-block split-K partial kernel for tp8_wqkv_a_m16 - +// iteration 4 architecture (A-only LDS staging). Same launch geometry as +// iteration 3 (one 64-thread wavefront per block, ZERO in-loop barriers, +// grid = (N/16) * split_k, split-major / ntile-minor blockIdx encoding, +// step-based 32-aligned K slicing) with one change on the load path: +// * each block stages its own A slice [16, slice_k] once at block start +// into dynamic LDS (extern __shared__): 16-byte int4 global loads +// (16-B-aligned: k multiple of 16, k_lo multiple of 32, col16 multiple +// of 16) + ds_write_b128 stores, row-major with padded row stride +// local_stride = slice_k + 16 bytes (16-B-aligned non-power-of-two +// bank skew so the 16 rows do not alias one LDS bank phase; validated +// down_proj rule). Rows are not contiguous in global (A row stride = +// k), so lane `chunk` owns 16 bytes of row chunk/chunks_per_row at +// column (chunk % chunks_per_row) * 16 - consecutive lanes walk +// consecutive 16-B chunks (coalesced), one wavefront iteration stages +// 64 chunks = 1 KiB; +// * the K loop reads the A fragment from LDS with the unchanged +// du_load_matrix_sync (smem pointer, ldm = local_stride in bytes), so +// A's ~8-12 scalar global loads + per-load vmcnt waits per step leave +// the loop (replaced by low-latency ds_read); B fragments are loaded +// from global exactly as iteration 3 (unchanged byte-load path, B is +// the later staging/packing round); +// * dynamic LDS = 16 * (max_slice_k + 16) bytes is uniform per launch +// (host computes max_slice_k = ceil(total_steps / split_k) * 32; for +// the default split_k=5 that is 16 * 848 = 13,568 B so four blocks +// fit per CU); each block uses its own slice_k stride; +// * int32 partial publish unchanged (du_store_matrix_sync, plane +// (split, ntile)). Exact bf16 equality is preserved: staged bytes are +// the same A data in the same row-major order, and the int32 +// accumulation order is unchanged (k-ascending per split, ascending +// split sum). +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16n16k32_sk_astage_partial_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ b, // [K, N] row-major int8 + int32_t* __restrict__ partials, // [split_k * ntiles][16][16] + int n, + int k, + int split_k) { + extern __shared__ __align__(16) signed char smem_a[]; + + const int lane = static_cast(threadIdx.x) & (kDummaWaveSize - 1); + const int ntiles = n / kDummaTileN; + const int ntile = static_cast(blockIdx.x) % ntiles; + const int split = static_cast(blockIdx.x) / ntiles; + const int n0 = ntile * kDummaTileN; + const int total_steps = k / kDummaTileK; + // Step-based 32-aligned slice boundaries (identical to iteration 3): + // slice s owns DUMMA steps [s*total_steps/split_k, (s+1)*total_steps/ + // split_k), tiling [0, K) exactly once ascending. + const int k_lo = + (static_cast(split) * total_steps / split_k) * kDummaTileK; + const int k_hi = + (static_cast(split + 1) * total_steps / split_k) * + kDummaTileK; + const int slice_k = k_hi - k_lo; // multiple of 32 + const int local_stride = slice_k + 16; // padded LDS row stride (bytes) + + // One-time A-slice stage: 16 rows x slice_k bytes, row-major into LDS. + const int chunks_per_row = slice_k / 16; + const int total_chunks = 16 * chunks_per_row; + const int iters = (total_chunks + kDummaWaveSize - 1) / kDummaWaveSize; + for (int it = 0; it < iters; ++it) { + const int chunk = it * kDummaWaveSize + lane; + if (chunk < total_chunks) { + const int row = chunk / chunks_per_row; + const int col16 = (chunk % chunks_per_row) * 16; + const int4 v = *reinterpret_cast( + a + static_cast(row) * k + k_lo + col16); + *reinterpret_cast(smem_a + row * local_stride + col16) = v; + } + } + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + for (int k0 = k_lo; k0 < k_hi; k0 += kDummaTileK) { + du::dumma::du_load_matrix_sync( + a_frag, + reinterpret_cast(smem_a) + (k0 - k_lo), + local_stride); + du::dumma::du_load_matrix_sync( + b_frag, b + static_cast(k0) * n + n0, n); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + du::dumma::du_store_matrix_sync( + partials + + (static_cast(split) * ntiles + ntile) * + (kDummaTileM * kDummaTileN), + acc_frag, kDummaTileN, du::dumma::mem_row_major); +} + +// gfx928 DUMMA INT8 cross-block split-K partial kernel for tp8_wqkv_a_m16 - +// iteration 5 architecture (A-only LDS staging + PACKED-B fragment slots). +// Same launch geometry, K slicing, A-slice LDS stage and partial publish as +// the iteration-4 astage kernel with one change on the B load path: +// * B fragments are read from the packed [n_tile][k_step][lane][8] +// fragment-slot layout (w8a8_pack_b_fragslot_kernel) with ONE aligned +// 8-byte (dwordx2) global load per lane per k32 step +// (load_packed_b_fragment), in the exact du_mma.hpp matrix_b row_major +// byte order - B's per-step scalar byte loads + fragment reassembly + +// per-load vmcnt cascade (8 vmem instructions and ~50 VALU per step in +// the iteration-4 exact ISA) leave the K loop, which becomes ds_read +// (A, from LDS) -> 1 dwordx2 (B, packed) -> v_mmac; +// * dynamic LDS is unchanged (16 x (max_slice_k + 16) B; default split_k=5 +// -> 13,568 B, four blocks/CU), int32 partial publish unchanged. Exact +// bf16 equality is preserved: every packed fragment feeds the same bytes +// to the same v_mmac in the same order as the iteration-4 direct loads, +// and the int32 accumulation order is unchanged (k-ascending per split, +// ascending split sum). +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16n16k32_sk_astage_packedb_partial_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ packed_b, // [n_tile][k_step][lane][8] + int32_t* __restrict__ partials, // [split_k * ntiles][16][16] + int n, + int k, + int split_k) { + extern __shared__ __align__(16) signed char smem_a[]; + + const int lane = static_cast(threadIdx.x) & (kDummaWaveSize - 1); + const int ntiles = n / kDummaTileN; + const int ntile = static_cast(blockIdx.x) % ntiles; + const int split = static_cast(blockIdx.x) / ntiles; + const int total_steps = k / kDummaTileK; + const int k_lo = + (static_cast(split) * total_steps / split_k) * kDummaTileK; + const int k_hi = + (static_cast(split + 1) * total_steps / split_k) * + kDummaTileK; + const int slice_k = k_hi - k_lo; // multiple of 32 + const int local_stride = slice_k + 16; // padded LDS row stride (bytes) + + // One-time A-slice stage: 16 rows x slice_k bytes, row-major into LDS. + const int chunks_per_row = slice_k / 16; + const int total_chunks = 16 * chunks_per_row; + const int iters = (total_chunks + kDummaWaveSize - 1) / kDummaWaveSize; + for (int it = 0; it < iters; ++it) { + const int chunk = it * kDummaWaveSize + lane; + if (chunk < total_chunks) { + const int row = chunk / chunks_per_row; + const int col16 = (chunk % chunks_per_row) * 16; + const int4 v = *reinterpret_cast( + a + static_cast(row) * k + k_lo + col16); + *reinterpret_cast(smem_a + row * local_stride + col16) = v; + } + } + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + for (int k0 = k_lo; k0 < k_hi; k0 += kDummaTileK) { + du::dumma::du_load_matrix_sync( + a_frag, + reinterpret_cast(smem_a) + (k0 - k_lo), + local_stride); + load_packed_b_fragment(b_frag, packed_b, ntile, total_steps, + k0 / kDummaTileK, lane); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + du::dumma::du_store_matrix_sync( + partials + + (static_cast(split) * ntiles + ntile) * + (kDummaTileM * kDummaTileN), + acc_frag, kDummaTileN, du::dumma::mem_row_major); +} + +// gfx928 DUMMA INT8 cross-block split-K partial kernel for tp8_wqkv_a_m16 - +// iteration 6 architecture (A+B LDS staging, stage-K=64 double buffer). +// Same launch geometry, K slicing and partial publish as iteration 5 (one +// 64-thread wavefront per block, ZERO barriers, grid = (N/16) * split_k, +// split-major / ntile-minor blockIdx, step-based 32-aligned K slices) with +// the staging family changed from A-only full-slice to A+B stage-K=64 +// double-buffered prefetch: +// * the K slice is processed in 2-step (K=64) chunks; dynamic LDS holds +// two ping-pong buffers of 2,560 B each: an A chunk region (16 rows x +// 64 cols, padded 96-B row stride) and a B chunk region (the packed +// fragment slots of the 2 steps, 8 B per lane per step, in the exact +// w8a8_pack_b_fragslot_kernel order); +// * per chunk each lane issues 3 cooperative global loads (1 x 16-B A +// dwordx4 + 2 x 8-B B dwordx2) into registers BEFORE the current +// chunk's 1-2 v_mmac_i32_16x16x32_i8, then commits them to the +// alternate LDS buffer AFTER the MMAC burst - the compiler's vmcnt wait +// lands at that ds_write, overlapping the cold once-read B latency with +// the current chunk's compute (the exact source is verified by the +// harness ISA extraction: next-chunk global_load_dwordx4/dwordx2 before +// the current v_mmac, vmcnt wait near the alternate-buffer +// ds_write_b128/b64); +// * the steady K loop then contains ZERO global loads: per k32 step it is +// ds_read2_b32 (A, from LDS) -> ds_read_b64 (B, from LDS) -> v_mmac; +// * dynamic LDS = 5,120 B per block (constant, independent of split_k, so +// the old split_k=1 64-KiB A-slice overflow fallback is no longer +// needed); at the default split_k=5 grid 480 = 4 x 120 CUs, four blocks +// per CU = 20,480 B <= 64 KiB. The last chunk of an odd-step slice +// stages/computes 1 step only (the A load still reads 64 cols, always +// inside the A row [0, K)); +// * int32 partial publish unchanged. Exact bf16 equality is preserved: +// every staged byte is the same data fed to the same v_mmac in the same +// order, and the int32 accumulation order is unchanged (k-ascending per +// split, ascending split sum). +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16n16k32_sk_astage_ab_dbuf_partial_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ packed_b, // [n_tile][k_step][lane][8] + int32_t* __restrict__ partials, // [split_k * ntiles][16][16] + int n, + int k, + int split_k) { + extern __shared__ __align__(16) signed char smem[]; + + const int lane = static_cast(threadIdx.x) & (kDummaWaveSize - 1); + const int ntiles = n / kDummaTileN; + const int ntile = static_cast(blockIdx.x) % ntiles; + const int split = static_cast(blockIdx.x) / ntiles; + const int total_steps = k / kDummaTileK; + const int k_lo = + (static_cast(split) * total_steps / split_k) * kDummaTileK; + const int k_hi = + (static_cast(split + 1) * total_steps / split_k) * + kDummaTileK; + const int slice_steps = (k_hi - k_lo) / kDummaTileK; + const int nchunks = (slice_steps + kStageSteps - 1) / kStageSteps; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Per-lane roles inside one chunk: A lane owns 16 B of row (lane>>2) at + // column (lane&3)*16 (4 consecutive lanes walk one 64-B A row, coalesced + // 16-B loads); B lane owns its 8-B fragment slot of each of the 2 steps. + const int a_row = lane >> 2; + const int a_col16 = (lane & 3) * 16; + const int64_t a_base = static_cast(k_lo); // + chunk*64 + col16 + // B slot base for the slice's FIRST global k32 step (k_lo/kDummaTileK): + // the packed [n_tile][k_step][lane][8] layout is indexed by the GLOBAL + // step (pack kernel: slot = ntile*k_steps + step, step in [0, K/32)), so + // the slice-relative chunk/step offsets below must be added on top of the + // split's global step offset - iteration-6 repair 2 (the original base + // pointed at global step 0, so every split with k_lo > 0 loaded B from + // the wrong global steps while A was staged correctly). + const int64_t b_slot_base = + static_cast(ntile) * total_steps * kPackedTileBytes + + (k_lo / kDummaTileK) * kPackedTileBytes + + lane * kPackedSlotBytes; // + step*kPackedTileBytes + + // Prologue: stage chunk 0 into buffer 0 (single wait before the loop). + { + signed char* buf = smem; + const int4 av = *reinterpret_cast( + a + static_cast(a_row) * k + a_base + a_col16); + *reinterpret_cast(buf + a_row * kStageAStride + a_col16) = av; + const uint64_t bv0 = + *reinterpret_cast(packed_b + b_slot_base); + *reinterpret_cast(buf + kStageABufBytes + + lane * kPackedSlotBytes) = bv0; + if (1 < slice_steps) { + const uint64_t bv1 = *reinterpret_cast( + packed_b + b_slot_base + kPackedTileBytes); + *reinterpret_cast(buf + kStageABufBytes + + kPackedTileBytes + + lane * kPackedSlotBytes) = bv1; + } + } + + for (int c = 0; c < nchunks; ++c) { + const signed char* buf = smem + (c & 1) * kStageBufBytes; + + // Prefetch chunk c+1 from global into registers (no wait yet). The + // values are consumed only by the alternate-buffer stores below, so the + // compiler issues these loads before (or at worst alongside) the MMAC + // burst and lands the vmcnt wait at the ds_write - the overlap that + // removes the iteration-5 per-step vmcnt(0)-before-mmac exposure. + int4 a_v; + uint64_t b_v0, b_v1; + const bool has_next = (c + 1 < nchunks); + if (has_next) { + const int nc = c + 1; + const int64_t a_off = + a_base + static_cast(nc) * kStageK + a_col16; + a_v = *reinterpret_cast( + a + static_cast(a_row) * k + a_off); + const int64_t boff = b_slot_base + + static_cast(nc) * kStageSteps * + kPackedTileBytes; + b_v0 = *reinterpret_cast(packed_b + boff); + if (nc * kStageSteps + 1 < slice_steps) { + b_v1 = *reinterpret_cast(packed_b + boff + + kPackedTileBytes); + } + } + + // Compute the current chunk from LDS: 1 or 2 k-ascending k32 steps, + // each ds_read (A) -> ds_read (B) -> v_mmac. nsteps must be clamped to + // the chunk width: slice_steps - s0 is the remaining-step count, which + // is > kStageSteps for the second-to-last chunk (e.g. 25 - 2*11 = 3 at + // split_k=5, 26 - 2*11 = 4 for an even slice), where the extra steps + // would read A/B bytes beyond the staged 64-column chunk (stale LDS). + // Only the kStageSteps steps staged in this chunk may be computed; the + // last chunk of an odd-step slice stages/computes 1 step only. + const int s0 = c * kStageSteps; + const int nsteps = (slice_steps - s0 > kStageSteps) + ? kStageSteps + : (slice_steps - s0); // 1 or 2 + const signed char* a_buf = buf; + const signed char* b_buf = buf + kStageABufBytes; + for (int j = 0; j < nsteps; ++j) { + du::dumma::du_load_matrix_sync(a_frag, a_buf + j * kDummaTileK, + kStageAStride); + const uint64_t bv = *reinterpret_cast( + b_buf + j * kPackedTileBytes + lane * kPackedSlotBytes); + __builtin_memcpy(b_frag.x, &bv, sizeof(bv)); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Commit the prefetched chunk into the alternate buffer (the compiler + // inserts the vmcnt wait here, after the MMAC burst above). + if (has_next) { + const int nc = c + 1; + signed char* nb = smem + ((nc & 1) * kStageBufBytes); + *reinterpret_cast(nb + a_row * kStageAStride + a_col16) = a_v; + const int64_t boff = b_slot_base + + static_cast(nc) * kStageSteps * + kPackedTileBytes; + *reinterpret_cast(nb + kStageABufBytes + + lane * kPackedSlotBytes) = b_v0; + if (nc * kStageSteps + 1 < slice_steps) { + *reinterpret_cast(nb + kStageABufBytes + + kPackedTileBytes + + lane * kPackedSlotBytes) = b_v1; + } + } + } + + du::dumma::du_store_matrix_sync( + partials + + (static_cast(split) * ntiles + ntile) * + (kDummaTileM * kDummaTileN), + acc_frag, kDummaTileN, du::dumma::mem_row_major); +} + +// gfx928 DUMMA INT8 cross-block split-K partial kernel for tp8_wqkv_a_m16 - +// iteration 7 architecture (SINGLE-buffered A+B LDS staging, stage-K=64) - +// the mandated single-vs-double buffering comparison arm. Identical to the +// accepted iteration-6 double-buffered kernel (same launch geometry, K +// slicing, per-chunk register prefetch of chunk c+1 before chunk c's MMAC +// burst, vmcnt wait placed at the post-MMAC ds_write, zero in-loop barriers, +// same packed-B fragment-slot layout and int32 partial publish) with ONE +// change: the two ping-pong 2,560-B buffers collapse into a SINGLE 2,560-B +// buffer, so +// * dynamic LDS per block drops 5,120 -> 2,560 B (per CU 20,480 -> +// 10,240 B at the default split_k=5 grid; occupancy is grid-limited at +// 480 blocks = 4 x 120 CUs, so residency is unchanged); +// * the chunk c+1 prefetch commits to the SAME LDS addresses chunk c was +// just read from - the write-after-read ordering is program order +// inside the single wave (the j-loop reads precede the commit stores), +// the compiler keeps the same load-before-MMAC / wait-at-ds_write +// pipeline, and the ds_write base is a fixed constant instead of the +// (c&1)*kStageBufBytes toggle; +// * barriers per K step stay 0 (one 64-lane wave per block, no cross-wave +// LDS producer/consumer; sync is s_waitcnt vmcnt/lgkmcnt only - the +// TP4 4-wave reference's ~1 barrier per K step does not apply here); +// * the last chunk of an odd-step slice still stages/computes 1 step only +// (nsteps clamp), and every A stage load stays inside the A row +// [0, K). Exact bf16 equality is preserved: the staged bytes are the +// same data fed to the same v_mmac in the same order, and the int32 +// accumulation order is unchanged (k-ascending per split, ascending +// split sum). Falsifiable gate vs the double-buffer best: if the +// ping-pong LDS buffer and its base toggle are NOT the mechanism behind +// the iteration-6 win (the register prefetch + wait placement is), this +// kernel lands within noise of 18.551 us; PMC discriminates via +// lds_bytes 5,120 -> 2,560 with unchanged vmem_read/lds_instruction +// counts and 0 barriers per K step in both. +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16n16k32_sk_astage_ab_sbuf_partial_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ packed_b, // [n_tile][k_step][lane][8] + int32_t* __restrict__ partials, // [split_k * ntiles][16][16] + int n, + int k, + int split_k) { + extern __shared__ __align__(16) signed char smem[]; + + const int lane = static_cast(threadIdx.x) & (kDummaWaveSize - 1); + const int ntiles = n / kDummaTileN; + const int ntile = static_cast(blockIdx.x) % ntiles; + const int split = static_cast(blockIdx.x) / ntiles; + const int total_steps = k / kDummaTileK; + const int k_lo = + (static_cast(split) * total_steps / split_k) * kDummaTileK; + const int k_hi = + (static_cast(split + 1) * total_steps / split_k) * + kDummaTileK; + const int slice_steps = (k_hi - k_lo) / kDummaTileK; + const int nchunks = (slice_steps + kStageSteps - 1) / kStageSteps; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Per-lane roles inside one chunk (unchanged from iteration 6): A lane + // owns 16 B of row (lane>>2) at column (lane&3)*16; B lane owns its 8-B + // fragment slot of each of the 2 steps. + const int a_row = lane >> 2; + const int a_col16 = (lane & 3) * 16; + const int64_t a_base = static_cast(k_lo); + // B slot base for the slice's FIRST GLOBAL k32 step (k_lo/kDummaTileK) - + // the packed [n_tile][k_step][lane][8] layout is indexed by the GLOBAL + // step (iteration-6 repair 2). + const int64_t b_slot_base = + static_cast(ntile) * total_steps * kPackedTileBytes + + (k_lo / kDummaTileK) * kPackedTileBytes + + lane * kPackedSlotBytes; // + step*kPackedTileBytes + + // Prologue: stage chunk 0 into the single buffer (the compiler inserts + // the vmcnt wait before the ds_write). + { + const int4 av = *reinterpret_cast( + a + static_cast(a_row) * k + a_base + a_col16); + *reinterpret_cast(smem + a_row * kStageAStride + a_col16) = av; + const uint64_t bv0 = + *reinterpret_cast(packed_b + b_slot_base); + *reinterpret_cast(smem + kStageABufBytes + + lane * kPackedSlotBytes) = bv0; + if (1 < slice_steps) { + const uint64_t bv1 = *reinterpret_cast( + packed_b + b_slot_base + kPackedTileBytes); + *reinterpret_cast(smem + kStageABufBytes + + kPackedTileBytes + + lane * kPackedSlotBytes) = bv1; + } + } + + for (int c = 0; c < nchunks; ++c) { + // Single buffer: the current chunk is always at the smem base. + const signed char* buf = smem; + + // Prefetch chunk c+1 from global into registers (no wait yet). The + // values are consumed only by the commit stores below, so the compiler + // issues these loads before (or at worst alongside) the MMAC burst and + // lands the vmcnt wait at the ds_write - the overlap that removes the + // iteration-5 per-step vmcnt(0)-before-mmac exposure. + int4 a_v; + uint64_t b_v0, b_v1; + const bool has_next = (c + 1 < nchunks); + if (has_next) { + const int nc = c + 1; + const int64_t a_off = + a_base + static_cast(nc) * kStageK + a_col16; + a_v = *reinterpret_cast( + a + static_cast(a_row) * k + a_off); + const int64_t boff = b_slot_base + + static_cast(nc) * kStageSteps * + kPackedTileBytes; + b_v0 = *reinterpret_cast(packed_b + boff); + if (nc * kStageSteps + 1 < slice_steps) { + b_v1 = *reinterpret_cast(packed_b + boff + + kPackedTileBytes); + } + } + + // Compute the current chunk from the single LDS buffer: 1 or 2 + // k-ascending k32 steps, each ds_read (A) -> ds_read (B) -> v_mmac + // (nsteps clamped to the staged chunk width, iteration-6 repair 1). + const int s0 = c * kStageSteps; + const int nsteps = (slice_steps - s0 > kStageSteps) + ? kStageSteps + : (slice_steps - s0); // 1 or 2 + const signed char* a_buf = buf; + const signed char* b_buf = buf + kStageABufBytes; + for (int j = 0; j < nsteps; ++j) { + du::dumma::du_load_matrix_sync(a_frag, a_buf + j * kDummaTileK, + kStageAStride); + const uint64_t bv = *reinterpret_cast( + b_buf + j * kPackedTileBytes + lane * kPackedSlotBytes); + __builtin_memcpy(b_frag.x, &bv, sizeof(bv)); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Commit the prefetched chunk into the single buffer (the compiler + // inserts the vmcnt wait here, after the MMAC burst above; the write + // follows the reads of the same buffer in program order). + if (has_next) { + const int nc = c + 1; + *reinterpret_cast(smem + a_row * kStageAStride + a_col16) = + a_v; + const int64_t boff = b_slot_base + + static_cast(nc) * kStageSteps * + kPackedTileBytes; + *reinterpret_cast(smem + kStageABufBytes + + lane * kPackedSlotBytes) = b_v0; + if (nc * kStageSteps + 1 < slice_steps) { + *reinterpret_cast(smem + kStageABufBytes + + kPackedTileBytes + + lane * kPackedSlotBytes) = b_v1; + } + } + } + + du::dumma::du_store_matrix_sync( + partials + + (static_cast(split) * ntiles + ntile) * + (kDummaTileM * kDummaTileN), + acc_frag, kDummaTileN, du::dumma::mem_row_major); +} + +// gfx928 DUMMA INT8 cross-block split-K partial kernel for tp8_wqkv_a_m16 - +// iteration 14 architecture (FUSED last-arrival combine tail): the +// iteration-7 single-buffered stage-K=64 kernel EXACTLY as measured +// (18.484 us median / 18.496 us p90 shadow; same launch geometry split_k=5 +// grid 480 one-wave zero-barrier blocks = 4 x 120 CUs, same per-chunk +// register prefetch of chunk c+1 before chunk c's MMAC burst with the vmcnt +// wait at the post-MMAC fixed-base ds_write, same single 2,560-B buffer, +// same packed-B fragment-slot layout, same K slicing and int32 partial +// publish) PLUS a fused combine/scale tail: after publishing its int32 +// plane, each block issues __threadfence() (device scope, orders the plane +// store), lane 0 does one atomicAdd on the ntile's monotonic arrival +// counter (counters live in the last counter_bytes of the caller +// workspace, zeroed once per distinct workspace on the first launcher call +// - the harness warmup runs OUTSIDE Graph capture - and advance by exactly +// split_k per complete launch/replay, so the (old % split_k == split_k-1) +// last-arriver test is valid forever without any in-Graph reset), and the +// LAST-ARRIVING block for the ntile (one 64-lane wave covers all 256 +// elements, 4 per lane, coalesced plane reads) sums the split_k planes in +// ascending split order (bit-identical int32 order to the iteration-3 +// combine kernel: k-ascending per split, ascending split sum), applies +// x_scale[row] * weight_scale[n0+col], and stores bf16 with the same +// expression. The separate 96-block x 256-thread combine kernel is REMOVED +// from the timed Graph - one kernel launch per replay. The tail's end-of- +// kernel s_barrier is the only barrier (one wave per block; accepted +// down_proj-style END-of-K combine barrier). Every staged byte and every +// v_mmac is unchanged, so partials are bit-identical and the only output- +// order change is where the ascending split sum happens; |full dot| = +// 4096*127*127 ~= 6.6e7 << 2^31. Falsifiable gate vs the 18.484 us shadow: +// if the ~2 us gap between the 16.48 us profiled partial phase and the +// Graph median is the second launch plus the combine kernel's L2-plane +// round trip (480 KB read + 48 KB write), removing them drops the median +// well below the 1% acceptance bar (predicted ~1-1.5 us); if the gap is +// actually graph dispatch overhead or profiler perturbation, the round +// lands flat; if the atomic/fence tail or its L2 traffic costs more than +// the saved launch, it lands mildly regressed - the comparison +// discriminates the mechanism either way. +// Iteration 15: the measured round landed essentially flat (18.382 us +// shadow median vs the 18.484 us iteration-7 shadow, +0.55%) - the ~2 us +// partial-to-Graph gap was the FUSED TAIL's own fence/atomic/tail cost, +// not the combine kernel's launch+L2 round trip (consistent with the TP4 +// lineage, where the atomic last-arriver fused finalize on the 4-wave +// design measured +1.52 us: 18.864 vs 17.344). This kernel is now a +// PRESERVED SYMBOL: the workspace path dispatches the iteration-15 +// TP4-validated 4-wave partial kernel + separate combine kernel instead. +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16n16k32_sk_astage_ab_sbuf_fused_partial_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ packed_b, // [n_tile][k_step][lane][8] + int32_t* __restrict__ partials, // [split_k * ntiles][16][16] + int32_t* __restrict__ counters, // [ntiles] monotonic arrivals + const float* __restrict__ x_scale, // [16, 1] fp32 + const float* __restrict__ weight_scale, // [N, 1] fp32 + hip_bfloat16* __restrict__ out, // [16, N] bf16 + int n, + int k, + int split_k) { + extern __shared__ __align__(16) signed char smem[]; + __shared__ int s_arrived; // tail arrival broadcast (one wave per block) + + const int lane = static_cast(threadIdx.x) & (kDummaWaveSize - 1); + const int ntiles = n / kDummaTileN; + const int ntile = static_cast(blockIdx.x) % ntiles; + const int split = static_cast(blockIdx.x) / ntiles; + const int total_steps = k / kDummaTileK; + const int k_lo = + (static_cast(split) * total_steps / split_k) * kDummaTileK; + const int k_hi = + (static_cast(split + 1) * total_steps / split_k) * + kDummaTileK; + const int slice_steps = (k_hi - k_lo) / kDummaTileK; + const int nchunks = (slice_steps + kStageSteps - 1) / kStageSteps; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Per-lane roles inside one chunk (unchanged from iteration 6/7): A lane + // owns 16 B of row (lane>>2) at column (lane&3)*16; B lane owns its 8-B + // fragment slot of each of the 2 steps. + const int a_row = lane >> 2; + const int a_col16 = (lane & 3) * 16; + const int64_t a_base = static_cast(k_lo); + // B slot base for the slice's FIRST GLOBAL k32 step (k_lo/kDummaTileK) - + // the packed [n_tile][k_step][lane][8] layout is indexed by the GLOBAL + // step (iteration-6 repair 2). + const int64_t b_slot_base = + static_cast(ntile) * total_steps * kPackedTileBytes + + (k_lo / kDummaTileK) * kPackedTileBytes + + lane * kPackedSlotBytes; // + step*kPackedTileBytes + + // Prologue: stage chunk 0 into the single buffer (the compiler inserts + // the vmcnt wait before the ds_write). + { + const int4 av = *reinterpret_cast( + a + static_cast(a_row) * k + a_base + a_col16); + *reinterpret_cast(smem + a_row * kStageAStride + a_col16) = av; + const uint64_t bv0 = + *reinterpret_cast(packed_b + b_slot_base); + *reinterpret_cast(smem + kStageABufBytes + + lane * kPackedSlotBytes) = bv0; + if (1 < slice_steps) { + const uint64_t bv1 = *reinterpret_cast( + packed_b + b_slot_base + kPackedTileBytes); + *reinterpret_cast(smem + kStageABufBytes + + kPackedTileBytes + + lane * kPackedSlotBytes) = bv1; + } + } + + for (int c = 0; c < nchunks; ++c) { + // Single buffer: the current chunk is always at the smem base. + const signed char* buf = smem; + + // Prefetch chunk c+1 from global into registers (no wait yet). The + // values are consumed only by the commit stores below, so the compiler + // issues these loads before (or at worst alongside) the MMAC burst and + // lands the vmcnt wait at the ds_write - the overlap that removes the + // iteration-5 per-step vmcnt(0)-before-mmac exposure. + int4 a_v; + uint64_t b_v0, b_v1; + const bool has_next = (c + 1 < nchunks); + if (has_next) { + const int nc = c + 1; + const int64_t a_off = + a_base + static_cast(nc) * kStageK + a_col16; + a_v = *reinterpret_cast( + a + static_cast(a_row) * k + a_off); + const int64_t boff = b_slot_base + + static_cast(nc) * kStageSteps * + kPackedTileBytes; + b_v0 = *reinterpret_cast(packed_b + boff); + if (nc * kStageSteps + 1 < slice_steps) { + b_v1 = *reinterpret_cast(packed_b + boff + + kPackedTileBytes); + } + } + + // Compute the current chunk from the single LDS buffer: 1 or 2 + // k-ascending k32 steps, each ds_read (A) -> ds_read (B) -> v_mmac + // (nsteps clamped to the staged chunk width, iteration-6 repair 1). + const int s0 = c * kStageSteps; + const int nsteps = (slice_steps - s0 > kStageSteps) + ? kStageSteps + : (slice_steps - s0); // 1 or 2 + const signed char* a_buf = buf; + const signed char* b_buf = buf + kStageABufBytes; + for (int j = 0; j < nsteps; ++j) { + du::dumma::du_load_matrix_sync(a_frag, a_buf + j * kDummaTileK, + kStageAStride); + const uint64_t bv = *reinterpret_cast( + b_buf + j * kPackedTileBytes + lane * kPackedSlotBytes); + __builtin_memcpy(b_frag.x, &bv, sizeof(bv)); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Commit the prefetched chunk into the single buffer (the compiler + // inserts the vmcnt wait here, after the MMAC burst above; the write + // follows the reads of the same buffer in program order). + if (has_next) { + const int nc = c + 1; + *reinterpret_cast(smem + a_row * kStageAStride + a_col16) = + a_v; + const int64_t boff = b_slot_base + + static_cast(nc) * kStageSteps * + kPackedTileBytes; + *reinterpret_cast(smem + kStageABufBytes + + lane * kPackedSlotBytes) = b_v0; + if (nc * kStageSteps + 1 < slice_steps) { + *reinterpret_cast(smem + kStageABufBytes + + kPackedTileBytes + + lane * kPackedSlotBytes) = b_v1; + } + } + } + + du::dumma::du_store_matrix_sync( + partials + + (static_cast(split) * ntiles + ntile) * + (kDummaTileM * kDummaTileN), + acc_frag, kDummaTileN, du::dumma::mem_row_major); + + // Iteration 14 fused tail: publish this block's partial (above), fence, + // then the last-arriving block for this ntile sums the split_k planes in + // ascending split order (bit-identical to the iteration-3 combine + // kernel), scales with the same expression, and stores bf16 - so the + // timed Graph needs only this ONE kernel per replay. + __threadfence(); + if (lane == 0) { + s_arrived = atomicAdd(&counters[ntile], 1); + } + __syncthreads(); + const int arrived = s_arrived; + if (arrived % split_k == split_k - 1) { + const int plane = ntile * (kDummaTileM * kDummaTileN); + const int n0 = ntile * kDummaTileN; + int32_t total[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + total[i] = 0; + } + for (int s = 0; s < split_k; ++s) { + const int32_t* plane_ptr = + partials + (static_cast(s) * ntiles) * + (kDummaTileM * kDummaTileN) + + plane; +#pragma unroll + for (int i = 0; i < 4; ++i) { + total[i] += plane_ptr[lane + kDummaWaveSize * i]; + } + } +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int e = lane + kDummaWaveSize * i; + const int row = e >> 4; + const int col = e & 15; + const float scaled = static_cast(total[i]) * x_scale[row] * + weight_scale[n0 + col]; + out[static_cast(row) * n + n0 + col] = hip_bfloat16(scaled); + } + } +} + +// gfx928 DUMMA INT8 cross-block split-K partial kernel for tp8_wqkv_a_m16 - +// iteration 16 (REPAIR of the iteration-15 candidate; the fast mapping is +// preserved byte-for-byte). Iteration 15 measured 60.158 us (single sample) +// with 24,555/24,576 mismatches; the SOLE defect was the B-staging global +// address: each stage chunk was read as one contiguous 4,096-B window of the +// packed [n_tile][k_step][lane][8] layout starting at +// (4*tile64*k_steps + k_lo/32)*512 + tid*16 + c*4,096. The packed layout is +// indexed by the GLOBAL step ((nt*k_steps + step)*512 + lane*8, each ntile's +// 128 steps form a contiguous 64 KiB run), so a 4-ntile x 2-step chunk is +// NOT contiguous: only wave 0 of chunk 0 read correct bytes, waves 1-3 read +// ntile 4*tile64's later steps, and every chunk c>=1 read step k_lo/32+8c +// instead of k_lo/32+2c. The repair addresses each thread's 16-B slot pair +// with the exact per-(ntile, step) packed offset (same one 16-B +// global_load_dwordx4 per thread per chunk and the same 16-B ds_write_b128 +// into the unchanged slot-mirror LDS layout at byte tid*16, so the fragment +// reads b_buf + (wave*2 + step)*512 + lane*8 hit exactly their own slot); +// the chunk stride becomes kStageSteps*kPackedTileBytes = 1,024 B. A +// staging, the double-buffer pipeline, both __syncthreads() per stage, the +// plane publish, the separate combine kernel, the dispatch and all +// fallbacks are byte-identical to the measured iteration-15 candidate. +// Iteration 15 architecture (TP4-validated 4-WAVE block-N=64 stage-K=64 +// DOUBLE-buffer design; the TP4 lineage measured 17.344 us median / 17.616 +// us p90 on this exact (M=16, K=4096, N=1536) local shape at split_k=10, +// grid 240 = exactly 2 x 120 CUs, 4 x wave64 per block, one 16x16 quadrant +// per wave sharing one staged A chunk). Differences from the TP4 winner are +// forced by this session's packed-B weight layout: B is staged from the +// packed [n_tile][k_step][lane][8] fragment slots (every lane issues one +// 16-B int4 global load = two adjacent 8-B slots of the same (ntile, step) +// pair and commits one 16-B ds_write_b128 into an LDS mirror of the slot +// layout, so the B fragment read is the same single ds_read_b64 per lane +// per k32 step as the accepted 1-wave kernels), instead of the TP4 raw +// [K,N] 16-B row loads; the A chunk is staged by wave 0 (64 lanes x 16 B = +// the full 16x64 raw A tile, row stride 64) exactly as the TP4 winner; the +// per-stage pipeline is identical (stage i+1's int4 loads issue before the +// current stage's 2 x v_mmac_i32_16x16x32_i8 per wave, the vmcnt wait +// lands at the alternate-buffer ds_write, and two __syncthreads() per +// stage - the pure-HIP form of the TP4 "raw first s_barrier + compiler- +// managed second barrier" sync mode; no raw asm per the round policy); +// partials publish to the unchanged [split][ntile][16][16] plane layout +// (per-wave 16x16 quadrant), and the separate 96-block x 256-thread +// combine kernel returns to the timed Graph (the TP4 lineage measured the +// atomic last-arriver fused finalize on this design at +1.52 us - 18.864 +// vs 17.344 - so the fused tail is NOT used here; the iteration-14 fused +// 1-wave kernel stays as a preserved symbol). split_k is runtime- +// selectable (default 10; env METAINFER_W8A8_SPLIT_K), sliced in +// stage-K=64 units non-uniformly (4 x 7 + 6 x 6 stages at split10; every +// boundary stage-aligned, k-ascending per split). Dynamic LDS = 10,240 B +// per block. Exact bf16: identical staged bytes fed to the same v_mmac in +// the same order; int32 accumulation order unchanged (k-ascending per +// split, ascending split sum). +__global__ __launch_bounds__(k4wBlockThreads) void +w8a8_dumma_m16n16k32_4wave_sk_astage_packedb_dbuf_partial_kernel( + const int8_t* __restrict__ a, // [16, K] row-major int8 + const int8_t* __restrict__ packed_b, // [n_tile][k_step][lane][8] + int32_t* __restrict__ partials, // [split_k * ntiles][16][16] + int n, + int k, + int split_k) { + extern __shared__ __align__(16) signed char smem[]; + const int tid = static_cast(threadIdx.x); + const int wave = tid >> 6; // 0..3 -> quadrant ntile64*4 + wave + const int lane = tid & (kDummaWaveSize - 1); + const int ntiles64 = n / (k4wBlockNTiles * kDummaTileN); // 24 + const int tile64 = static_cast(blockIdx.x) % ntiles64; + const int split = static_cast(blockIdx.x) / ntiles64; + // Non-uniform stage-based K slicing (stage-K=64 units, TP4 formula): + // split10 -> 4 x 7 + 6 x 6 stages; every boundary stage-aligned. + const int total_stages = k / kStageK; // 64 + const int base_stages = total_stages / split_k; + const int extra_stages = total_stages - base_stages * split_k; + const int begin_stage = split * base_stages + + (split < extra_stages ? split : extra_stages); + const int split_stages = base_stages + (split < extra_stages ? 1 : 0); + const int k_lo = begin_stage * kStageK; + + signed char* a_tile = smem; // [2][1024] + signed char* b_tile = smem + 2 * k4wStageABytes; // [2][4096] + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Staging roles: A - wave 0 (tid < 64) loads 16 B per lane, raw [16][64] + // row stride 64; B - every thread loads one 16-B int4 = two adjacent 8-B + // packed slots (same (ntile, step) pair, adjacent lanes), committed as + // one 16-B ds_write_b128 to the slot-mirror LDS layout at the same byte + // offset (each lane's later 8-B fragment read hits exactly its own slot). + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) * 16; + const int64_t a_base = static_cast(k_lo); + // B per-thread roles: tid maps to one 16-B slot pair inside exactly ONE + // (ntile, step) pair: pair index p = tid >> 5 (p>>1 = ntile within the + // block's 4-tile group, p&1 = step within the K=64 chunk), 16-B + // sub-offset (tid&31)*16. The packed [n_tile][k_step][lane][8] layout is + // indexed by the GLOBAL step: byte ((nt*k_steps + step)*64 + lane)*8, so + // the pair (w, s) of chunk c sits at (4*tile64 + w)*k_steps + k_lo/32 + + // 2c + s - the 4-ntile x 2-step chunk is NOT one contiguous 4,096-B + // global run (iteration-16 repair: the iteration-15 code read a + // contiguous window from (4*tile64*k_steps + k_lo/32)*512 + tid*16 + c* + // 4,096, which is only correct for the first 1,024 B; waves 1-3 consumed + // ntile 4*tile64's later steps and every chunk c>=1 read step k_lo/32+8c + // instead of k_lo/32+2c - the sole correctness defect of iteration 15). + const int b_pair = tid >> 5; // 0..7 + const int b_w = b_pair >> 1; // ntile index within the block (0..3) + const int b_s = b_pair & 1; // step within the chunk (0..1) + const int64_t b_off = // this thread's 16-B slot pair of chunk 0 + (static_cast(tile64 * k4wBlockNTiles + b_w) * + (k / kDummaTileK) + + (k_lo / kDummaTileK) + b_s) * + kPackedTileBytes + + (tid & 31) * 16; // 16-B aligned + + // Prologue: stage chunk 0 into buffer 0, then make it visible to all + // four waves. + if (tid < kDummaWaveSize) { + const int4 av = *reinterpret_cast( + a + static_cast(a_row) * k + a_base + a_col16); + *reinterpret_cast(a_tile + a_row * kStageK + a_col16) = av; + } + const int4 bv = *reinterpret_cast(packed_b + b_off); + *reinterpret_cast(b_tile + tid * 16) = bv; + __syncthreads(); + + for (int c = 0; c < split_stages; ++c) { + const int cur = c & 1; + const int nxt = cur ^ 1; + const bool has_next = (c + 1 < split_stages); + + // Prefetch chunk c+1 from global into registers (no wait yet; the + // compiler issues these before the MMAC burst and lands the vmcnt wait + // at the alternate-buffer ds_write below - the TP4-validated overlap). + int4 a_v{}; + int4 b_v{}; + if (has_next) { + if (tid < kDummaWaveSize) { + a_v = *reinterpret_cast( + a + static_cast(a_row) * k + a_base + + static_cast(c + 1) * kStageK + a_col16); + } + // Chunk c+1 = the same (ntile, step-in-chunk) pair two GLOBAL steps + // later: +kStageSteps*kPackedTileBytes = 1,024 B per chunk (NOT the + // 4,096-B contiguous-window stride of the iteration-15 bug). + b_v = *reinterpret_cast( + packed_b + b_off + static_cast(c + 1) * kStageSteps * + kPackedTileBytes); + } + + // Consume the current chunk: 2 k-ascending k32 steps per wave, each + // ds_read (A, shared) -> ds_read (B, own 8-B slot) -> v_mmac. + const signed char* a_buf = a_tile + cur * k4wStageABytes; + const signed char* b_buf = b_tile + cur * k4wStageBBytes; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kDummaTileK) { + du::dumma::du_load_matrix_sync(a_frag, a_buf + kk, kStageK); + const uint64_t bval = *reinterpret_cast( + b_buf + (wave * kStageSteps + (kk >> 5)) * kPackedTileBytes + + lane * kPackedSlotBytes); + __builtin_memcpy(b_frag.x, &bval, sizeof(bval)); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // All four waves have consumed buffer cur (pure-HIP form of the TP4 + // raw first s_barrier). + __syncthreads(); + + // Commit the prefetched chunk into the alternate buffer (the compiler + // inserts the vmcnt wait before these stores, after the MMAC burst). + if (has_next) { + if (tid < kDummaWaveSize) { + *reinterpret_cast(a_tile + nxt * k4wStageABytes + + a_row * kStageK + a_col16) = a_v; + } + *reinterpret_cast(b_tile + nxt * k4wStageBBytes + tid * 16) = + b_v; + } + // The committed chunk is visible to all four waves. + __syncthreads(); + } + + const int plane_ntile = tile64 * k4wBlockNTiles + wave; + du::dumma::du_store_matrix_sync( + partials + + (static_cast(split) * (n / kDummaTileN) + plane_ntile) * + (kDummaTileM * kDummaTileN), + acc_frag, kDummaTileN, du::dumma::mem_row_major); +} + +// gfx928 DUMMA INT8 combine + scale kernel for tp8_wqkv_a_m16 - iteration +// 3: sums the split_k workspace int32 planes per output element in +// ascending split order (bit-identical int32 to the serial k-ascending +// accumulation: integer addition is exact and |full dot| << 2^31), +// applies x_scale[m] * weight_scale[n], stores bf16. grid = N/16 blocks x +// 256 threads (4 wavefronts), one output element per thread: thread t owns +// element (row = t >> 4, col = t & 15) of its tile, so each wavefront +// covers 4 rows x all 16 columns with coalesced 16-wide plane reads and +// bf16 stores; no barriers, no LDS. +__global__ __launch_bounds__(4 * kDummaWaveSize) void +w8a8_dumma_m16n16k32_combine_kernel( + const int32_t* __restrict__ partials, // [split_k * ntiles][16][16] + const float* __restrict__ x_scale, // [16, 1] fp32 + const float* __restrict__ weight_scale, // [N, 1] fp32 + hip_bfloat16* __restrict__ out, // [16, N] bf16 + int n, + int split_k) { + const int ntiles = n / kDummaTileN; + const int tid = static_cast(threadIdx.x); + const int row = tid >> 4; // 0..15 + const int col = tid & 15; // 0..15 + const int ntile = static_cast(blockIdx.x); + const int n0 = ntile * kDummaTileN; + const int plane = (ntile << 8); // ntile * 16 * 16 int32 elements + + int32_t total = 0; +#pragma unroll + for (int s = 0; s < split_k; ++s) { + total += partials[static_cast(s) * ntiles * (kDummaTileM * + kDummaTileN) + + plane + row * kDummaTileN + col]; + } + const float scaled = static_cast(total) * x_scale[row] * + weight_scale[n0 + col]; + out[static_cast(row) * n + n0 + col] = hip_bfloat16(scaled); +} + +// Generic elementwise device-to-device copy (used by the identity +// pack_weight bootstrap and valid for every (K, N)). +template +__global__ void w8a8_device_copy_kernel(const T* __restrict__ src, + T* __restrict__ dst, + int64_t numel) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < numel) { + dst[i] = src[i]; + } +} + +// Iteration 5: one-time device permutation of the raw logical [K, N] +// row-major weight into the packed [n_tile][k_step][lane][8] B fragment-slot +// layout for the exact (K, N) == (4096, 1536) shape. One thread owns one +// 8-byte lane slot; slot ((ntile * (k/32) + step) * 64 + lane) * 8 + i gets +// logical weight[step*32 + (lane>>4)*8 + i][ntile*16 + (lane&15)]. Byte +// count is unchanged (K*N), so packed_weight keeps the same numel and +// graph-stable address; runs outside the timed region and Graph capture. +__global__ __launch_bounds__(kCopyBlockThreads) void +w8a8_pack_b_fragslot_kernel(const int8_t* __restrict__ raw, // [K, N] + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t tid = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int k_steps = k / kDummaTileK; + const int64_t total = static_cast(n / kDummaTileN) * k_steps * + kDummaWaveSize; // 786,432 lane slots + if (tid >= total) { + return; + } + const int lane = static_cast(tid & (kDummaWaveSize - 1)); + const int64_t slot = tid >> 6; // ntile * k_steps + step + const int step = static_cast(slot % k_steps); + const int ntile = static_cast(slot / k_steps); + const int nn = lane & 15; + const int k8 = (lane >> 4) << 3; + const int8_t* src = + raw + (static_cast(step) * kDummaTileK + k8) * n + + ntile * kDummaTileN + nn; + int8_t* dst = packed + tid * kPackedSlotBytes; +#pragma unroll + for (int i = 0; i < kPackedSlotBytes; ++i) { + dst[i] = src[static_cast(i) * n]; + } +} + +} // namespace + +// Host launcher for torch.ops.zth_w8a8.gemm_out. Runs only on the caller's +// stream; no allocations, no synchronization, no autotuning, no packing. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + + // Exact-shape-guarded DUMMA fast path for the assigned shape + // tp8_wqkv_a_m16: (M=16, N=1536, K=4096). Iteration 15 architecture: + // the TP4-validated 4-WAVE block-N=64 stage-K=64 DOUBLE-buffered + // cross-block split-K partial kernel (256 threads = 4 wave64, one 16x16 + // quadrant per wave sharing one staged A chunk; B staged from the packed + // [n_tile][k_step][lane][8] fragment-slot layout) + the separate + // 96-block x 256-thread combine/scale kernel, both launched on the + // caller stream inside the timed Graph. Default split_k=10 -> grid = + // 24*10 = 240 blocks = exactly 2 x 120 CUs (the TP4-validated occupancy + // for this exact shape: 17.344 us median / 17.616 us p90 at TP4). + // Dynamic LDS = 10,240 B per block (2 x 5,120-B ping-pong buffers), + // independent of split_k. Int32 partial planes live in the caller- + // provided workspace: plane (split, ntile) is 16x16 int32 at + // partials[(split*ntiles + ntile) * 256], so split_k*96 planes total. + if (m == 16 && n == 1536 && k == 4096) { + const int ntiles = n / kDummaTileN; // 96 (16x16 tiles) + const int ntiles64 = n / (k4wBlockNTiles * kDummaTileN); // 24 + const int64_t plane_bytes = + static_cast(kDummaTileM) * kDummaTileN * sizeof(int32_t); + if (workspace != nullptr && workspace_bytes >= plane_bytes) { + // split_k selection: METAINFER_W8A8_SPLIT_K override for trusted + // occupancy-probe candidates, default 10 (the TP4-validated + // non-uniform split for this shape: grid 240 = exactly 2 x 120 CUs + // with 24 block-N=64 tiles); clamped to [1, min(16, K/32)] and + // workspace capacity. The staged-chunk design keeps dynamic LDS at + // 10,240 B per block for every split_k (no 64 KiB A-slice concern). + int split_k = 10; + if (const char* env = getenv("METAINFER_W8A8_SPLIT_K")) { + const long parsed = strtol(env, nullptr, 10); + if (parsed >= 1 && parsed <= 16) { + split_k = static_cast(parsed); + } + } + const int64_t capacity = + workspace_bytes / plane_bytes / static_cast(ntiles); + if (static_cast(split_k) > capacity) { + split_k = static_cast(capacity); + } + const int64_t stage_capacity = k / kStageK; + if (static_cast(split_k) > stage_capacity) { + split_k = static_cast(stage_capacity); + } + if (split_k < 1) { + split_k = 1; + } + + int32_t* partials = static_cast(workspace); + // Iteration 15: the TP4-validated 4-wave block-N=64 stage-K=64 + // double-buffer partial kernel (default split_k=10 -> grid 240 = + // exactly 2 x 120 CUs; 10,240 B dynamic LDS/block) + the separate + // 96-block x 256-thread combine kernel, both inside the timed Graph. + // The TP4 lineage measured the atomic last-arriver fused finalize on + // this design at +1.52 us (18.864 vs 17.344), so the fused tail is + // NOT used here; the iteration-14 fused 1-wave kernel stays as a + // preserved symbol. Exact bf16: same staged bytes, same int32 + // accumulation order (k-ascending per split, ascending split sum). + const unsigned grid = static_cast(ntiles64 * split_k); + hipLaunchKernelGGL( + w8a8_dumma_m16n16k32_4wave_sk_astage_packedb_dbuf_partial_kernel, + dim3(grid), + dim3(k4wBlockThreads), + static_cast(k4wStageTotalBytes), + stream, + a, + b, + partials, + n, + k, + split_k); + hipLaunchKernelGGL( + w8a8_dumma_m16n16k32_combine_kernel, + dim3(static_cast(ntiles)), + dim3(4 * kDummaWaveSize), + 0, + stream, + partials, + x_scale, + weight_scale, + reinterpret_cast(out), + n, + split_k); + return; + } + // Caller provided no usable workspace: keep the accepted iteration-2 + // workspace-free in-block split-K=2 kernel (now with the packed-B + // fragment-slot load) so the fast path stays correct for such callers + // (the fixed-API contract always provides capacity for 16 planes, so + // the harness takes the split-K path above). + const unsigned grid = static_cast(n / kDummaTileN); + hipLaunchKernelGGL( + w8a8_dumma_m16n16k32_2wave_sk2_packedb_kernel, + dim3(grid), + dim3(2 * kDummaWaveSize), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast(out), + n, + k); + return; + } + + // Generic scalar fallback for every other (m, n, k) (including the paired + // M=2 shape with the same (N, K)). + const int64_t total = static_cast(m) * n; + const unsigned grid = static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + dim3(grid), + dim3(kScalarBlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast(out), + m, + n, + k); +} + +// Host launcher for the optional torch.ops.zth_w8a8.pack_weight. For the +// exact (K, N) == (4096, 1536) weight it performs the iteration-5 one-time +// device permutation into the packed [n_tile][k_step][lane][8] B fragment +// slots (byte count unchanged: 6,291,456 B); for every other shape it stays +// the identity device-to-device copy. packed_weight_scale[n] = +// weight_scale[n] always. Out of the timed region and out of Graph capture. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_numel = static_cast(k) * n; + if (weight_numel > 0) { + if (k == kTargetK && n == kTargetN) { + const int64_t slots = + static_cast(n / kDummaTileN) * (k / kDummaTileK) * + kDummaWaveSize; // 786,432 lane slots of 8 B + const unsigned grid = static_cast( + (slots + kCopyBlockThreads - 1) / kCopyBlockThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_pack_b_fragslot_kernel), + dim3(grid), + dim3(kCopyBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + const unsigned grid = static_cast( + (weight_numel + kCopyBlockThreads - 1) / kCopyBlockThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_device_copy_kernel), + dim3(grid), + dim3(kCopyBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_numel); + } + } + if (n > 0) { + const unsigned grid = static_cast( + (n + kCopyBlockThreads - 1) / kCopyBlockThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_device_copy_kernel), + dim3(grid), + dim3(kCopyBlockThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + static_cast(n)); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/fused_qkv_a_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/fused_qkv_a_proj.hip new file mode 100644 index 00000000..a7e83daf --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/fused_qkv_a_proj.hip @@ -0,0 +1,1889 @@ +// @@variant shape=glm_tp8_fused_qkv_a_proj_m16 commit=bf306c40c48713a81efbd7ace430c9cf21cbb429 added=2026-08-31 +// median_us=41.97 p90_us=42.09 +// source=glm5-2-dsh-tp8-m16-test1-cb4d262c +// INT8 W8A8 GEMM kernel for Hygon K500SM_AI / gfx928 (worker_0). +// +// Assigned shape: +// glm_tp8_fused_qkv_a_proj_m16 : M=16, N=2624, K=6144 +// +// Iteration 0 (correctness-first scalar bootstrap): one thread computes one +// output element with a plain int32 dot product. Generic fallback, still +// used for every shape that is NOT the exact assigned (m,n,k) below, +// including the paired M=2 API shapes with the same (N,K). +// +// Iteration 1 (DUMMA m16n16k32 bootstrap): the exact assigned shape ran the +// minimal gfx928 INT8 DUMMA kernel -- one 64-lane wavefront per 16x16 output +// tile (grid = N/16 = 164 one-wave blocks, no cross-wave barrier, no LDS +// staging, no split-K), direct global fragment loads via du_load_matrix_sync, +// explicit int32 accumulation over the full K=6144 (192 m16n16k32 steps), +// then a direct register epilogue that applies x_scale[m] * +// packed_weight_scale[n] and stores bf16. Measured median 205.41 us +// (0.48x vs the fixed 97.99 us Triton baseline): the 164-block grid leaves +// ~1.37 one-wave blocks per CU, below the two-blocks-per-CU latency-hiding +// target, so each wave's direct fragment-load vmcnt stalls are serialized +// with no co-resident wave to overlap them. +// +// Iteration 2 (architecture round, split-K occupancy probe): keep the +// zero-barrier one-wave block, but split K across SPLIT_K=4 blocks, giving a +// finer one-wave zero-barrier grid of 164*4 = 656 independent blocks (5.47 +// blocks/CU >= 2 blocks/CU target) with each MMA chain shortened to +// K/SPLIT_K = 48 k32 steps. Each block publishes its 16x16 int32 partial +// plane into the preallocated workspace ([4][164][256] int32 = 671,744 B, +// within the 16-plane API budget), and a tiny same-stream combine kernel +// (164 one-wave blocks) sums the planes in ascending split order -- exact +// int32 accumulation, bit-identical to the reference -- then applies the +// x_scale * weight_scale bf16 epilogue. Both launches stay inside the timed +// Graph replay (split-K combine is part of the operator wall). +// +// Iteration 3 (rejected occupancy probe, source reverted): an env-selectable +// SPLIT_K sweep defaulted to 2 measured 184.38 us median (0.53x baseline) -- +// worse than SPLIT_K=4 -- so the three-point curve {1,2,4} = {205.41, 184.38, +// 138.56} us shows 5.47 blocks/CU co-residency plus the 48-step chains is the +// good operating point and 2.73 blocks/CU regresses. +// +// Iteration 4 (architecture/pipeline round, multi-N-tile reuse): the +// iteration-2 partial kernel's exact code object still serializes 16 +// global_load_ubyte fragment loads with a full per-load s_waitcnt vmcnt +// ladder (vmcnt(15)..vmcnt(0)) before every single v_mmac (48 chains per +// block at SPLIT_K=4, zero LDS), so the residual limiter is the load-path +// latency ladder, not split count. This round reuses the A fragment across +// two adjacent 16x16 N-tiles in the same one-wave zero-barrier block: each +// k32 step issues one shared A fragment load plus two B fragment loads (24 +// byte loads total) and then two v_mmacs, halving the per-MMA load-ladder +// stalls while keeping the proven 656-block grid (82 N-tile pairs x 8 K +// splits = 5.47 one-wave blocks/CU) and shortening each block's MMA chain to +// K/SPLIT_K = 24 k32 steps. Reused bytes: the per-block A slice +// (A[0:16, split-slice], 12,288 B at SPLIT_K=8) is loaded once per k32 step +// and shared by both 16x16 N-tiles (A total reads halve from 16.12 MB to +// 8.06 MB, L2-hot); every B byte is still read exactly once by its owning +// tile (cold 16.12 MB stream, no B reuse). Workspace layout (iteration 4: +// [8][164][256] int32; reordered to tile-major [164][8][256] in iteration +// 18 so a tile's 8 split planes are one contiguous 8-KB region) totals +// 1,343,488 B, within the API's 16-plane budget +// (allocate_workspace yields 2,686,976 B for this shape). The combine kernel +// (SPLIT_K=8, 164 one-wave blocks) is unchanged apart from its template +// instantiation: ascending split sum keeps the int32 accumulation +// bit-identical to the k-ascending reference. +// +// Iteration 5 (HIP-only packed-weight round): the iteration-4 partial +// kernel's exact code object still compiles every fragment load of a k32 +// step to one global_load_ubyte per byte (8 A + 16 B) with a progressive +// s_waitcnt vmcnt ladder and byte reassembly before the two v_mmacs -- the +// library's int8 du_load_matrix_sync scalarizes to 8 per-lane byte loads +// regardless of layout. This round replaces the B data path with the +// lineage-validated n-major pack: launch_pack_w8a8_weight transposes the +// exact (k,n)==(6144,2624) weight once, outside the timed region and out of +// Graph capture, into packed[n*K + k] = raw[k*N + n] (same byte count, same +// buffer, graph-stable addresses; the identity copy is retained for every +// other (K,N)). In the partial kernel the B fragments are declared col_major +// against the packed buffer (du_mma.hpp lane mapping: lane l holds the 8 +// CONSECUTIVE k bytes of column n0+(l&15) at k-offset ((l>>4)*8)), and the +// row-major A fragment already maps to 8 consecutive bytes of one row, so +// each of the three fragments of a k32 step is fetched with ONE aligned 8-B +// vector load per lane (int2 + __builtin_memcpy into the fragment storage, +// the same register-to-register fill the library loader produces): 3 +// global_load_dwordx2 per step instead of 24 global_load_ubyte (47,232 vs +// 377,856 vmem_read instructions per replay, -87.5%) and one vmcnt drain per +// step instead of the per-byte ladder. Geometry (grid (82,8) = 656 one-wave +// blocks, two tiles per block, SPLIT_K=8), the zero-barrier zero-LDS +// structure, the combine kernel, the workspace planes and the exact int32 +// accumulation order are all unchanged, so outputs stay bit-identical (0 +// mismatches). The generic scalar fallback decodes the n-major pack for +// (n,k)==(2624,6144) (keeps the paired M=2 API shape byte-exact). +// +// Iteration 6 (pipeline round, B-only LDS staging): the iteration-5 code +// object's steady-state loop still serializes, per k32 step, 3 +// global_load_dwordx2 + s_waitcnt vmcnt(1) + v_mmac + s_waitcnt vmcnt(0) + +// v_mmac with the next step's loads issued only AFTER both MMAs (24 cold-B +// latency exposures per block, no cross-step overlap), while the L2/VMEM +// evidence shows B is the cold once-read stream (16.12 MB/replay, ~all L2 +// misses / HBM bytes) and A is a 98 KiB L2-resident stream re-read 82x. +// This round stages ONLY B through LDS: each 2-k32-step chunk's two B tiles +// are fetched with one aligned 16-B cooperative load per lane per tile +// (global_load_dwordx4, 64 contiguous packed B bytes per row, was 32 B), +// drained ONCE per chunk (12 drains/block, was 24), written to a 2,304-B +// bank-skewed (stride 72) LDS buffer, and re-read per fragment with 8-B +// ds_read_b64 while the next chunk's loads are issued into fresh registers +// before the current chunk is consumed (structure only; overlap is verified +// from the exact ISA after this round). A stays direct (L2-hot, 8-B int2 +// loads). vmem_read 47,232 -> 31,488 per replay (-33%); HBM read unchanged +// at the compulsory floor (~17.66 MB); LDS +32.2 MB/replay (16.12 MB +// write + 16.12 MB read); 5.47 one-wave blocks/CU preserved (LDS 2,304 B -> +// 28 blocks/CU, VGPR ~40-46 -> >= 10 waves/CU). Fragment values are +// byte-identical to iteration 5, so partials stay bit-identical (0 +// mismatches) and Graph replay with changed contents must pass. +// +// Iteration 7 (pipeline round, B-only LDS DOUBLE buffering, mandated by the +// pipeline decision rule: K=6144 >= 1024, L2 hit 40.6% < 70%, doubled LDS +// 4,608 B < 48 KiB): the iteration-6 exact code object keeps the next +// chunk's B loads in flight across the back edge but drains the next chunk's +// A loads at the BOTTOM of every body (vmcnt(1)/vmcnt(0) before the rotation +// v_movs at 0x56E8/0x56F0) -- the register rotation forces a per-chunk vmcnt +// drain on the A path (zero cross-chunk overlap for A, ~1 chunk for cold B). +// This round replaces the single LDS plane + rotation pipeline with TWO +// 2,304-B LDS planes (ping-pong) and TWO register sets alternating by chunk +// parity, each chunk's four loads (2 dwordx4 B + 2 dwordx2 A) issued TWO +// chunks ahead and consumed in place by the same-parity stage: no rotation +// copies, no bottom-of-body drains, one vmcnt drain per chunk at its own +// staging write, barriers per K step = 0 (one wave; parity replaces the +// cross-wave barrier). Fragment values are unchanged, so partials stay +// bit-identical and Graph replay with changed contents must pass. +// +// Iteration 8 (HIP-only occupancy probe, REJECTED and reverted): SPLIT_K=12 +// (984 one-wave blocks = 2.05 waves/SIMD) regressed to 50.46 us median (vs +// the iteration-7 best 47.59 us), falsifying per-SIMD co-residency as the +// binding limiter: the residual is per-block serial cost, not the number of +// co-resident waves. +// +// Iteration 9 (HIP-only pipeline round, stage width 64 K -> 128 K): the +// iteration-7 partial kernel still pays one near-full cold-B vmcnt drain per +// 64-K chunk (12 chunks/block) with only ~1 pass of work between issue and +// drain. This round WIDENS each staged chunk to FOUR k32 steps (128 K) while +// keeping the exact double-buffer ping-pong structure (two 136-stride LDS +// planes, two parity register sets, two-chunk prefetch, 3 passes of 2 +// chunks, grid (82,8) = 656 one-wave blocks, zero barriers): drains halve +// 12 -> 6, the in-flight window doubles (each set now carries 2 x 32 B of B +// + 32 B of A per lane), LDS becomes 8,704 B/block (7 blocks/CU headroom vs +// the scheduled 5.47), VGPR ~58 -> ~80-95 (no spills expected), vmem_read +// stays 31,488/replay, LDS instructions stay 31,488/replay, HBM read stays +// at the compulsory ~17.66 MB floor. B fragment VALUES are byte-identical +// (same packed bytes staged, same 8-B ds_read_b64 fills), the k-ascending +// int32 accumulation order is unchanged, so partials and the combine output +// stay bit-identical (0 mismatches) and Graph replay with changed contents +// must pass. +// +// Iteration 10 (HIP-only pipeline round, stage width 128 K -> 192 K): the +// exact iteration-9 code object (digest b766fc17..., 96 VGPR/32 SGPR/8,704 B +// LDS/0 scratch, profiled 37.44 us) emits ONE full progressive vmcnt drain +// per PASS at the pass top (s_waitcnt vmcnt(7)..vmcnt(0) ladder at 0x571C- +// 0x574C covering all 16 loads issued in the previous pass; the B staging +// writes and the A capture after it are wait-free) -- so the real per-block +// drain-event count is 3 (passes), not 6 (chunks), each exposing the cold-B +// round trip of 2 x 128-K chunks against an issue-to-drain window of only +// ~1 pass (~16 MMAs + 8 ds_read2 + 8 stagings + 16 loads). This round +// WIDENS each staged chunk to SIX k32 steps (192 K) while keeping the exact +// double-buffer ping-pong structure (two 200-stride LDS planes, two parity +// register sets, two-chunk prefetch, 2 passes of 2 chunks, grid (82,8) = +// 656 one-wave blocks, zero barriers): drain events drop 3 -> 2 per block +// and each event now covers ~50% more in-flight bytes (2 x 192-K chunks = +// 12,288 B of B + 6,144 B of A) against an issue-to-drain window of one +// full pass (24 MMAs + 12 ds_read2 + 12 stagings + 24 loads). LDS becomes +// s_b[2][2][16][200] = 12,800 B/block (stride 200 = 192 + 8: 50 dwords = 18 +// mod 32 -> row bank bases 18*n mod 32 for n = 0..15, the same 16-distinct- +// base conflict-free pattern as iteration 9's 2*n mod 32; the plane offset +// 6,400 B = 1,600 dwords is 0 mod 32, so both planes keep the same pattern). +// LDS 12,800 B -> 5 blocks/CU of headroom vs the scheduled 5.47 (5 slots +// still >= the 4 SIMDs per CU, so co-execution is unchanged; the 56-block +// shortfall only delays block START by at most one block-launch latency on +// those CUs); VGPR ~96 -> ~110-130 (in-flight set registers 48 -> 72; no +// spills expected, 65,536/(130*64) ~= 7.9 waves/CU >= 5.47). vmem_read stays +// 31,488/replay (4 chunks x 6 dwordx4 B + 6 dwordx2 A = 48 loads/block, same +// total), LDS instructions stay 31,488/replay (the same staged bytes and the +// same 8-B ds_read fills), LDS traffic unchanged, HBM read unchanged at the +// compulsory ~17.66 MB floor. Fragment VALUES are byte-identical (the staged +// rows hold the same packed bytes; the per-step 8-B ds_read_b64 fills are +// unchanged), the k-ascending int32 accumulation order is unchanged, so the +// partial planes and the combine output stay bit-identical (0 mismatches) +// and Graph replay with changed contents must pass. +// +// Iteration 13 (combine latency consolidation): iteration 12's single-drain +// partial kernel landed at 43.64 us median (43.81 -> 43.64, +0.38%), +// falsifying the cold-B drain chain as the residual limiter: the operator is +// at the once-read-B memory-pipe floor (PMC: 19.09 MB of traffic per replay +// at ~436 GB/s derived HBM bandwidth; partial kernel profiled 41.92 us, +// combine kernel profiled 4.48 us with 31.4% L2 hit rate -- the 16.1-MB B +// stream thrashes L2 ahead of the combine). The exact combine ISA shows the +// four weight_scale elements loaded as FOUR serialized 4-B loads, each +// waited to completion before its multiply and its own 2-B store: four +// round trips per lane. This round consolidates the combine epilogue: ONE +// aligned 16-B float4 scale load (issued before the 8 plane reads) + ONE +// 8-B packed bf16 store per lane. The partial kernel, its 656-block +// geometry, the workspace, and the per-element arithmetic (int32 ascending +// split sum, fp32 left-to-right scale multiplies, RNE bf16) are unchanged, +// so the combine output stays bit-identical (0 mismatches) and Graph replay +// with changed contents must pass. +// +// Iteration 18 (workspace plane order [tile][split]): the accepted best is +// 43.00 us median and the iteration-14 PMC shows the operator is at the +// once-read-B DRAM floor: 19.08 MB of traffic per replay (17.65 MB read + +// 1.43 MB write) at 443.7 GB/s derived; partial kernel profiled 36.32 us +// (16.27 MB read + 1.34 MB write), combine kernel profiled 3.52 us +// (1.38 MB read, 31.4% L2 hit rate -- the 16.12-MB B stream thrashes L2, so +// ~0.95 MB of the combine's plane re-reads are DRAM-cold). The fusion +// attempts (iteration 15: x-major grid + last-arriver tail, 48.24 us; +// iteration 16: split-major grid + last-arriver tail, 43.96 us) both lost to +// the two-launch structure: the tail's plane re-reads are DRAM-cold under +// the x-major spread, and the tile-pair-major B order perturbs the main +// loop; no cache-policy hint is available (raw inline asm banned; the DTK's +// __ldcs/__ldcg are plain dereferences, unimplemented for int8/int4). The +// remaining exposed cost inside the combine kernel is the DRAM locality of +// its plane reads: with the [split][tile] layout each of the 164 combine +// blocks reads its tile's eight 1-KB planes at 164-KB strides -- 1,312 +// scattered 1-KB cold streams. CHANGE: swap the plane order to [tile][split] +// (plane (tile, split) at (tile*SPLIT_K + split)*256 int32), so each combine +// block reads its tile's eight planes as ONE contiguous 8-KB region: 164 +// sequential 8-KB streams instead of 1,312 scattered 1-KB streams (8x fewer +// concurrent streams, sequential-line access). The partial kernel's stores +// carry byte-identical values (same planes, same mem_row_major fragment +// stores; only the base index changes -- each block still writes two +// line-complete 1-KB planes), the ascending split-sum order and the +// per-element fp32 scale arithmetic are unchanged, so the combine output +// stays bit-identical (0 mismatches) and Graph replay with changed contents +// must pass. Exact shape guards, the M=2 scalar fallback, the n-major pack, +// the generic scalar fallback and the workspace guard (1,343,488 B, within +// the API's 16-plane budget) are untouched. Resource deltas: none (no +// register/LDS/grid changes; combine vmem_read stays 1,640/replay, vmem_write +// 164/replay; partial vmem_write stays 5,248/replay). Falsifiable +// prediction: if the combine's scattered cold 1-KB reads are part of its +// ~3.5 us wall, the contiguous 8-KB streams cut its DRAM service time and +// the median drops from 43.00 toward ~42.0-42.6 us with p90 below the +// current 43.07; a flat ~43.0-43.6 median falsifies it (the combine is +// launch/ramp-bound or already at the machine DRAM rate -- the residual is +// the once-read-B floor plus the combine launch, and the next round attacks +// the launch gap or re-tests byte reduction); a median above ~44.5 would +// show the store-address scatter hurt the partial kernel's write path and +// the round must be reverted. +// +// Contract (see int8_w8a8_gemm_api.py): +// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) +// * x_scale[m, 0] * weight_scale[n, 0] ) +// x_q : int8 [M, K] row-major +// weight : int8 [K, N] row-major logical; the packed_weight buffer +// holds packed[n*K + k] = weight[k*N + n] (n-major transpose) +// for the exact (k,n)==(6144,2624) and the identity [K, N] +// copy for every other (K, N) (iteration 5) +// x_scale : fp32 [M, 1] +// weight_scale: fp32 [N, 1] +// out : bf16 [M, N] row-major, caller-provided +// workspace : uint8, caller-allocated before Graph capture; holds the +// [164][8][256] int32 split-K partial planes (tile-major, +// iteration 18: a tile's 8 split planes are one contiguous +// 8-KB region) for the exact shape (1,343,488 B, within the +// API 16-plane budget) +// +// The timed operator (launch_w8a8_gemm) performs no allocation, compilation, +// autotuning, packing, host/device synchronization, or default-stream +// launch: it only validates the exact shape guard and launches the DUMMA +// kernel (or the generic scalar fallback) on the caller-provided stream. + +#include +#include +#include + +#include + +namespace { + +// gfx928 wavefront = 64 lanes; block sizes must be multiples of 64. +constexpr int kScalarBlockThreads = 128; +constexpr int kCopyBlockThreads = 256; + +// DUMMA INT8 primitive: m16n16k32, int8 x int8 -> int32. +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; +constexpr int kDummaBlockThreads = 64; // one wavefront per output tile + +// --------------------------------------------------------------------------- +// Generic scalar INT8 dot-product GEMM (fallback for every shape). +// +// One thread computes one output element out[m, n]. The full K loop runs in +// int32 (exact for every assigned K; max K=6144 -> |dot| <= 6144*127*127 +// ~= 9.9e7 << 2^31), then the two fp32 scales are applied and the result is +// stored as bf16. +// +// Thread t maps to (row = t / n, col = t % n): adjacent lanes land on +// adjacent addresses in the fastest-changing N dimension, so the per-K-row +// B reads (b[k*n + col]) stay lane-coalesced and the A reads are broadcast +// within a row. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_scalar_gemm_kernel(const int8_t* __restrict__ a, // [M, K] + const int8_t* __restrict__ b, // [K, N] + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [M, N] + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t t = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (t >= total) { + return; + } + const int row = static_cast(t / n); + const int col = static_cast(t - static_cast(row) * n); + + const int8_t* a_row = a + static_cast(row) * k; + // Iteration 5: for the exact (n, k) == (2624, 6144) the packed_weight + // buffer holds the n-major transpose packed[col * k + kk] = raw[kk * n + + // col] (produced once out of the timed region), so the fallback decodes + // that layout there (keeps the paired M=2 API shape with the same (N, K) + // byte-exact); every other (k, n) keeps the identity copy (raw row-major + // [K, N]). + const bool packed_nmajor = (n == 2624 && k == 6144); + const int8_t* b_ptr = packed_nmajor ? b + static_cast(col) * k + : b + col; + const int b_stride = packed_nmajor ? 1 : n; + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_ptr[0]); + b_ptr += b_stride; + } + + // Same arithmetic order as the harness reference + // (dot.to(fp32) * x_scale * weight_scale.T): fp32 multiply left-to-right, + // then round-to-nearest-even bf16. + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[t] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Iteration 2 + iteration 4 + iteration 5: split-K partial + combine +// kernels (exact assigned shape M=16, N=2624, K=6144). +// +// Partial kernel geometry (iteration 4): one 64-lane wavefront per (two +// adjacent 16x16 N-tiles, K split) pair; grid = (N/32, SPLIT_K) = (82, 8) -> +// 656 one-wave zero-barrier blocks, each computing two 16x16 tiles over its +// own contiguous K slice (K/SPLIT_K = 768 = 24 k32 steps, k-ascending within +// the slice). The A fragment of each k32 step is loaded once and reused by +// both tiles (A slice bytes are the reused bytes); B fragments are loaded +// per tile, each B byte exactly once. No LDS, no __syncthreads, no atomics: +// every block writes only its own two 256-int32 planes, so the partial grid +// is deterministic and Graph-safe. Partial planes are published with +// du_store_matrix_sync (mem_row_major), the canonical accumulator-fragment +// store matching the verified gfx928 ownership (row = lane & 15, +// col = (lane >> 4) + 4*i). +// +// Iteration 5 (packed-weight round): the B fragments are now declared +// col_major against the packed [N, K] n-major weight +// (packed[n*K + k] = raw[k*N + n], produced once out of the timed region; +// see launch_pack_w8a8_weight). With the du_mma.hpp lane mappings, every +// lane's fragment bytes are 8 CONSECUTIVE bytes of one row/column, so the +// three fragments of a k32 step are fetched with one aligned 8-B vector +// load per lane each (int2 + __builtin_memcpy into the fragment storage -- +// the same register-to-register fill the library loader produces): 3 +// global_load_dwordx2 per step instead of 24 global_load_ubyte with the +// per-byte vmcnt ladder. Fragment VALUES are identical to the iteration-4 +// library loads (row-major A: 8 bytes of row (lane&15) at k-offset +// ((lane>>4)*8); col_major B: 8 k bytes of column n0+(lane&15) at the same +// k-offset), so the k-ascending int32 accumulation and the partial planes +// stay bit-identical. +// +// Iteration 6 (pipeline round): the B data path now stages through LDS (see +// the kernel body): per 2-k32-step chunk, each of the two B tiles is fetched +// with ONE aligned 16-B cooperative load per lane (int4 / dwordx4, 64 +// contiguous packed bytes per row), drained once, written to s_b (stride +// 72), and the per-step 8-B fragments are re-read from LDS with ds_read_b64 +// while the next chunk's loads are already in flight. A stays direct (8-B +// int2, L2-hot). Fragment VALUES are unchanged, so partials stay +// bit-identical. +// +// Iteration 7 (pipeline round, double buffering): the single LDS plane and +// the rotation-copy pipeline are replaced by TWO ping-pong LDS planes +// (4,608 B) and TWO parity-alternating register sets with a two-chunk +// prefetch (see the kernel body): every load stays in flight across the +// back edge and is drained only at its own staging write. Fragment VALUES +// are unchanged, so partials stay bit-identical. +// +// Combine kernel geometry (unchanged from iteration 2, SPLIT_K=8): one +// 64-lane wavefront per 16x16 output tile (164 one-wave blocks, zero +// barriers); lane l owns the four contiguous row-major elements e = 4*l + i +// (row = e >> 4, col = e & 15), sums them across the SPLIT_K planes in +// ascending split order (exact int32 accumulation), applies +// x_scale[row] * weight_scale[col], and stores bf16. +// +// Iteration 13 (combine latency consolidation): the exact iteration-12 +// code object's combine kernel (164 blocks, 48 VGPR, profiled 4.48 us) +// issues the 8 plane dwordx4 loads up front, but loads the four +// weight_scale elements with four SEPARATE 4-B loads, each waited to +// completion (vmcnt(0)) before its multiply and its own 2-B bf16 store -- +// four serialized round trips per lane. This round replaces them with ONE +// aligned 16-B float4 scale load (issued before the plane reads) and ONE +// 8-B packed store per lane; per-element arithmetic (int32 ascending sum, +// fp32 left-to-right scale multiplies, RNE bf16) is byte-for-byte +// unchanged, so the combine output stays bit-identical. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kDummaBlockThreads) void +w8a8_dumma_m16n16k32_splitk_partial_kernel( + const int8_t* __restrict__ a, // [16, K] row-major (logical A) + const int8_t* __restrict__ b, // packed [N, K] n-major (iter 5) + int* __restrict__ partials, // [N/16][SPLIT_K][256] int32 + // (tile-major, iteration 18) + int n, + int k) { + constexpr int kNTiles = 2; // adjacent 16x16 N-tiles per block + const int tile0 = static_cast(blockIdx.x) * kNTiles; // 0..N/16-2 + const int split = static_cast(blockIdx.y); // 0..SPLIT_K-1 + const int n0 = tile0 * kDummaTileN; + const int k_len = k / SPLIT_K; // K % SPLIT_K == 0 (exact shape guard) + const int k_lo = split * k_len; + const int lane = static_cast(threadIdx.x); // 0..63 + + // Iteration 7 (pipeline round, B-only LDS DOUBLE buffering): the exact + // iteration-6 code object (source digest f04e11e9..., 41 VGPR, single + // 2,304-B LDS plane) keeps the next chunk's B loads in flight across the + // back edge, but drains the next chunk's A loads at the BOTTOM of every + // body (s_waitcnt vmcnt(1)/vmcnt(0) at 0x56E8/0x56F0 immediately before + // the rotation v_movs v24->v22/v26->v20): the register rotation forces a + // per-chunk vmcnt drain on the A path, so A (L2-hot) is re-exposed to L2 + // latency 12x per block with ZERO cross-chunk overlap, and the cold B + // stream gets only ~1 chunk of hiding. This round double-buffers the B + // staging: TWO 2,304-B LDS planes (4,608 B/block total; 64 KiB/4,608 B = + // 14 blocks/CU of LDS headroom vs the scheduled 5.47) and TWO register + // sets alternating by chunk parity, each chunk's four loads (2 dwordx4 B + // + 2 dwordx2 A) issued TWO chunks ahead and consumed IN PLACE by the + // matching parity stage -- there are no rotation copies, so no bottom-of- + // body vmcnt drains: every load stays in flight across the back edge and + // the single vmcnt drain per chunk lands only at its own LDS staging + // write. A now gets ~2 chunks of latency hiding (was ~0), B ~2 (was ~1). + // Row stride 72 = 64 + 8 keeps every ds_read_b64/ds_write_b64 address 8-B + // aligned AND gives 18*n mod 32 distinct row bank bases for n = 0..15 (no + // LDS bank conflicts on either the fragment reads or the b64 staging + // writes; the ping-pong plane offset 2,304 B = 576 dwords is 0 mod 32, so + // both planes keep the same conflict-free bank pattern). One wavefront per + // block: no __syncthreads, no barriers -- the ping-pong parity replaces + // the cross-wave barrier a multi-wave block would need. + // + // Iteration 9 (HIP-only pipeline round, stage width 64 K -> 128 K): the + // iteration-8 occupancy probe (SPLIT_K=12, 984 blocks = 2.05 waves/SIMD) + // REGRESSED to 50.46 us median (vs the iteration-7 best 47.59), falsifying + // per-SIMD co-residency as the limiter and pointing back at the per-block + // serial cost: each 64-K chunk still pays one near-full cold-B vmcnt drain + // (12 chunks/block), and the issue-to-drain window is only ~1 pass of work + // (~8 MMAs + 4 ds_read2 + 2 stagings) against a ~600+ cycle HBM cold + // latency. This round keeps the winning 656-block geometry and the exact + // double-buffer ping-pong structure but WIDENS each staged chunk to FOUR + // k32 steps (128 K): per chunk each B tile is fetched with TWO aligned + // 16-B cooperative loads per lane (global_load_dwordx4 at b_wide and + // b_wide+64; 128 CONTIGUOUS packed B bytes per row, was 64), the drain + // count halves (12 -> 6), and the in-flight window doubles (each reload + // set now carries 2 x 32 B of B + 32 B of A per lane). LDS planes become + // s_b[2][2][16][136] = 8,704 B/block (stride 136 = 128 + 8: 34 dwords = 2 + // mod 32 -> row bank bases 2*n mod 32, still 16 distinct bases for n = + // 0..15, so reads (ds_read2_b64 pairs at +0/+32 and +64/+96) and staging + // writes (ds_write2_b64 at b_wide and b_wide+64) stay conflict-free; the + // plane offset 4,352 B = 1,088 dwords is 0 mod 32, so both planes keep the + // same pattern). LDS 8,704 B -> 7 blocks/CU of headroom vs the scheduled + // 5.47; VGPR ~58 -> ~80-95 (in-flight set registers 24 -> 48; no spills + // expected, 65,536/(95*64) ~= 10 waves/CU >= 5.47). vmem_read stays + // 31,488/replay (6 chunks x 4 dwordx4 B + 4 dwordx2 A = 48 loads/block, + // same total), LDS instructions stay 31,488/replay (24 ds_write2 + 24 + // ds_read2 per block), LDS traffic unchanged (every B byte still staged + // exactly once), HBM read unchanged at the compulsory ~17.66 MB floor. + // Fragment VALUES are byte-identical (the staged rows hold the same packed + // bytes; the per-step 8-B ds_read_b64 fills are unchanged), the k-ascending + // int32 accumulation order is unchanged, so the partial planes and the + // combine output stay bit-identical (0 mismatches) and Graph replay with + // changed contents must pass. + // + // Iteration 10 (stage width 128 K -> 192 K): the exact iteration-9 code + // object emits one full progressive vmcnt ladder per PASS at the pass top + // (vmcnt(7)..vmcnt(0) at 0x571C-0x574C, covering all 16 loads issued in the + // previous pass; the B staging writes and the A capture after it are + // wait-free), so the real drain-event count is 3 per block (passes), each + // exposing the cold-B round trip of 2 x 128-K chunks. This round WIDENS + // each staged chunk to SIX k32 steps (192 K) with the same double-buffer + // ping-pong structure: per chunk each B tile is fetched with THREE aligned + // 16-B cooperative loads per lane (global_load_dwordx4 at b_wide, + // b_wide+64, b_wide+128; 192 CONTIGUOUS packed B bytes per row, was 128), + // drain events drop 3 -> 2 per block, and each pass now carries 24 loads + // (2 chunks x 6 dwordx4 B + 6 dwordx2 A) against an issue-to-drain window + // of one full pass (24 MMAs + 12 ds_read2 + 12 stagings + 24 loads). LDS + // planes become s_b[2][2][16][200] = 12,800 B/block (stride 200 = 192 + 8: + // 50 dwords = 18 mod 32 -> row bank bases 18*n mod 32, still 16 distinct + // bases for n = 0..15, so reads (ds_read2_b64 at +0/+32/+64/+96/+128/+160) + // and staging writes (ds_write2_b64 at b_wide, b_wide+64, b_wide+128) stay + // conflict-free; the plane offset 6,400 B = 1,600 dwords is 0 mod 32, so + // both planes keep the same pattern). LDS 12,800 B -> 5 blocks/CU of + // headroom vs the scheduled 5.47 (5 slots still >= the 4 SIMDs per CU); + // VGPR ~96 -> ~110-130 (in-flight set registers 48 -> 72; no spills + // expected, 65,536/(130*64) ~= 7.9 waves/CU >= 5.47). vmem_read stays + // 31,488/replay (4 chunks x 12 loads = 48 loads/block, same total), LDS + // instructions stay 31,488/replay (24 ds_write2 + 24 ds_read2 per block), + // LDS traffic unchanged (every B byte still staged exactly once), HBM read + // unchanged at the compulsory ~17.66 MB floor. Fragment VALUES are + // byte-identical (the staged rows hold the same packed bytes; the per-step + // 8-B ds_read_b64 fills are unchanged), the k-ascending int32 accumulation + // order is unchanged, so the partial planes and the combine output stay + // bit-identical (0 mismatches) and Graph replay with changed contents must + // pass. + alignas(16) __shared__ int8_t s_b[2][kNTiles][kDummaTileN][200]; + + du::dumma::DUFragment + a_frag0; + du::dumma::DUFragment + a_frag1; + du::dumma::DUFragment + a_frag2; + du::dumma::DUFragment + a_frag3; + du::dumma::DUFragment + a_frag4; + du::dumma::DUFragment + a_frag5; + du::dumma::DUFragment + b_frag0; + du::dumma::DUFragment + b_frag1; + du::dumma::DUFragment + acc0_frag; + du::dumma::DUFragment + acc1_frag; + du::dumma::du_fill_fragment(acc0_frag, 0); + du::dumma::du_fill_fragment(acc1_frag, 0); + + // Lane roles are fixed by the m16n16k32 fragment layouts in du_mma.hpp + // (see the iteration-5 note): row-major A lane l holds 8 consecutive bytes + // of row (l&15) at k-offset ((l>>4)*8); col_major B over the packed buffer + // holds 8 consecutive k bytes of column n0+(l&15) at the same k-offset. + // Iteration 6 widens the B LOADS to 16 B per lane (int4 / dwordx4): lane + // (row, c) fetches packed bytes [c*16, c*16+16) of its column across a + // 2-step (64-K) chunk, i.e. 64 CONTIGUOUS packed B bytes per row per + // instruction (was 32 B per row per step), stages them into LDS, and re- + // reads the per-step 8-B fragment from LDS -- so one global drain covers + // the whole 2-step chunk and the MMA path is decoupled from global latency. + const int a_row = lane & 0xf; + const int a_col = (lane >> 4) << 3; // {0, 8, 16, 24} (A fragment offset) + const int b_row = lane & 0xf; // n index within the 16-col tile + const int b_col = (lane >> 4) << 3; // {0, 8, 16, 24} (B fragment offset) + const int b_wide = (lane >> 4) << 4; // {0, 16, 32, 48} (B 16-B load offset) + const int8_t* a_base = + a + static_cast(a_row) * k + k_lo + a_col; + // B chunk bases are 16-B aligned: k = 6144 is a multiple of 16, n0 and + // k_lo are multiples of 16/64, and b_wide is a multiple of 16. + const int8_t* b0_base = + b + static_cast(n0 + b_row) * k + k_lo; + const int8_t* b1_base = b0_base + static_cast(kDummaTileN) * k; + + constexpr int kChunkK = 6 * kDummaTileK; // 192 K per staged chunk (iter 10) + + // Iteration 12 (single cold-B drain per block): the exact iteration-10 + // code object (source digest 6ca5190e..., 130 VGPR/12,800 B LDS/0 scratch, + // compile_cache_key 34ee6cfd...) emits TWO cold-B vmcnt drains per block: + // (1) the prologue ladder s_waitcnt vmcnt(13)..vmcnt(0) at 0x575C-0x57A4 + // covering chunks 0+1 with a ZERO issue-to-drain window, and (2) the + // pass-0 bottom ladder s_waitcnt vmcnt(5)..vmcnt(0) at 0x5AD8-0x5B00 + // covering the 24 reloaded pass-1 loads (chunks 2+3) before the 24 + // rotation v_movs -- so every block still serially exposes TWO cold-B + // round trips. This round issues ALL FOUR chunks' B loads (24 + // global_load_dwordx4) plus the chunks 0/1 A loads in the prologue, so + // the FIRST staging consumer drains the whole K-slice with ONE vmcnt + // ladder: one cold-B round trip per block, the pass-1 staging (chunks + // 2+3) becomes wait-free, and the 24 rotation v_movs disappear. The loop + // is UNROLLED into two explicit passes because the B register sets are no + // longer reloaded in place (the exact-shape launcher guard guarantees + // k_len = 1536 at SPLIT_K = 4, i.e. chunks == 8 -- the only launch + // configuration of this kernel; iteration 20 adds the second 4-chunk + // body below). A stays L2-pipelined exactly as in + // iteration 10: the chunk 2/3 A loads are issued inside pass 0 (after the + // chunk 0/1 captures, reusing the dead set-0/1 A registers) and are + // consumed by the pass-1 captures one full pass later. Resource deltas: + // in-flight load registers 72 -> 120 (24 dwordx4 B = 96 regs + 12 int2 A + // = 24 regs at the prologue), expected arch_vgpr ~136 -> ~150-175 + // (65,536/(175*64) ~= 5.85 waves/CU >= the LDS-capped 5 slots, no spills + // at < 256 VGPR, scratch must stay 0), LDS + // stays 12,800 B (5 blocks/CU, unchanged), grid stays (82,8) = 656 + // one-wave zero-barrier blocks. vmem_read stays 31,488/replay (the same + // 48 loads/block: 24 dwordx4 B + 24 dwordx2 A -- 36 issued in the + // prologue, 12 A issued in pass 0), LDS instructions stay 31,488/replay + // (24 ds_write2 + 24 ds_read2 per block), LDS traffic unchanged (every B + // byte still staged exactly once), HBM read unchanged at the compulsory + // ~17.66 MB floor. Exactness is preserved: the staged LDS rows hold + // byte-for-byte the same packed B bytes (same planes, same b_wide / + // b_wide+64 / b_wide+128 staging addresses, same 8-B ds_read_b64 fragment + // fills), the A fragments are unchanged, and the k-ascending int32 + // accumulation order within each split slice (chunks 0,1,2,3) is + // unchanged, so the partial planes and the combine output stay + // bit-identical (0 mismatches) and Graph replay with changed contents + // must pass. Falsifiable prediction: if the residual limiter is the + // per-block serial cold-B drain chain (the mechanism behind iterations 9 + // and 10), removing the second cold-B round trip should land the median + // in roughly 42.0-43.3 us with p90 below the current 43.92 (combine + // ~2.5-3 us unchanged); a flat ~43.8 us median falsifies the drain-chain + // hypothesis (the residual would be the memory-pipe / launch+combine + // floor, and the next round would attack the combine or the B request + // partition); a median above ~45 us would show register-pressure or + // scheduling loss (e.g. the compiler spilling the 96 in-flight B + // registers or re-inserting a per-pass wait at the pass-1 staging). + // + // Iteration 20 (SPLIT_K 8 -> 4 byte-reduction round): the operator is at + // the once-read-B floor -- the counter-derived whole-operator HBM rate is + // ~448 GB/s and the partial kernel alone moves 16.22 MB read + 1.34 MB + // write in ~36.5 us (~480 GB/s effective), so the iteration-18/19 + // falsification branches' remaining lever is BYTE REDUCTION: the workspace + // plane traffic (1.34 MB write + 1.31 MB combine re-read at SPLIT_K=8) + // halves at SPLIT_K=4. The partial kernel keeps the exact iteration-12 + // single-drain-per-group structure but covers EIGHT 192-K chunks (k_len = + // 1536): body 1 (chunks 0-3) is the accepted iteration-12 code path + // verbatim; body 2 (chunks 4-7) re-issues its 24 dwordx4 B loads plus the + // chunks 4/5 A loads into the body-1 dead registers (same 96-B + 48-B live + // load sets -> arch_vgpr stays ~120-140, LDS stays 12,800 B, no spills) + // and drains them with ONE ladder at its chunk-4 staging -- the block's + // second cold-B round trip, but device-wide the drain count is unchanged + // (328 blocks x 2 = 656 vs the accepted 656 blocks x 1), and drain count + // was measured ~flat at the memory ceiling (iteration 10: 3 drains 43.81 + // us vs iteration 12: 1 drain 43.64 us). Grid becomes (82,4) = 328 one-wave + // zero-barrier blocks = 2.73 blocks/CU, a SINGLE wave (all resident), vs + // the old 656-block 1.31-wave grid. The combine kernel is templated on + // SPLIT_K and needs NO change: 164 one-wave blocks each sum the four + // planes in ascending split order with the fused x_scale * weight_scale + // bf16 epilogue. Exactness: the int32 accumulation stays k-ascending + // within each split slice (chunks 0..7 in order, 48 k32 steps) and the + // combine sums s=0..3 ascending, so the output stays bit-identical (0 + // mismatches) and Graph replay with changed contents must pass. The + // [tile][split] workspace layout (iteration 18), the n-major pack + // (iteration 5), the exact shape guards, the M=2 scalar fallback and the + // two-launch Graph structure are untouched (the pack and the scalar + // fallback do not depend on SPLIT_K). Resource deltas: partial vmem_read + // 31,488 -> 23,616/replay (48 loads/block x 328), vmem_write 5,248 -> + // 2,624/replay (plane bytes halved), LDS instructions 31,488 -> 15,744/ + // replay, LDS stays 12,800 B/block, scratch stays 0; combine vmem_read + // 1,640 -> ~984/replay, vmem_write stays 164/replay; HBM traffic drops + // 19,080,640 -> ~17,771,008 B per replay (partial read 16.22 MB compulsory + // + partial write 0.66 MB + combine read 0.66 MB + out 0.08 MB). + // Falsifiable prediction: if the machine's effective HBM rate holds at + // ~480 GB/s for the smaller grid, the 1.31-MB plane-traffic cut lands the + // median in roughly 39.5-40.8 us with p90 below the current 42.61 (partial + // toward ~35.2 us profiled, combine toward ~2.5 us, launch gap unchanged); + // a flat ~42.5 us median falsifies it (the smaller grid or the second + // per-block drain costs as much as the byte cut -- the machine's effective + // rate drops with grid size -- and the next round reverts or tries + // SPLIT_K=2 / the launch gap); a median above ~44.5 us shows the second + // zero-window drain serialized per-block latency or a register/scheduling + // loss and the round must be reverted. + + // ---- Prologue: all four chunks' B loads (24 dwordx4) + chunks 0/1 A. + // The first staging consumer below drains the whole set with ONE vmcnt + // ladder: a single cold-B round trip per block (was two). ---- + int4 b0s0a = *reinterpret_cast(b0_base + b_wide); + int4 b0s0b = *reinterpret_cast(b0_base + b_wide + 64); + int4 b0s0c = *reinterpret_cast(b0_base + b_wide + 128); + int4 b1s0a = *reinterpret_cast(b1_base + b_wide); + int4 b1s0b = *reinterpret_cast(b1_base + b_wide + 64); + int4 b1s0c = *reinterpret_cast(b1_base + b_wide + 128); + int4 b0s1a = *reinterpret_cast(b0_base + kChunkK + b_wide); + int4 b0s1b = + *reinterpret_cast(b0_base + kChunkK + b_wide + 64); + int4 b0s1c = + *reinterpret_cast(b0_base + kChunkK + b_wide + 128); + int4 b1s1a = *reinterpret_cast(b1_base + kChunkK + b_wide); + int4 b1s1b = + *reinterpret_cast(b1_base + kChunkK + b_wide + 64); + int4 b1s1c = + *reinterpret_cast(b1_base + kChunkK + b_wide + 128); + int4 b0s2a = *reinterpret_cast(b0_base + 2 * kChunkK + b_wide); + int4 b0s2b = + *reinterpret_cast(b0_base + 2 * kChunkK + b_wide + 64); + int4 b0s2c = + *reinterpret_cast(b0_base + 2 * kChunkK + b_wide + 128); + int4 b1s2a = *reinterpret_cast(b1_base + 2 * kChunkK + b_wide); + int4 b1s2b = + *reinterpret_cast(b1_base + 2 * kChunkK + b_wide + 64); + int4 b1s2c = + *reinterpret_cast(b1_base + 2 * kChunkK + b_wide + 128); + int4 b0s3a = *reinterpret_cast(b0_base + 3 * kChunkK + b_wide); + int4 b0s3b = + *reinterpret_cast(b0_base + 3 * kChunkK + b_wide + 64); + int4 b0s3c = + *reinterpret_cast(b0_base + 3 * kChunkK + b_wide + 128); + int4 b1s3a = *reinterpret_cast(b1_base + 3 * kChunkK + b_wide); + int4 b1s3b = + *reinterpret_cast(b1_base + 3 * kChunkK + b_wide + 64); + int4 b1s3c = + *reinterpret_cast(b1_base + 3 * kChunkK + b_wide + 128); + int2 a0s0 = *reinterpret_cast(a_base); + int2 a1s0 = *reinterpret_cast(a_base + kDummaTileK); + int2 a2s0 = *reinterpret_cast(a_base + 2 * kDummaTileK); + int2 a3s0 = *reinterpret_cast(a_base + 3 * kDummaTileK); + int2 a4s0 = *reinterpret_cast(a_base + 4 * kDummaTileK); + int2 a5s0 = *reinterpret_cast(a_base + 5 * kDummaTileK); + int2 a0s1 = *reinterpret_cast(a_base + kChunkK); + int2 a1s1 = + *reinterpret_cast(a_base + kChunkK + kDummaTileK); + int2 a2s1 = + *reinterpret_cast(a_base + kChunkK + 2 * kDummaTileK); + int2 a3s1 = + *reinterpret_cast(a_base + kChunkK + 3 * kDummaTileK); + int2 a4s1 = + *reinterpret_cast(a_base + kChunkK + 4 * kDummaTileK); + int2 a5s1 = + *reinterpret_cast(a_base + kChunkK + 5 * kDummaTileK); + + // ---- Pass 0: chunks 0 and 1. Each B tile's 48 B per lane (three dwordx4 + // halves at b_wide, b_wide+64, b_wide+128) land in the 192-B row: six + // ds_write2_b64 per chunk, all conflict-free (row bank bases 18*n mod 32). + { + // Stage chunk 0 (set 0) into LDS plane 0 -- the single cold-B drain of + // the whole block lands here (first consumer of the 36 prologue loads). + { + int2* d0 = reinterpret_cast(&s_b[0][0][b_row][b_wide]); + int2* d1 = reinterpret_cast(&s_b[0][1][b_row][b_wide]); + d0[0] = *reinterpret_cast(&b0s0a); + d0[1] = *reinterpret_cast( + reinterpret_cast(&b0s0a) + 2); + d1[0] = *reinterpret_cast(&b1s0a); + d1[1] = *reinterpret_cast( + reinterpret_cast(&b1s0a) + 2); + int2* d0b = reinterpret_cast(&s_b[0][0][b_row][b_wide + 64]); + int2* d1b = reinterpret_cast(&s_b[0][1][b_row][b_wide + 64]); + d0b[0] = *reinterpret_cast(&b0s0b); + d0b[1] = *reinterpret_cast( + reinterpret_cast(&b0s0b) + 2); + d1b[0] = *reinterpret_cast(&b1s0b); + d1b[1] = *reinterpret_cast( + reinterpret_cast(&b1s0b) + 2); + int2* d0c = reinterpret_cast(&s_b[0][0][b_row][b_wide + 128]); + int2* d1c = reinterpret_cast(&s_b[0][1][b_row][b_wide + 128]); + d0c[0] = *reinterpret_cast(&b0s0c); + d0c[1] = *reinterpret_cast( + reinterpret_cast(&b0s0c) + 2); + d1c[0] = *reinterpret_cast(&b1s0c); + d1c[1] = *reinterpret_cast( + reinterpret_cast(&b1s0c) + 2); + } + // Capture chunk 0's six A fragments (steps 0..5) BEFORE set 0's A + // registers are overwritten below by the chunk-2 A issue; these captures + // are the first consumers of the chunks 0/1 A loads (L2-hot, complete + // during the single cold-B drain). + __builtin_memcpy(a_frag0.x, &a0s0, 8); + __builtin_memcpy(a_frag1.x, &a1s0, 8); + __builtin_memcpy(a_frag2.x, &a2s0, 8); + __builtin_memcpy(a_frag3.x, &a3s0, 8); + __builtin_memcpy(a_frag4.x, &a4s0, 8); + __builtin_memcpy(a_frag5.x, &a5s0, 8); + // Issue chunk 2's A loads (L2-hot) into the now-dead set-0 A registers; + // consumed by the pass-1 capture one full pass later. + a0s0 = *reinterpret_cast(a_base + 2 * kChunkK); + a1s0 = *reinterpret_cast( + a_base + 2 * kChunkK + kDummaTileK); + a2s0 = *reinterpret_cast( + a_base + 2 * kChunkK + 2 * kDummaTileK); + a3s0 = *reinterpret_cast( + a_base + 2 * kChunkK + 3 * kDummaTileK); + a4s0 = *reinterpret_cast( + a_base + 2 * kChunkK + 4 * kDummaTileK); + a5s0 = *reinterpret_cast( + a_base + 2 * kChunkK + 5 * kDummaTileK); + // Consume chunk 0 from LDS plane 0: six k32 steps covering k offsets + // [0, 32), [32, 64), [64, 96), [96, 128), [128, 160), [160, 192) of the + // chunk. All twelve fragment reads issue up front so the LDS latency of + // the whole chunk is hidden behind the twelve-MMA burst. + { + const int2 rb00 = + *reinterpret_cast(&s_b[0][0][b_row][b_col]); + const int2 rb10 = + *reinterpret_cast(&s_b[0][1][b_row][b_col]); + const int2 rb01 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + kDummaTileK]); + const int2 rb11 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + kDummaTileK]); + const int2 rb02 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 2 * kDummaTileK]); + const int2 rb12 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 2 * kDummaTileK]); + const int2 rb03 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 3 * kDummaTileK]); + const int2 rb13 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 3 * kDummaTileK]); + const int2 rb04 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 4 * kDummaTileK]); + const int2 rb14 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 4 * kDummaTileK]); + const int2 rb05 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 5 * kDummaTileK]); + const int2 rb15 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 5 * kDummaTileK]); + __builtin_memcpy(b_frag0.x, &rb00, 8); + __builtin_memcpy(b_frag1.x, &rb10, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag0, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag0, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb01, 8); + __builtin_memcpy(b_frag1.x, &rb11, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag1, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag1, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb02, 8); + __builtin_memcpy(b_frag1.x, &rb12, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag2, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag2, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb03, 8); + __builtin_memcpy(b_frag1.x, &rb13, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag3, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag3, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb04, 8); + __builtin_memcpy(b_frag1.x, &rb14, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag4, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag4, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb05, 8); + __builtin_memcpy(b_frag1.x, &rb15, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag5, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag5, b_frag1, acc1_frag); + } + } + { + // Stage chunk 1 (set 1) into LDS plane 1 (wait-free: already drained + // with chunk 0 by the single prologue ladder). + { + int2* d0 = reinterpret_cast(&s_b[1][0][b_row][b_wide]); + int2* d1 = reinterpret_cast(&s_b[1][1][b_row][b_wide]); + d0[0] = *reinterpret_cast(&b0s1a); + d0[1] = *reinterpret_cast( + reinterpret_cast(&b0s1a) + 2); + d1[0] = *reinterpret_cast(&b1s1a); + d1[1] = *reinterpret_cast( + reinterpret_cast(&b1s1a) + 2); + int2* d0b = reinterpret_cast(&s_b[1][0][b_row][b_wide + 64]); + int2* d1b = reinterpret_cast(&s_b[1][1][b_row][b_wide + 64]); + d0b[0] = *reinterpret_cast(&b0s1b); + d0b[1] = *reinterpret_cast( + reinterpret_cast(&b0s1b) + 2); + d1b[0] = *reinterpret_cast(&b1s1b); + d1b[1] = *reinterpret_cast( + reinterpret_cast(&b1s1b) + 2); + int2* d0c = reinterpret_cast(&s_b[1][0][b_row][b_wide + 128]); + int2* d1c = reinterpret_cast(&s_b[1][1][b_row][b_wide + 128]); + d0c[0] = *reinterpret_cast(&b0s1c); + d0c[1] = *reinterpret_cast( + reinterpret_cast(&b0s1c) + 2); + d1c[0] = *reinterpret_cast(&b1s1c); + d1c[1] = *reinterpret_cast( + reinterpret_cast(&b1s1c) + 2); + } + // Capture chunk 1's six A fragments before the chunk-3 A issue below. + __builtin_memcpy(a_frag0.x, &a0s1, 8); + __builtin_memcpy(a_frag1.x, &a1s1, 8); + __builtin_memcpy(a_frag2.x, &a2s1, 8); + __builtin_memcpy(a_frag3.x, &a3s1, 8); + __builtin_memcpy(a_frag4.x, &a4s1, 8); + __builtin_memcpy(a_frag5.x, &a5s1, 8); + // Issue chunk 3's A loads (L2-hot) into the now-dead set-1 A registers. + a0s1 = *reinterpret_cast(a_base + 3 * kChunkK); + a1s1 = *reinterpret_cast( + a_base + 3 * kChunkK + kDummaTileK); + a2s1 = *reinterpret_cast( + a_base + 3 * kChunkK + 2 * kDummaTileK); + a3s1 = *reinterpret_cast( + a_base + 3 * kChunkK + 3 * kDummaTileK); + a4s1 = *reinterpret_cast( + a_base + 3 * kChunkK + 4 * kDummaTileK); + a5s1 = *reinterpret_cast( + a_base + 3 * kChunkK + 5 * kDummaTileK); + // Consume chunk 1 from LDS plane 1. + { + const int2 rb00 = + *reinterpret_cast(&s_b[1][0][b_row][b_col]); + const int2 rb10 = + *reinterpret_cast(&s_b[1][1][b_row][b_col]); + const int2 rb01 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + kDummaTileK]); + const int2 rb11 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + kDummaTileK]); + const int2 rb02 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 2 * kDummaTileK]); + const int2 rb12 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 2 * kDummaTileK]); + const int2 rb03 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 3 * kDummaTileK]); + const int2 rb13 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 3 * kDummaTileK]); + const int2 rb04 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 4 * kDummaTileK]); + const int2 rb14 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 4 * kDummaTileK]); + const int2 rb05 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 5 * kDummaTileK]); + const int2 rb15 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 5 * kDummaTileK]); + __builtin_memcpy(b_frag0.x, &rb00, 8); + __builtin_memcpy(b_frag1.x, &rb10, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag0, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag0, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb01, 8); + __builtin_memcpy(b_frag1.x, &rb11, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag1, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag1, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb02, 8); + __builtin_memcpy(b_frag1.x, &rb12, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag2, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag2, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb03, 8); + __builtin_memcpy(b_frag1.x, &rb13, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag3, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag3, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb04, 8); + __builtin_memcpy(b_frag1.x, &rb14, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag4, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag4, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb05, 8); + __builtin_memcpy(b_frag1.x, &rb15, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag5, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag5, b_frag1, acc1_frag); + } + } + + // ---- Pass 1: chunks 2 and 3. Both stagings are wait-free (their B loads + // were drained by the single prologue ladder); the only waits here are + // the L2 A captures of the chunk 2/3 A loads issued in pass 0. ---- + { + // Stage chunk 2 (set 2) into LDS plane 0. + { + int2* d0 = reinterpret_cast(&s_b[0][0][b_row][b_wide]); + int2* d1 = reinterpret_cast(&s_b[0][1][b_row][b_wide]); + d0[0] = *reinterpret_cast(&b0s2a); + d0[1] = *reinterpret_cast( + reinterpret_cast(&b0s2a) + 2); + d1[0] = *reinterpret_cast(&b1s2a); + d1[1] = *reinterpret_cast( + reinterpret_cast(&b1s2a) + 2); + int2* d0b = reinterpret_cast(&s_b[0][0][b_row][b_wide + 64]); + int2* d1b = reinterpret_cast(&s_b[0][1][b_row][b_wide + 64]); + d0b[0] = *reinterpret_cast(&b0s2b); + d0b[1] = *reinterpret_cast( + reinterpret_cast(&b0s2b) + 2); + d1b[0] = *reinterpret_cast(&b1s2b); + d1b[1] = *reinterpret_cast( + reinterpret_cast(&b1s2b) + 2); + int2* d0c = reinterpret_cast(&s_b[0][0][b_row][b_wide + 128]); + int2* d1c = reinterpret_cast(&s_b[0][1][b_row][b_wide + 128]); + d0c[0] = *reinterpret_cast(&b0s2c); + d0c[1] = *reinterpret_cast( + reinterpret_cast(&b0s2c) + 2); + d1c[0] = *reinterpret_cast(&b1s2c); + d1c[1] = *reinterpret_cast( + reinterpret_cast(&b1s2c) + 2); + } + // Capture chunk 2's six A fragments (issued in pass 0; L2-hot). + __builtin_memcpy(a_frag0.x, &a0s0, 8); + __builtin_memcpy(a_frag1.x, &a1s0, 8); + __builtin_memcpy(a_frag2.x, &a2s0, 8); + __builtin_memcpy(a_frag3.x, &a3s0, 8); + __builtin_memcpy(a_frag4.x, &a4s0, 8); + __builtin_memcpy(a_frag5.x, &a5s0, 8); + // Consume chunk 2 from LDS plane 0. + { + const int2 rb00 = + *reinterpret_cast(&s_b[0][0][b_row][b_col]); + const int2 rb10 = + *reinterpret_cast(&s_b[0][1][b_row][b_col]); + const int2 rb01 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + kDummaTileK]); + const int2 rb11 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + kDummaTileK]); + const int2 rb02 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 2 * kDummaTileK]); + const int2 rb12 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 2 * kDummaTileK]); + const int2 rb03 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 3 * kDummaTileK]); + const int2 rb13 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 3 * kDummaTileK]); + const int2 rb04 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 4 * kDummaTileK]); + const int2 rb14 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 4 * kDummaTileK]); + const int2 rb05 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 5 * kDummaTileK]); + const int2 rb15 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 5 * kDummaTileK]); + __builtin_memcpy(b_frag0.x, &rb00, 8); + __builtin_memcpy(b_frag1.x, &rb10, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag0, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag0, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb01, 8); + __builtin_memcpy(b_frag1.x, &rb11, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag1, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag1, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb02, 8); + __builtin_memcpy(b_frag1.x, &rb12, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag2, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag2, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb03, 8); + __builtin_memcpy(b_frag1.x, &rb13, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag3, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag3, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb04, 8); + __builtin_memcpy(b_frag1.x, &rb14, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag4, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag4, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb05, 8); + __builtin_memcpy(b_frag1.x, &rb15, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag5, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag5, b_frag1, acc1_frag); + } + } + { + // Stage chunk 3 (set 3) into LDS plane 1 (wait-free). + { + int2* d0 = reinterpret_cast(&s_b[1][0][b_row][b_wide]); + int2* d1 = reinterpret_cast(&s_b[1][1][b_row][b_wide]); + d0[0] = *reinterpret_cast(&b0s3a); + d0[1] = *reinterpret_cast( + reinterpret_cast(&b0s3a) + 2); + d1[0] = *reinterpret_cast(&b1s3a); + d1[1] = *reinterpret_cast( + reinterpret_cast(&b1s3a) + 2); + int2* d0b = reinterpret_cast(&s_b[1][0][b_row][b_wide + 64]); + int2* d1b = reinterpret_cast(&s_b[1][1][b_row][b_wide + 64]); + d0b[0] = *reinterpret_cast(&b0s3b); + d0b[1] = *reinterpret_cast( + reinterpret_cast(&b0s3b) + 2); + d1b[0] = *reinterpret_cast(&b1s3b); + d1b[1] = *reinterpret_cast( + reinterpret_cast(&b1s3b) + 2); + int2* d0c = reinterpret_cast(&s_b[1][0][b_row][b_wide + 128]); + int2* d1c = reinterpret_cast(&s_b[1][1][b_row][b_wide + 128]); + d0c[0] = *reinterpret_cast(&b0s3c); + d0c[1] = *reinterpret_cast( + reinterpret_cast(&b0s3c) + 2); + d1c[0] = *reinterpret_cast(&b1s3c); + d1c[1] = *reinterpret_cast( + reinterpret_cast(&b1s3c) + 2); + } + // Capture chunk 3's six A fragments (issued in pass 0; L2-hot). + __builtin_memcpy(a_frag0.x, &a0s1, 8); + __builtin_memcpy(a_frag1.x, &a1s1, 8); + __builtin_memcpy(a_frag2.x, &a2s1, 8); + __builtin_memcpy(a_frag3.x, &a3s1, 8); + __builtin_memcpy(a_frag4.x, &a4s1, 8); + __builtin_memcpy(a_frag5.x, &a5s1, 8); + // Consume chunk 3 from LDS plane 1. + { + const int2 rb00 = + *reinterpret_cast(&s_b[1][0][b_row][b_col]); + const int2 rb10 = + *reinterpret_cast(&s_b[1][1][b_row][b_col]); + const int2 rb01 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + kDummaTileK]); + const int2 rb11 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + kDummaTileK]); + const int2 rb02 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 2 * kDummaTileK]); + const int2 rb12 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 2 * kDummaTileK]); + const int2 rb03 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 3 * kDummaTileK]); + const int2 rb13 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 3 * kDummaTileK]); + const int2 rb04 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 4 * kDummaTileK]); + const int2 rb14 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 4 * kDummaTileK]); + const int2 rb05 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 5 * kDummaTileK]); + const int2 rb15 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 5 * kDummaTileK]); + __builtin_memcpy(b_frag0.x, &rb00, 8); + __builtin_memcpy(b_frag1.x, &rb10, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag0, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag0, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb01, 8); + __builtin_memcpy(b_frag1.x, &rb11, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag1, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag1, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb02, 8); + __builtin_memcpy(b_frag1.x, &rb12, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag2, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag2, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb03, 8); + __builtin_memcpy(b_frag1.x, &rb13, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag3, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag3, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb04, 8); + __builtin_memcpy(b_frag1.x, &rb14, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag4, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag4, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb05, 8); + __builtin_memcpy(b_frag1.x, &rb15, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag5, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag5, b_frag1, acc1_frag); + } + } + + // ---- Iteration 20, body 2 (chunks 4-7, slice offsets 768..1535): the + // body-1 B/A registers are dead after the pass-1 captures, so re-issue the + // 24 dwordx4 B loads plus the chunks 4/5 A loads in place; the chunk-4 + // staging below drains them with ONE ladder (the block's second cold-B + // round trip; 328 blocks x 2 = 656 drains device-wide, the same total as + // the accepted 656-block kernel). All offsets stay 16-B aligned: 4*kChunkK + // = 768 = 16*48, max offset 7*192 + 128 + 15 = 1487 < 1536. ---- + b0s0a = *reinterpret_cast(b0_base + 4 * kChunkK + b_wide); + b0s0b = *reinterpret_cast(b0_base + 4 * kChunkK + b_wide + 64); + b0s0c = *reinterpret_cast(b0_base + 4 * kChunkK + b_wide + 128); + b1s0a = *reinterpret_cast(b1_base + 4 * kChunkK + b_wide); + b1s0b = *reinterpret_cast(b1_base + 4 * kChunkK + b_wide + 64); + b1s0c = *reinterpret_cast(b1_base + 4 * kChunkK + b_wide + 128); + b0s1a = *reinterpret_cast(b0_base + 5 * kChunkK + b_wide); + b0s1b = *reinterpret_cast(b0_base + 5 * kChunkK + b_wide + 64); + b0s1c = *reinterpret_cast(b0_base + 5 * kChunkK + b_wide + 128); + b1s1a = *reinterpret_cast(b1_base + 5 * kChunkK + b_wide); + b1s1b = *reinterpret_cast(b1_base + 5 * kChunkK + b_wide + 64); + b1s1c = *reinterpret_cast(b1_base + 5 * kChunkK + b_wide + 128); + b0s2a = *reinterpret_cast(b0_base + 6 * kChunkK + b_wide); + b0s2b = *reinterpret_cast(b0_base + 6 * kChunkK + b_wide + 64); + b0s2c = *reinterpret_cast(b0_base + 6 * kChunkK + b_wide + 128); + b1s2a = *reinterpret_cast(b1_base + 6 * kChunkK + b_wide); + b1s2b = *reinterpret_cast(b1_base + 6 * kChunkK + b_wide + 64); + b1s2c = *reinterpret_cast(b1_base + 6 * kChunkK + b_wide + 128); + b0s3a = *reinterpret_cast(b0_base + 7 * kChunkK + b_wide); + b0s3b = *reinterpret_cast(b0_base + 7 * kChunkK + b_wide + 64); + b0s3c = *reinterpret_cast(b0_base + 7 * kChunkK + b_wide + 128); + b1s3a = *reinterpret_cast(b1_base + 7 * kChunkK + b_wide); + b1s3b = *reinterpret_cast(b1_base + 7 * kChunkK + b_wide + 64); + b1s3c = *reinterpret_cast(b1_base + 7 * kChunkK + b_wide + 128); + a0s0 = *reinterpret_cast(a_base + 4 * kChunkK); + a1s0 = *reinterpret_cast(a_base + 4 * kChunkK + kDummaTileK); + a2s0 = *reinterpret_cast(a_base + 4 * kChunkK + 2 * kDummaTileK); + a3s0 = *reinterpret_cast(a_base + 4 * kChunkK + 3 * kDummaTileK); + a4s0 = *reinterpret_cast(a_base + 4 * kChunkK + 4 * kDummaTileK); + a5s0 = *reinterpret_cast(a_base + 4 * kChunkK + 5 * kDummaTileK); + a0s1 = *reinterpret_cast(a_base + 5 * kChunkK); + a1s1 = *reinterpret_cast(a_base + 5 * kChunkK + kDummaTileK); + a2s1 = *reinterpret_cast(a_base + 5 * kChunkK + 2 * kDummaTileK); + a3s1 = *reinterpret_cast(a_base + 5 * kChunkK + 3 * kDummaTileK); + a4s1 = *reinterpret_cast(a_base + 5 * kChunkK + 4 * kDummaTileK); + a5s1 = *reinterpret_cast(a_base + 5 * kChunkK + 5 * kDummaTileK); + + // ---- Pass 2: chunks 4 and 5. The chunk-4 staging is the block's second + // cold-B drain; chunk 5 is wait-free. A rotation identical to body 1. ---- + { + // Stage chunk 4 (set 0) into LDS plane 0. + { + int2* d0 = reinterpret_cast(&s_b[0][0][b_row][b_wide]); + int2* d1 = reinterpret_cast(&s_b[0][1][b_row][b_wide]); + d0[0] = *reinterpret_cast(&b0s0a); + d0[1] = *reinterpret_cast( + reinterpret_cast(&b0s0a) + 2); + d1[0] = *reinterpret_cast(&b1s0a); + d1[1] = *reinterpret_cast( + reinterpret_cast(&b1s0a) + 2); + int2* d0b = reinterpret_cast(&s_b[0][0][b_row][b_wide + 64]); + int2* d1b = reinterpret_cast(&s_b[0][1][b_row][b_wide + 64]); + d0b[0] = *reinterpret_cast(&b0s0b); + d0b[1] = *reinterpret_cast( + reinterpret_cast(&b0s0b) + 2); + d1b[0] = *reinterpret_cast(&b1s0b); + d1b[1] = *reinterpret_cast( + reinterpret_cast(&b1s0b) + 2); + int2* d0c = reinterpret_cast(&s_b[0][0][b_row][b_wide + 128]); + int2* d1c = reinterpret_cast(&s_b[0][1][b_row][b_wide + 128]); + d0c[0] = *reinterpret_cast(&b0s0c); + d0c[1] = *reinterpret_cast( + reinterpret_cast(&b0s0c) + 2); + d1c[0] = *reinterpret_cast(&b1s0c); + d1c[1] = *reinterpret_cast( + reinterpret_cast(&b1s0c) + 2); + } + // Capture chunk 4's six A fragments BEFORE set 0's A registers are + // overwritten below by the chunk-6 A issue. + __builtin_memcpy(a_frag0.x, &a0s0, 8); + __builtin_memcpy(a_frag1.x, &a1s0, 8); + __builtin_memcpy(a_frag2.x, &a2s0, 8); + __builtin_memcpy(a_frag3.x, &a3s0, 8); + __builtin_memcpy(a_frag4.x, &a4s0, 8); + __builtin_memcpy(a_frag5.x, &a5s0, 8); + // Issue chunk 6's A loads (L2-hot) into the now-dead set-0 A registers; + // consumed by the pass-3 capture one full pass later. + a0s0 = *reinterpret_cast(a_base + 6 * kChunkK); + a1s0 = *reinterpret_cast( + a_base + 6 * kChunkK + kDummaTileK); + a2s0 = *reinterpret_cast( + a_base + 6 * kChunkK + 2 * kDummaTileK); + a3s0 = *reinterpret_cast( + a_base + 6 * kChunkK + 3 * kDummaTileK); + a4s0 = *reinterpret_cast( + a_base + 6 * kChunkK + 4 * kDummaTileK); + a5s0 = *reinterpret_cast( + a_base + 6 * kChunkK + 5 * kDummaTileK); + // Consume chunk 4 from LDS plane 0: six k32 steps covering slice k + // offsets [768, 800), [800, 832), [832, 864), [864, 896), [896, 928), + // [928, 960). + { + const int2 rb00 = + *reinterpret_cast(&s_b[0][0][b_row][b_col]); + const int2 rb10 = + *reinterpret_cast(&s_b[0][1][b_row][b_col]); + const int2 rb01 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + kDummaTileK]); + const int2 rb11 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + kDummaTileK]); + const int2 rb02 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 2 * kDummaTileK]); + const int2 rb12 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 2 * kDummaTileK]); + const int2 rb03 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 3 * kDummaTileK]); + const int2 rb13 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 3 * kDummaTileK]); + const int2 rb04 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 4 * kDummaTileK]); + const int2 rb14 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 4 * kDummaTileK]); + const int2 rb05 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 5 * kDummaTileK]); + const int2 rb15 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 5 * kDummaTileK]); + __builtin_memcpy(b_frag0.x, &rb00, 8); + __builtin_memcpy(b_frag1.x, &rb10, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag0, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag0, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb01, 8); + __builtin_memcpy(b_frag1.x, &rb11, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag1, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag1, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb02, 8); + __builtin_memcpy(b_frag1.x, &rb12, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag2, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag2, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb03, 8); + __builtin_memcpy(b_frag1.x, &rb13, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag3, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag3, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb04, 8); + __builtin_memcpy(b_frag1.x, &rb14, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag4, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag4, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb05, 8); + __builtin_memcpy(b_frag1.x, &rb15, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag5, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag5, b_frag1, acc1_frag); + } + } + { + // Stage chunk 5 (set 1) into LDS plane 1 (wait-free: drained with + // chunk 4 by the body-2 ladder). + { + int2* d0 = reinterpret_cast(&s_b[1][0][b_row][b_wide]); + int2* d1 = reinterpret_cast(&s_b[1][1][b_row][b_wide]); + d0[0] = *reinterpret_cast(&b0s1a); + d0[1] = *reinterpret_cast( + reinterpret_cast(&b0s1a) + 2); + d1[0] = *reinterpret_cast(&b1s1a); + d1[1] = *reinterpret_cast( + reinterpret_cast(&b1s1a) + 2); + int2* d0b = reinterpret_cast(&s_b[1][0][b_row][b_wide + 64]); + int2* d1b = reinterpret_cast(&s_b[1][1][b_row][b_wide + 64]); + d0b[0] = *reinterpret_cast(&b0s1b); + d0b[1] = *reinterpret_cast( + reinterpret_cast(&b0s1b) + 2); + d1b[0] = *reinterpret_cast(&b1s1b); + d1b[1] = *reinterpret_cast( + reinterpret_cast(&b1s1b) + 2); + int2* d0c = reinterpret_cast(&s_b[1][0][b_row][b_wide + 128]); + int2* d1c = reinterpret_cast(&s_b[1][1][b_row][b_wide + 128]); + d0c[0] = *reinterpret_cast(&b0s1c); + d0c[1] = *reinterpret_cast( + reinterpret_cast(&b0s1c) + 2); + d1c[0] = *reinterpret_cast(&b1s1c); + d1c[1] = *reinterpret_cast( + reinterpret_cast(&b1s1c) + 2); + } + // Capture chunk 5's six A fragments before the chunk-7 A issue below. + __builtin_memcpy(a_frag0.x, &a0s1, 8); + __builtin_memcpy(a_frag1.x, &a1s1, 8); + __builtin_memcpy(a_frag2.x, &a2s1, 8); + __builtin_memcpy(a_frag3.x, &a3s1, 8); + __builtin_memcpy(a_frag4.x, &a4s1, 8); + __builtin_memcpy(a_frag5.x, &a5s1, 8); + // Issue chunk 7's A loads (L2-hot) into the now-dead set-1 A registers. + a0s1 = *reinterpret_cast(a_base + 7 * kChunkK); + a1s1 = *reinterpret_cast( + a_base + 7 * kChunkK + kDummaTileK); + a2s1 = *reinterpret_cast( + a_base + 7 * kChunkK + 2 * kDummaTileK); + a3s1 = *reinterpret_cast( + a_base + 7 * kChunkK + 3 * kDummaTileK); + a4s1 = *reinterpret_cast( + a_base + 7 * kChunkK + 4 * kDummaTileK); + a5s1 = *reinterpret_cast( + a_base + 7 * kChunkK + 5 * kDummaTileK); + // Consume chunk 5 from LDS plane 1: slice k offsets [960, 1152). + { + const int2 rb00 = + *reinterpret_cast(&s_b[1][0][b_row][b_col]); + const int2 rb10 = + *reinterpret_cast(&s_b[1][1][b_row][b_col]); + const int2 rb01 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + kDummaTileK]); + const int2 rb11 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + kDummaTileK]); + const int2 rb02 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 2 * kDummaTileK]); + const int2 rb12 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 2 * kDummaTileK]); + const int2 rb03 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 3 * kDummaTileK]); + const int2 rb13 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 3 * kDummaTileK]); + const int2 rb04 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 4 * kDummaTileK]); + const int2 rb14 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 4 * kDummaTileK]); + const int2 rb05 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 5 * kDummaTileK]); + const int2 rb15 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 5 * kDummaTileK]); + __builtin_memcpy(b_frag0.x, &rb00, 8); + __builtin_memcpy(b_frag1.x, &rb10, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag0, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag0, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb01, 8); + __builtin_memcpy(b_frag1.x, &rb11, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag1, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag1, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb02, 8); + __builtin_memcpy(b_frag1.x, &rb12, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag2, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag2, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb03, 8); + __builtin_memcpy(b_frag1.x, &rb13, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag3, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag3, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb04, 8); + __builtin_memcpy(b_frag1.x, &rb14, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag4, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag4, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb05, 8); + __builtin_memcpy(b_frag1.x, &rb15, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag5, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag5, b_frag1, acc1_frag); + } + } + + // ---- Pass 3: chunks 6 and 7. Both stagings are wait-free (their B loads + // were drained by the body-2 ladder at the chunk-4 staging); the only + // waits here are the L2 A captures of the chunk 6/7 A loads issued in + // pass 2. ---- + { + // Stage chunk 6 (set 2) into LDS plane 0. + { + int2* d0 = reinterpret_cast(&s_b[0][0][b_row][b_wide]); + int2* d1 = reinterpret_cast(&s_b[0][1][b_row][b_wide]); + d0[0] = *reinterpret_cast(&b0s2a); + d0[1] = *reinterpret_cast( + reinterpret_cast(&b0s2a) + 2); + d1[0] = *reinterpret_cast(&b1s2a); + d1[1] = *reinterpret_cast( + reinterpret_cast(&b1s2a) + 2); + int2* d0b = reinterpret_cast(&s_b[0][0][b_row][b_wide + 64]); + int2* d1b = reinterpret_cast(&s_b[0][1][b_row][b_wide + 64]); + d0b[0] = *reinterpret_cast(&b0s2b); + d0b[1] = *reinterpret_cast( + reinterpret_cast(&b0s2b) + 2); + d1b[0] = *reinterpret_cast(&b1s2b); + d1b[1] = *reinterpret_cast( + reinterpret_cast(&b1s2b) + 2); + int2* d0c = reinterpret_cast(&s_b[0][0][b_row][b_wide + 128]); + int2* d1c = reinterpret_cast(&s_b[0][1][b_row][b_wide + 128]); + d0c[0] = *reinterpret_cast(&b0s2c); + d0c[1] = *reinterpret_cast( + reinterpret_cast(&b0s2c) + 2); + d1c[0] = *reinterpret_cast(&b1s2c); + d1c[1] = *reinterpret_cast( + reinterpret_cast(&b1s2c) + 2); + } + // Capture chunk 6's six A fragments (issued in pass 2; L2-hot). + __builtin_memcpy(a_frag0.x, &a0s0, 8); + __builtin_memcpy(a_frag1.x, &a1s0, 8); + __builtin_memcpy(a_frag2.x, &a2s0, 8); + __builtin_memcpy(a_frag3.x, &a3s0, 8); + __builtin_memcpy(a_frag4.x, &a4s0, 8); + __builtin_memcpy(a_frag5.x, &a5s0, 8); + // Consume chunk 6 from LDS plane 0: slice k offsets [1152, 1344). + { + const int2 rb00 = + *reinterpret_cast(&s_b[0][0][b_row][b_col]); + const int2 rb10 = + *reinterpret_cast(&s_b[0][1][b_row][b_col]); + const int2 rb01 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + kDummaTileK]); + const int2 rb11 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + kDummaTileK]); + const int2 rb02 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 2 * kDummaTileK]); + const int2 rb12 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 2 * kDummaTileK]); + const int2 rb03 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 3 * kDummaTileK]); + const int2 rb13 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 3 * kDummaTileK]); + const int2 rb04 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 4 * kDummaTileK]); + const int2 rb14 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 4 * kDummaTileK]); + const int2 rb05 = *reinterpret_cast( + &s_b[0][0][b_row][b_col + 5 * kDummaTileK]); + const int2 rb15 = *reinterpret_cast( + &s_b[0][1][b_row][b_col + 5 * kDummaTileK]); + __builtin_memcpy(b_frag0.x, &rb00, 8); + __builtin_memcpy(b_frag1.x, &rb10, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag0, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag0, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb01, 8); + __builtin_memcpy(b_frag1.x, &rb11, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag1, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag1, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb02, 8); + __builtin_memcpy(b_frag1.x, &rb12, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag2, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag2, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb03, 8); + __builtin_memcpy(b_frag1.x, &rb13, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag3, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag3, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb04, 8); + __builtin_memcpy(b_frag1.x, &rb14, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag4, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag4, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb05, 8); + __builtin_memcpy(b_frag1.x, &rb15, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag5, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag5, b_frag1, acc1_frag); + } + } + { + // Stage chunk 7 (set 3) into LDS plane 1 (wait-free). + { + int2* d0 = reinterpret_cast(&s_b[1][0][b_row][b_wide]); + int2* d1 = reinterpret_cast(&s_b[1][1][b_row][b_wide]); + d0[0] = *reinterpret_cast(&b0s3a); + d0[1] = *reinterpret_cast( + reinterpret_cast(&b0s3a) + 2); + d1[0] = *reinterpret_cast(&b1s3a); + d1[1] = *reinterpret_cast( + reinterpret_cast(&b1s3a) + 2); + int2* d0b = reinterpret_cast(&s_b[1][0][b_row][b_wide + 64]); + int2* d1b = reinterpret_cast(&s_b[1][1][b_row][b_wide + 64]); + d0b[0] = *reinterpret_cast(&b0s3b); + d0b[1] = *reinterpret_cast( + reinterpret_cast(&b0s3b) + 2); + d1b[0] = *reinterpret_cast(&b1s3b); + d1b[1] = *reinterpret_cast( + reinterpret_cast(&b1s3b) + 2); + int2* d0c = reinterpret_cast(&s_b[1][0][b_row][b_wide + 128]); + int2* d1c = reinterpret_cast(&s_b[1][1][b_row][b_wide + 128]); + d0c[0] = *reinterpret_cast(&b0s3c); + d0c[1] = *reinterpret_cast( + reinterpret_cast(&b0s3c) + 2); + d1c[0] = *reinterpret_cast(&b1s3c); + d1c[1] = *reinterpret_cast( + reinterpret_cast(&b1s3c) + 2); + } + // Capture chunk 7's six A fragments (issued in pass 2; L2-hot). + __builtin_memcpy(a_frag0.x, &a0s1, 8); + __builtin_memcpy(a_frag1.x, &a1s1, 8); + __builtin_memcpy(a_frag2.x, &a2s1, 8); + __builtin_memcpy(a_frag3.x, &a3s1, 8); + __builtin_memcpy(a_frag4.x, &a4s1, 8); + __builtin_memcpy(a_frag5.x, &a5s1, 8); + // Consume chunk 7 from LDS plane 1: slice k offsets [1344, 1536). + { + const int2 rb00 = + *reinterpret_cast(&s_b[1][0][b_row][b_col]); + const int2 rb10 = + *reinterpret_cast(&s_b[1][1][b_row][b_col]); + const int2 rb01 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + kDummaTileK]); + const int2 rb11 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + kDummaTileK]); + const int2 rb02 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 2 * kDummaTileK]); + const int2 rb12 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 2 * kDummaTileK]); + const int2 rb03 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 3 * kDummaTileK]); + const int2 rb13 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 3 * kDummaTileK]); + const int2 rb04 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 4 * kDummaTileK]); + const int2 rb14 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 4 * kDummaTileK]); + const int2 rb05 = *reinterpret_cast( + &s_b[1][0][b_row][b_col + 5 * kDummaTileK]); + const int2 rb15 = *reinterpret_cast( + &s_b[1][1][b_row][b_col + 5 * kDummaTileK]); + __builtin_memcpy(b_frag0.x, &rb00, 8); + __builtin_memcpy(b_frag1.x, &rb10, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag0, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag0, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb01, 8); + __builtin_memcpy(b_frag1.x, &rb11, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag1, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag1, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb02, 8); + __builtin_memcpy(b_frag1.x, &rb12, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag2, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag2, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb03, 8); + __builtin_memcpy(b_frag1.x, &rb13, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag3, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag3, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb04, 8); + __builtin_memcpy(b_frag1.x, &rb14, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag4, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag4, b_frag1, acc1_frag); + __builtin_memcpy(b_frag0.x, &rb05, 8); + __builtin_memcpy(b_frag1.x, &rb15, 8); + du::dumma::du_mma_sync(acc0_frag, a_frag5, b_frag0, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag5, b_frag1, acc1_frag); + } + } + + // Publish both 16x16 int32 partial planes (2 x 256 int32 = 2 KiB/block). + // Iteration 18 (workspace plane order [tile][split]): tile t's split-s + // plane lives at (t*SPLIT_K + s)*256 int32, so the four split planes of + // an output tile form ONE contiguous 4-KB region that the combine kernel + // reads as a sequential stream (164 streams instead of 1,312 scattered + // 1-KB streams). The stored bytes are identical; only the base index + // changes, and each block still writes two line-complete 1-KB planes. + du::dumma::du_store_matrix_sync( + partials + (tile0 * SPLIT_K + split) * (kDummaTileM * kDummaTileN), + acc0_frag, static_cast(kDummaTileN), + du::dumma::mem_row_major); + du::dumma::du_store_matrix_sync( + partials + ((tile0 + 1) * SPLIT_K + split) * (kDummaTileM * kDummaTileN), + acc1_frag, static_cast(kDummaTileN), + du::dumma::mem_row_major); +} + +template +__global__ __launch_bounds__(kDummaBlockThreads) void +w8a8_dumma_m16n16k32_splitk_combine_kernel( + const int* __restrict__ partials, // [N/16][SPLIT_K][256] int32 + // (tile-major, iteration 18) + const float* __restrict__ x_scale, // [16] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [16, N] + int n) { + const int tile = static_cast(blockIdx.x); + const int n0 = tile * kDummaTileN; + const int lane = static_cast(threadIdx.x); + // Iteration 18: tile-major plane order -- plane (tile, split) lives at + // (tile*SPLIT_K + split)*256, so this tile's four split planes are one + // contiguous 4-KB region (164 sequential 4-KB streams instead of 1,312 + // scattered 1-KB streams at 164-KB strides). + const int* p = + partials + tile * SPLIT_K * (kDummaTileM * kDummaTileN) + lane * 4; + + // Iteration 13 (combine latency consolidation): the exact iteration-12 + // code object's combine kernel issues the 8 plane loads up front (one + // round trip) but then loads the four weight_scale elements with four + // SEPARATE 4-B loads, each waited to completion (s_waitcnt vmcnt(0)) + // before its multiply and its own 2-B bf16 store -- four serialized + // round trips plus four store instructions per lane (weight_scale is + // only 10.5 KiB, but the 16.1-MB B stream thrashes L2 ahead of the + // combine, whose l2 hit rate is 31.4%, so most of those round trips are + // DRAM-cold). This round fetches the four contiguous scales (col_base is + // a multiple of 4 -> 16-B aligned) with ONE aligned dwordx4 load issued + // before the plane reads below, and packs the four 2-B bf16 outputs into + // ONE 8-B store per lane (byte offset (row*n + col_base)*2 is a multiple + // of 8: n is a multiple of 8 and col_base is a multiple of 4). The + // per-element arithmetic is byte-for-byte unchanged -- int32 ascending + // split sum, then float(acc)*xs*ws (fp32 left-to-right), then + // round-to-nearest-even bf16 -- only the load grouping and the store + // width change, so the combine output stays bit-identical (0 mismatches) + // and Graph replay with changed contents must pass. + const int e0 = lane * 4; + const int row = e0 >> 4; // 0..15 + const int col0 = e0 & 15; // 0,4,8,12 + const int col_base = n0 + col0; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col_base); + // Lane l owns row-major elements e = 4*l + i (i = 0..3). Sum the SPLIT_K + // planes in ascending split order: exact int32 accumulation, bit-identical + // to the k-ascending full-K reference dot. + int acc0 = 0, acc1 = 0, acc2 = 0, acc3 = 0; +#pragma unroll + for (int s = 0; s < SPLIT_K; ++s) { + const int* ps = p + s * (kDummaTileM * kDummaTileN); + acc0 += ps[0]; + acc1 += ps[1]; + acc2 += ps[2]; + acc3 += ps[3]; + } + + // DTK note: __float2bfloat16 returns __hip_bfloat16 (amd_hip_bf16.h), which + // is a distinct type from hip_bfloat16 (amd_hip_bfloat16.h, explicit float + // ctor only) -- keep these locals as __hip_bfloat16 so the RNE bits are + // stored directly (same bits the iter-12 out[i] = __float2bfloat16(...) + // assignment path wrote) with the single 8-B packed store. + const __hip_bfloat16 out0 = + __float2bfloat16(static_cast(acc0) * xs * ws.x); + const __hip_bfloat16 out1 = + __float2bfloat16(static_cast(acc1) * xs * ws.y); + const __hip_bfloat16 out2 = + __float2bfloat16(static_cast(acc2) * xs * ws.z); + const __hip_bfloat16 out3 = + __float2bfloat16(static_cast(acc3) * xs * ws.w); + const __hip_bfloat16 packed[4] = {out0, out1, out2, out3}; + __builtin_memcpy(&out[row * n + col_base], packed, sizeof(packed)); +} + +// Shared scalar launch helper: works for any (m, n, k) satisfying the API +// contract, so it doubles as the generic fallback. +void launch_scalar_gemm(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + int m, + int n, + int k, + hipStream_t stream) { + const int64_t total = static_cast(m) * n; + const int64_t blocks = + (total + kScalarBlockThreads - 1) / kScalarBlockThreads; + hipLaunchKernelGGL(w8a8_scalar_gemm_kernel, + dim3(static_cast(blocks)), + dim3(kScalarBlockThreads), 0, stream, a, b, x_scale, + weight_scale, reinterpret_cast(out), m, + n, k); +} + +// Identity device-to-device copy; generic pack_weight fallback. +template +__global__ __launch_bounds__(kCopyBlockThreads) void +w8a8_identity_copy_kernel( + const T* __restrict__ src, + T* __restrict__ dst, + int count) { + const int i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { + dst[i] = src[i]; + } +} + +// Exact-shape n-major pack (iteration 5): for (K, N) == (6144, 2624) the +// weight is transposed once, outside the timed region and out of Graph +// capture, into packed[n * K + k] = raw[k * N + n] so each lane's 8-byte +// col_major B fragment is contiguous and every m16n16k32 B fragment is a +// single 8-B vector load in the partial kernel. Same byte count and buffer +// as the identity pack, so captured addresses are unchanged. One thread per +// output byte; runs once during weight prep. +__global__ __launch_bounds__(kCopyBlockThreads) void +w8a8_pack_nmajor_bytes_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t total = static_cast(k) * n; + const int64_t stride = + static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; idx += stride) { + const int col = static_cast(idx / k); + const int kk = static_cast(idx - static_cast(col) * k); + packed[idx] = raw[static_cast(kk) * n + col]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Host launch symbols consumed by csrc/bindings.cpp (TORCH_LIBRARY zth_w8a8). +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // Exact assigned shape: glm_tp8_fused_qkv_a_proj_m16 (M=16, N=2624, + // K=6144). Iteration 4 architecture: multi-N-tile reuse -- partial grid + // (82, 8) = 656 one-wave zero-barrier blocks (5.47 blocks/CU, same as the + // iteration-2 SPLIT_K=4 winner), each block computing two adjacent 16x16 + // N-tiles over its K slice (K/8 = 768 = 24 k32 steps) with one shared A + // fragment load per step feeding two MMAs, then a same-stream combine + // kernel (164 one-wave blocks) summing the eight int32 planes in ascending + // split order with the fused x_scale * weight_scale bf16 epilogue. + // Iteration 20: SPLIT_K 8 -> 4 (byte-reduction round) -- the partial grid + // becomes (82, 4) = 328 one-wave zero-barrier blocks (2.73 blocks/CU, a + // single wave with all blocks resident) over K/4 = 1536 = 48 k32 steps, + // and the combine sums the four int32 planes in ascending split order. + // Iteration 5: `b` is the n-major packed weight packed[n*K + k] = + // raw[k*N + n] (launch_pack_w8a8_weight transposes it once out of the + // timed region for (k,n)==(6144,2624)); the partial kernel loads A and + // both B fragments with one aligned 8-B vector load per lane (int2 + + // __builtin_memcpy) instead of the library's per-byte loaders. + // Iteration 6: the partial kernel stages ONLY the cold B stream through a + // 2,304-B LDS buffer (16-B dwordx4 cooperative loads, one vmcnt drain per + // 2-k32-step chunk, ds_read_b64 fragment re-reads, next-chunk loads issued + // before consumption); A stays direct and L2-hot. B fragment values are + // byte-identical, so the int32 partial planes and the combine output stay + // bit-identical; geometry, workspace and Graph structure are unchanged. + // Iteration 7: the partial kernel DOUBLE-buffers the B staging -- two + // ping-pong 2,304-B LDS planes (4,608 B total, < 48 KiB) and two + // parity-alternating register sets with a two-chunk prefetch, no rotation + // copies, zero barriers per K step (one wave per block) -- so every load + // stays in flight across the back edge and the single vmcnt drain per + // chunk lands only at its own staging write. B fragment values are still + // byte-identical, so partials and the combine output stay bit-identical; + // geometry, workspace and Graph structure are unchanged. + // Iteration 9: the partial kernel WIDENS each staged chunk to FOUR k32 + // steps (128 K, 6 chunks per block at SPLIT_K=8) with the same double- + // buffer ping-pong structure -- two 136-stride LDS planes (8,704 B total, + // < 48 KiB), two parity register sets carrying 2 x 32 B of B + 32 B of A + // per lane, two-chunk prefetch, zero barriers per K step -- halving the + // per-chunk vmcnt drains and doubling the in-flight window. vmem_read + // stays 31,488/replay, LDS instructions stay 31,488/replay, HBM read + // stays at the compulsory floor; B fragment values are still + // byte-identical, so partials and the combine output stay bit-identical; + // geometry, workspace and Graph structure are unchanged. + // Iteration 10: the partial kernel WIDENS each staged chunk to + // SIX k32 steps (192 K, 4 chunks per block at SPLIT_K=8, 2 passes of 2) + // with the same double-buffer ping-pong structure -- two 200-stride LDS + // planes (12,800 B total, < 48 KiB), two parity register sets carrying + // 2 x 48 B of B + 48 B of A per lane, two-chunk prefetch, zero barriers + // per K step -- cutting the per-pass vmcnt drain events from 3 to 2 per + // block with ~50% more in-flight bytes per event. vmem_read stays + // 31,488/replay, LDS instructions stay 31,488/replay, HBM read stays at + // the compulsory floor; B fragment values are still byte-identical, so + // partials and the combine output stay bit-identical; geometry, workspace + // and Graph structure are unchanged. + // Iteration 12 (current partial-kernel architecture): the partial kernel + // issues ALL FOUR chunks' B loads (24 dwordx4) plus the chunks 0/1 A + // loads in the prologue and unrolls the 2-pass loop, so the first staging + // consumer drains the whole K-slice with ONE vmcnt ladder: one cold-B + // round trip per block (was two), wait-free pass-1 staging, no rotation + // v_movs. A stays L2-pipelined (chunk 2/3 A issued in pass 0, captured in + // pass 1). Same 656-block geometry, same 12,800-B LDS planes, same staged + // bytes and k-ascending int32 order -> bit-identical partials and combine + // output; vmem_read stays 31,488/replay, LDS instructions stay + // 31,488/replay. + // Iteration 13 (combine latency consolidation): the partial kernel is + // UNTOUCHED. The combine kernel (164 one-wave blocks) now fetches the + // four contiguous weight_scale elements of each lane with ONE aligned + // 16-B load (was four serialized 4-B loads, each waited to completion + // before its multiply) and stores the four 2-B bf16 outputs with ONE 8-B + // store per lane (was four). Per-element arithmetic is byte-for-byte + // unchanged -> the combine output stays bit-identical; grid, workspace + // and Graph structure (partial + combine launches on the caller stream) + // are unchanged. + // K % 4 == 0, N % 32 == 0 hold for this exact shape. The workspace layout + // is [164][4][256] int32 (tile-major since iteration 18) = 671,744 B + // (iteration 20: SPLIT_K 8 -> 4 halves the plane traffic), guaranteed by + // the API's 16-plane split budget for this shape + // (allocate_workspace yields 2,686,976 B); the + // guard below falls back to the generic scalar path if the caller-provided + // workspace is smaller. The shape guard stays exact, so paired M=2 shapes + // with the same (N, K) keep taking the generic scalar fallback (which + // decodes the n-major pack for (n,k)==(2624,6144)). + if (m == kDummaTileM && n == 2624 && k == 6144) { + constexpr int kSplitK = 4; + constexpr int kTilesPerBlock = 2; // adjacent 16x16 N-tiles per block + constexpr int kTiles = 2624 / kDummaTileN; // 164 + constexpr int kBlockTiles = kTiles / kTilesPerBlock; // 82 + constexpr int kPlaneInts = kDummaTileM * kDummaTileN; // 256 + constexpr int64_t kPartialBytes = + static_cast(kSplitK) * kTiles * kPlaneInts * + static_cast(sizeof(int)); // 671,744 + if (workspace != nullptr && workspace_bytes >= kPartialBytes) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_m16n16k32_splitk_partial_kernel), + dim3(static_cast(kBlockTiles), + static_cast(kSplitK)), + dim3(kDummaBlockThreads), + 0, // no dynamic shared memory + stream, + a, + b, + reinterpret_cast(workspace), + n, + k); + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_m16n16k32_splitk_combine_kernel), + dim3(static_cast(kTiles)), + dim3(kDummaBlockThreads), + 0, // no dynamic shared memory + stream, + reinterpret_cast(workspace), + x_scale, + weight_scale, + reinterpret_cast(out), + n); + return; + } + // Workspace too small for split-K=8: fall through to the generic scalar + // fallback (still exact for this shape). + } + + // Generic scalar fallback for every unmatched (m, n, k). + launch_scalar_gemm(a, b, x_scale, weight_scale, out, m, n, k, stream); +} + +extern "C" void launch_pack_w8a8_weight(const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Iteration 5: the exact (k, n) == (6144, 2624) weight is packed once, + // outside the timed region and out of Graph capture, into the n-major + // layout packed[n * K + k] = raw[k * N + n] (same byte count and buffer + // addresses, so captured pointers stay valid; the partial kernel reads it + // with col_major B fragments, one 8-B vector load per lane). Every other + // (K, N) keeps the identity copy (raw row-major [K, N]). + const int elem_count = k * n; + if (k == 6144 && n == 2624) { + const int64_t pack_blocks = + (static_cast(elem_count) + kCopyBlockThreads - 1) / + kCopyBlockThreads; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_pack_nmajor_bytes_kernel), + dim3(static_cast(pack_blocks)), + dim3(kCopyBlockThreads), + 0, // no dynamic shared memory + stream, + raw_weight, + packed_weight, + k, + n); + } else { + const int elem_blocks = + (elem_count + kCopyBlockThreads - 1) / kCopyBlockThreads; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + dim3(static_cast(elem_blocks)), + dim3(kCopyBlockThreads), + 0, // no dynamic shared memory + stream, + raw_weight, + packed_weight, + elem_count); + } + + const int scale_blocks = + (n + kCopyBlockThreads - 1) / kCopyBlockThreads; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + dim3(static_cast(scale_blocks)), + dim3(kCopyBlockThreads), + 0, // no dynamic shared memory + stream, + weight_scale, + packed_weight_scale, + n); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/kv_b_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/kv_b_proj.hip new file mode 100644 index 00000000..33985f73 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/kv_b_proj.hip @@ -0,0 +1,1100 @@ +// @@variant shape=glm_tp8_kv_b_proj_m16 commit=a150485d1a63d39ca9251b13d3da304013431efd added=2026-08-31 +// median_us=8.607 p90_us=9.766 +// source=glm5-2-dsh-tp8-m16-test1-cb4d262c +// csrc/w8a8_gemm_hip.hip +// +// Worker_1 bootstrap (iteration 1) for the GLM5.2 TP8 M=16 decode shapes on +// Hygon K500SM_AI / gfx928 (wavefront = 64, 64 KiB LDS per CU): +// +// glm_tp8_q_b_proj_m16 : M=16, N=2048, K=2048 +// glm_tp8_kv_b_proj_m16 : M=16, N=3584, K=512 +// +// Logical operation (frozen contract, see int8_w8a8_gemm_api.py): +// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) +// * x_scale[m, 0] * weight_scale[n, 0] ) +// +// Strategy: assigned M=16 shapes run gfx928 DUMMA INT8 fast paths +// (m16n16k32, exact int32 accumulation). q_b (K=2048, N=2048) runs the +// grid-level split-K partial kernel: one 64-thread wavefront per (tile, +// split) block, the 32-aligned B K slice is staged once into LDS with +// vectorized 16-B bulk loads, one staging barrier, then a zero-barrier K +// loop whose B fragments come from LDS and whose A fragments are loaded +// straight from the L2-resident row-major A (8 contiguous bytes per lane per +// k32 step, 2-deep software pipeline; iteration 7 deepened the pipeline +// because the exact iteration-5 code object showed the 1-deep source +// degrading to an issue-mid-body / wait-same-body schedule that still +// exposed most of the L2 latency per step, and iteration 6 proved the LDS +// side of the body is not the critical path). Iteration 8 (HIP-only +// occupancy round): SPLIT_K 3 -> 8 (128 x 8 = 1,024 blocks = 8.53 blocks/CU +// = 2.13 waves/SIMD on the 120-CU device, uniform 256-K slices, LDS 11,264 -> +// 4,096 B/block -> 16 resident blocks/CU); each block publishes its exact +// int32 partial to a caller workspace plane and a separate combine+scale +// kernel (in the timed Graph) sums the split planes in ascending order and +// emits scaled bf16. Iteration 9's batched-staging HIP change (all four +// global_load_dwordx4 -> ds_write_b128 chains collapsed into one L2 round +// trip, verified in its exact code object) regressed 21.191 -> 22.033 us, so +// the staging latency is hidden by the 2.13 co-resident waves/SIMD and the +// partial kernel is at its HIP floor. Iteration 10 (this code) rewrites the +// combine+scale kernel as a single-pass int4-vectorized wavefront (one drain +// instead of four serialized load->wait->add->store chains; vmem_read 4,736 +// -> ~1,280 per replay) with bit-identical int32 sums and float rounding. +// kv_b (K=512, N=3584) runs the fused in-block split-K=2 staged kernel +// (iteration 3): two 64-lane wavefronts per block (blockDim 128, grid = +// N/16 = 224 = 3.73 waves/CU), A/B K-slices staged once into LDS with 16-B +// vector loads (padded 528 stride), each wave computes one exact +// k-ascending int32 partial over its uniform 256-K half (8 m16n16k32 steps, +// zero-vmcnt ds_read_b64 fragments), wave 0 publishes its partial to a +// per-block LDS plane, one END-of-K barrier, wave 1 adds s=0 then s=1 and +// emits scaled bf16 straight from registers. The (k, n) == (512, 3584) +// weight is the one-time n-major pack P[n*K+kk] = W[kk*N+n]; the generic +// scalar fallback decodes it via b_layout == 2 (paired M=2 stays exact). +// No caller workspace and no combine dispatch: one kernel launch per +// replay. +// Iteration 4 (this code): the kv_b fused split-K=2 kernel replaces the A/B +// LDS staging round trip with BOUNDED REGISTER PREFETCH (16 aligned 8-B +// uint2 fragment loads per lane per wave in two 4-step bursts, wave-1 +// scale operands prefetched into registers at kernel top, LDS reduced to +// the 1 KiB s_part plane, one barrier left). Geometry (224 blocks x 128 +// threads = 3.73 waves/CU), split, combine and epilogue expression are +// unchanged -> outputs stay bit-identical to iteration 3. +// Iteration 17 (this code): one-time tile-major B pack for (k, n) == +// (2048, 2048) (P[(t*K + kk)*16 + col] = B[kk][t*16 + col], out of the timed +// region and the Graph). The exact current-best code object + PMC show the +// partial kernel is memory-service-bound: the B staging loop's 64-lane 16-B +// chunks are scattered over 64 different rows' 128-B sectors (262,144 sector +// requests/replay, 16 B used per request) and each sector is shared by 8 +// tile-blocks whose accesses are spread over the whole kernel, so evicted +// sectors are re-fetched (PMC l2_misses 83,350 ~= 2.5x the 32,768 compulsory +// B sectors; ~10.7 MB/replay at ~660 GB/s effective). The pack makes every +// staging iteration one contiguous 1,024-B chunk (8 fully-consumed sectors, +// one consumer block, no re-fetch possible): B sector requests drop 262,144 +// -> 32,768 and B DRAM traffic drops toward the 4.19 MB compulsory set. The +// staged bytes, lds_b, the single barrier, the hoisted uniform A path, the +// k-ascending v_mmac body, the exact int32 partials, the combine kernel, the +// 1,024-block grid and the 2-dispatch Graph are unchanged -> outputs stay +// bit-identical to iteration 11 (0 mismatches). The (n, k) == (2048, 2048) +// generic scalar fallback (paired M=2 validation shape) decodes the pack via +// the b_layout flag; every other shape keeps the identity pack and the +// logical decode. +// Iteration 18 (HIP-only consolidation; final conditional inline-asm round, +// raw asm still gated): the uniform-path B staging loop's 4 chains are +// batched into one in-flight group (4 named int4 loads back-to-back, one +// drain, 4 ds_write_b128) - the exact current-best code object shows the +// rolled loop still executes one load in flight per block (4 serialized +// cold-DRAM round trips), and the pack made the 4 chunks contiguous so the +// batch is DRAM-row-friendly (iteration 9's pre-pack scatter regression does +// not apply). Named registers avoid iteration 14's scratch spill. Staged +// bytes, barrier, body, partials and combine are unchanged -> outputs stay +// bit-identical to iteration 17. +// The generic scalar kernel remains the fallback for every unmatched (m, n, +// k), including the paired M=2 API shapes that share (N, K) with the +// assigned M=16 shapes. +// +// Graph-safety: gemm_out only launches kernels on the caller-provided stream +// (PyTorch's current HIP stream); it performs no allocation, no compilation, +// no autotuning, no packing, and no host/device synchronization, and touches +// only the caller-provided `out` and `workspace` (split-K partial planes). + +#include +#include +#include + +#include +#include + +namespace { + +// blockDim must be a multiple of the gfx928 wavefront size (64). +constexpr int kScalarBlockThreads = 256; +constexpr int kPackBlockThreads = 256; + +// gfx928 DUMMA INT8 tile geometry for the M=16 decode fast path: the +// installed DTK exposes m16n16k32 signed-char fragments with int32 +// accumulation, and one 64-thread wavefront owns one independent tile. +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; +constexpr int kDummaWaveSize = 64; + +using namespace du::dumma; + +// One thread -> one output element. Adjacent threads map to adjacent columns +// (fastest-changing N dimension), so both the B row loads and the bf16 stores +// are coalesced; the A row load is broadcast across the threads of the same +// row. +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k, + int b_layout) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + + const int8_t* a_row = a + static_cast(row) * k; + // Packed-layout decode flags (b_layout): + // 0: logical [K][N] row-major (column stride n); + // 1: (k, n) == (2048, 2048) tile-major pack + // P[(t*K + kk)*16 + col_tile] = raw[kk][t*16 + col_tile] (t = col >> 4, + // col_tile = col & 15), so the logical column stride n becomes a 16-B + // stride in kk; + // 2: (k, n) == (512, 3584) n-major pack P[n*K + kk] = raw[kk*N + n], so + // the per-column bytes are contiguous in kk (stride 1). + const int8_t* b_col; + int64_t b_stride; + if (b_layout == 2) { + b_col = b + static_cast(col) * k; + b_stride = 1; + } else if (b_layout == 1) { + b_col = b + (static_cast(col >> 4) * k * 16 + (col & 15)); + b_stride = 16; + } else { + b_col = b + col; + b_stride = static_cast(n); + } + + // Exact int32 dot product over the complete K loop. The maximum assigned K + // is 2048, so |dot| <= 2048 * 128 * 128 = 2^25 stays far below INT32_MAX. + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * b_stride]); + } + + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Iteration 3: grid-level split-K partial + combine (CU-aligned probe). +// +// Iteration 1 (one wavefront per 16x16 tile, direct global fragment loads) +// measured 63.695 us and iteration 2 (in-block split-K=2 with the same direct +// byte loads) regressed to 69.771 us: the 16 global_load_ubyte + vmcnt chain +// per m16n16k32 step (PMC: 131,712 vmem_reads) exposes full memory latency +// with ~1.07 blocks/CU and no in-loop overlap, i.e. the direct fragment +// loader is the poison, not the grid size. Iteration 3 replaced the q_b +// (16, 2048, 2048) path with a grid-level split-K partial kernel that stages +// each 32-aligned K slice into LDS once with vectorized 16-B bulk loads +// (one barrier), then runs a zero-barrier LDS-only K loop; every block is a +// single wavefront and the grid is 128 tiles x SPLIT_K blocks (SPLIT_K=3 -> +// 384 blocks = 3.2 blocks/CU, a non-power-of-two CU-aligned occupancy probe +// from the trusted set). Exact int32 partials go to caller workspace planes +// and the combine+scale kernel runs in the timed Graph. +// +// Iteration 4 (bounded staging prefetch, kStageDepth=4) regressed to 43.876 +// us while keeping LDS at 22,784 B/block (2 resident blocks/CU), so the +// staging vmem chain was NOT the lever: the LDS-resident footprint itself is. +// +// Iteration 5 (this code): B-only staging. The exact iteration-3 code object +// shows the K loop's A fragment is 8 contiguous bytes per lane (one +// ds_read2_b32 in the staged kernel; 8 consecutive global_load_ubyte in the +// direct kernel) while the B fragment is 8 scattered 16-B-strided byte loads +// (16-way LDS bank conflicts). A is therefore moved out of LDS entirely: +// per k32 step each lane issues ONE 8-B global load from the L2-resident +// row-major A (1-deep software pipeline hides the L2 latency behind the +// current step's LDS loads + mma), the 11,520-B A tile disappears, LDS drops +// from 22,784 B to 11,264 B per block, and the resident blocks/CU rise from 2 +// to 5. B keeps the exact iteration-3 staged layout and byte loads. +// +// Iteration 8 (this code): HIP-only occupancy round. PMC/ISA evidence pins +// the binding occupancy limiter: 24 VGPR / 32 SGPR / 11,264 B LDS / 0 scratch +// -> LDS is the limiter at 5 resident blocks/CU (64 KiB / 11.25 KiB), while +// the grid (384 = 3.2 blocks/CU = 0.8 waves/SIMD) leaves residency half +// empty and SQ_WAIT_INST_LDS = 36/replay (~0.09/block) plus the exact loop +// body's s_waitcnt vmcnt(0) right before v_mmac (and the staging loop's 11 +// serialized global_load_dwordx4 -> vmcnt(0) -> ds_write_b128 chains) show +// the shaders stall on VMEM latency, not LDS. SPLIT_K is raised 3 -> 8: the +// trusted occupancy-probe split with uniform 256-K slices (64 k32 steps / 8 +// = 0 remainder), grid 128 x 8 = 1,024 blocks = 8.53 blocks/CU = 2.13 +// waves/SIMD, LDS 11,264 -> 4,096 B -> 16 resident blocks/CU (residency +// never binds, zero queueing, tail factor 1.055 vs 1.25). HBM is unchanged: +// every B byte is still staged exactly once (4.19 MB/replay) and the A +// request volume is invariant in S (each 32-KB L2-hot A is re-read per +// block, 4.19 MB/replay for any S); only the partial-plane writes grow +// 393,216 -> 1,048,576 B (8 planes <= the API 16-plane workspace for this +// shape). +// --------------------------------------------------------------------------- + +// Compile-time non-uniform, 32-aligned K slicing for the guarded K=2048 +// shape (64 k32 steps split into SPLIT_K contiguous slices; the remainder +// steps go to the first slices so every boundary stays a multiple of the +// DUMMA tile K=32, keeping each partial an exact k-ascending int32 sum). +template +struct KSplit2048 { + static constexpr int kSteps = 2048 / kDummaTileK; // 64 + static constexpr int kBase = kSteps / SPLIT_K; + static constexpr int kRem = kSteps % SPLIT_K; + static constexpr int kMaxSliceSteps = kBase + (kRem > 0 ? 1 : 0); + static constexpr int kMaxSlice = kMaxSliceSteps * kDummaTileK; + // Closed form (no C++14 constexpr loops): slice i contributes + // kBase + (i < kRem ? 1 : 0) k32 steps. + static constexpr int begin(int s) { + return (s * kBase + (s < kRem ? s : kRem)) * kDummaTileK; + } + static constexpr int slice(int s) { + return (kBase + (s < kRem ? 1 : 0)) * kDummaTileK; + } +}; + +template +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16_splitk_partial_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int* __restrict__ partials, + int n, + int k) { + using Split = KSplit2048; + // B-only staging (iteration 5): only the B tile is staged into LDS. The A + // fragment of this compiler/DUMMA header (verified in the exact iteration-3 + // code object, profiles/.../iteration4/current-best-isa) is 8 CONTIGUOUS + // bytes per lane: the staged path reads them with one ds_read2_b32 at + // (lane&15)*720 + (lane>>4)*8, and the direct kernel reads them with 8 + // consecutive global_load_ubyte at (lane&15)*k + (lane>>4)*8. A therefore + // needs no LDS round trip: one 8-B global load per lane per k32 step from + // the L2-resident row-major A reproduces the fragment exactly. Dropping the + // 11,520-B A tile halves the LDS footprint per block (22,784 -> 11,264 B) + // and raises the resident blocks/CU from 2 to 5 (64 KiB LDS / 11.25 KiB), + // while the A load latency is hidden by a 2-deep software pipeline in the + // K loop (iteration 7). The B tile keeps the exact iteration-3 layout and + // byte loads. Iteration 8 (occupancy round) raises SPLIT_K to 8 via the + // launcher only: kSlice becomes a uniform 256, LDS drops to 4,096 B/block + // and the grid becomes 1,024 blocks; the body below is unchanged. + constexpr int kLdsStrideB = kDummaTileN; // 16 + + const int lane = static_cast(threadIdx.x); // 0..63 + const int split = static_cast(blockIdx.x) % SPLIT_K; + const int tile = static_cast(blockIdx.x) / SPLIT_K; + const int n0 = tile * kDummaTileN; + const int kBeg = Split::begin(split); + const int kSlice = Split::slice(split); + + __shared__ __align__(16) int8_t lds_b[Split::kMaxSlice][kLdsStrideB]; + + // A fragment lane mapping (verified against the iteration-3 code object): + // lane l owns the 8 contiguous bytes of A row (l & 15) starting at column + // ((l >> 4) << 3) inside the k32 window, with x[e] = byte at address + e. + // Global row-major A (row stride k) yields one 8-B (int2) load per lane per + // step, always 8-B aligned (k, kBeg, kDummaTileK are all multiples of 8). + const int a_row = lane & 15; + const int a_col0 = (lane >> 4) << 3; // 0, 8, 16, 24 + const int8_t* a_lane = + a + static_cast(a_row) * k + kBeg + a_col0; + + // Iteration 11 (HIP-only consolidation): the exact iteration-10 code object + // (profiles/.../iteration11/current-best-isa, source digest + // 5187c0a5a308822b9914dae2c086c979f0fcda9093fa20009c201b19597794e1) + // compiles the 2-deep A rotation below into a schedule where the load for + // step i+2 issues at the TAIL of body i (global_load_dwordx2 at 0x4374) and + // the body opens with a full s_waitcnt vmcnt(0) (0x4324) that drains ALL + // outstanding loads, so the load consumed by body i's mma has only ~1 body + // (~60-120 cycles of issue) of lookahead against an L2 round trip of + // ~250-350 cycles: every one of the 8 bodies still exposes most of one + // A-load round trip (~1.6-2.4k cycles of the ~3.2-3.9k cycle per-block + // path), and the ~2.1 co-resident waves/SIMD stall in phase (264,192 VALU + // over 18.24 us profiled ~= 2% issue busy). Uniform slices (Split::kRem == + // 0, true for SPLIT_K=8) therefore take a compile-time hoisted path: all 8 + // int2 A loads issue back-to-back BEFORE the B staging loop, so the + // staging chains' serialized DRAM waits (4 x global_load_dwordx4 -> + // vmcnt(0) -> ds_write_b128 at 0x4190-0x41C0) give every A load + // ~2.4-3.6k cycles of slack, and the K body is fully unrolled so the + // compiler can overlap the independent B-fragment LDS assembly of later + // steps with the serial v_mmac chain. The B-only staging, lds_b layout, + // k-ascending mma order and exact int32 accumulation are unchanged, so the + // partials stay bit-identical to iteration 10 (0 mismatches; Graph replay + // with changed contents must pass). The generic kRem>0 path keeps the + // original rotation unchanged. + // a_buf is function-scoped so both branches can reference it; only the + // uniform (kRem == 0) branch fills it, and only that branch's unrolled body + // reads it (kMaxSliceSteps >= kBase covers every SPLIT_K instantiation). + int2 a_buf[Split::kMaxSliceSteps]; + if (Split::kRem == 0) { +#pragma unroll + for (int i = 0; i < Split::kBase; ++i) { + a_buf[i] = *reinterpret_cast(a_lane + i * kDummaTileK); + } + } + + // Stage only B: chunk c covers 16 contiguous bytes = exactly one tile row. + // Iteration 17 (HIP-only consolidation round; final conditional inline-asm + // round, raw asm still gated: plateau=false, phase=hip_only, + // raw_inline_asm_allowed=false, skill_allowed=false): the (k, n) == + // (2048, 2048) weight is now the one-time tile-major pack + // P[(t*K + kk)*16 + col] = B[kk][t*16 + col] (t = tile index), so lane c + // reads its 16 B from packed + ((t*K + kBeg + c)*16). The 64 lanes' chunks + // of one staging iteration are CONTIGUOUS (one 1,024-B chunk = 8 + // fully-consumed 128-B sectors) instead of 64 scattered 16-B requests + // spread over 64 different rows' sectors (row stride n = 2048), and each + // sector has exactly ONE consumer block (no cross-tile sharing, so an + // evicted sector is never re-fetched: PMC l2_misses 83,350 ~= 2.5x the + // 32,768 compulsory B sectors). The staged bytes are identical + // (B[kBeg+c][n0..n0+16)), so lds_b, the single staging barrier, the hoisted + // uniform A path, the unrolled k-ascending v_mmac body and the exact int32 + // partials are unchanged (bit-identical to iteration 11). + // Iteration 18 (HIP-only consolidation round; final conditional inline-asm + // round, raw asm still gated: plateau=false, phase=hip_only, + // raw_inline_asm_allowed=false, skill_allowed=false): the exact + // current-best code object (profiles/.../iteration18/current-best-isa, + // source digest 8b6bcaa866f26094d84762db45bd259a3a6492f900b42461e60c3f4264 + // 94853c) compiles the uniform-path staging loop below into a ROLLED loop + // with ONE global load in flight per block (global_load_dwordx4 -> + // s_waitcnt vmcnt(0) -> ds_write_b128 at 0x4A08-0x4A20), so every block + // pays 4 SERIALIZED cold-DRAM round trips for its 4 KiB slice while the + // 2.13 co-resident waves/SIMD stall in phase. Pre-pack, iteration 9 batched + // the same chains and regressed (+4%) because the 4 chains were 256 + // SCATTERED sector requests (row thrash) and the co-resident waves covered + // the drains; post-pack each chain is 8 contiguous sectors and the 4 chains + // tile the slice CONTIGUOUSLY (4 x 1,024 B = one 4,096-B group), so issuing + // all four per-lane chunk loads back-to-back (one in-flight group, one + // drain) collapses 4 serialized round trips toward ~1 while keeping the + // DRAM row-buffer-friendly. The four chunks are NAMED locals, not a + // runtime-indexed array: iteration 14's `int4 b_chunk[kChunksPerLane]` + // under a runtime `c < kSlice` trip bound spilled to scratch + // (private_segment_fixed_size 80) and regressed to 36.675 us; named + // registers with compile-time addressing keep the loads in VGPR (arch_vgpr + // 40 -> ~44-56; occupancy stays grid-limited at 8.53 blocks/CU, LDS 16 + // resident blocks/CU unchanged, no scratch). The staged bytes, lds_b + // addresses, single barrier, hoisted uniform A path, unrolled k-ascending + // v_mmac body and exact int32 partials are unchanged, so outputs stay + // bit-identical to iteration 17 (0 mismatches; Graph replay with changed + // contents must pass). The generic kRem>0 path keeps the rolled loop + // (runtime slice sizes); every other uniform instantiation falls back to it. + const int64_t b_tile_base = + static_cast(n0 / kDummaTileN) * k; + if (Split::kRem == 0) { + // kMaxSlice is constexpr and equals kSlice on this branch (kRem == 0). + constexpr int kChunks = Split::kMaxSlice / kDummaWaveSize; // 4 for SPLIT_K=8 + if (kChunks == 4) { + const int8_t* b_lane = + b + (b_tile_base + kBeg) * kDummaTileN + lane * kDummaTileN; + const int4 b0 = *reinterpret_cast(b_lane); + const int4 b1 = *reinterpret_cast( + b_lane + kDummaWaveSize * kDummaTileN); + const int4 b2 = *reinterpret_cast( + b_lane + 2 * kDummaWaveSize * kDummaTileN); + const int4 b3 = *reinterpret_cast( + b_lane + 3 * kDummaWaveSize * kDummaTileN); + *reinterpret_cast(&lds_b[lane][0]) = b0; + *reinterpret_cast(&lds_b[lane + kDummaWaveSize][0]) = b1; + *reinterpret_cast(&lds_b[lane + 2 * kDummaWaveSize][0]) = b2; + *reinterpret_cast(&lds_b[lane + 3 * kDummaWaveSize][0]) = b3; + } else { + for (int c = lane; c < kSlice; c += kDummaWaveSize) { + *reinterpret_cast(&lds_b[c][0]) = + *reinterpret_cast( + b + (b_tile_base + (kBeg + c)) * kDummaTileN); + } + } + } else { + for (int c = lane; c < kSlice; c += kDummaWaveSize) { + *reinterpret_cast(&lds_b[c][0]) = + *reinterpret_cast( + b + (b_tile_base + (kBeg + c)) * kDummaTileN); + } + } + __syncthreads(); + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + if (Split::kRem == 0) { + // Hoisted uniform path: a_buf holds all kSteps int2 A fragments, loaded + // before the B staging (every load has >= the staging chains' serialized + // DRAM waits of slack). The fully unrolled body keeps the exact + // k-ascending mma order of the generic path. +#pragma unroll + for (int i = 0; i < Split::kBase; ++i) { + const int8_t* a_bytes = reinterpret_cast(&a_buf[i]); +#pragma unroll + for (int e = 0; e < 8; ++e) { + a_frag.x[e] = a_bytes[e]; + } + const int k0 = i * kDummaTileK; + du_load_matrix_sync(b_frag, lds_b[0] + k0 * kLdsStrideB, kLdsStrideB); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + } else { + // Zero-barrier K loop: exact k-ascending int32 accumulation over + // [kBeg, kBeg + kSlice). B fragments come from LDS (unchanged layout). A + // fragments come from global with a 2-deep software pipeline (iteration + // 7): three int2 values (steps i, i+1, i+2) rotate through registers and + // the load for step i+3 issues at the top of step i's body, so every A + // load has ~1.5 step bodies (~400-550 cycles) to arrive in instead of the + // ~half-body window the 1-deep source degraded to in the exact + // iteration-5 code object (flat_load issued mid-body at 0x42C0, + // s_waitcnt vmcnt(0) at the END of the SAME body at 0x4380 before the + // rotation movs -> each step still exposes most of the ~350-450-cycle L2 + // latency). Iteration 6 proved the LDS side of the body is NOT the + // critical path (one ds_read_b64 replacing the 8 conflicted ds_read_u8 + + // ~12 VALU shuffles + lgkmcnt chain changed the median by 0.03%), so the + // per-step serial cost is the A-load wait; this rotation attacks exactly + // that. steps is runtime (e.g. 22/21/21 for SPLIT_K=3), so the loop is + // not unrolled and the rotation is explicit; the guards keep every load + // inside the A slice [kBeg, kBeg + kSlice). + const int steps = kSlice / kDummaTileK; + int2 a_cur = *reinterpret_cast(a_lane); // step 0 + int2 a_next = a_cur; + if (1 < steps) { + a_next = + *reinterpret_cast(a_lane + kDummaTileK); // step 1 + } + for (int i = 0; i < steps; ++i) { + int2 a_nxt2 = a_next; + if (i + 2 < steps) { + a_nxt2 = *reinterpret_cast( + a_lane + (i + 2) * kDummaTileK); + } + const int8_t* a_bytes = reinterpret_cast(&a_cur); +#pragma unroll + for (int e = 0; e < 8; ++e) { + a_frag.x[e] = a_bytes[e]; + } + const int k0 = i * kDummaTileK; + du_load_matrix_sync(b_frag, lds_b[0] + k0 * kLdsStrideB, kLdsStrideB); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + a_cur = a_next; + a_next = a_nxt2; + } + } + + // Publish the exact int32 partial to caller workspace plane `split` + // ([SPLIT_K][M][N] row-major). Every plane element is overwritten on every + // launch, so no workspace clear is needed and the combine kernel below is + // stream-ordered after this store. + int* plane = + partials + static_cast(split) * kDummaTileM * n + n0; + du_store_matrix_sync(plane, acc_frag, n, mem_row_major); +} + +// Combine+scale kernel: one wavefront per 16x16 output tile; sums the +// SPLIT_K int32 planes in ascending split order (bit-identical to the scalar +// k-ascending reference), applies x_scale * weight_scale in the scalar +// epilogue order, and stores bf16. Runs in the timed Graph after the partial +// kernel on the same stream. +// +// Iteration 10 (HIP-only consolidation round): the exact iteration-8 code +// object compiled the previous 4-iteration scalar loop (idx = lane, lane+64, +// lane+128, lane+192; 8 global_load_dword per iteration, one per split +// plane, at 0x464C-0x4694 of the combine symbol) into four SERIALIZED +// load -> wait-drain -> add -> store chains: each iteration's loads issue +// only after the previous iteration's bf16 store, so a single wave pays ~4 x +// (one DRAM round trip + add chain) while the 128-block grid delivers only +// 1.07 waves/CU (no second co-resident wave to overlap the drains; PMC: +// combine 5.12 us profiled, 4,736 vmem_read and 512 vmem_write per replay). +// This rewrite is ONE pass: lane l owns the contiguous 16-B chunk (row = +// l>>2, col4 = (l&3)<<2, 4 consecutive columns) of the tile, issues all +// SPLIT_K int4 plane loads back-to-back (a single drain, ~4x fewer memory +// instructions), accumulates the SAME ascending-s int32 sum per element, +// applies the same scalar epilogue expression order (float(dot) * x_scale * +// weight_scale), and stores the 4 bf16 as two dword stores. Per-element +// values, int32 accumulation order, and float rounding are identical to the +// previous combine, so the outputs stay bit-identical (0 mismatches) and the +// Graph structure (partial then combine on the caller stream) is unchanged. +template +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16_splitk_combine_kernel( + const int* __restrict__ partials, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n) { + const int lane = static_cast(threadIdx.x); + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + const int64_t plane = static_cast(kDummaTileM) * n; + + // lane -> one contiguous 16-B chunk of the 16x16 tile: row = lane/4, first + // column = (lane%4)*4. 64 lanes x 4 elements = 256 = the full tile, exactly + // once. All chunk addresses are 16-B aligned (n, n0, col4 and plane are all + // multiples of 4 int32). + const int row = lane >> 2; + const int col4 = (lane & 3) << 2; + const int* chunk = + partials + static_cast(row) * n + n0 + col4; + + int32_t dot[4] = {0, 0, 0, 0}; +#pragma unroll + for (int s = 0; s < SPLIT_K; ++s) { + const int4 p = *reinterpret_cast(chunk + s * plane); + dot[0] += p.x; + dot[1] += p.y; + dot[2] += p.z; + dot[3] += p.w; + } + + const float xs = x_scale[row]; + union { + float f[4]; + float4 v; + } ws; + ws.v = *reinterpret_cast(weight_scale + n0 + col4); + union { + __hip_bfloat16 b[4]; + uint32_t u[2]; + } packed; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float scaled = static_cast(dot[e]) * xs * ws.f[e]; + packed.b[e] = __float2bfloat16(scaled); + } + uint32_t* out32 = reinterpret_cast( + out + static_cast(row) * n + n0 + col4); + out32[0] = packed.u[0]; + out32[1] = packed.u[1]; +} + +// --------------------------------------------------------------------------- +// kv_b (m, n, k) == (16, 3584, 512) fused in-block split-K=2 register +// prefetch kernel (iteration 4: bounded register/LDS prefetch). +// +// Iteration 3 accepted the fused in-block split-K=2 staged kernel (two +// 64-lane wavefronts per block, blockDim 128, grid = N/16 = 224 = 3.73 +// waves/CU; A/B K-halves staged once into LDS with 16-B vector loads, one +// staging barrier, zero-vmcnt ds_read_b64 fragments, wave-0 publish to a +// per-block LDS plane, one END-of-K barrier, wave-1 ascending combine and +// scaled-bf16 epilogue) at 11.290 us median / 11.558 us p90 - stable, but +// ~1.5 us above the iteration-1 one-wave staged geometry's 9.765 us median +// (rejected only on environmental ~50 us p90 spikes). The exact iteration-3 +// code object shows the staged fragment path is that extra cost: 8 +// global_load_dwordx4 -> 8 serialized s_waitcnt vmcnt(N) + ds_write_b128 +// staging chains, one staging s_barrier, 8 ds_read2_b64 with lgkmcnt waits +// feeding 8 v_mmac, and a wave-1 epilogue that issues x_scale/weight_scale +// only at the tail and stalls per element (s_waitcnt vmcnt(1) before every +// epilogue mul) - plus 53,760 LDS bank conflicts and 8,512 lds_instructions +// per replay (PMC iteration4). +// +// This round replaces the LDS staging round trip with BOUNDED REGISTER +// PREFETCH while keeping the accepted geometry, split, combine, epilogue +// expression and grid byte-for-byte: each lane loads exactly its own 8-B A +// fragment (a[lane&15][w*256 + s*32 + (lane>>4)*8 + i], 8-B aligned in the +// row-major A) and its own 8-B B fragment (8 contiguous bytes of one packed +// n-major row P[n0 + lane&15][...]) directly from global as 16 uint2 loads +// per lane per wave, issued in two 4-step bursts (burst 1 in flight while +// burst 0's MMAC chain runs; the depth-1 rotation the hy3 TP4 o_proj +// register-transport lineage validated). No A/B LDS staging, no staging +// barrier, zero ds_read/ds_write in the K loop; LDS drops to the 1 KiB +// s_part plane only (17,920 -> 1,024 B/block). The wave-1 scale operands +// (x_scale[row], weight_scale[n0 + col] x4) are prefetched into registers +// at kernel top so their L2 latency overlaps the whole K loop and the +// epilogue never waits on vmcnt. Fragment bytes, per-wave k-ascending int32 +// order, the ascending split sum, and the epilogue float expression are +// identical to iteration 3 -> outputs stay bit-identical (0 mismatches +// expected). Reused bytes: B pack bytes are compulsory (each byte read by +// exactly one block, one K-half); A (8 KiB) is re-read per block from L2 +// (unchanged); scales per tile. No caller workspace, no combine dispatch, +// one kernel launch per replay (unchanged). Falsifiable claim: if LDS +// staging, the staging barrier and bank conflicts were the iteration-3 +// cost, median/p90 improve below 11.290/11.558 us; if the global +// fragment-load latency dominates instead, the change is flat and the +// split-combine tail or the load path is the floor. +// --------------------------------------------------------------------------- + +// See the section comment. blockDim = 128 = two 64-lane wavefronts; each +// wavefront owns one uniform 256-K half (8 k32 steps) and the per-block LDS +// combine keeps the int32 order bit-identical to the reference. +__global__ __launch_bounds__(2 * kDummaWaveSize) void +w8a8_dumma_m16_k512_sk2_reg_kernel( + const int8_t* __restrict__ a, // [16, K] row-major + const int8_t* __restrict__ b, // [N, K] n-major pack P[n*K+kk]=W[kk*N+n] + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n, + int k) { + // Exact assigned shape: K = 512 = 16 * 32. The compile-time step count is + // required for the fully unrolled per-wave K body. + constexpr int kExactKSteps = 512 / kDummaTileK; // 16 + if (k != kDummaTileK * kExactKSteps) { + return; + } + const int tid = static_cast(threadIdx.x); // 0..127 + const int w = tid >> 6; // wavefront 0 / 1 + const int lane = tid & 63; + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + // Wave w's uniform K half starts at w * 256 (k32 steps 0..7 / 8..15). + const int koff = w * (kExactKSteps / 2) * kDummaTileK; + + __shared__ __align__(16) int32_t s_part[kDummaTileM * kDummaTileN]; + + // Fragment lane layout (the verified iteration-1/3 layout, byte-identical): + // row = lane & 15, col = (lane >> 4) * 8, a.x[i] = a[row][koff + s*32 + + // col + i] and b.x[i] = P[n0 + row][koff + s*32 + col + i], i = 0..7. Both + // byte offsets are multiples of 8 (row-major A rows are K=512 B; the pack + // rows are K=512 B), so each fragment is ONE aligned 8-B (uint2) global + // load straight from the source buffers - no LDS staging, no cross-lane + // sharing, no staging barrier. + const int row = lane & 15; + const int col = (lane >> 4) * 8; + const int64_t a_off = static_cast(row) * k + koff + col; + const int64_t b_off = static_cast(n0 + row) * k + koff + col; + + // Bounded register prefetch: the wave's 8 A/B fragment pairs (16 x 8 B per + // lane) load in two 4-step bursts; burst 1 is issued before burst 0's MMAC + // chain so its L2/DRAM latency overlaps the first 4 v_mmac (the depth-1 + // rotation the hy3 TP4 o_proj register-transport lineage validated). + // Expect the exact code object to show 16 global_load_dwordx2 with <= 2 + // vmcnt drains before the 8-mmac chain and zero ds instructions in the K + // body. + const uint2 a0 = *reinterpret_cast(a + a_off); + const uint2 b0 = *reinterpret_cast(b + b_off); + const uint2 a1 = *reinterpret_cast(a + a_off + 32); + const uint2 b1 = *reinterpret_cast(b + b_off + 32); + const uint2 a2 = *reinterpret_cast(a + a_off + 64); + const uint2 b2 = *reinterpret_cast(b + b_off + 64); + const uint2 a3 = *reinterpret_cast(a + a_off + 96); + const uint2 b3 = *reinterpret_cast(b + b_off + 96); + const uint2 a4 = *reinterpret_cast(a + a_off + 128); + const uint2 b4 = *reinterpret_cast(b + b_off + 128); + const uint2 a5 = *reinterpret_cast(a + a_off + 160); + const uint2 b5 = *reinterpret_cast(b + b_off + 160); + const uint2 a6 = *reinterpret_cast(a + a_off + 192); + const uint2 b6 = *reinterpret_cast(b + b_off + 192); + const uint2 a7 = *reinterpret_cast(a + a_off + 224); + const uint2 b7 = *reinterpret_cast(b + b_off + 224); + + // Scale operands prefetched at kernel top (wave 1 only): consumed after + // the END-of-K barrier, so their L2 latency overlaps the whole K body and + // the epilogue has no vmcnt wait (the iteration-3 ISA stalled per element + // with s_waitcnt vmcnt(1) before every epilogue mul). + float xs = 0.0f; + float ws0 = 0.0f; + float ws1 = 0.0f; + float ws2 = 0.0f; + float ws3 = 0.0f; + if (w == 1) { + const int ecol4 = lane >> 4; // 0..3 + xs = x_scale[row]; + ws0 = weight_scale[n0 + ecol4 + 0]; + ws1 = weight_scale[n0 + ecol4 + 4]; + ws2 = weight_scale[n0 + ecol4 + 8]; + ws3 = weight_scale[n0 + ecol4 + 12]; + } + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + // Zero-LDS, zero-barrier K body: 8 exact k-ascending m16n16k32 steps over + // the wave's uniform 256-K half, each fragment pair from the registers + // loaded above (the library lane layout x[i] is reproduced byte-identically + // by the 8-B memcpys, so the int32 accumulation order is unchanged). + memcpy(a_frag.x, &a0, sizeof(a0)); + memcpy(b_frag.x, &b0, sizeof(b0)); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + memcpy(a_frag.x, &a1, sizeof(a1)); + memcpy(b_frag.x, &b1, sizeof(b1)); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + memcpy(a_frag.x, &a2, sizeof(a2)); + memcpy(b_frag.x, &b2, sizeof(b2)); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + memcpy(a_frag.x, &a3, sizeof(a3)); + memcpy(b_frag.x, &b3, sizeof(b3)); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + memcpy(a_frag.x, &a4, sizeof(a4)); + memcpy(b_frag.x, &b4, sizeof(b4)); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + memcpy(a_frag.x, &a5, sizeof(a5)); + memcpy(b_frag.x, &b5, sizeof(b5)); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + memcpy(a_frag.x, &a6, sizeof(a6)); + memcpy(b_frag.x, &b6, sizeof(b6)); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + memcpy(a_frag.x, &a7, sizeof(a7)); + memcpy(b_frag.x, &b7, sizeof(b7)); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + + // Fused per-block combine (unchanged from iteration 3): wave 0 publishes + // its [0, 256) partial to s_part through the library store, one END-of-K + // barrier, then wave 1 adds split 0 + split 1 per element in ascending + // split order (identical int32 expression to the scalar k-ascending + // reference) and emits scaled bf16 straight from registers with the + // verified gfx928 accumulator ownership: row = lane&15, col_mod4 = + // lane>>4, x[i] maps to column col_mod4 + 4*i. + if (w == 0) { + du_store_matrix_sync(s_part, acc_frag, kDummaTileN, mem_row_major); + } + __syncthreads(); + if (w == 1) { + const int erow = lane & 15; + const int ecol4 = lane >> 4; // 0..3 + const int32_t dot0 = + s_part[erow * kDummaTileN + ecol4 + 0] + acc_frag.x[0]; + out[static_cast(erow) * n + n0 + ecol4 + 0] = __float2bfloat16( + static_cast(dot0) * xs * ws0); + const int32_t dot1 = + s_part[erow * kDummaTileN + ecol4 + 4] + acc_frag.x[1]; + out[static_cast(erow) * n + n0 + ecol4 + 4] = __float2bfloat16( + static_cast(dot1) * xs * ws1); + const int32_t dot2 = + s_part[erow * kDummaTileN + ecol4 + 8] + acc_frag.x[2]; + out[static_cast(erow) * n + n0 + ecol4 + 8] = __float2bfloat16( + static_cast(dot2) * xs * ws2); + const int32_t dot3 = + s_part[erow * kDummaTileN + ecol4 + 12] + acc_frag.x[3]; + out[static_cast(erow) * n + n0 + ecol4 + 12] = __float2bfloat16( + static_cast(dot3) * xs * ws3); + } +} + +// One launch per replay: grid = N/16 two-wave blocks, no caller workspace, +// no combine dispatch. Requires K = 512 and N divisible by 16 (guaranteed +// by the exact-shape guard that reaches this launcher). +void launch_dumma_m16_k512_sk2_reg( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + __hip_bfloat16* out, + int n, + int k, + hipStream_t stream) { + const dim3 grid(static_cast(n / kDummaTileN)); + const dim3 block(2 * kDummaWaveSize); + hipLaunchKernelGGL( + w8a8_dumma_m16_k512_sk2_reg_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out, + n, + k); +} + +// Identity device-to-device copy used by the bootstrap pack op. Works for any +// element type and any (K, N): unmatched shapes keep this generic fallback. +template +__global__ void w8a8_identity_copy_kernel( + const T* __restrict__ src, + T* __restrict__ dst, + int64_t numel) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < numel) { + dst[i] = src[i]; + } +} + +// One-time tile-major B pack for the (k, n) == (2048, 2048) q_b weight, +// outside the timed region and the Graph: P[(t*K + k)*16 + col] = +// raw[k][t*16 + col], with t = n-tile index (0..N/16-1). The 16 B of one +// (tile, k-row) stay contiguous and the 64 rows of a staging iteration become +// one contiguous 1,024-B chunk, so the split-K partial kernel stages fully +// consumed 128-B sectors instead of 64 scattered 16-B requests (iteration 17; +// see the partial kernel's staging comment). Same 4,194,304 B, same bytes, so +// the packed buffer stays graph-stable and the API contract is unchanged; the +// (n, k) == (2048, 2048) generic scalar fallback decodes this layout via the +// b_layout flag. Every other (k, n) keeps the identity pack. +__global__ void w8a8_pack_b_tilemajor_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int t = static_cast(blockIdx.x); // n-tile 0..n/16-1 + // First k-row of this thread's 16-row chunk: threadIdx.x in 0..127 -> + // chunks [0,16), [16,32), ..., [2032,2048), covering every k-row exactly + // once (128 threads x 16 rows = the full K=2048). + const int k16 = static_cast(threadIdx.x) * 16; + const int64_t src = + static_cast(k16) * n + static_cast(t) * 16; + const int64_t dst = (static_cast(t) * k + k16) * 16; +#pragma unroll + for (int r = 0; r < 16; ++r) { + *reinterpret_cast(packed + dst + r * 16) = + *reinterpret_cast(raw + src + r * n); + } +} + +// One-time n-major transpose pack for the (k, n) == (512, 3584) kv_b weight, +// outside the timed region and the Graph: P[n_idx * k + kk] = raw[kk * n + +// n_idx]. Packed rows are full K columns of W, so every lane's 8-B B fragment +// is 8 contiguous bytes of one packed row (the fused staged kernel reads it +// as one ds_read_b64). One thread per (n_idx, 16-B K chunk): the 16 source +// bytes are 16-B-strided in W (byte loads), but the 16 packed bytes are +// contiguous (two 8-B stores). Same 1,835,008 B, same graph-stable buffer; +// the (n, k) == (3584, 512) generic scalar fallback (paired M=2 validation +// shape) decodes this layout via b_layout == 2. +__global__ void w8a8_pack_b_nmajor_kernel( + const int8_t* __restrict__ w, + int8_t* __restrict__ packed, + int n, + int k) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t chunks = static_cast(n) * (k / 16); + if (idx >= chunks) { + return; + } + const int k16 = static_cast(idx % (k / 16)); + const int n_idx = static_cast(idx / (k / 16)); + const int8_t* src = w + static_cast(k16) * 16 * n + n_idx; + int8_t* dst = packed + static_cast(n_idx) * k + k16 * 16; + uint64_t lo = 0; + uint64_t hi = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) { + lo |= static_cast(static_cast(src[i * n])) << (8 * i); + hi |= static_cast(static_cast(src[(i + 8) * n])) + << (8 * i); + } + *reinterpret_cast(dst) = lo; + *reinterpret_cast(dst + 8) = hi; +} + +void launch_scalar_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + __hip_bfloat16* out, + int m, + int n, + int k, + int b_layout, + hipStream_t stream) { + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads)); + const dim3 block(kScalarBlockThreads); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out, + m, + n, + k, + b_layout); +} + +// Grid-level split-K partial kernel + separate combine+scale kernel (both in +// the timed Graph, on the caller stream). One wavefront per (tile, split) +// block; exact int32 partials live in the caller workspace planes, so the +// partial kernel needs no cross-block synchronization and the combine kernel +// sums the planes in ascending split order. Requires K = 2048 and N divisible +// by 16 (guaranteed by the exact-shape guard that reaches this launcher). +template +void launch_dumma_m16_splitk( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + __hip_bfloat16* out, + void* workspace, + int n, + int k, + hipStream_t stream) { + auto* partials = reinterpret_cast(workspace); + const dim3 partial_grid( + static_cast(n / kDummaTileN * SPLIT_K)); + const dim3 block(kDummaWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk_partial_kernel), + partial_grid, + block, + 0, + stream, + a, + b, + partials, + n, + k); + const dim3 combine_grid(static_cast(n / kDummaTileN)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk_combine_kernel), + combine_grid, + block, + 0, + stream, + partials, + x_scale, + weight_scale, + out, + n); +} + +} // namespace + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // The workspace holds the split-K int32 partial planes for the q_b + // grid-level split path; scalar/tail paths ignore it. + auto* out_bf16 = reinterpret_cast<__hip_bfloat16*>(out); + + // Assigned M=16 shapes run the DUMMA fast paths. q_b (K=2048, N=2048) runs + // the grid-level split-K=8 partial kernel + combine kernel (iteration 8: + // 1,024 blocks = 8.53 blocks/CU = 2.13 waves/SIMD, uniform 256-K slices, + // B-only staging at 4,096 B LDS/block = 16 resident blocks/CU, A fragments + // loaded directly from the L2-resident row-major A, workspace int32 + // partials, combine in the timed Graph). Iteration 17: the (k, n) == + // (2048, 2048) weight is the one-time tile-major pack P[(t*K+kk)*16+col] = + // B[kk][t*16+col] (see w8a8_pack_b_tilemajor_kernel), so the partial + // kernel's staging reads contiguous fully-consumed sectors; the scalar + // fallback for this (n, k) decodes the pack via b_layout == 1. kv_b + // (K=512, N=3584) runs the fused in-block split-K=2 staged kernel + // (iteration 3): two waves per block, per-block LDS combine, no caller + // workspace, no combine dispatch, one launch per replay. + // Each guard pins the exact M=16 so the paired M=2 API shape with the same + // (N, K) still reaches the generic scalar fallback below. + if (m == 16 && n == 2048 && k == 2048) { + // 8 int32 partial planes of [16][2048]; the caller workspace contract + // allocates 16 planes for this shape (2,097,152 B), so the guard is + // defensive only. + constexpr int64_t kRequiredWorkspace = + 8LL * 16 * 2048 * static_cast(sizeof(int32_t)); + if (workspace != nullptr && workspace_bytes >= kRequiredWorkspace) { + launch_dumma_m16_splitk<8>( + a, b, x_scale, weight_scale, out_bf16, workspace, n, k, stream); + } else { + // Defensive-only: the packed B must be decoded, so the packed scalar + // decode is always correct here. + launch_scalar_gemm( + a, b, x_scale, weight_scale, out_bf16, m, n, k, + /*b_layout=*/1, stream); + } + return; + } + if (m == 16 && n == 3584 && k == 512) { + // Iteration 4 (bounded register/LDS prefetch): two 64-lane wavefronts + // per block (blockDim 128), grid = N/16 = 224 blocks = 3.73 waves/CU; + // each lane prefetches its own 8-B A fragment (row-major A) and 8-B B + // fragment (n-major pack) per k32 step into registers (16 aligned uint2 + // loads per lane per wave in two 4-step bursts), computes its exact + // k-ascending int32 partial over its uniform 256-K half (8 m16n16k32 + // steps, zero LDS in the K body), wave 0 publishes to the per-block LDS + // plane, one END-of-K barrier, wave 1 adds s=0 then s=1 and emits scaled + // bf16 from registers (scale operands prefetched at kernel top). The + // (k, n) == (512, 3584) weight is the one-time n-major pack (b_layout + // == 2 in the scalar fallback). No caller workspace and no combine + // dispatch: one kernel launch per replay. + launch_dumma_m16_k512_sk2_reg( + a, b, x_scale, weight_scale, out_bf16, n, k, stream); + return; + } + + // Generic scalar fallback for every unmatched (m, n, k). The (n, k) == + // (2048, 2048) pair (including the paired M=2 validation shape) decodes + // the tile-major pack (b_layout == 1) and the (n, k) == (3584, 512) pair + // decodes the n-major pack (b_layout == 2); every other shape reads the + // logical layout. + launch_scalar_gemm( + a, b, x_scale, weight_scale, out_bf16, m, n, k, + (n == 2048 && k == 2048) ? 1 : ((n == 3584 && k == 512) ? 2 : 0), + stream); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // One-time packing outside the timed region and the Graph. The (k, n) == + // (2048, 2048) q_b weight is repacked tile-major + // (P[(t*K+kk)*16+col] = raw[kk][t*16+col], iteration 17); the (k, n) == + // (512, 3584) kv_b weight is repacked n-major + // (P[n*K+kk] = raw[kk*N+n], iteration 3, for the fused staged kernel's + // contiguous B rows). The matching GEMM interpretations live in the + // kernel staging and the b_layout scalar decode. Every other (K, N) shape + // keeps the identity copy so the logical [K, N] row-major layout is + // preserved. + if (k == 2048 && n == 2048) { + const dim3 pack_grid(static_cast(n / 16)); + hipLaunchKernelGGL( + w8a8_pack_b_tilemajor_kernel, + pack_grid, + dim3(128), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else if (k == 512 && n == 3584) { + const int64_t total_chunks = static_cast(n) * (k / 16); + const dim3 pack_grid(static_cast( + (total_chunks + kPackBlockThreads - 1) / kPackBlockThreads)); + hipLaunchKernelGGL( + w8a8_pack_b_nmajor_kernel, + pack_grid, + dim3(kPackBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + n, + k); + } else { + const int64_t weight_numel = static_cast(k) * n; + const dim3 weight_grid(static_cast( + (weight_numel + kPackBlockThreads - 1) / kPackBlockThreads)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + weight_grid, + dim3(kPackBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_numel); + } + + const int64_t scale_numel = static_cast(n); + const dim3 scale_grid(static_cast( + (scale_numel + kPackBlockThreads - 1) / kPackBlockThreads)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + scale_grid, + dim3(kPackBlockThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + scale_numel); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/o_proj.hip new file mode 100644 index 00000000..3ec28b39 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/o_proj.hip @@ -0,0 +1,559 @@ +// @@variant shape=glm_tp8_o_proj_m16 commit=c1cfe1975c2fcb8d4bc23c86e8026afae225a4c4 added=2026-08-31 +// median_us=24 p90_us=24.02 +// source=glm5-2-dsh-tp8-m16-test1-cb4d262c +// MetaInfer INT8 W8A8 GEMM for Hygon K500SM_AI / gfx928. +// +// Iteration 10 (worker_2): HIP-only request-granularity round on the accepted +// iteration-6 baseline (the harness restored the accepted source digest +// 15eae0f3... after iteration 9's rejection). The occupancy sweep (iterations +// 7, 8: split 3 -> 27.02 us, split 1 -> 29.76 us) and the per-wait MLP probe +// (iteration 9: k128 stages, 4 KiB/wait -> 26.46 us) all falsified their +// models, leaving the 768-stream / 2 KiB-in-flight / 16-wait-per-wave +// configuration the measured memory-system sweet spot. The one untested +// dimension at that sweet spot is per-instruction request granularity: the +// iteration-6 steady state issues FOUR global_load_dwordx2 per k64 step (2 A +// halves + 2 B halves of 512 B each per wave), and iteration 9's prediction +// named the 16-B true-vector packed-B variant as the next attack if deeper +// per-wait MLP stayed flat/regressed (it regressed). Change: repack B into +// packed16[n_tile][k64_step][lane][16] -- each lane's two k32 fragment slots +// (lower half bytes 0..7, upper half bytes 8..15) side by side, a pure +// same-byte-count permutation -- so each wave's 1-KiB B slice per k64 step is +// ONE aligned 16-B-per-lane global_load_dwordx4 instead of two 8-B dwordx2 +// loads. Stream count (768), per-wait in-flight bytes (2 KiB), wait count +// (16/wave), A transport, geometry, split, LDS/barrier, global bytes and the +// exact k-ascending int32 order are all unchanged. +// +// Iteration 10 repair 2: the first iteration-10 draft read the packed-B slot +// with the k64-step index missing its *kWaveSize stride (b_slots[ps0+ps] +// instead of b_slots[(ps0+ps)*kWaveSize]), so every k64 step after the first +// loaded a slot 16 B away instead of 1 KiB away and corrupted ~all outputs +// (98211/98304 mismatches). The 16-B slot layout, the one-dwordx4-per-step +// load and the k64 pipeline are all preserved; the wave's slot anchor is now +// b_base + ps0*kWaveSize + lane with in-loop index ps*kWaveSize. +// +// Iteration 6 (worker_2, physical GPU 2): HIP-only pipeline round on the +// assigned shape glm_tp8_o_proj_m16 (M=16, N=6144, K=2048). The staging-family +// decision is DIRECT (no LDS staging for A or B), argued from the iteration-5 +// L2/VMEM PMC evidence: l2_misses 201,111 ~= 196,608 unique B 64-B lines +// (12,582,912 B = the whole once-read B stream) so B has ZERO reuse and an +// LDS round trip for it would only re-add the ds_write/ds_read + lgkmcnt +// chain that iteration 4 priced at 45.03 us vs the 26.83 us direct kernel; +// A's HBM footprint is ~88,960 B/replay (~0.7% of traffic, 32 KiB L2-resident +// working set) so A-only staging saves no HBM bytes and cannot remove a wait +// from the step chain (A and B loads share one vmcnt). The measured per-wave +// budget is ~28,475 cycles over 32 k32 steps (~890/step): the kernel is +// serialized per-step global-latency-wait bound (in-flight ~1 KiB/wave at +// each wait), not issue-bound and not HBM-bound (479 GB/s logical << peak). +// Change: keep the exact iteration-5 geometry, packed layout, bytes and +// int32 accumulation, but restructure the K loop from 32 one-wait k32 steps +// into 16 k64 pipeline steps -- each step = TWO m16n16k32 halves, all four +// per-lane 8-B cooperative loads (2 A + 2 B) issued together and consumed by +// the step's two serialized du_mma_sync calls, so ONE vmcnt(0) wait serves +// 2 x K; the per-wave wait count halves 32 -> 16 and the bytes in flight at +// each wait point double (2 KiB/wave), with depth-1 k64 prefetch (step j+1's +// four loads issued before step j's two-MMAC burst). +// +// History: iteration 1 established the minimal native DUMMA tile (scalar +// bootstrap 251 us -> 85.02 us median; 384 one-wave blocks, 3.2 waves/CU); +// iteration 2 (accepted, 50.16 us) doubled resident waves to 6.4/CU via +// in-block split-K=2 (384 blocks x 128 thr); iteration 3 (rejected, 51.50 us) +// probed grid-side split-K=5 at 16.0 one-wave blocks/CU and was flat -- extra +// co-resident waves are NOT the lever; iteration 4 (accepted, 45.03 us) added +// bounded depth-1 register prefetch of the global loads, which removed the +// two serialized global-latency waits per step but left the per-step LDS +// fragment-read chain (8 ds_read_u8 B byte reads at stride kBStride + packing +// VALU + ~3 ds_write, 436,224 LDS bank conflicts) as the dominant per-wave +// cost; iteration 5 (accepted, 26.83 us) implemented the packed-layout arm of +// the mandate: launch_pack_w8a8_weight permutes the exact (k,n)==(2048,6144) +// weight ONCE, outside the timed region and out of Graph capture, into the B +// fragment-slot layout packed[n_tile][k_step][lane][8] (same byte count K*N +// as the logical weight; identity copy retained for every other (K,N)); the +// DUMMA kernel then consumes A and B directly from global memory as two +// aligned 8-B vector loads per lane per k32 step (register-only K-loop +// transport with depth-1 prefetch), eliminating the LDS staging round trip, +// the 8 ds_read_u8 byte reads, the packing VALU, and every in-loop lgkmcnt +// wait. +// +// Logical contract (see int8_w8a8_gemm_api.py): +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// * x_scale[m, 0] * weight_scale[n, 0] ) +// +// This file provides the two stable host launch symbols: +// launch_w8a8_gemm(...) -- timed operator, graph-safe +// launch_pack_w8a8_weight(...) -- optional out-of-timed-region weight prep +// +// Iteration-5 round strategy (packed-B fragment-slot transport): +// * the fragment lane mapping is taken from the exact DTK du_mma.hpp +// loaders, which the accepted iteration-2/4 kernels call with the SAME +// m16n16k32 signed-char fragment type that du_mma_sync consumes (verified +// against the iteration-4 code object's ds_read address arithmetic): +// matrix_a row_major: row = lane & 15, col = (lane >> 4) << 3, +// x[i] = p[row*ldm + col + i] (8 consecutive B); +// matrix_b row_major: row = lane & 15, col = (lane >> 4) << 3, +// x[i] = p[(col + i)*ldm + row]; +// so each lane's B fragment bytes for one k32 step are the 8 k-ascending +// bytes raw[(32*s + 8*(lane>>4) + i) * n + 16*t + (lane&15)], i=0..7; +// * pack kernel: one thread per (n_tile, k_step, lane) writes the 8 raw +// bytes (strided by n) into one aligned 8-B slot of the same-size packed +// tensor (permutation, byte count unchanged); runs once per weight, +// outside the timed GEMM and outside Graph capture; +// * DUMMA kernel: grid = N/16 = 384 blocks x 128 thr = 2 wavefronts, +// in-block split-K=2 (wave w owns k32 steps [32w, 32w+32) = K=1024 each) +// -- the exact proven iteration-2/4 geometry and resident-wave count +// (768 waves = 6.4/CU, all 120 CUs at 3.2 blocks/CU) is frozen; per k32 +// step each lane issues ONE aligned 8-B A load (logical row-major A: +// a + (lane&15)*k + 32*s + 8*(lane>>4), L2-hot 32 KiB, reused by all 384 +// blocks) and ONE aligned 8-B packed-B load (fully coalesced: 64 lanes x +// 8 B = 512 B contiguous per step), with depth-1 prefetch of the NEXT +// stage's eight loads (4 steps x 2) issued before the current 4-MMAC +// burst (#pragma-unrolled K=128 stage), so the compiler-inserted vmcnt(0) +// wait lands at the next stage's first fragment fill and no LDS hop / +// lgkmcnt wait exists in the K loop; +// * split combine unchanged: wave 0 publishes its 16x16 int32 partial into +// the single LDS c_part plane, one END-of-K __syncthreads(), wave 1 adds +// both halves (order-independent: no partial can overflow int32, max +// |dot| = 2048*127*127 << 2^31) and emits the scaled bf16 output from +// registers -- bit-identical to the exact k-ascending int32 reference; +// * the launch is guarded by the exact (m,n,k)=(16,6144,2048) shape; every +// other shape -- including the paired M=2 API shape with the same (N,K) -- +// keeps the scalar generic fallback, which now decodes the packed layout +// elementwise for (n,k)==(6144,2048) and reads the logical row-major +// weight for every other (K,N); +// * pack_weight stays an identity device-to-device copy for unmatched +// (K,N); the DUMMA kernel consumes packed_weight for the assigned shape. + +// Known-good include order for this DTK: HIP runtime first, then +// hip_bfloat16.h, then du_mma.h (du_mma.h is not self-contained when it is +// included before the HIP runtime headers). +#include +#include +#include + +#include + +namespace { + +// gfx928 wavefront is 64 lanes; every block dimension must be a multiple of +// 64. 128 threads = 2 wavefronts keeps the scalar bootstrap simple while +// still giving a straightforward grid over M*N. +constexpr int kGemmThreads = 128; +constexpr int kPackThreads = 256; + +// Assigned-shape DUMMA constants. +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; // gfx928 INT8 DUMMA K depth +constexpr int kWaveSize = 64; + +// Iteration-2 launch geometry: two 64-lane wavefronts per block, each owning +// a contiguous K-half of the block's single 16x16 N tile (in-block split-K=2, +// 128 threads/block, no cross-wave barrier inside the K loop). +constexpr int kWavesPerTile = 2; +constexpr int kTileThreads = kWavesPerTile * kWaveSize; + +// K=64 pipeline step: two m16n16k32 halves per wait point (iteration 6). +constexpr int kPipeK = 2 * kTileK; // 64 +constexpr int kPipeHalves = kPipeK / kTileK; // 2 halves per pipeline step + +// 16-B packed-B fragment slot (iteration 10): a lane's TWO adjacent k32 +// fragment 8-B slots for one k64 pipeline step, side by side (lower k32 half +// in .lo, upper in .hi), so the wave's whole 1-KiB B slice per step is ONE +// aligned 16-B-per-lane global load (global_load_dwordx4) instead of two +// 512-B dwordx2 loads. Pure same-byte-count permutation of the packed tensor. +struct alignas(16) BSlot { + uint64_t lo; + uint64_t hi; +}; + +// Round-to-nearest-even float -> bf16 bit pattern. +// +// Bit-compatible with the trusted exact reference, which converts the scaled +// fp32 value with PyTorch's BFloat16 conversion (c10 BFloat16 +// round_to_bfloat16: add 0x7fff + lsb, then truncate). Timed values are +// finite (small int8 dots scaled by <= 1.0), so no inf/nan special casing is +// needed to match the reference; denormals round identically. +__device__ __forceinline__ uint16_t bf16_bits_from_float(float f) { + union { + float f32; + uint32_t u32; + } cvt; + cvt.f32 = f; + const uint32_t lsb = (cvt.u32 >> 16) & 1u; + cvt.u32 += 0x7fffu + lsb; + return static_cast(cvt.u32 >> 16); +} + +// One-time (out-of-timed-region, out-of-Graph) pack for the exact assigned +// weight (k,n) == (2048, 6144): permute the logical [K,N] int8 layout into +// packed[n_tile][k64_step][lane][16] B fragment slots (iteration 10), where +// for n_tile t (0..383), k64_step ps (0..31) and lane l (0..63), bytes +// 0..7 (lower k32 half) and 8..15 (upper k32 half) are +// byte 8*h+i = raw[(64*ps + 32*h + 8*(l>>4) + i) * n + 16*t + (l&15)] +// -- exactly the bytes du_load_matrix_sync assigned to lane l's B fragment +// for the two m16n16k32 steps of the k64 pipeline step in the accepted +// iteration-2/4/6 kernels (du_mma.hpp matrix_b row_major: row = lane & 15, +// col = (lane >> 4) << 3, x[i] = p[(col + i)*ldm + row]), k-ascending in +// memory order. One thread per (t, ps, l) writes one aligned 16-B slot; byte +// count unchanged (K*N), so the packed tensor is a same-size permutation of +// the iteration-5/6 packed[n_tile][k_step][lane][8] layout (the two 8-B slots +// for k32 steps 2*ps and 2*ps+1 of the same lane become one contiguous +// 16-B slot). The identity copy remains the generic fallback for every other +// (K, N). +__global__ __launch_bounds__(kPackThreads) void w8a8_pack_o_proj_tp8_m16_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t k64_steps = (k / kTileK) / kPipeHalves; // 32 for k=2048 + const int64_t n_tiles = n / kTileN; // 384 for n=6144 + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + const int64_t total = n_tiles * k64_steps * kWaveSize; + if (idx >= total) { + return; + } + const int64_t t = idx / (k64_steps * kWaveSize); // n_tile + const int64_t rem = idx % (k64_steps * kWaveSize); + const int64_t ps = rem / kWaveSize; // k64 step + const int lane = static_cast(rem % kWaveSize); + const int row = lane & 15; // B column within the tile + const int col = (lane >> 4) << 3; // k offset within each k32 half + const int8_t* src = + raw + (ps * kPipeK + col) * n + t * kTileN + row; + uint64_t lo = 0; + uint64_t hi = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) { + lo |= (static_cast(static_cast(src[i * n]))) << (8 * i); + hi |= (static_cast(static_cast(src[(kTileK + i) * n]))) + << (8 * i); + } + BSlot slot; + slot.lo = lo; + slot.hi = hi; + __builtin_memcpy(packed + idx * 16, &slot, sizeof(BSlot)); +} + +// 16x16x32 DUMMA tile with 2 wavefronts, in-block split-K=2, packed-B +// fragment-slot register-only transport and a k64 pipelined K loop for the +// assigned M=16 decode shape (iteration-6 pipeline round; iteration 10 moves +// the B transport to one aligned 16-B-per-lane load per k64 step). +// +// Launch geometry: grid = N/16 = 384 blocks, 128 threads (2 wavefronts) per +// block, one 16x16 output tile per block, so 384*2 = 768 independent waves +// are resident across the 120 CUs (6.4 waves/CU) -- identical to the accepted +// iteration-2/4/5 geometry. Each wave owns a contiguous K-half (K=1024 = 16 +// k64 pipeline steps) of the tile: +// * per k64 step each lane issues THREE aligned cooperative loads -- A +// halves h=0,1 (logical row-major A: row = lane&15, col = 8*(lane>>4), +// 8 consecutive bytes per half, du_mma.hpp matrix_a row_major fragment +// mapping) as two 8-B dwordx2, and the packed-B k64 slot +// (packed[n_tile][k64_step][lane][16]; 64 lanes x 16 B = 1 KiB coalesced +// per step, ONE global_load_dwordx4 instead of the two 512-B dwordx2 +// loads of iteration 6; .lo = lower k32 half, .hi = upper) -- and the +// step's two serialized m16n16k32 du_mma_sync calls (k-ascending: lower +// half first) consume all three loads, so ONE compiler-inserted +// s_waitcnt vmcnt(0) serves 2 x K (2 KiB in flight per wave at the wait, +// unchanged byte count and stream count vs iteration 6); +// * depth-1 k64 prefetch: step j+1's three loads are issued BEFORE step +// j's two-MMAC burst, so the wait lands at the next fill -- no LDS hop, +// no lgkmcnt wait anywhere in the K loop; the per-wave wait count stays +// 16 (the iteration-6 sweet spot; iteration 9's 8-wait/4-KiB probe +// regressed 26.46 us and is NOT repeated); +// * split combine unchanged: wave 0 publishes its 16x16 int32 partial into +// one LDS plane, one END-of-K __syncthreads(), wave 1 adds its own +// partial (k-ascending halves; no partial sum can overflow int32, so any +// split-sum order is bit-identical to the exact k-ascending reference) +// and performs the direct-register scale/bf16 epilogue. +__global__ __launch_bounds__(kTileThreads) void +w8a8_gemm_m16_dumma_packedb_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n, + int k) { + const int wave = static_cast(threadIdx.x) / kWaveSize; + const int lane = static_cast(threadIdx.x) % kWaveSize; + const int n_tile = static_cast(blockIdx.x); + + __shared__ __align__(16) int32_t c_part[kTileM * kTileN]; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Per-lane fragment ownership is loop-invariant (du_mma.hpp matrix_a / + // matrix_b row_major loaders): row = lane & 15, col = (lane >> 4) << 3. + const int a_row = lane & 15; + const int a_col = (lane >> 4) << 3; + + const int64_t k_steps = k / kTileK; // 64 + const int64_t s_per_wave = k_steps / kWavesPerTile; // 32 k32 steps per wave + const int64_t p_steps = s_per_wave / kPipeHalves; // 16 k64 steps per wave + const int64_t s0 = wave * s_per_wave; // first step of wave + const int64_t k64_steps = k_steps / kPipeHalves; // 32 + const int64_t ps0 = wave * p_steps; // first k64 step of wave + + // Packed-B region of this block's n_tile: [n_tile * k64_steps + ps][lane][16] + // (one aligned 16-B BSlot per lane per k64 step, iteration 10). The packed + // slot index for (ps_global, lane) is b_base + ps_global * kWaveSize + lane + // (pack kernel: idx = t*k64_steps*kWaveSize + ps*kWaveSize + lane). Repair 2: + // the iteration-10 change dropped the *kWaveSize stride on the ps index and + // read slots 16 B apart instead of 1 KiB apart after the first k64 step, + // corrupting every output (98211/98304 mismatches). Anchor this wave's slice + // at ps0 * kWaveSize so the in-loop index is just ps * kWaveSize. + const int64_t b_base = n_tile * k64_steps * kWaveSize + + ps0 * kWaveSize + lane; // in 16-B slots + const BSlot* b_slots = reinterpret_cast(packed_b) + b_base; + + // Logical row-major A: lane's 8 bytes at a + a_row*k + 32*s + a_col. + const uint64_t* a_slots = + reinterpret_cast(a + a_row * k + a_col); + + // Prologue: k64 step 0 (k32 steps s0, s0+1) lands directly in registers + // (one cold-start DRAM latency per wave, amortized over 16 steps; the two + // waves' cold starts overlap). B arrives as one 16-B slot per lane. + uint64_t ca0 = a_slots[(s0 * kTileK) / 8]; + uint64_t ca1 = a_slots[((s0 + 1) * kTileK) / 8]; + BSlot cb; + __builtin_memcpy(&cb, &b_slots[0], sizeof(BSlot)); + uint64_t cb0 = cb.lo; + uint64_t cb1 = cb.hi; + + for (int64_t ps = 0; ps < p_steps; ++ps) { + // Issue the NEXT k64 step's global loads (2 A dwordx2 + ONE B dwordx4 = + // 2 KiB per wave in flight per wait point -- same bytes and stream count + // as iteration 6, one instruction fewer per step) BEFORE this step's + // two-MMAC burst so DRAM/L2 latency overlaps the burst; the + // compiler-inserted s_waitcnt vmcnt(0) lands at the next iteration's + // first fragment fill -- no LDS hop. + uint64_t na0, na1; + BSlot nb; + if (ps + 1 < p_steps) { + const int64_t step = s0 + (ps + 1) * kPipeHalves; + na0 = a_slots[(step * kTileK) / 8]; + na1 = a_slots[((step + 1) * kTileK) / 8]; + __builtin_memcpy(&nb, &b_slots[(ps + 1) * kWaveSize], sizeof(BSlot)); + } + // Step-ps MMAC pair straight from registers, k-ascending (lower k32 half + // first), in the exact a_frag.x / b_frag.x byte order du_mma_sync consumes + // (one 64-bit value per lane). All three loads were issued together, so + // the single vmcnt(0) wait that precedes this fill serves the whole k64 + // step (2 x K per wait point). + __builtin_memcpy(a_frag.x, &ca0, 8); + __builtin_memcpy(b_frag.x, &cb0, 8); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + __builtin_memcpy(a_frag.x, &ca1, 8); + __builtin_memcpy(b_frag.x, &cb1, 8); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + // Rotate: the prefetched k64 step ps+1 becomes the current step. + if (ps + 1 < p_steps) { + ca0 = na0; + ca1 = na1; + cb0 = nb.lo; + cb1 = nb.hi; + } + } + + // Split combine + direct fragment epilogue (verified gfx928 INT8 m16n16k32 + // accumulator ownership: row = lane&15, col_mod4 = lane>>4, x[i] -> col + // col_mod4+4*i). + const int row = lane & 15; + const int col_mod4 = lane >> 4; + if (wave == 0) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + c_part[row * kTileN + col_mod4 + 4 * i] = acc_frag.x[i]; + } + } + __syncthreads(); // single END-of-K barrier (cheap, order-independent) + if (wave == 1) { + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int out_col = n_tile * kTileN + col_mod4 + 4 * i; + const int32_t total = + acc_frag.x[i] + c_part[row * kTileN + col_mod4 + 4 * i]; + const float scaled = + static_cast(total) * xs * weight_scale[out_col]; + out[row * n + out_col] = __float2bfloat16(scaled); + } + } +} + +// Scalar INT8 W8A8 GEMM: one thread per output element. +// +// a: [M, K] int8, row-major (x_q) +// b: [K, N] int8, row-major (packed_weight; identity-packed in bootstrap) +// x_scale: [M] fp32 +// weight_scale: [N] fp32 +// out: [M, N] bf16, row-major +// +// Linear index idx = row * n + col, so consecutive lanes own consecutive N +// addresses. The full K dot is accumulated exactly in int32 (k-ascending, +// matching the exact int64 reference sum), then scaled in contract order. +// For the exact (n,k) == (6144, 2048) pair the weight tensor is the packed +// [n_tile][k64_step][lane][16] fragment-slot layout (see +// w8a8_pack_o_proj_tp8_m16_kernel), which this fallback decodes elementwise +// so the paired M=2 API shape with the same (N,K) stays byte-exact; every +// other (K,N) reads the logical row-major weight. +__global__ __launch_bounds__(kGemmThreads) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t stride = + static_cast(gridDim.x) * static_cast(blockDim.x); + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + idx < total; idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = a + static_cast(row) * k; + const bool packed_b = (n == 6144 && k == 2048); + int32_t acc = 0; + if (packed_b) { + // Decode packed[n_tile][k64_step][lane][16] (iteration 10): element + // (kk, col) is byte 8*h+j of the 16-B slot for n_tile = col>>4, + // k64_step = kk>>6, lane = (((kk&63)>>5)*32 + ((kk&31)>>3))*16 + + // (col&15), where bytes 0..7 are the lower k32 half and bytes 8..15 the + // upper half (a pure permutation of the iteration-5/6 packed layout, + // byte count unchanged). + const int64_t t = col >> 4; + for (int kk = 0; kk < k; ++kk) { + const int64_t ps = kk >> 6; + const int64_t h = (kk >> 5) & 1; + const int64_t c = (kk & 31) >> 3; + const int64_t j = kk & 7; + const int64_t slot = + (((t * (k / kPipeK) + ps) * kWaveSize) + (c << 4) + (col & 15)) * + 16 + + (h << 3) + j; + acc += static_cast(a_row[kk]) * + static_cast(b[slot]); + } + } else { + const int8_t* b_col = b + col; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + } + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + reinterpret_cast(out)[idx] = bf16_bits_from_float(scaled); + } +} + +// Generic elementwise device-to-device copy used by the identity pack_weight +// bootstrap. Handles every (K, N) through a grid-stride loop, so it stays +// valid as the generic fallback for unmatched (K, N) in later rounds. +template +__global__ __launch_bounds__(kPackThreads) void w8a8_identity_pack_kernel( + const T* __restrict__ src, + T* __restrict__ dst, + int64_t count) { + const int64_t stride = + static_cast(gridDim.x) * static_cast(blockDim.x); + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + i < count; i += stride) { + dst[i] = src[i]; + } +} + +} // namespace + +// Timed operator entry point. Graph-safe: launches only on `stream`; no +// allocation, compilation, autotuning, packing, or host/device sync. +// +// The assigned shape (m,n,k) == (16, 6144, 2048) is dispatched to the +// packed-B two-wave in-block split-K=2 DUMMA kernel (384 blocks x 128 +// threads, one 16x16 tile per block, register-only transport). Every other +// (m, n, k) -- including the paired M=2 API shape with the same (N, K) -- +// takes the scalar generic fallback (which decodes the packed layout for +// (n,k) == (6144, 2048)). +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + if (m == 16 && n == 6144 && k == 2048) { + const unsigned blocks = static_cast(n / kTileN); + hipLaunchKernelGGL( + w8a8_gemm_m16_dumma_packedb_kernel, dim3(blocks), dim3(kTileThreads), 0, + stream, a, b, x_scale, weight_scale, reinterpret_cast(out), + n, k); + return; + } + const int64_t total = static_cast(m) * n; + const unsigned blocks = + static_cast((total + kGemmThreads - 1) / kGemmThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, dim3(blocks), dim3(kGemmThreads), 0, stream, + a, b, x_scale, weight_scale, reinterpret_cast(out), m, + n, k); +} + +// Optional out-of-timed-region weight prep. For the exact (k, n) == +// (2048, 6144) o_proj weight this permutes into the packed-B fragment-slot +// layout (one-time, untimed, outside Graph capture); every other (K, N) +// keeps the generic identity device-to-device copy. Weight scales are always +// copied through (identity [N] fp32). +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + if (k == 2048 && n == 6144) { + const int64_t total = + static_cast(n / kTileN) * ((k / kTileK) / kPipeHalves) * + kWaveSize; // one thread per (n_tile, k64_step, lane) 16-B slot + const unsigned blocks = static_cast( + (total + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_pack_o_proj_tp8_m16_kernel), dim3(blocks), + dim3(kPackThreads), 0, stream, raw_weight, packed_weight, k, n); + } else { + const int64_t weight_elems = static_cast(k) * n; + const unsigned weight_blocks = static_cast( + (weight_elems + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_pack_kernel), dim3(weight_blocks), + dim3(kPackThreads), 0, stream, raw_weight, packed_weight, weight_elems); + } + + const int64_t scale_elems = n; + const unsigned scale_blocks = static_cast( + (scale_elems + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_pack_kernel), dim3(scale_blocks), + dim3(kPackThreads), 0, stream, weight_scale, packed_weight_scale, + scale_elems); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/q_b_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/q_b_proj.hip new file mode 100644 index 00000000..8d888111 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/q_b_proj.hip @@ -0,0 +1,853 @@ +// @@variant shape=glm_tp8_q_b_proj_m16 commit=8a9a28e409277a6f266a59e0c053a3d2bb516f12 added=2026-08-31 +// median_us=14.12 p90_us=14.14 +// source=glm5-2-dsh-tp8-m16-test1-cb4d262c +// csrc/w8a8_gemm_hip.hip +// +// Worker_1 bootstrap (iteration 1) for the GLM5.2 TP8 M=16 decode shapes on +// Hygon K500SM_AI / gfx928 (wavefront = 64, 64 KiB LDS per CU): +// +// glm_tp8_q_b_proj_m16 : M=16, N=2048, K=2048 +// glm_tp8_kv_b_proj_m16 : M=16, N=3584, K=512 +// +// Logical operation (frozen contract, see int8_w8a8_gemm_api.py): +// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) +// * x_scale[m, 0] * weight_scale[n, 0] ) +// +// Strategy: assigned M=16 shapes run gfx928 DUMMA INT8 fast paths +// (m16n16k32, exact int32 accumulation). q_b (K=2048, N=2048) runs the +// grid-level split-K partial kernel: one 64-thread wavefront per (tile, +// split) block, the 32-aligned B K slice is staged once into LDS with +// vectorized 16-B bulk loads, one staging barrier, then a zero-barrier K +// loop whose B fragments come from LDS and whose A fragments are loaded +// straight from the L2-resident row-major A (8 contiguous bytes per lane per +// k32 step, 2-deep software pipeline; iteration 7 deepened the pipeline +// because the exact iteration-5 code object showed the 1-deep source +// degrading to an issue-mid-body / wait-same-body schedule that still +// exposed most of the L2 latency per step, and iteration 6 proved the LDS +// side of the body is not the critical path). Iteration 8 (HIP-only +// occupancy round): SPLIT_K 3 -> 8 (128 x 8 = 1,024 blocks = 8.53 blocks/CU +// = 2.13 waves/SIMD on the 120-CU device, uniform 256-K slices, LDS 11,264 -> +// 4,096 B/block -> 16 resident blocks/CU); each block publishes its exact +// int32 partial to a caller workspace plane and a separate combine+scale +// kernel (in the timed Graph) sums the split planes in ascending order and +// emits scaled bf16. Iteration 9's batched-staging HIP change (all four +// global_load_dwordx4 -> ds_write_b128 chains collapsed into one L2 round +// trip, verified in its exact code object) regressed 21.191 -> 22.033 us, so +// the staging latency is hidden by the 2.13 co-resident waves/SIMD and the +// partial kernel is at its HIP floor. Iteration 10 (this code) rewrites the +// combine+scale kernel as a single-pass int4-vectorized wavefront (one drain +// instead of four serialized load->wait->add->store chains; vmem_read 4,736 +// -> ~1,280 per replay) with bit-identical int32 sums and float rounding. +// kv_b (K=512, N=3584) keeps the one-wave direct kernel. +// Iteration 17 (this code): one-time tile-major B pack for (k, n) == +// (2048, 2048) (P[(t*K + kk)*16 + col] = B[kk][t*16 + col], out of the timed +// region and the Graph). The exact current-best code object + PMC show the +// partial kernel is memory-service-bound: the B staging loop's 64-lane 16-B +// chunks are scattered over 64 different rows' 128-B sectors (262,144 sector +// requests/replay, 16 B used per request) and each sector is shared by 8 +// tile-blocks whose accesses are spread over the whole kernel, so evicted +// sectors are re-fetched (PMC l2_misses 83,350 ~= 2.5x the 32,768 compulsory +// B sectors; ~10.7 MB/replay at ~660 GB/s effective). The pack makes every +// staging iteration one contiguous 1,024-B chunk (8 fully-consumed sectors, +// one consumer block, no re-fetch possible): B sector requests drop 262,144 +// -> 32,768 and B DRAM traffic drops toward the 4.19 MB compulsory set. The +// staged bytes, lds_b, the single barrier, the hoisted uniform A path, the +// k-ascending v_mmac body, the exact int32 partials, the combine kernel, the +// 1,024-block grid and the 2-dispatch Graph are unchanged -> outputs stay +// bit-identical to iteration 11 (0 mismatches). The (n, k) == (2048, 2048) +// generic scalar fallback (paired M=2 validation shape) decodes the pack via +// the b_tilemajor flag; every other shape keeps the identity pack and the +// logical decode. +// Iteration 18 (HIP-only consolidation; final conditional inline-asm round, +// raw asm still gated): the uniform-path B staging loop's 4 chains are +// batched into one in-flight group (4 named int4 loads back-to-back, one +// drain, 4 ds_write_b128) - the exact current-best code object shows the +// rolled loop still executes one load in flight per block (4 serialized +// cold-DRAM round trips), and the pack made the 4 chunks contiguous so the +// batch is DRAM-row-friendly (iteration 9's pre-pack scatter regression does +// not apply). Named registers avoid iteration 14's scratch spill. Staged +// bytes, barrier, body, partials and combine are unchanged -> outputs stay +// bit-identical to iteration 17. +// The generic scalar kernel remains the fallback for every unmatched (m, n, +// k), including the paired M=2 API shapes that share (N, K) with the +// assigned M=16 shapes. +// +// Graph-safety: gemm_out only launches kernels on the caller-provided stream +// (PyTorch's current HIP stream); it performs no allocation, no compilation, +// no autotuning, no packing, and no host/device synchronization, and touches +// only the caller-provided `out` and `workspace` (split-K partial planes). + +#include +#include +#include + +#include + +namespace { + +// blockDim must be a multiple of the gfx928 wavefront size (64). +constexpr int kScalarBlockThreads = 256; +constexpr int kPackBlockThreads = 256; + +// gfx928 DUMMA INT8 tile geometry for the M=16 decode fast path: the +// installed DTK exposes m16n16k32 signed-char fragments with int32 +// accumulation, and one 64-thread wavefront owns one independent tile. +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; +constexpr int kDummaWaveSize = 64; + +using namespace du::dumma; + +// One thread -> one output element. Adjacent threads map to adjacent columns +// (fastest-changing N dimension), so both the B row loads and the bf16 stores +// are coalesced; the A row load is broadcast across the threads of the same +// row. +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k, + int b_tilemajor) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + + const int8_t* a_row = a + static_cast(row) * k; + // For (k, n) == (2048, 2048) the B buffer is the one-time tile-major pack + // P[(t*K + kk)*16 + col_tile] = raw[kk][t*16 + col_tile] (t = col >> 4, + // col_tile = col & 15), so the logical [K][N] column stride n becomes a + // 16-B stride in kk; every other shape keeps the logical row-major decode. + const int8_t* b_col = b_tilemajor + ? b + (static_cast(col >> 4) * k * 16 + (col & 15)) + : b + col; + const int64_t b_stride = b_tilemajor ? 16 : static_cast(n); + + // Exact int32 dot product over the complete K loop. The maximum assigned K + // is 2048, so |dot| <= 2048 * 128 * 128 = 2^25 stays far below INT32_MAX. + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * b_stride]); + } + + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// Minimal M=16 DUMMA fast path: one wavefront (blockDim=64) computes one +// 16x16 output tile with a 16x16x32 signed-char fragment pair and exact int32 +// accumulation, so there is no cross-wave barrier and no LDS staging in the K +// loop. A and B fragments are loaded directly from the caller row-major +// global tensors (A leading dimension K, B leading dimension N). The epilogue +// materializes the accumulator fragment through du_store_matrix_sync into a +// per-block LDS tile, applies the two float scales in the same order as the +// scalar path, and stores bf16. +__global__ __launch_bounds__(kDummaWaveSize) void w8a8_dumma_m16_direct_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n, + int k) { + constexpr int kTileElems = kDummaTileM * kDummaTileN; // 256 + const int lane = static_cast(threadIdx.x); // 0..63 + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + __shared__ __align__(16) int32_t acc_tile[kTileElems]; + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + +#pragma unroll 4 + for (int k0 = 0; k0 < k; k0 += kDummaTileK) { + du_load_matrix_sync(a_frag, a + k0, k); + du_load_matrix_sync( + b_frag, b + static_cast(k0) * n + n0, n); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + du_store_matrix_sync(acc_tile, acc_frag, kDummaTileN, mem_row_major); + __syncthreads(); + +#pragma unroll + for (int linear = lane; linear < kTileElems; linear += kDummaWaveSize) { + const int row = linear / kDummaTileN; + const int col = linear - row * kDummaTileN; + const float scaled = static_cast(acc_tile[linear]) * + x_scale[row] * weight_scale[n0 + col]; + out[row * n + n0 + col] = __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// Iteration 3: grid-level split-K partial + combine (CU-aligned probe). +// +// Iteration 1 (one wavefront per 16x16 tile, direct global fragment loads) +// measured 63.695 us and iteration 2 (in-block split-K=2 with the same direct +// byte loads) regressed to 69.771 us: the 16 global_load_ubyte + vmcnt chain +// per m16n16k32 step (PMC: 131,712 vmem_reads) exposes full memory latency +// with ~1.07 blocks/CU and no in-loop overlap, i.e. the direct fragment +// loader is the poison, not the grid size. Iteration 3 replaced the q_b +// (16, 2048, 2048) path with a grid-level split-K partial kernel that stages +// each 32-aligned K slice into LDS once with vectorized 16-B bulk loads +// (one barrier), then runs a zero-barrier LDS-only K loop; every block is a +// single wavefront and the grid is 128 tiles x SPLIT_K blocks (SPLIT_K=3 -> +// 384 blocks = 3.2 blocks/CU, a non-power-of-two CU-aligned occupancy probe +// from the trusted set). Exact int32 partials go to caller workspace planes +// and the combine+scale kernel runs in the timed Graph. +// +// Iteration 4 (bounded staging prefetch, kStageDepth=4) regressed to 43.876 +// us while keeping LDS at 22,784 B/block (2 resident blocks/CU), so the +// staging vmem chain was NOT the lever: the LDS-resident footprint itself is. +// +// Iteration 5 (this code): B-only staging. The exact iteration-3 code object +// shows the K loop's A fragment is 8 contiguous bytes per lane (one +// ds_read2_b32 in the staged kernel; 8 consecutive global_load_ubyte in the +// direct kernel) while the B fragment is 8 scattered 16-B-strided byte loads +// (16-way LDS bank conflicts). A is therefore moved out of LDS entirely: +// per k32 step each lane issues ONE 8-B global load from the L2-resident +// row-major A (1-deep software pipeline hides the L2 latency behind the +// current step's LDS loads + mma), the 11,520-B A tile disappears, LDS drops +// from 22,784 B to 11,264 B per block, and the resident blocks/CU rise from 2 +// to 5. B keeps the exact iteration-3 staged layout and byte loads. +// +// Iteration 8 (this code): HIP-only occupancy round. PMC/ISA evidence pins +// the binding occupancy limiter: 24 VGPR / 32 SGPR / 11,264 B LDS / 0 scratch +// -> LDS is the limiter at 5 resident blocks/CU (64 KiB / 11.25 KiB), while +// the grid (384 = 3.2 blocks/CU = 0.8 waves/SIMD) leaves residency half +// empty and SQ_WAIT_INST_LDS = 36/replay (~0.09/block) plus the exact loop +// body's s_waitcnt vmcnt(0) right before v_mmac (and the staging loop's 11 +// serialized global_load_dwordx4 -> vmcnt(0) -> ds_write_b128 chains) show +// the shaders stall on VMEM latency, not LDS. SPLIT_K is raised 3 -> 8: the +// trusted occupancy-probe split with uniform 256-K slices (64 k32 steps / 8 +// = 0 remainder), grid 128 x 8 = 1,024 blocks = 8.53 blocks/CU = 2.13 +// waves/SIMD, LDS 11,264 -> 4,096 B -> 16 resident blocks/CU (residency +// never binds, zero queueing, tail factor 1.055 vs 1.25). HBM is unchanged: +// every B byte is still staged exactly once (4.19 MB/replay) and the A +// request volume is invariant in S (each 32-KB L2-hot A is re-read per +// block, 4.19 MB/replay for any S); only the partial-plane writes grow +// 393,216 -> 1,048,576 B (8 planes <= the API 16-plane workspace for this +// shape). +// --------------------------------------------------------------------------- + +// Compile-time non-uniform, 32-aligned K slicing for the guarded K=2048 +// shape (64 k32 steps split into SPLIT_K contiguous slices; the remainder +// steps go to the first slices so every boundary stays a multiple of the +// DUMMA tile K=32, keeping each partial an exact k-ascending int32 sum). +template +struct KSplit2048 { + static constexpr int kSteps = 2048 / kDummaTileK; // 64 + static constexpr int kBase = kSteps / SPLIT_K; + static constexpr int kRem = kSteps % SPLIT_K; + static constexpr int kMaxSliceSteps = kBase + (kRem > 0 ? 1 : 0); + static constexpr int kMaxSlice = kMaxSliceSteps * kDummaTileK; + // Closed form (no C++14 constexpr loops): slice i contributes + // kBase + (i < kRem ? 1 : 0) k32 steps. + static constexpr int begin(int s) { + return (s * kBase + (s < kRem ? s : kRem)) * kDummaTileK; + } + static constexpr int slice(int s) { + return (kBase + (s < kRem ? 1 : 0)) * kDummaTileK; + } +}; + +template +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16_splitk_partial_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int* __restrict__ partials, + int n, + int k) { + using Split = KSplit2048; + // B-only staging (iteration 5): only the B tile is staged into LDS. The A + // fragment of this compiler/DUMMA header (verified in the exact iteration-3 + // code object, profiles/.../iteration4/current-best-isa) is 8 CONTIGUOUS + // bytes per lane: the staged path reads them with one ds_read2_b32 at + // (lane&15)*720 + (lane>>4)*8, and the direct kernel reads them with 8 + // consecutive global_load_ubyte at (lane&15)*k + (lane>>4)*8. A therefore + // needs no LDS round trip: one 8-B global load per lane per k32 step from + // the L2-resident row-major A reproduces the fragment exactly. Dropping the + // 11,520-B A tile halves the LDS footprint per block (22,784 -> 11,264 B) + // and raises the resident blocks/CU from 2 to 5 (64 KiB LDS / 11.25 KiB), + // while the A load latency is hidden by a 2-deep software pipeline in the + // K loop (iteration 7). The B tile keeps the exact iteration-3 layout and + // byte loads. Iteration 8 (occupancy round) raises SPLIT_K to 8 via the + // launcher only: kSlice becomes a uniform 256, LDS drops to 4,096 B/block + // and the grid becomes 1,024 blocks; the body below is unchanged. + constexpr int kLdsStrideB = kDummaTileN; // 16 + + const int lane = static_cast(threadIdx.x); // 0..63 + const int split = static_cast(blockIdx.x) % SPLIT_K; + const int tile = static_cast(blockIdx.x) / SPLIT_K; + const int n0 = tile * kDummaTileN; + const int kBeg = Split::begin(split); + const int kSlice = Split::slice(split); + + __shared__ __align__(16) int8_t lds_b[Split::kMaxSlice][kLdsStrideB]; + + // A fragment lane mapping (verified against the iteration-3 code object): + // lane l owns the 8 contiguous bytes of A row (l & 15) starting at column + // ((l >> 4) << 3) inside the k32 window, with x[e] = byte at address + e. + // Global row-major A (row stride k) yields one 8-B (int2) load per lane per + // step, always 8-B aligned (k, kBeg, kDummaTileK are all multiples of 8). + const int a_row = lane & 15; + const int a_col0 = (lane >> 4) << 3; // 0, 8, 16, 24 + const int8_t* a_lane = + a + static_cast(a_row) * k + kBeg + a_col0; + + // Iteration 11 (HIP-only consolidation): the exact iteration-10 code object + // (profiles/.../iteration11/current-best-isa, source digest + // 5187c0a5a308822b9914dae2c086c979f0fcda9093fa20009c201b19597794e1) + // compiles the 2-deep A rotation below into a schedule where the load for + // step i+2 issues at the TAIL of body i (global_load_dwordx2 at 0x4374) and + // the body opens with a full s_waitcnt vmcnt(0) (0x4324) that drains ALL + // outstanding loads, so the load consumed by body i's mma has only ~1 body + // (~60-120 cycles of issue) of lookahead against an L2 round trip of + // ~250-350 cycles: every one of the 8 bodies still exposes most of one + // A-load round trip (~1.6-2.4k cycles of the ~3.2-3.9k cycle per-block + // path), and the ~2.1 co-resident waves/SIMD stall in phase (264,192 VALU + // over 18.24 us profiled ~= 2% issue busy). Uniform slices (Split::kRem == + // 0, true for SPLIT_K=8) therefore take a compile-time hoisted path: all 8 + // int2 A loads issue back-to-back BEFORE the B staging loop, so the + // staging chains' serialized DRAM waits (4 x global_load_dwordx4 -> + // vmcnt(0) -> ds_write_b128 at 0x4190-0x41C0) give every A load + // ~2.4-3.6k cycles of slack, and the K body is fully unrolled so the + // compiler can overlap the independent B-fragment LDS assembly of later + // steps with the serial v_mmac chain. The B-only staging, lds_b layout, + // k-ascending mma order and exact int32 accumulation are unchanged, so the + // partials stay bit-identical to iteration 10 (0 mismatches; Graph replay + // with changed contents must pass). The generic kRem>0 path keeps the + // original rotation unchanged. + // a_buf is function-scoped so both branches can reference it; only the + // uniform (kRem == 0) branch fills it, and only that branch's unrolled body + // reads it (kMaxSliceSteps >= kBase covers every SPLIT_K instantiation). + int2 a_buf[Split::kMaxSliceSteps]; + if (Split::kRem == 0) { +#pragma unroll + for (int i = 0; i < Split::kBase; ++i) { + a_buf[i] = *reinterpret_cast(a_lane + i * kDummaTileK); + } + } + + // Stage only B: chunk c covers 16 contiguous bytes = exactly one tile row. + // Iteration 17 (HIP-only consolidation round; final conditional inline-asm + // round, raw asm still gated: plateau=false, phase=hip_only, + // raw_inline_asm_allowed=false, skill_allowed=false): the (k, n) == + // (2048, 2048) weight is now the one-time tile-major pack + // P[(t*K + kk)*16 + col] = B[kk][t*16 + col] (t = tile index), so lane c + // reads its 16 B from packed + ((t*K + kBeg + c)*16). The 64 lanes' chunks + // of one staging iteration are CONTIGUOUS (one 1,024-B chunk = 8 + // fully-consumed 128-B sectors) instead of 64 scattered 16-B requests + // spread over 64 different rows' sectors (row stride n = 2048), and each + // sector has exactly ONE consumer block (no cross-tile sharing, so an + // evicted sector is never re-fetched: PMC l2_misses 83,350 ~= 2.5x the + // 32,768 compulsory B sectors). The staged bytes are identical + // (B[kBeg+c][n0..n0+16)), so lds_b, the single staging barrier, the hoisted + // uniform A path, the unrolled k-ascending v_mmac body and the exact int32 + // partials are unchanged (bit-identical to iteration 11). + // Iteration 18 (HIP-only consolidation round; final conditional inline-asm + // round, raw asm still gated: plateau=false, phase=hip_only, + // raw_inline_asm_allowed=false, skill_allowed=false): the exact + // current-best code object (profiles/.../iteration18/current-best-isa, + // source digest 8b6bcaa866f26094d84762db45bd259a3a6492f900b42461e60c3f4264 + // 94853c) compiles the uniform-path staging loop below into a ROLLED loop + // with ONE global load in flight per block (global_load_dwordx4 -> + // s_waitcnt vmcnt(0) -> ds_write_b128 at 0x4A08-0x4A20), so every block + // pays 4 SERIALIZED cold-DRAM round trips for its 4 KiB slice while the + // 2.13 co-resident waves/SIMD stall in phase. Pre-pack, iteration 9 batched + // the same chains and regressed (+4%) because the 4 chains were 256 + // SCATTERED sector requests (row thrash) and the co-resident waves covered + // the drains; post-pack each chain is 8 contiguous sectors and the 4 chains + // tile the slice CONTIGUOUSLY (4 x 1,024 B = one 4,096-B group), so issuing + // all four per-lane chunk loads back-to-back (one in-flight group, one + // drain) collapses 4 serialized round trips toward ~1 while keeping the + // DRAM row-buffer-friendly. The four chunks are NAMED locals, not a + // runtime-indexed array: iteration 14's `int4 b_chunk[kChunksPerLane]` + // under a runtime `c < kSlice` trip bound spilled to scratch + // (private_segment_fixed_size 80) and regressed to 36.675 us; named + // registers with compile-time addressing keep the loads in VGPR (arch_vgpr + // 40 -> ~44-56; occupancy stays grid-limited at 8.53 blocks/CU, LDS 16 + // resident blocks/CU unchanged, no scratch). The staged bytes, lds_b + // addresses, single barrier, hoisted uniform A path, unrolled k-ascending + // v_mmac body and exact int32 partials are unchanged, so outputs stay + // bit-identical to iteration 17 (0 mismatches; Graph replay with changed + // contents must pass). The generic kRem>0 path keeps the rolled loop + // (runtime slice sizes); every other uniform instantiation falls back to it. + const int64_t b_tile_base = + static_cast(n0 / kDummaTileN) * k; + if (Split::kRem == 0) { + // kMaxSlice is constexpr and equals kSlice on this branch (kRem == 0). + constexpr int kChunks = Split::kMaxSlice / kDummaWaveSize; // 4 for SPLIT_K=8 + if (kChunks == 4) { + const int8_t* b_lane = + b + (b_tile_base + kBeg) * kDummaTileN + lane * kDummaTileN; + const int4 b0 = *reinterpret_cast(b_lane); + const int4 b1 = *reinterpret_cast( + b_lane + kDummaWaveSize * kDummaTileN); + const int4 b2 = *reinterpret_cast( + b_lane + 2 * kDummaWaveSize * kDummaTileN); + const int4 b3 = *reinterpret_cast( + b_lane + 3 * kDummaWaveSize * kDummaTileN); + *reinterpret_cast(&lds_b[lane][0]) = b0; + *reinterpret_cast(&lds_b[lane + kDummaWaveSize][0]) = b1; + *reinterpret_cast(&lds_b[lane + 2 * kDummaWaveSize][0]) = b2; + *reinterpret_cast(&lds_b[lane + 3 * kDummaWaveSize][0]) = b3; + } else { + for (int c = lane; c < kSlice; c += kDummaWaveSize) { + *reinterpret_cast(&lds_b[c][0]) = + *reinterpret_cast( + b + (b_tile_base + (kBeg + c)) * kDummaTileN); + } + } + } else { + for (int c = lane; c < kSlice; c += kDummaWaveSize) { + *reinterpret_cast(&lds_b[c][0]) = + *reinterpret_cast( + b + (b_tile_base + (kBeg + c)) * kDummaTileN); + } + } + __syncthreads(); + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + if (Split::kRem == 0) { + // Hoisted uniform path: a_buf holds all kSteps int2 A fragments, loaded + // before the B staging (every load has >= the staging chains' serialized + // DRAM waits of slack). The fully unrolled body keeps the exact + // k-ascending mma order of the generic path. +#pragma unroll + for (int i = 0; i < Split::kBase; ++i) { + const int8_t* a_bytes = reinterpret_cast(&a_buf[i]); +#pragma unroll + for (int e = 0; e < 8; ++e) { + a_frag.x[e] = a_bytes[e]; + } + const int k0 = i * kDummaTileK; + du_load_matrix_sync(b_frag, lds_b[0] + k0 * kLdsStrideB, kLdsStrideB); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + } else { + // Zero-barrier K loop: exact k-ascending int32 accumulation over + // [kBeg, kBeg + kSlice). B fragments come from LDS (unchanged layout). A + // fragments come from global with a 2-deep software pipeline (iteration + // 7): three int2 values (steps i, i+1, i+2) rotate through registers and + // the load for step i+3 issues at the top of step i's body, so every A + // load has ~1.5 step bodies (~400-550 cycles) to arrive in instead of the + // ~half-body window the 1-deep source degraded to in the exact + // iteration-5 code object (flat_load issued mid-body at 0x42C0, + // s_waitcnt vmcnt(0) at the END of the SAME body at 0x4380 before the + // rotation movs -> each step still exposes most of the ~350-450-cycle L2 + // latency). Iteration 6 proved the LDS side of the body is NOT the + // critical path (one ds_read_b64 replacing the 8 conflicted ds_read_u8 + + // ~12 VALU shuffles + lgkmcnt chain changed the median by 0.03%), so the + // per-step serial cost is the A-load wait; this rotation attacks exactly + // that. steps is runtime (e.g. 22/21/21 for SPLIT_K=3), so the loop is + // not unrolled and the rotation is explicit; the guards keep every load + // inside the A slice [kBeg, kBeg + kSlice). + const int steps = kSlice / kDummaTileK; + int2 a_cur = *reinterpret_cast(a_lane); // step 0 + int2 a_next = a_cur; + if (1 < steps) { + a_next = + *reinterpret_cast(a_lane + kDummaTileK); // step 1 + } + for (int i = 0; i < steps; ++i) { + int2 a_nxt2 = a_next; + if (i + 2 < steps) { + a_nxt2 = *reinterpret_cast( + a_lane + (i + 2) * kDummaTileK); + } + const int8_t* a_bytes = reinterpret_cast(&a_cur); +#pragma unroll + for (int e = 0; e < 8; ++e) { + a_frag.x[e] = a_bytes[e]; + } + const int k0 = i * kDummaTileK; + du_load_matrix_sync(b_frag, lds_b[0] + k0 * kLdsStrideB, kLdsStrideB); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + a_cur = a_next; + a_next = a_nxt2; + } + } + + // Publish the exact int32 partial to caller workspace plane `split` + // ([SPLIT_K][M][N] row-major). Every plane element is overwritten on every + // launch, so no workspace clear is needed and the combine kernel below is + // stream-ordered after this store. + int* plane = + partials + static_cast(split) * kDummaTileM * n + n0; + du_store_matrix_sync(plane, acc_frag, n, mem_row_major); +} + +// Combine+scale kernel: one wavefront per 16x16 output tile; sums the +// SPLIT_K int32 planes in ascending split order (bit-identical to the scalar +// k-ascending reference), applies x_scale * weight_scale in the scalar +// epilogue order, and stores bf16. Runs in the timed Graph after the partial +// kernel on the same stream. +// +// Iteration 10 (HIP-only consolidation round): the exact iteration-8 code +// object compiled the previous 4-iteration scalar loop (idx = lane, lane+64, +// lane+128, lane+192; 8 global_load_dword per iteration, one per split +// plane, at 0x464C-0x4694 of the combine symbol) into four SERIALIZED +// load -> wait-drain -> add -> store chains: each iteration's loads issue +// only after the previous iteration's bf16 store, so a single wave pays ~4 x +// (one DRAM round trip + add chain) while the 128-block grid delivers only +// 1.07 waves/CU (no second co-resident wave to overlap the drains; PMC: +// combine 5.12 us profiled, 4,736 vmem_read and 512 vmem_write per replay). +// This rewrite is ONE pass: lane l owns the contiguous 16-B chunk (row = +// l>>2, col4 = (l&3)<<2, 4 consecutive columns) of the tile, issues all +// SPLIT_K int4 plane loads back-to-back (a single drain, ~4x fewer memory +// instructions), accumulates the SAME ascending-s int32 sum per element, +// applies the same scalar epilogue expression order (float(dot) * x_scale * +// weight_scale), and stores the 4 bf16 as two dword stores. Per-element +// values, int32 accumulation order, and float rounding are identical to the +// previous combine, so the outputs stay bit-identical (0 mismatches) and the +// Graph structure (partial then combine on the caller stream) is unchanged. +template +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_dumma_m16_splitk_combine_kernel( + const int* __restrict__ partials, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n) { + const int lane = static_cast(threadIdx.x); + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + const int64_t plane = static_cast(kDummaTileM) * n; + + // lane -> one contiguous 16-B chunk of the 16x16 tile: row = lane/4, first + // column = (lane%4)*4. 64 lanes x 4 elements = 256 = the full tile, exactly + // once. All chunk addresses are 16-B aligned (n, n0, col4 and plane are all + // multiples of 4 int32). + const int row = lane >> 2; + const int col4 = (lane & 3) << 2; + const int* chunk = + partials + static_cast(row) * n + n0 + col4; + + int32_t dot[4] = {0, 0, 0, 0}; +#pragma unroll + for (int s = 0; s < SPLIT_K; ++s) { + const int4 p = *reinterpret_cast(chunk + s * plane); + dot[0] += p.x; + dot[1] += p.y; + dot[2] += p.z; + dot[3] += p.w; + } + + const float xs = x_scale[row]; + union { + float f[4]; + float4 v; + } ws; + ws.v = *reinterpret_cast(weight_scale + n0 + col4); + union { + __hip_bfloat16 b[4]; + uint32_t u[2]; + } packed; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float scaled = static_cast(dot[e]) * xs * ws.f[e]; + packed.b[e] = __float2bfloat16(scaled); + } + uint32_t* out32 = reinterpret_cast( + out + static_cast(row) * n + n0 + col4); + out32[0] = packed.u[0]; + out32[1] = packed.u[1]; +} + +// Identity device-to-device copy used by the bootstrap pack op. Works for any +// element type and any (K, N): unmatched shapes keep this generic fallback. +template +__global__ void w8a8_identity_copy_kernel( + const T* __restrict__ src, + T* __restrict__ dst, + int64_t numel) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < numel) { + dst[i] = src[i]; + } +} + +// One-time tile-major B pack for the (k, n) == (2048, 2048) q_b weight, +// outside the timed region and the Graph: P[(t*K + k)*16 + col] = +// raw[k][t*16 + col], with t = n-tile index (0..N/16-1). The 16 B of one +// (tile, k-row) stay contiguous and the 64 rows of a staging iteration become +// one contiguous 1,024-B chunk, so the split-K partial kernel stages fully +// consumed 128-B sectors instead of 64 scattered 16-B requests (iteration 17; +// see the partial kernel's staging comment). Same 4,194,304 B, same bytes, so +// the packed buffer stays graph-stable and the API contract is unchanged; the +// (n, k) == (2048, 2048) generic scalar fallback decodes this layout via the +// b_tilemajor flag. Every other (k, n) keeps the identity pack. +__global__ void w8a8_pack_b_tilemajor_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int t = static_cast(blockIdx.x); // n-tile 0..n/16-1 + // First k-row of this thread's 16-row chunk: threadIdx.x in 0..127 -> + // chunks [0,16), [16,32), ..., [2032,2048), covering every k-row exactly + // once (128 threads x 16 rows = the full K=2048). + const int k16 = static_cast(threadIdx.x) * 16; + const int64_t src = + static_cast(k16) * n + static_cast(t) * 16; + const int64_t dst = (static_cast(t) * k + k16) * 16; +#pragma unroll + for (int r = 0; r < 16; ++r) { + *reinterpret_cast(packed + dst + r * 16) = + *reinterpret_cast(raw + src + r * n); + } +} + +void launch_scalar_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + __hip_bfloat16* out, + int m, + int n, + int k, + int b_tilemajor, + hipStream_t stream) { + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads)); + const dim3 block(kScalarBlockThreads); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out, + m, + n, + k, + b_tilemajor); +} + +// One wavefront per 16x16 output tile. N must be divisible by 16 and K by 32 +// (guaranteed by the exact-shape guards that reach this launcher). +void launch_dumma_m16_direct( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + __hip_bfloat16* out, + int n, + int k, + hipStream_t stream) { + const dim3 grid(static_cast(n / kDummaTileN)); + const dim3 block(kDummaWaveSize); + hipLaunchKernelGGL( + w8a8_dumma_m16_direct_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out, + n, + k); +} + +// Grid-level split-K partial kernel + separate combine+scale kernel (both in +// the timed Graph, on the caller stream). One wavefront per (tile, split) +// block; exact int32 partials live in the caller workspace planes, so the +// partial kernel needs no cross-block synchronization and the combine kernel +// sums the planes in ascending split order. Requires K = 2048 and N divisible +// by 16 (guaranteed by the exact-shape guard that reaches this launcher). +template +void launch_dumma_m16_splitk( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + __hip_bfloat16* out, + void* workspace, + int n, + int k, + hipStream_t stream) { + auto* partials = reinterpret_cast(workspace); + const dim3 partial_grid( + static_cast(n / kDummaTileN * SPLIT_K)); + const dim3 block(kDummaWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk_partial_kernel), + partial_grid, + block, + 0, + stream, + a, + b, + partials, + n, + k); + const dim3 combine_grid(static_cast(n / kDummaTileN)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk_combine_kernel), + combine_grid, + block, + 0, + stream, + partials, + x_scale, + weight_scale, + out, + n); +} + +} // namespace + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // The workspace holds the split-K int32 partial planes for the q_b + // grid-level split path; scalar/tail paths ignore it. + auto* out_bf16 = reinterpret_cast<__hip_bfloat16*>(out); + + // Assigned M=16 shapes run the DUMMA fast paths. q_b (K=2048, N=2048) runs + // the grid-level split-K=8 partial kernel + combine kernel (iteration 8: + // 1,024 blocks = 8.53 blocks/CU = 2.13 waves/SIMD, uniform 256-K slices, + // B-only staging at 4,096 B LDS/block = 16 resident blocks/CU, A fragments + // loaded directly from the L2-resident row-major A, workspace int32 + // partials, combine in the timed Graph). Iteration 17: the (k, n) == + // (2048, 2048) weight is the one-time tile-major pack P[(t*K+kk)*16+col] = + // B[kk][t*16+col] (see w8a8_pack_b_tilemajor_kernel), so the partial + // kernel's staging reads contiguous fully-consumed sectors; the scalar + // fallback for this (n, k) decodes the pack via b_tilemajor. kv_b + // (K=512, N=3584) keeps the one-wave direct kernel and its identity pack. + // Each guard pins the exact M=16 so the paired M=2 API shape with the same + // (N, K) still reaches the generic scalar fallback below. + if (m == 16 && n == 2048 && k == 2048) { + // 8 int32 partial planes of [16][2048]; the caller workspace contract + // allocates 16 planes for this shape (2,097,152 B), so the guard is + // defensive only. + constexpr int64_t kRequiredWorkspace = + 8LL * 16 * 2048 * static_cast(sizeof(int32_t)); + if (workspace != nullptr && workspace_bytes >= kRequiredWorkspace) { + launch_dumma_m16_splitk<8>( + a, b, x_scale, weight_scale, out_bf16, workspace, n, k, stream); + } else { + // Defensive-only: the packed B must be decoded, so the direct kernel + // (which reads the logical [K][N] layout) is not usable here; the + // packed scalar decode is always correct. + launch_scalar_gemm( + a, b, x_scale, weight_scale, out_bf16, m, n, k, + /*b_tilemajor=*/1, stream); + } + return; + } + if (m == 16 && n == 3584 && k == 512) { + launch_dumma_m16_direct( + a, b, x_scale, weight_scale, out_bf16, n, k, stream); + return; + } + + // Generic scalar fallback for every unmatched (m, n, k). The (n, k) == + // (2048, 2048) pair (including the paired M=2 validation shape) decodes + // the tile-major pack; every other shape reads the logical layout. + launch_scalar_gemm( + a, b, x_scale, weight_scale, out_bf16, m, n, k, + (n == 2048 && k == 2048) ? 1 : 0, stream); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // One-time packing outside the timed region and the Graph. The (k, n) == + // (2048, 2048) q_b weight is repacked tile-major + // (P[(t*K+kk)*16+col] = raw[kk][t*16+col], iteration 17); the matching + // GEMM interpretation lives in the split-K partial staging and the + // b_tilemajor scalar decode. Every other (K, N) shape keeps the identity + // copy so the logical [K, N] row-major layout is preserved. + if (k == 2048 && n == 2048) { + const dim3 pack_grid(static_cast(n / 16)); + hipLaunchKernelGGL( + w8a8_pack_b_tilemajor_kernel, + pack_grid, + dim3(128), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + const int64_t weight_numel = static_cast(k) * n; + const dim3 weight_grid(static_cast( + (weight_numel + kPackBlockThreads - 1) / kPackBlockThreads)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + weight_grid, + dim3(kPackBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_numel); + } + + const int64_t scale_numel = static_cast(n); + const dim3 scale_grid(static_cast( + (scale_numel + kPackBlockThreads - 1) / kPackBlockThreads)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + scale_grid, + dim3(kPackBlockThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + scale_numel); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_down_proj.hip new file mode 100644 index 00000000..fdf64fd2 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_down_proj.hip @@ -0,0 +1,1472 @@ +// @@variant shape=glm_tp8_shared_down_proj_m16 commit=42d49afdb24af21a5480786c88ca11a84d57f44a added=2026-08-31 +// median_us=8.239 p90_us=9.122 +// source=glm5-2-dsh-tp8-m16-test1-cb4d262c +// W8A8 INT8 GEMM for Hygon K500SM_AI / gfx928 (worker_3). +// +// Exact-shape split-K DUMMA specialization for +// glm_tp8_shared_gate_up_proj_m16 (M=16, N=512, K=6144) established before +// the generic scalar fallback: +// +// glm_tp8_shared_gate_up_proj_m16 : (M=16, N=512, K=6144) +// glm_tp8_shared_down_proj_m16 : (M=16, N=6144, K=256) +// +// Iteration 1 measured the minimal gfx928 INT8 m16n16k32 DUMMA tile (one +// 64-thread wavefront per 16x16 output tile, grid = N/16 = 32 blocks) at +// 173.6 us median: PMC showed grbm_count == grbm_gui_active with one +// wavefront per CU and 88 of 120 CUs idle, i.e. a latency-bound launch +// geometry where the 192-step K loop exposes its full load->vmcnt->mmac +// chain with no co-resident wave to hide it. +// +// Iteration 2 (architecture round) split K=6144 into 8 uniform 768-element +// slices (24 m16n16k32 steps each) so grid = 32 N-tiles x 8 splits = 256 +// independent one-wave zero-barrier blocks = 2.13 blocks/CU, reaching the +// two-blocks-per-CU latency-hiding target with all 120 CUs covered. Each +// block kept the minimal zero-barrier DUMMA K loop over its own slice and +// stored its int32 partial plane to the caller workspace (8 x 16 x 512 x 4 = +// 256 KiB, within the 16-plane contract budget); a separate 128-thread +// combine kernel (second dispatch) summed the 8 planes in ascending split +// order (bit-exact int32 accumulation order, identical to the unsplit +// kernel) with the fused scale/bf16 epilogue: 76.9 us median (1.29x over the +// 99.3 us Triton Graph baseline). +// +// Iteration 3 (architecture round, this source): the fresh PMC showed the +// two-launch wall is dominated by the combine dispatch, not the GEMM -- +// w8a8_dumma_m16_sk8_combine_kernel (grid = 1 block, 128 threads on one CU, +// only 256 KiB of L2-hot planes) profiles at 47.2 us vs 39.4 us for all 256 +// partial blocks (86.6 us profiled aggregate vs 76.9 us unprofiled +// operator): a one-CU serial dependent-load-chain cliff worth roughly half +// the wall. The decision: keep the accepted split-K=8 one-wave zero-barrier +// grid and remove the one-CU combine dispatch by fusing the combine into the +// partial kernel as a validated qkv-style last-arrival tail. Every block +// ends with __threadfence() + atomicAdd on a monotonic per-N-tile counter in +// the workspace tail; the 8th arriver of each tile (32 of 256 blocks) sums +// the tile's 8 planes in ascending split order (bit-exact int32, identical +// to the unsplit kernel) and applies the fused scale/bf16 epilogue. +// Counters are monotonic ((arrived & 7) == 7 fires exactly once per tile per +// replay), so no reset is needed between Graph replays; one dispatch per +// replay, still Graph-safe. +// +// Repair 1 (in-round, iteration 3) suspected uninitialized counters and added +// a hipStreamIsCapturing-gated one-time hipMemsetAsync; the exact check still +// failed with the identical 7168/8192 mismatch signature, so that diagnosis +// was wrong. +// +// Repair 2 (in-round, iteration 3, this source): the real defect is that the +// fused tail let ALL 64 lanes of every block execute the arrival atomicAdd +// (64 increments per block = 512 per tile per replay), so (arrived & 7) == 7 +// fired ~64 times per tile -- the first fires inside the first-arriving +// block, before its 7 siblings had stored their planes, and the premature +// combiner summed unwritten zero planes: exact check failed with +// first_mismatch (0,0) actual=0.0 and 28/32 tiles (7168/8192 elements) wrong, +// the identical signature already documented by the validated hy3 TP8 gate_up +// lineage. Fix (mirroring that lineage): exactly ONE arrival per block -- +// lane 0 fences + atomicAdds + fences and publishes the result through an LDS +// flag with barriers; only the last arriver of each tile runs the combine. +// The one-time counter zero is now a workspace-pointer-static-guarded +// hipMemsetAsync (the validated qkv/gate_up pattern; no capture-status +// query), issued on the caller stream by the first launch with the workspace +// (the pre-capture eager warmup), so the timed replay keeps one dispatch and +// the monotonic counters need no reset between Graph replays. +// +// Iteration 4 (architecture round, this source): A-only staging. The fresh +// iteration-4 PMC/ISA of the accepted iteration-3 kernel shows the K-loop +// body is the documented decode poison: 16 global_load_ubyte per lane per +// step (8 A + 8 B bytes), a 16-deep s_waitcnt vmcnt cascade, ~20 byte-OR +// reassembly VALU, and ~30 address-carry VALU -- ~64 instructions per step +// with the whole A side of the load chain inside the mmac critical path. +// The DTK du_mma.hpp implementation confirms the m16n16k32 row_major +// matrix_a fragment is 8 CONSECUTIVE bytes at (row = lane&15, k = +// (lane>>4)*8 + {0..7}) -- vectorizable -- while the matrix_b fragment's 8 +// bytes are strided by the leading dimension (not vectorizable without +// packing, out of scope). Decision: keep the accepted split-K=8 geometry +// (grid = 256 blocks = 2.13 blocks/CU, all 120 CUs covered, fused +// last-arrival combine tail) and stage each block's 16x768 A slice into LDS +// once with vectorized 16-B loads (12 global_load_dwordx4 + 12 +// ds_write_b128 per lane, one __syncthreads), then read every k-step's A +// fragment from LDS as one 8-B ds_read_b64 (x[0..7] = A[row][k0 + +// (lane>>4)*8 + {0..7}], bit-identical to du_load_matrix_sync). B keeps +// its direct global path. A bytes reused: the 12 KiB slice is loaded from +// global exactly once per block and each byte is read 24x from LDS (once +// per k-step); cross-block A/B reuse via L2 is unchanged. No workspace, +// plane, counter, or combine-tail change; exact int32 accumulation order +// (k-ascending within a split, ascending split sum) is preserved because +// the mmac consumes identical fragment bytes. +// +// Iteration 5 (packed-weight round, this source): the fresh iteration-5 +// PMC of the accepted iteration-4 code object (digest 7d5673b9..., vmem_read +// 53152, valu_instructions 255872) shows the K-loop body is still half +// decode poison: per lane per k-step B contributes 8 global_load_ubyte, ~15 +// of the 16-deep s_waitcnt vmcnt cascade, ~20 byte-OR reassembly VALU and +// ~30 v_add/v_addc address-carry VALU -- the whole B side of the load chain +// inside the mmac critical path, because the m16n16k32 row_major matrix_b +// fragment's 8 bytes are strided by the leading dimension N and cannot be +// vectorized in the raw [K][N] layout (the iteration-4 round staged A only, +// per its mandate). Decision (mandated option: packed layout): pack B once, +// outside the timed region and outside Graph capture, into the n-major +// transpose packed[n*K + kk] = raw[kk*n + n] for the exact (k, n) == +// (6144, 512) pair (byte count unchanged, so captured addresses stay +// valid). In the packed layout the B fragment is 8 CONTIGUOUS bytes at +// packed[(n0 + lane&15)*k + k0 + (lane>>4)*8 + {0..7}] -- the exact byte +// pattern du_load_matrix_sync(b_frag, b + n0*k + k0, k) reads -- so each +// k-step's B fragment becomes ONE aligned global_load_dwordx2 per lane with +// no byte-OR reassembly and no per-byte waitcnt cascade; the A-side LDS +// stage (iteration 4) is unchanged. This is the exact load-path pattern +// validated on the K=6144 m16n16k32 gate_up lineage (MiniMax TP8 gate_up, +// accepted iterations 5-16: n-major pack + one 8-byte dwordx2 load per +// fragment). The generic scalar fallback decodes the same n-major pack +// whenever (n, k) == (512, 6144), so the paired M=2 API shape with the same +// (N, K) stays correct through the packed buffer; every other (K, N) keeps +// the identity pack and the raw row-major decode, and the exact int32 +// accumulation order (k-ascending within a split, ascending split sum) is +// unchanged because the mmac consumes identical fragment bytes. +// +// Iteration 7 (pipeline round): double-buffered B transport. The fresh +// iteration-5 ISA (digest d2ffb4f6...) shows the single-buffered K-loop +// body is exactly global_load_dwordx2 (B) -> ds_read_b64 (A) -> +// s_waitcnt vmcnt(0) lgkmcnt(0) -> v_mmac per step, with no unroll, so the +// next B load cannot issue until the branch returns to the loop head: every +// per-block B load (a zero-reuse HBM stream -- each (tile, split) block +// reads a disjoint 12 KiB B slice, l2_misses 55733 ~ the compulsory 3.57 +// MiB traffic) exposes its full ~590-cycle HBM round trip on the critical +// path. Mandated this round: compare single buffering with double +// buffering (K=6144 >= 1024, L2 hit 55.11% < 70%, doubled LDS budget +// 25,096 B < 48 KiB all hold). Chosen double buffering: depth-1 register +// prefetch of the packed B fragment (the o_proj-validated pattern -- next +// step's loads issued before the current mmac, rotated in after, compiler +// wait lands at the next fill): step i+1's aligned global_load_dwordx2 is +// issued BEFORE step i's ds_read_b64 + v_mmac, so its HBM latency overlaps +// the previous step's compute instead of sitting on the critical path; the +// final step is peeled so the prefetch address never leaves the block's K +// slice (the packed buffer ends exactly at 512*6144 B; an unconditional +// k0+32 read on the last step would be out of bounds). No second LDS +// buffer: the LDS-fed B path regressed twice (MiniMax TP8 gate_up 25.9 us; +// our iteration 6 at 19.624 us vs 16.778 us), so packed-global dwordx2 B +// remains the transport. LDS per block stays 12,544 + 4 B (doubled budget +// 25,096 B < 48 KiB; 2.13 blocks/CU unchanged), barriers per K step stay 0 +// (only the single pre-loop staging barrier and the two tail barriers, +// unchanged from iteration 5), and the mmac consumes bit-identical fragment +// bytes in k-ascending order, so the output stays bit-identical to the CPU +// reference (0 mismatches, exact check). +// +// Iteration 9 (MLP round, this source): grouped depth-8 register prefetch. +// The iteration-8 occupancy probe (split-K 15 -> 480 one-wave blocks = exactly +// 4 blocks/CU, integer-balanced) REGRESSED to 21.54 us median, so co-resident +// wave count is not the lever: the kernel is DRAM-latency-bound on per-wave +// in-flight B bytes, not wave-starved (l2_misses ~55.7k = the compulsory 3.57 +// MiB A+B traffic; each block still serially waits ~24 ~440-540-cycle HBM +// round trips). This round raises memory-level parallelism INSIDE the wave +// while keeping the exact 256-block split-K=8 grid and its contiguous 12 KiB +// per-block B slices (no DRAM scatter): the 24-step K loop is regrouped into +// 3 x 8-step groups; before each group's 8 mmacs consume the previously +// loaded B fragments, the NEXT group's 8 packed-B dwordx2 loads are issued +// back-to-back (8 independent loads in flight per lane = 4 KiB per wave), so +// ONE vmcnt fill covers 8 steps and each group's loads have the whole +// preceding 8-mmac burst (~8x the iteration-7 single-step window) to expire +// off the critical path. Grid, split-K=8, A LDS staging, packed-B transport, +// zero in-loop barriers, workspace, and the fused last-arrival tail are +// untouched; the final group is peeled (no prefetch) so the prefetch address +// never leaves the block's [k_start, k_start + k_per_split) slice, and the +// mmacs consume bit-identical fragment bytes in k-ascending order, so the +// output stays bit-identical to the CPU reference (0 mismatches, exact +// check). +// +// Iteration 11 (layout consolidation, this source): fragment-slot B pack. +// The iteration-10 direct-A probe REGRESSED to 15.06 us median: its ISA shows +// the compiler could not hoist the 48 global A+B dwordx2 loads (register +// pressure) and fell back to a depth-2 rotation with `s_waitcnt vmcnt(8)` +// before EVERY mmac, re-exposing a full global round trip per step -- so the +// A-side LDS staging stays (the 12 KiB stage is served by ~30-40-cycle +// ds_read2_b64 with 2-step lookahead, not by global loads). The accepted +// iteration-9 kernel's remaining B-side inefficiency is LAYOUT, not MLP: +// the 24 packed-B dwordx2 loads are fully hoisted (all in flight from the +// prologue, per-step vmcnt waits mostly satisfied), but the iteration-5 +// n-major pack packed[n*K+kk] places a block's 16 rows x 768 B B-slice in 16 +// chunks 6,144 B apart, so each 64-lane load instruction touches 16 strided +// 64-B lines (8 B used each) and the 256 concurrent blocks stream 4,096 +// strided DRAM row streams -- the kernel delivers only ~250-270 GB/s +// (counter-derived) where the o_proj-validated fragment-slot pack +// (packed[n_tile][k_step][lane][8]) carried ~461 GB/s on this part. This +// round re-packs the exact (k, n) == (6144, 512) pair once (outside the +// timed region, same buffer bytes) into b2[(tile*(k/32) + step)*512 + row*32 +// + kk]: one 512-B slot per m16n16k32 step, so each lane still reads its 8 +// CONSECUTIVE fragment bytes as ONE aligned global_load_dwordx2 (the exact +// du_load_matrix_sync fragment bytes, k-ascending order, bit-exact int32), +// a full 64-lane load instruction covers 512 CONTIGUOUS bytes (8 full 64-B +// lines, 100% line utilization), and each block's 24-slot slice is one +// contiguous 12 KiB stream -- the 256 blocks now stream the 3 MiB pack in +// 256 sequential DRAM row runs. Grid (256), split-K=8, A LDS staging, +// zero in-loop barriers, the b8 grouped rotation, the workspace, and the +// fused last-arrival tail are untouched; the scalar fallback decodes the +// same fragment-slot pack for (n, k) == (512, 6144) so the paired M=2 API +// shape stays correct through the same packed buffer; l2_misses stay the +// compulsory ~55.6k (same bytes, no over-fetch) while the B second-half +// line L2 hits disappear (l2_hits drop), and the output stays bit-identical +// to the CPU reference (0 mismatches, exact check). +// +// Iteration 20 (stream-shape probe, this source): split-K 8 -> 6. The +// measured split-K/stream lattice for this shape is {blocks x per-wave KiB +// slice -> median}: 64 x 48 -> 13.64 us (iter 15), 128 x 24 -> 12.124 us +// (iter 13), 256 x 12 -> 11.736 us (accepted, iter 11), 480 x 4 -> 21.54 us +// (iter 8). Two readings survive every probe so far: (a) the counter-derived +// delivery rate rises with the number of concurrent per-wave streams +// (260 -> 287 -> 302 GB/s at 64 -> 128 -> 256 streams), with diminishing +// elasticity; (b) there is a steep DRAM-efficiency cliff below ~12-KiB +// per-stream length (4-KiB slices collapse to ~165 GB/s even at 480 streams +// and 4 blocks/CU). Those two readings leave exactly ONE untested lattice +// point that keeps every stream >= 12 KiB with an integer uniform split: +// 192 blocks x 16-KiB slices = split-K 6 (6144/6 = 1024 = 32 k32 steps per +// slice; 32 N-tiles x 6 = 192 blocks = 1.6 blocks/CU, fully resident at +// every plausible occupancy >= 2/CU, so no queued tail at any occupancy). +// 6 is deliberately OUTSIDE the round's trusted set {2,3,4,7,8,10,11,15}: +// the set's untested members are either length-cliffed (10/11 would need +// 8.7-9.6-KiB slices), underfilling (3 -> 96 blocks < 120 CUs), or require +// ragged non-constexpr splits (7/10/11: 192 is not divisible by 7/10/11), +// while 6 divides 192 evenly and keeps the exact 8-step grouped structure +// (32 = 4 x 8). The decode skill's non-power-of-two sweep family includes +// 6 explicitly, and the validated wqkv_a winner used split-K 6. Mechanism: +// the 3.55-MiB compulsory traffic, the 3-MiB A L2 re-reads (192 x 16 KiB == +// 256 x 12 KiB), the 512-B per-instruction B request granularity, the +// fragment-slot pack, the grouped depth-8 rotation, the zero-barrier +// single-wave loop, the plane store, and the fused last-arrival tail are +// all unchanged -- only the per-wave B stream grows 12 -> 16 KiB and the +// concurrent stream count falls 256 -> 192. Falsifiable: median stays +// 11.736 +/- 0.3 if ~302 GB/s is a count- and length-independent DRAM-mix +// floor for this 3.55-MiB traffic (the session's second plateau-eligible +// flat candidate after iteration 19), drops toward ~11.5-11.65 if the +// longer 16-KiB rows lift delivery to ~305-310 GB/s, or regresses to +// ~12.0-12.1 if count-dominance dominates (falsifying the +// length-compensation model). Outputs stay bit-identical to the CPU +// reference: the pack is split-agnostic (pure (tile, k-step) slots), each +// split's 32 steps are consumed in k-ascending order, and the fused tail +// sums the 6 planes in ascending split order (int32 addition commutative, +// every partial sum within int32 range as before), firing on +// (arrived % 6) == 5 -- the monotonic modulo fire lands exactly once per +// tile per replay on the true 6th arrival. +// +// The scalar kernel stays the fallback for every unmatched (m, n, k), +// including the paired M=2 API shape with the same (N, K) and the +// down_proj shape. +// +// Down_proj iteration 1 (this source): minimal DUMMA bootstrap for +// glm_tp8_shared_down_proj_m16 (M=16, N=6144, K=256). Until this round the +// down_proj shape ran the generic scalar fallback (grid 768 x 128 thr = +// 98304 threads, one output element per thread, 45.0 us median vs the +// 18.177 us fixed Triton Graph baseline) -- the primary kernel in the fresh +// PMC was w8a8_gemm_scalar_kernel (grid_blocks 768, workgroup 128, ~5.5 +// VALU per byte of B traffic). The seed's DUMMA kernel covers only the +// gate_up (N=512, K=6144) pair, so this round establishes the mandated +// minimal gfx928 INT8 m16n16k32 DUMMA tile for the down_proj shape: ONE +// 64-thread wavefront per block, one 16x16 N-tile per block, grid = N/16 = +// 384 blocks (3.2 blocks/CU on 120 CUs), explicit int32 accumulation, zero +// cross-wave barriers (single wave -> the compiler's lgkmcnt orders the +// LDS write-before-read; no __syncthreads at all). Transport (the +// lineage-validated decode recipe): the whole K=256 slice is staged ONCE +// into 272 = 256+16 bank-skewed LDS with cooperative 16-B int4 loads (4 per +// lane per matrix); A is read as row_major matrix_a fragments (8 contiguous +// bytes per lane) from the logical [16][256] slice, B as col_major +// matrix_b fragments (8 contiguous k bytes per lane) from a NEW one-time +// packed [N,K] n-major transpose P[n*K+k] = W[k*N+n] for the exact +// (k,n) == (256, 6144) pair (packed outside the timed region by +// launch_pack_w8a8_weight; the generic scalar fallback decodes the same +// n-major pack for (n,k) == (6144,256) so the paired M=2 API shape with +// the same (N,K) stays correct). Each k-step is one ds_read_b64 per +// fragment (no byte-load cascade, no waitcnt-per-byte), and the epilogue +// uses the verified gfx928 accumulator ownership (row = lane&15, +// col_mod4 = lane>>4, acc_frag.x[i] -> column col_mod4 + 4*i) with the +// fused x_scale/weight_scale/bf16 store straight from registers; no +// workspace use (no split-K). Exact int32 accumulation order is +// k-ascending over the full K, bit-identical to the scalar reference +// (max |acc| = 256*127*127 = 4,129,024 << 2^31). +// +// Math contract: +// out[m, n] = bf16( int32_dot(a[m, :], b[:, n]) * x_scale[m] * weight_scale[n] ) +// The int8 dot product is accumulated exactly in int32 (the largest assigned +// K=6144 keeps every partial sum within int32 range: +// 6144 * 127 * 127 = 99,058,176 < 2^31), so the result is bit-identical to +// the CPU int64 reference regardless of accumulation order. + +// Include the installed headers in this known-good order: the DTK's +// du_mma.h is not self-contained when included before the HIP runtime +// headers. +#include +#include +#include + +#include + +namespace { + +constexpr int kScalarBlockThreads = 128; +constexpr int kPackBlockThreads = 256; + +// DUMMA tile constants for the exact-shape specialization. gfx928 INT8 +// tensor cores execute m16n16k32 (signed char x signed char -> int32) with +// wavefront=64; each 16x16 output tile is owned by exactly one wavefront. +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; +constexpr int kDummaBlockThreads = 64; // one 64-thread wavefront per block +constexpr int kDummaSk2BlockThreads = 128; // down_proj iter 2: 2 waves/block +// down_proj iter 8 (multi-N-tile amortization): 4 waves/block, each block +// computes TWO 16x16 N tiles (N=32/block) so the shared A stage, the +// cooperative staging round trip and the combine barrier amortize over 2 +// tiles while keeping the accepted 768-wave / 1.6-waves-per-SIMD class. +constexpr int kDummaSk2N32BlockThreads = 256; // down_proj iter 8: 4 waves/block + +// Split-K geometry for the exact-shape specialization (iteration 2; +// iteration 20: 8 -> 6). K=6144 / 32 = 192 m16n16k32 steps, split into 6 +// uniform 32-step slices; grid = (N/16) * kDummaSplitK = 32 * 6 = 192 blocks +// = 1.6 blocks/CU on the 120-CU device (fully resident at every occupancy +// >= 2/CU, no queued tail), each slice a contiguous 16-KiB fragment-slot +// run. 6 divides 192 evenly, so every K boundary stays aligned to the +// DUMMA k=32 unit with zero wasted work. Probe rationale (iteration 20): +// the measured stream lattice (64 x 48 KiB -> 13.64 us, 128 x 24 KiB -> +// 12.124 us, 256 x 12 KiB -> 11.736 us, 480 x 4 KiB -> 21.54 us) reads +// monotone in concurrent-stream count (260 -> 287 -> 302 GB/s) with a steep +// length cliff below ~12 KiB; 192 x 16 KiB is the only untested lattice +// point that keeps every stream >= 12 KiB with an integer uniform split. +constexpr int kDummaSplitK = 6; +constexpr int kDummaCombineThreads = 128; + +// Iteration-4 A-only staging constants (iteration 20: k_per_split 768 -> +// 1024 with split-K 6). The exact-shape guard pins k=6144, so each split +// slice is k_per_split = 6144/6 = 1024 int8 columns. The A stage is the +// 16x1024 slice padded to a 1040-column row stride: 1040 is 16-aligned +// (staging stores compile to ds_write_b128) and 8-aligned (the per-k-step +// A-fragment read compiles to one ds_read_b64); the +16 skew bounds LDS +// bank conflicts to 2-way. 16 * 1040 = 16,640 B LDS per block; with ~2-3 +// blocks/CU that is ~33-50 KiB of the 64 KiB/CU LDS, so the accepted +// occupancy class is unchanged (2/CU = 33,280 B < 64 KiB still holds). +constexpr int kPerSplitK = kDummaTileK * 32; // 1024 = 6144 / 6 +constexpr int kAStageStride = kPerSplitK + 16; // 1040 + +// --------------------------------------------------------------------------- +// Generic scalar W8A8 GEMM fallback. +// +// One thread computes exactly one output element out[row, col]. Threads are +// mapped linear = row * n + col, so adjacent lanes sit on adjacent addresses +// in the fastest-changing N dimension: B[kk, col] loads coalesce across a +// wavefront and A[row, kk] broadcasts within a row. blockDim is a multiple +// of the gfx928 wavefront size 64. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear % n); + + const int8_t* __restrict__ a_row = a + static_cast(row) * k; + // The exact gate_up (n, k) == (512, 6144) uses the iteration-11 + // fragment-slot packed weight b2[(tile*(k/32) + step)*512 + row*32 + kk] + // (produced once by launch_pack_w8a8_weight outside the timed region); the + // exact down_proj (n, k) == (6144, 256) uses the iteration-1 n-major + // packed weight P[n*K + kk] = raw[kk*n + n] (produced once by the same + // pack launch); all other shapes keep the identity [K, N] row-major + // layout. Decode all three so the generic fallback -- including the + // paired M=2 API shape with the same (N, K) -- stays correct through the + // same packed pointer. + const bool fragslot_packed = (n == 512 && k == 6144); + const bool nmajor_packed = (n == 6144 && k == 256); + const int8_t* __restrict__ b_col = b + col; + int64_t b_stride = static_cast(n); // identity layout per-kk stride + if (fragslot_packed) { + // b2[(col>>4)*(k>>5)*512 + (col&15)*32 + (kk>>5)*512 + (kk&31)] + b_col = b + static_cast(col >> 4) * ((k >> 5) * 512) + + static_cast(col & 15) * 32; + b_stride = 0; // addressed via the per-kk slot formula below + } else if (nmajor_packed) { + // P[col*K + kk] = raw[kk*n + col]: 8-byte-aligned contiguous k bytes. + b_col = b + static_cast(col) * k; + b_stride = 1; + } + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + const int64_t b_off = fragslot_packed + ? (static_cast(kk >> 5) * 512 + (kk & 31)) + : kk * b_stride; + acc += static_cast(a_row[kk]) * + static_cast(b_col[b_off]); + } + + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Exact-shape split-K DUMMA specialization for glm_tp8_shared_gate_up_proj_m16: +// (M=16, N=512, K=6144). Guarded by an exact (m, n, k) match in +// launch_w8a8_gemm so paired M=2 API shapes and every other shape keep the +// scalar fallback. +// +// Architecture round (iteration 2): the iteration-1 grid of 32 one-wave +// blocks left 88 of 120 CUs idle and each 192-step K loop ran with no +// co-resident wave (PMC: grbm_count == grbm_gui_active, 173.6 us). Splitting +// K into 8 uniform 768-element slices turns the launch into 256 independent +// zero-barrier one-wave blocks (2.13 blocks/CU), so while one block's +// fragment loads sit on vmcnt, the co-resident block issues its mmac chain. +// Each block keeps the minimal gfx928 INT8 m16n16k32 DUMMA tile with explicit +// int32 accumulation over its own k-slice (split-major block order keeps the +// first 32 blocks on one 384 KiB B slice: L2-hot), then stores the int32 +// partial plane to the workspace using the verified gfx928 accumulator lane +// ownership (row = lane & 15, col_mod4 = lane >> 4, acc_frag.x[i] -> columns +// col_mod4 + 4*i). +// +// Iteration 3 fused combine (this source): the standalone 128-thread combine +// dispatch profiled at 47.2 us (grid = 1 block on one CU) -- the largest +// single component of the 76.9 us wall. The combine is now a last-arrival +// tail inside this kernel: after the plane store each block publishes one +// monotonic arrival on counters[tile] (tile = the block's N-tile); the block +// whose arrival makes (arrived & 7) == 7 is the last of the tile's 8 splits, +// so it sums the 8 planes in ascending split order (bit-exact int32 +// accumulation order, identical to the unsplit kernel) and applies the fused +// scale/bf16 epilogue. Repair 2: exactly ONE arrival per block (lane 0 fence +// + atomicAdd + LDS flag + barriers; the round-3 draft's all-64-lane atomicAdd +// fired the combine prematurely, before sibling planes were stored -- see the +// tail comment). Only 32 of 256 blocks run the combine; counters live in the +// workspace tail, are zeroed once per workspace by launch_w8a8_gemm, and +// never need a reset (monotonic across Graph replays). Single Graph-safe +// dispatch per replay. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDummaBlockThreads) void +w8a8_dumma_m16_sk8_partial_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int32_t* __restrict__ partials, // [kDummaSplitK][M][N] int32 planes + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + __hip_bfloat16* __restrict__ out, // [M][N] + int32_t* __restrict__ counters, // [N/16] monotonic arrival counters + int n, + int k) { + const int lane = static_cast(threadIdx.x); + const int n_tiles = n / kDummaTileN; + const int split = static_cast(blockIdx.x) / n_tiles; + const int tile = static_cast(blockIdx.x) % n_tiles; + const int n0 = tile * kDummaTileN; + + const int k_per_split = k / kDummaSplitK; // 1024, a multiple of 32 + const int k_start = split * k_per_split; + + // Iteration 4 (A-only staging): stage this block's 16 x k_per_split A + // slice into LDS once with vectorized 16-B loads, then serve every + // k-step's A fragment from LDS as one 8-B read. The DTK du_mma.hpp + // row_major matrix_a loader maps lane l to 8 CONSECUTIVE bytes + // A[lane&15][k0 + ((lane>>4)<<3) + {0..7}], so the staged fragment bytes + // are bit-identical to du_load_matrix_sync while the accepted kernel's 8 + // global ubyte loads + waitcnt cascade + byte-OR reassembly for A leave + // the per-step critical path (see the header comment for the iteration-4 + // ISA evidence). B (iteration 5) is the n-major packed weight + // (packed[n*K+kk] = raw[kk*n+n] for (k, n) == (6144, 512), produced once + // by launch_pack_w8a8_weight outside the timed region), whose fragment is + // 8 CONTIGUOUS bytes at packed[(n0 + lane&15)*k + k0 + (lane>>4)*8 + + // {0..7}], loaded below as one aligned global_load_dwordx2 per lane + // (bit-identical to du_load_matrix_sync on the packed layout). + __shared__ int8_t s_a[kDummaTileM * kAStageStride]; + { + // 16 x 1024 B = 16,384 B = 1024 x 16-B groups; 64 lanes x 16 = 1024. + constexpr int kGroupsPerRow = kPerSplitK / 16; // 64 + constexpr int kStageVecs = + (kDummaTileM * kGroupsPerRow) / kDummaBlockThreads; // 16 + static_assert(kDummaTileM * kGroupsPerRow == + kDummaBlockThreads * kStageVecs, + "A slice must divide evenly across the staging lanes"); +#pragma unroll + for (int t = 0; t < kStageVecs; ++t) { + const int g = lane + kDummaBlockThreads * t; + const int row = g / kGroupsPerRow; + const int c16 = (g % kGroupsPerRow) * 16; + // 16-aligned: a base, k_start (multiple of 1024), row*k (k=6144) and + // c16 are all multiples of 16, so this is one global_load_dwordx4. + const int4 v = + *reinterpret_cast(a + k_start + row * k + c16); + // 16-aligned in LDS: kAStageStride = 1040 is a multiple of 16, so this + // is one ds_write_b128. + *reinterpret_cast(&s_a[row * kAStageStride + c16]) = v; + } + } + __syncthreads(); // every lane reads rows staged by other lanes below + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + const int a_row = lane & (kDummaTileM - 1); + const int a_kq = (lane >> 4) << 3; // 8 consecutive k bytes per lane + // Iteration 5 (packed weight): the same lane bits select the B fragment + // of the n-major packed weight -- n = lane & 15, k offset (lane>>4)*8 -- + // so b_row and b_kq coincide with a_row / a_kq. + const int b_row = lane & (kDummaTileM - 1); + const int b_kq = (lane >> 4) << 3; + // Iteration 11 (layout consolidation, this source): the per-step packed-B + // address is loop-invariant except the step index, so precompute its base + // here. The packed weight is now the o_proj-validated fragment-slot + // layout b2[(tile*(k/32) + step)*512 + row*32 + kk] (one 512-B slot per + // m16n16k32 step): lane l still owns the 8 CONSECUTIVE bytes + // B[n0 + lane&15][k0 + (lane>>4)*8 + {0..7}] -- the exact + // du_load_matrix_sync fragment -- but a full 64-lane load instruction now + // covers 512 CONTIGUOUS bytes (8 full 64-B lines, 100% line utilization) + // instead of 16 rows x 8 B strided 6,144 B apart under the iteration-5 + // n-major pack. Each block's 24-slot slice is one contiguous 12 KiB + // stream (the contiguity the iteration-9 hypothesis assumed), so the 256 + // blocks stream the 3 MiB pack in 256 sequential DRAM row runs instead of + // 4,096 strided row streams -- the same fragment-slot transport that + // carried the validated o_proj kernel at ~461 GB/s effective read + // bandwidth on this part (vs ~250-270 GB/s counter-derived for this + // kernel's strided n-major pack). Current and prefetch addresses stay + // 8-byte aligned (b_base and kBFragStep are multiples of 8). + constexpr int kBFragStep = kDummaTileN * kDummaTileK; // 512 B per k32 slot + const int64_t b_base = + (static_cast(tile) * (k >> 5) + + split * (k_per_split / kDummaTileK)) * + kBFragStep + + static_cast(b_row) * kDummaTileK + b_kq; + + // Iteration 7 (pipeline round): double-buffered B transport. The + // iteration-5 K-loop body (one global_load_dwordx2 (B) + one ds_read_b64 + // (A) + s_waitcnt vmcnt(0) lgkmcnt(0) + one v_mmac per step, no unroll) + // exposes the full ~590-cycle HBM round trip of every B load on the + // critical path because the next load cannot issue until the branch + // returns to the loop head, and B is a zero-reuse HBM stream (disjoint + // 12 KiB slice per (tile, split) block). This round is the mandated + // single-vs-double buffering comparison. Double buffering here = one + // register-resident B fragment in flight (depth-1 software prefetch, the + // o_proj-validated pattern): step i+1's packed dwordx2 is issued BEFORE + // step i's ds_read + mmac and rotated in after, so its HBM latency + // overlaps the previous step's compute and the compiler places the vmcnt + // wait at the next fill point. LDS stays 12,544 + 4 B per block (no + // second LDS buffer: the LDS-fed B path regressed twice -- MiniMax TP8 + // gate_up 25.9 us and our iteration 6 at 19.624 us -- while packed-global + // dwordx2 B is the validated transport), so the doubled LDS budget + // (25,096 B) stays below 48 KiB and the 2.13 blocks/CU geometry is + // unchanged. Barriers per K step: 0 -- the only __syncthreads are the + // single pre-loop A-staging barrier and the two tail barriers, unchanged + // from iteration 5. The final step is peeled so the prefetch address + // never leaves the block's [k_start, k_start + k_per_split) slice (the + // packed buffer ends exactly at 512*6144 B; an unconditional k0+32 read + // on the last step would read past its end). The mmac consumes + // bit-identical fragment bytes in k-ascending order, so the output stays + // bit-identical to the CPU reference (0 mismatches, exact check). + // Iteration 9 (MLP round, this source): grouped depth-8 register prefetch. + // The iteration-8 occupancy probe (split-K 15 -> 480 one-wave blocks = + // exactly 4 blocks/CU, integer-balanced, no tail wave) REGRESSED to 21.54 + // us median, so co-resident wave count is NOT the lever: this kernel is + // DRAM-latency-bound on per-wave in-flight B bytes, not wave-starved + // (l2_misses ~55.7k = the compulsory 3.57 MiB A+B traffic; every block + // still serially waits ~24 ~440-540-cycle HBM round trips). This round + // raises memory-level parallelism INSIDE the wave while keeping the exact + // 256-block split-K=8 grid and its contiguous 12 KiB per-block B slices + // (no DRAM scatter): the 24-step K loop is regrouped into 3 x 8-step + // groups; before each group's 8 mmacs consume the previously loaded B + // fragments, the NEXT group's 8 packed-B dwordx2 loads are issued + // back-to-back (8 independent loads in flight per lane = 4 KiB per wave), + // so ONE vmcnt fill covers 8 steps and each group's loads have the whole + // preceding 8-mmac burst (~8x the iteration-7 single-step window) to expire + // off the critical path. Grid, split-K=8, A LDS staging, packed-B + // transport, zero in-loop barriers, workspace, and the fused last-arrival + // tail are untouched; the final group is peeled (no prefetch) so the + // prefetch address never leaves the block's [k_start, k_start + + // k_per_split) slice (the packed buffer ends exactly at 512*6144 B), and + // the mmacs consume bit-identical fragment bytes in k-ascending order, so + // the output stays bit-identical to the CPU reference (0 mismatches, + // exact check). + constexpr int kUnroll = 8; + constexpr int kGroupK = kUnroll * kDummaTileK; // 256 = 8 k32 steps + static_assert(kPerSplitK % kGroupK == 0, + "32 steps must split evenly into 8-step groups"); + uint64_t b8[kUnroll]; + // b_first is the base of this block's 32-slot (16 KiB) contiguous + // fragment-slot slice; each slot is kBFragStep = 512 B. Declared at + // function scope (repair 2: the draft scoped it inside the prologue block, + // so the grouped loop below could not see it) so the prologue, the grouped + // loop and the peeled tail all reuse the same base. + const int64_t b_first = b_base; + { + // Prologue: issue group 0 (the split slice's first 8 fragment slots) so + // its round trip is in flight before the first group fill. +#pragma unroll + for (int i = 0; i < kUnroll; ++i) { + b8[i] = *reinterpret_cast(b + b_first + i * kBFragStep); + } + } + int k0 = k_start; + int64_t b_off = 0; // B byte offset of the current 8-step group's first slot + // Three full groups: issue the NEXT group's 8 B loads before this group's + // 8 mmacs (each next-group load is in flight during the whole 8-mmac + // burst) and rotate the registers in after; the compiler places the vmcnt + // fill at the group head (the iteration-7 wait placement on this shape). + for (int g = 0; g < kPerSplitK / kGroupK - 1; ++g) { + uint64_t b8_next[kUnroll]; +#pragma unroll + for (int i = 0; i < kUnroll; ++i) { + // 8-aligned: b_first, b_off, kUnroll*kBFragStep, i*kBFragStep and b_kq + // are all multiples of 8, so each of these is one global_load_dwordx2 + // over 8 consecutive bytes of one 512-B fragment slot (8 full 64-B + // lines per full 64-lane load instruction). + b8_next[i] = *reinterpret_cast( + b + b_first + b_off + kUnroll * kBFragStep + i * kBFragStep); + } +#pragma unroll + for (int i = 0; i < kUnroll; ++i) { + // A fragment from LDS: one 8-B read reproduces the exact + // du_load_matrix_sync bytes (x[0..7] = A[row][k0 + a_kq + {0..7}]); + // the offset is 8-aligned (kAStageStride and a_kq are multiples of 8 + // and (k0 - k_start) + i*kDummaTileK is a multiple of 32), so this is + // one ds_read_b64. + const uint64_t a8 = *reinterpret_cast( + &s_a[a_row * kAStageStride + (k0 - k_start) + i * kDummaTileK + + a_kq]); + __builtin_memcpy(a_frag.x, &a8, sizeof(a8)); + // B fragment from the iteration-11 fragment-slot pack: 8 contiguous + // bytes at b2[(tile*(k/32) + 24*split + s)*512 + b_row*32 + b_kq] for + // step s of the block's split slice (i = 0..7) -- the exact bytes + // du_load_matrix_sync(b_frag, b + n0*k + k0, k) reads from the packed + // layout, delivered as ONE aligned global_load_dwordx2 per lane + // instead of 8 strided global_load_ubyte + waitcnt cascade + byte-OR + // reassembly. + __builtin_memcpy(b_frag.x, &b8[i], sizeof(b8[i])); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } +#pragma unroll + for (int i = 0; i < kUnroll; ++i) { + b8[i] = b8_next[i]; + } + k0 += kGroupK; + b_off += kUnroll * kBFragStep; // next 8-step group's first slot + } + { + // Peeled final group (no prefetch): consumes the last group loaded by the + // loop's final issue (b8 holds the split slice's slots 24 .. 31). +#pragma unroll + for (int i = 0; i < kUnroll; ++i) { + const uint64_t a8 = *reinterpret_cast( + &s_a[a_row * kAStageStride + (k0 - k_start) + i * kDummaTileK + + a_kq]); + __builtin_memcpy(a_frag.x, &a8, sizeof(a8)); + __builtin_memcpy(b_frag.x, &b8[i], sizeof(b8[i])); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + } + + // Register-only partial store with the verified gfx928 accumulator + // ownership: lane l owns row = l & 15 and columns col_mod4 + 4*i with + // col_mod4 = l >> 4, so acc_frag.x[i] belongs to output column + // n0 + col_mod4 + 4*i of row (l & 15). + const int plane = kDummaTileM * n; // stride between split planes + int32_t* plane_p = partials + split * plane; + const int row = lane & (kDummaTileM - 1); + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = n0 + col_mod4 + 4 * i; + plane_p[row * n + col] = acc_frag.x[i]; + } + + // Fused last-arrival combine tail (iteration 3; repair 2: exactly ONE + // arrival per block). The iteration-3 draft let all 64 lanes of every + // block execute the arrival atomicAdd (64 increments per block = 512 per + // tile per replay), so (arrived & 7) == 7 fired ~64 times per tile -- the + // first fires inside the first-arriving block, before its 7 siblings had + // stored their planes, and the premature combiner summed unwritten zero + // planes (exact check: 7168/8192 mismatched, first_mismatch (0,0) + // actual=0.0 -- the identical signature documented by the validated hy3 + // TP8 gate_up lineage). Fix: exactly one arrival per block -- lane 0 + // fences (release), atomicAdds, fences (acquire) and publishes the result + // through LDS; the rest of the block waits at a barrier, and only the last + // arriver of each tile (32 of 192 blocks) runs the combine. The monotonic + // modulo fire (arrived % kDummaSplitK) == kDummaSplitK-1 then lands + // exactly once per tile per replay, on the true 6th arrival (iteration 20: + // split-K 6 is non-power-of-two, so the power-of-two mask of iteration 3 + // becomes the generic modulo), after all 6 planes of the tile are visible; + // the counters are zeroed once per + // workspace by launch_w8a8_gemm (repair 2: static-pointer-guarded + // hipMemsetAsync, the validated qkv/gate_up pattern) and stay monotonic + // across Graph replays. + __syncthreads(); + __shared__ int s_is_last; + if (lane == 0) { + __threadfence(); // release: this block's plane stores visible to the + // observer of this block's arrival atomic + const int arrived = atomicAdd(&counters[tile], 1); + __threadfence(); // acquire: reads below see every sibling plane store + // that was released before its arrival atomic + // Iteration 20: kDummaSplitK = 6 is non-power-of-two, so the modulo-6 + // fire replaces the power-of-two mask of iteration 3; the monotonic + // counters still land the fire exactly once per tile per replay on the + // true 6th arrival (counters zeroed once per workspace, monotonic + // across Graph replays). + s_is_last = + ((arrived % kDummaSplitK) == (kDummaSplitK - 1)) ? 1 : 0; + } + __syncthreads(); + if (s_is_last == 0) { + return; + } + { + // Last arriver of this tile: sum the kDummaSplitK planes in ascending + // split order -- bit-exact int32 accumulation order, identical to the + // unsplit kernel -- and fuse the scale/bf16 epilogue. + const int e_local = lane * 4; // 64 lanes x 4 = 256 tile elements + const int tile_row = e_local >> 4; + const int col0 = n0 + (e_local & 15); + int32_t v[kDummaSplitK][4]; +#pragma unroll + for (int s = 0; s < kDummaSplitK; ++s) { + const int32_t* p = partials + s * plane + tile_row * n + col0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + v[s][i] = p[i]; + } + } + const float xs = x_scale[tile_row]; + __hip_bfloat16* orow = out + tile_row * n + col0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kDummaSplitK; ++s) { + acc += v[s][i]; + } + const float scaled = + static_cast(acc) * xs * weight_scale[col0 + i]; + orow[i] = __float2bfloat16(scaled); + } + } +} + +// Superseded by the iteration-3 fused last-arrival tail inside +// w8a8_dumma_m16_sk8_partial_kernel (single dispatch per replay); retained +// only for reference. The iteration-2 version of this kernel (grid = 1 +// block, 128 threads on one CU, serial per-element 8-plane dependent load +// chain) profiled at 47.2 us -- the largest single component of the 76.9 us +// operator wall. +__global__ __launch_bounds__(kDummaCombineThreads) void +w8a8_dumma_m16_sk8_combine_kernel( + const int32_t* __restrict__ partials, // [kDummaSplitK][M][N] int32 + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + __hip_bfloat16* __restrict__ out, // [M][N] + int n) { + const int total = kDummaTileM * n; // 8192 output elements + const int plane = kDummaTileM * n; // stride between split planes + for (int e = static_cast(threadIdx.x); e < total; + e += kDummaCombineThreads) { + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kDummaSplitK; ++s) { + acc += partials[s * plane + e]; + } + const int row = e / n; + const int col = e - row * n; + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[e] = __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// Down_proj iteration 1: minimal gfx928 INT8 m16n16k32 DUMMA tile for +// glm_tp8_shared_down_proj_m16 (M=16, N=6144, K=256), replacing the scalar +// fallback (45.0 us median) that ran this shape until now: one 64-thread +// wavefront per block, one 16x16 output tile per block, grid = N/16 = 384 +// blocks = 3.2 blocks/CU on the 120-CU device, zero-barrier, 8 k-ascending +// m16n16k32 steps over the LDS-staged K=256 slice. Accepted at 15.96 us +// median (1.14x over the 18.177 us Triton Graph baseline). +// +// Down_proj iteration 2 (architecture round, this source): the fresh PMC of +// the accepted iteration-1 code object (profiled_duration 4.96 us, gpu_active +// 100%, grid 384 x 64 threads, 48 VGPR) shows the launch has 384 waves on +// 480 SIMDs -- only 0.8 waves/SIMD, so most SIMDs run their one dependent +// staging -> ds_read -> mmac chain with no co-resident wave to hide it. +// Decision: in-block split-K=2, the first trusted occupancy-probe candidate +// -- 128 threads = two independent wavefronts per block, each staging and +// computing its own K=128 half (4 m16n16k32 steps) with zero cross-wave +// dependency in the K loop: no barrier before it (each wave reads only the +// LDS bytes it staged itself, and the compiler's lgkmcnt orders write-before- +// read within the wavefront). Grid stays 384 blocks = 3.2 blocks/CU (all +// 120 CUs covered), 768 waves total = 1.6 waves/SIMD: every SIMD busy, MMAC +// chains hidden. Each wave then publishes its int32 partial to a lane-major +// LDS slot (4 contiguous int32 per lane: one ds_write_b128, conflict-free), +// one __syncthreads, and wave 0 sums the two partials in ascending split +// order (bit-exact int32, identical to the scalar reference: integer +// addition is associative, and each partial <= 128*127*127 = 2,064,512 while +// the final sum stays 4,129,024 << 2^31) and runs the fused scale/bf16 +// epilogue. LDS rises to 8,704 + 2,048 = 10,752 B/block -> 6 blocks/CU +// capacity, still >= the 3.2 needed; ~48-64 VGPR x 64 lanes -> no register +// cliff. No workspace use (no split-K planes), single Graph-safe dispatch +// per replay. +// +// Data path (the lineage-validated decode recipe for K=256 down_proj): +// 1. Each wave stages its own K=128 half ONCE into LDS with cooperative +// 16-byte int4 loads: 2 global_load_dwordx4 per lane per matrix +// (lane -> (row = lane>>2, kk16 = (lane&3)<<4); 16 rows x 128 k per +// wave, split offset wave*128), all 4 loads issued back-to-back so +// their round trips overlap, then 2 ds_write_b128 per matrix. +// 2. A: logical row-major [16][256] slice staged at the 272 = 256+16 +// bank-skewed row stride; each k-step's matrix_a row_major fragment is +// the 8 CONTIGUOUS bytes A[row = lane&15][kk + (lane>>4)*8 + {0..7}] +// -- one ds_read_b64 per lane per step, bit-identical to +// du_load_matrix_sync. +// 3. B: the packed [N,K] n-major layout P[n*K+k] = W[k*N+n] for the exact +// (k,n) == (256,6144) pair (produced once out-of-timed-region by +// launch_pack_w8a8_weight; same buffer bytes, so captured addresses +// stay valid), staged at the same 272-stride; each k-step's matrix_b +// col_major fragment is the 8 CONTIGUOUS bytes +// P[(n0 + lane&15)][kk + (lane>>4)*8 + {0..7}] -- one ds_read_b64 per +// lane per step (du_load_matrix_sync col_major reads p[row*ldm + col + +// {0..7}] with row = lane&15, col = (lane>>4)*8, exactly these bytes). +// 4. 4 k-ascending m16n16k32 steps per wave over its own half: +// du_load_matrix_sync(a/b) + du_mma_sync(acc, a, b, acc) with explicit +// du_fill_fragment(acc, 0). +// 5. Combine + direct register epilogue with the verified gfx928 +// accumulator ownership (row = lane & 15, col_mod4 = lane >> 4, +// acc_frag.x[i] -> output column col_mod4 + 4*i): p0 + p1 (ascending +// split order) * x_scale[row] * weight_scale[col] -> bf16, four 2-byte +// stores per lane. No workspace use. +// +// Exact int32 accumulation order stays k-ascending over the full K=256 +// (half 0 ascending, then half 1 ascending), bit-identical to the scalar +// reference (max |acc| = 256*127*127 = 4,129,024 << 2^31, so the dot is +// exact regardless of order). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDummaSk2BlockThreads) void +w8a8_dumma_m16_n16_k256_sk2_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, // packed [N,K] n-major for this shape + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n, + int k) { + constexpr int kTileM = 16; + constexpr int kTileN = 16; + constexpr int kTileK = 32; + constexpr int kK = 256; + constexpr int kSplit = 2; + constexpr int kPerSplit = kK / kSplit; // 128 + constexpr int kLdsStride = kK + 16; // 272 = 256 + 16 bank skew + const int wave = static_cast(threadIdx.x) >> 6; // 0..1 (two waves) + const int lane = static_cast(threadIdx.x) & 63; + const int n0 = static_cast(blockIdx.x) * kTileN; + + __shared__ __align__(16) int8_t lds_a[kTileM * kLdsStride]; // 4,352 B + __shared__ __align__(16) int8_t lds_b[kTileN * kLdsStride]; // 4,352 B + __shared__ __align__(16) int32_t s_part[kSplit][64 * 4]; // 2,048 B + + // Stage this wave's K=128 half: lane -> (row = lane>>2, kk16 = (lane&3) + // <<4); i covers the two 64-k quarters of the half, so 16 rows x 128 k = + // 2,048 B per matrix per wave, each (row, chunk) written exactly once. + // All 4 global loads issue back-to-back (one wait covers all), then the + // 4 ds_write_b128 land. Each wave reads only its own staged half in the + // K loop below, so there is no cross-wave data dependency and no barrier: + // the compiler orders LDS write-before-read with s_waitcnt lgkmcnt (the + // single-wave pattern validated in iteration 1). + const int row = lane >> 2; + const int kk16 = (lane & 3) << 4; + const int k_off = wave * kPerSplit; + int4 va[2], vb[2]; +#pragma unroll + for (int i = 0; i < 2; ++i) { + // 16-aligned: a base, row*k (k=256), k_off (128), i*64 and kk16 are all + // multiples of 16, so each of these is one global_load_dwordx4. + va[i] = *reinterpret_cast(a + row * k + k_off + i * 64 + kk16); + vb[i] = *reinterpret_cast( + b + static_cast(n0 + row) * k + k_off + i * 64 + kk16); + } +#pragma unroll + for (int i = 0; i < 2; ++i) { + // 16-aligned in LDS: kLdsStride = 272 is a multiple of 16, so this is + // one ds_write_b128. + *reinterpret_cast(lds_a + row * kLdsStride + k_off + i * 64 + kk16) = + va[i]; + *reinterpret_cast(lds_b + row * kLdsStride + k_off + i * 64 + kk16) = + vb[i]; + } + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // 4 k-ascending m16n16k32 steps over this wave's K=128 half. No barrier: + // every ds_read touches bytes this wave staged itself. +#pragma unroll + for (int kk = 0; kk < kPerSplit; kk += kTileK) { + du::dumma::du_load_matrix_sync(a_frag, lds_a + k_off + kk, kLdsStride); + du::dumma::du_load_matrix_sync(b_frag, lds_b + k_off + kk, kLdsStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Publish this wave's int32 partial in lane-major order: 4 CONTIGUOUS + // int32 per lane at s_part[wave][lane*4 .. lane*4+3] (16-byte aligned -> + // one ds_write_b128, conflict-free). Ascending split order: wave 0 is + // the k=0..127 half, wave 1 the k=128..255 half. + *reinterpret_cast(&s_part[wave][lane * 4]) = + make_int4(acc_frag.x[0], acc_frag.x[1], acc_frag.x[2], acc_frag.x[3]); + + __syncthreads(); // wave 0's combine must observe wave 1's partial + + // Wave 0 sums the two partials in ascending split order (bit-exact int32 + // accumulation, identical to the scalar reference) and runs the direct + // register epilogue with the verified gfx928 accumulator ownership: + // lane l owns row = l & 15 and columns col_mod4 + 4*i with col_mod4 = + // l >> 4, so partial element i belongs to output column n0 + col_mod4 + + // 4*i of row (l & 15). + if (wave == 0) { + const int4 p0 = *reinterpret_cast(&s_part[0][lane * 4]); + const int4 p1 = *reinterpret_cast(&s_part[1][lane * 4]); + const int32_t v0[4] = {p0.x, p0.y, p0.z, p0.w}; + const int32_t v1[4] = {p1.x, p1.y, p1.z, p1.w}; + const int orow = lane & (kTileM - 1); + const int col_mod4 = lane >> 4; + const float xs = x_scale[orow]; + __hip_bfloat16* op = out + static_cast(orow) * n + n0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = col_mod4 + 4 * i; + // col is the LOCAL column within this block's 16-wide tile; the + // output element (orow, n0 + col) must be scaled by the GLOBAL weight + // scale weight_scale[n0 + col]. The iteration-1 draft indexed the + // local col (correct only for tile 0, n0 == 0), which the exact check + // caught as first_mismatch (m=0, n=16) with 97616/98304 mismatches. + const float scaled = + static_cast(v0[i] + v1[i]) * xs * weight_scale[n0 + col]; + op[col] = __float2bfloat16(scaled); + } + } +} + +// --------------------------------------------------------------------------- +// Down_proj iteration 8 (multi-N-tile amortization round): TWO 16x16 N tiles +// per block (N=32/block, grid = N/32 = 192 blocks) with four 64-thread +// wavefronts (256 threads/block). Each wave still computes exactly one 16x16 +// tile with in-block split-K=2 (wave w: tile t = w>>1, k-half h = w&1), so +// the per-wave K loop is byte-for-byte the accepted iteration-2 loop (4 +// k-ascending m16n16k32 steps over its K=128 half, one ds_read_b64 per +// fragment per step) and the int32 accumulation order (k-ascending within +// each half, ascending split sum) is bit-identical to the scalar reference. +// What amortizes over the two tiles: +// 1. A staging: the accepted kernel stages the L2-hot 4 KiB A slice ONCE +// per 16-N-column tile (2 dwordx4/lane/wave). Here the same 16x256 +// stage (272 = 256+16 bank-skewed stride) is shared by BOTH tiles: wave +// w stages only its 16x64 quarter k in [64w, 64w+64) (ONE dwordx4 + +// one ds_write_b128 per lane), and all four waves read the full stage. +// Per 2 tiles the A staging round trip drops from 2 (one per block +// today) to 1 cooperative one -> 3 staging dwordx4/lane per wave vs 4. +// The cross-wave A reads are ordered by ONE top __syncthreads right +// after the staging writes (the accepted kernel's zero-barrier prologue +// relied on each wave reading only its own staged half; that property +// is intentionally traded for the 2x A amortization). +// 2. Combine/epilogue: one __syncthreads covers both tiles' partial +// publish/read pairs (waves 0,2 publish s_part[tile][0], waves 1,3 +// publish s_part[tile][1]; after the barrier odd waves 1,3 each sum +// their tile's two partials in ascending split order and run the fused +// x_scale/weight_scale/bf16 epilogue with the verified gfx928 +// accumulator ownership (row = lane & 15, col_mod4 = lane >> 4), 4 +// bf16 stores per lane per tile) -> one barrier + one x_scale load per +// two tiles instead of per tile. +// Geometry/occupancy: 192 blocks x 256 threads = 768 waves = 1.6 waves/SIMD +// on 120 CUs -- the measured winning parallelism class (beat 0.8 waves/SIMD +// iteration 1: 15.96 us and 3.2 waves/SIMD iteration 3: 11.13 us); LDS +// 4,352 (A) + 2 x 4,352 (B) + 2 x 2,048 (partials) = 17,152 B/block -> 3 +// blocks/CU x 4 waves = 12 waves/CU co-resident, identical to the accepted +// 6 blocks/CU x 2 waves; ~40-56 VGPR x 256 threads x 3 blocks = well under +// the 65,536 VGPR/CU budget. B stays the cold once-read [N,K] n-major pack +// P[n*K+k] (same bytes, same per-row 272-stride stage, same per-wave 2 +// dwordx4/lane staging as iteration 2 -- the tile's 16 rows are just +// n0+16t+row). No workspace use, single Graph-safe dispatch per replay; +// exact shape guard (m==16 && n==6144 && k==256) and the generic scalar +// fallback (which still decodes the same n-major pack for the paired M=2 +// shape) untouched. +// Falsifiable PMC predictions vs the accepted code object (grid 384, 10,752 +// LDS instr, 4,992 vmem_read, 49,152 lds_bank_conflicts, 5.92 us): +// grid -> 192; vmem_read_instructions 4,992 -> ~4,224 (A staging dwordx4 +// halves: 768 waves x 2 -> 768 waves x 1 = -768, epilogue unchanged); +// lds_instructions 10,752 -> ~9,216 (staging ds_write_b128 4 -> 3 per +// wave: -768 waves x 1, K-loop reads and partials unchanged); +// lds_bank_conflicts ~49,152 UNCHANGED (identical per-wave fragment read +// pattern and count: 768 waves x 8 reads x 8-way); lds_wait follows +// lds_instructions; VGPR <= ~56; profiled_duration < 5.92 us. +// Accepted iff the fresh operator median < 10.9307 us with p90 <= 11.0731 +// us (official bests); a flat ~10.9-11.1 class rejects the mechanism and +// pins the wall on the per-wave staging round trip + fixed replay overhead. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDummaSk2N32BlockThreads) void +w8a8_dumma_m16_n32_k256_sk2_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, // packed [N,K] n-major for this shape + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n, + int k) { + constexpr int kTileM = 16; + constexpr int kTileN = 16; + constexpr int kTileK = 32; + constexpr int kK = 256; + constexpr int kSplit = 2; + constexpr int kPerSplit = kK / kSplit; // 128 + constexpr int kLdsStride = kK + 16; // 272 = 256 + 16 bank skew + constexpr int kTilesPerBlock = 2; + const int wave = static_cast(threadIdx.x) >> 6; // 0..3 (four waves) + const int lane = static_cast(threadIdx.x) & 63; + const int tile = wave >> 1; // 0..1: which of the block's two N tiles + const int half = wave & 1; // 0..1: this wave's K=128 half + const int n0 = static_cast(blockIdx.x) * (kTilesPerBlock * kTileN); + + __shared__ __align__(16) int8_t lds_a[kTileM * kLdsStride]; // 4,352 B + __shared__ __align__(16) int8_t lds_b[kTilesPerBlock][kTileN * kLdsStride]; + // // 2 x 4,352 B + __shared__ __align__(16) int32_t s_part[kTilesPerBlock][kSplit][64 * 4]; + // // 2 x 2,048 B + + // Cooperative A staging: wave w writes the 16x64 quarter k in + // [64w, 64w+64) of the shared 16x256 stage (lane -> (row = lane>>2, + // kk16 = (lane&3)<<4), one 16-aligned global_load_dwordx4 + one + // ds_write_b128 per lane). B staging unchanged from iteration 2 except + // the tile row base: this wave's tile t, its K=128 half (2 dwordx4 + + // 2 ds_write_b128 per lane). All 3 loads issue back-to-back (one vmcnt + // wait covers the whole staging round trip). + const int row = lane >> 2; + const int kk16 = (lane & 3) << 4; + const int4 va = *reinterpret_cast( + a + row * k + wave * 64 + kk16); + int4 vb[2]; +#pragma unroll + for (int i = 0; i < 2; ++i) { + vb[i] = *reinterpret_cast( + b + static_cast(n0 + tile * kTileN + row) * k + + half * kPerSplit + i * 64 + kk16); + } + *reinterpret_cast(lds_a + row * kLdsStride + wave * 64 + kk16) = va; +#pragma unroll + for (int i = 0; i < 2; ++i) { + *reinterpret_cast(lds_b[tile] + row * kLdsStride + + half * kPerSplit + i * 64 + kk16) = vb[i]; + } + + // Every wave reads A quarters staged by OTHER waves (wave 0/2 read + // quarters 0+1, wave 1/3 read quarters 2+3), so the K loop must wait for + // all four waves' staging writes: one top barrier (the accepted kernel's + // barrier-free prologue no longer applies by design). + __syncthreads(); + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // 4 k-ascending m16n16k32 steps over this wave's K=128 half: A fragment + // from the shared stage (lds_a + half*128), B fragment from this wave's + // tile stage (lds_b[tile] + half*128). Identical read pattern to the + // accepted iteration-2 loop. + const int k_off = half * kPerSplit; +#pragma unroll + for (int kk = 0; kk < kPerSplit; kk += kTileK) { + du::dumma::du_load_matrix_sync(a_frag, lds_a + k_off + kk, kLdsStride); + du::dumma::du_load_matrix_sync(b_frag, lds_b[tile] + k_off + kk, + kLdsStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Publish this wave's int32 partial (lane-major, 4 contiguous int32 = + // one conflict-free ds_write_b128): s_part[tile][half][lane*4]. Split 0 + // is the k=0..127 half, split 1 the k=128..255 half (ascending order). + *reinterpret_cast(&s_part[tile][half][lane * 4]) = + make_int4(acc_frag.x[0], acc_frag.x[1], acc_frag.x[2], acc_frag.x[3]); + + __syncthreads(); // publish visibility before the per-tile combine + + // Odd waves (1 for tile 0, 3 for tile 1) sum their tile's two partials in + // ascending split order (bit-exact int32, identical to the scalar + // reference) and run the direct register epilogue: row = lane & 15, + // col_mod4 = lane >> 4, output columns tile_n0 + col_mod4 + 4*i, scaled by + // x_scale[row] * weight_scale[tile_n0 + col] (GLOBAL column, iteration-1 + // lesson) -> 4 bf16 stores per lane. + if (half == 1) { + const int4 p0 = *reinterpret_cast(&s_part[tile][0][lane * 4]); + const int4 p1 = *reinterpret_cast(&s_part[tile][1][lane * 4]); + const int32_t v0[4] = {p0.x, p0.y, p0.z, p0.w}; + const int32_t v1[4] = {p1.x, p1.y, p1.z, p1.w}; + const int orow = lane & (kTileM - 1); + const int col_mod4 = lane >> 4; + const float xs = x_scale[orow]; + const int tile_n0 = n0 + tile * kTileN; + __hip_bfloat16* op = out + static_cast(orow) * n + tile_n0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = col_mod4 + 4 * i; + const float scaled = + static_cast(v0[i] + v1[i]) * xs * weight_scale[tile_n0 + col]; + op[col] = __float2bfloat16(scaled); + } + } +} + +// --------------------------------------------------------------------------- +// Device-to-device packing kernels (optional pack_weight op, outside the +// timed region). Iteration 5 introduced the packed weight for the exact +// gate_up (k, n) == (6144, 512) pair so every m16n16k32 B fragment is 8 +// contiguous bytes; iteration 11 re-packs that pair into the o_proj-validated +// fragment-slot layout (w8a8_pack_fragslot_i8_kernel) so each full 64-lane +// load instruction reads 512 contiguous bytes (see the kernel comment); every +// other (K, N) keeps the identity copy of raw_weight [K, N] int8 and +// weight_scale [N, 1] fp32, the generic fallback layout. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kPackBlockThreads) void +w8a8_pack_identity_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { + dst[i] = src[i]; + } +} + +__global__ __launch_bounds__(kPackBlockThreads) void +w8a8_pack_identity_scale_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { + dst[i] = src[i]; + } +} + +// Fragment-slot pack (iteration 11) for the exact gate_up (k, n) == +// (6144, 512) pair: dst[(tile*(k/32) + step)*512 + row*32 + kk] = +// src[kk*n + tile*16 + row], with tile = n_idx/16, row = n_idx%16, +// step = k_idx/32, kk = k_idx%32. Runs once, outside the timed region and +// outside Graph capture, into the same caller-owned packed_weight buffer +// (byte count unchanged, so captured addresses stay valid). Each +// m16n16k32 B-fragment lane still owns 8 CONTIGUOUS bytes at +// b2[(tile*(k/32) + 24*split + s)*512 + (lane&15)*32 + (lane>>4)*8 + {0..7}] +// (one 512-B slot per k32 step, see w8a8_dumma_m16_sk8_partial_kernel), so +// each fragment stays ONE aligned global_load_dwordx2 per lane -- while a +// full 64-lane load instruction now covers 512 contiguous bytes (8 full +// 64-B lines, 100% line utilization) and each block's 24-slot slice is one +// contiguous 12 KiB stream instead of 16 rows x 768 B strided 6,144 B +// apart under the iteration-5 n-major pack. This is the o_proj-validated +// fragment-slot transport that carried ~461 GB/s effective read bandwidth on +// this part. The identity copy remains the fallback for every other (K, N). +__global__ __launch_bounds__(kPackBlockThreads) void +w8a8_pack_fragslot_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count, + int n, + int k) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { + const int n_idx = static_cast(i / k); + const int k_idx = static_cast(i - static_cast(n_idx) * k); + const int tile = n_idx >> 4; // n_idx / kDummaTileN (16) + const int row = n_idx & 15; // n_idx % kDummaTileN + const int step = k_idx >> 5; // k_idx / kDummaTileK (32) + const int kk = k_idx & 31; // k_idx % kDummaTileK + const int64_t dst_idx = + (static_cast(tile) * (k >> 5) + step) * 512 + + static_cast(row) * 32 + kk; + dst[dst_idx] = src[static_cast(k_idx) * n + n_idx]; + } +} + +// N-major pack (down_proj iteration 1) for the exact (k, n) == (256, 6144) +// pair: dst[n*K + k] = src[k*n + n] -- the [N,K] n-major transpose. Runs +// once, outside the timed region and outside Graph capture, into the same +// caller-owned packed_weight buffer (byte count unchanged, so captured +// addresses stay valid). In this layout every matrix_b col_major fragment +// is 8 CONTIGUOUS k bytes at dst[(n0 + lane&15)*K + kk + (lane>>4)*8 + +// {0..7}], which is what the down_proj DUMMA kernel stages into LDS and the +// scalar fallback decodes for (n, k) == (6144, 256). 64(k) x 64(n) tile +// per block via an LDS transpose (grid = (K/64, N/64)), the lineage- +// validated gate_up n-major pack pattern. The identity copy remains the +// fallback for every other (K, N). +__global__ __launch_bounds__(kPackBlockThreads) void +w8a8_pack_nmajor_i8_kernel( + const int8_t* __restrict__ src, // W[K][N] row-major + int8_t* __restrict__ dst, // P[N][K] n-major + int k, + int n) { + constexpr int kPackTile = 64; + constexpr int kPackStride = 80; // 64 + 16 bank skew + __shared__ __align__(16) int8_t lds[kPackTile * kPackStride]; + const int k0 = static_cast(blockIdx.x) * kPackTile; + const int n0 = static_cast(blockIdx.y) * kPackTile; + const int tid = static_cast(threadIdx.x); + + // Load: W row k0+kk, 16 consecutive n columns -> LDS row kk. + const int kk = tid >> 2; // 0..63 + const int nn16 = (tid & 3) * 16; // 0,16,32,48 + *reinterpret_cast(&lds[kk * kPackStride + nn16]) = + *reinterpret_cast(src + (k0 + kk) * n + n0 + nn16); + __syncthreads(); + + // Store: P[n0+nn][k0 + kk16*16 .. +15] is 16 contiguous k bytes. + const int nn = tid & 63; // 0..63 + const int kk16 = tid >> 6; // 0..3 + int4 outv; + int8_t* op = reinterpret_cast(&outv); +#pragma unroll + for (int i = 0; i < 16; ++i) { + op[i] = lds[(kk16 * 16 + i) * kPackStride + nn]; + } + *reinterpret_cast(dst + static_cast(n0 + nn) * k + k0 + + kk16 * 16) = outv; +} + +} // namespace + +// Optional out-of-timed-region weight pack (stable host symbol used by +// csrc/bindings.cpp). Bootstrap: identity device-to-device copy. +// Iteration 5 packed the exact gate_up (k, n) == (6144, 512) pair into the +// n-major transpose packed[n*K+kk] = raw[kk*n+n]; iteration 11 re-packs that +// pair into the o_proj-validated fragment-slot layout +// b2[(tile*(k/32) + step)*512 + row*32 + kk] (one contiguous 512-B slot per +// m16n16k32 step; byte count unchanged); down_proj iteration 1 additionally +// packs the exact (k, n) == (256, 6144) pair into the [N,K] n-major +// transpose P[n*K+kk] = raw[kk*n+n] (w8a8_pack_nmajor_i8_kernel, the +// lineage-validated gate_up n-major pack pattern) so every matrix_b +// col_major fragment is 8 contiguous k bytes; every other (K, N) keeps the +// identity copy. The matching GEMM interpretation is selected by the exact +// (m, n, k) guard in launch_w8a8_gemm and by the scalar fallback's +// (n, k) == (512, 6144) / (n, k) == (6144, 256) decodes. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_count = static_cast(k) * n; + const int weight_grid = static_cast( + (weight_count + kPackBlockThreads - 1) / kPackBlockThreads); + if (k == 6144 && n == 512) { + hipLaunchKernelGGL( + w8a8_pack_fragslot_i8_kernel, + dim3(weight_grid), + dim3(kPackBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count, + n, + k); + } else if (k == 256 && n == 6144) { + // (256 % 64 == 0 && 6144 % 64 == 0) by the exact-pair guard. + hipLaunchKernelGGL( + w8a8_pack_nmajor_i8_kernel, + dim3(static_cast(k / 64), static_cast(n / 64)), + dim3(kPackBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + hipLaunchKernelGGL( + w8a8_pack_identity_kernel, + dim3(weight_grid), + dim3(kPackBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count); + } + + const int64_t scale_count = static_cast(n); + const int scale_grid = static_cast( + (scale_count + kPackBlockThreads - 1) / kPackBlockThreads); + hipLaunchKernelGGL( + w8a8_pack_identity_scale_kernel, + dim3(scale_grid), + dim3(kPackBlockThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + scale_count); +} + +// Timed W8A8 GEMM entry point (stable host symbol used by csrc/bindings.cpp). +// Graph-safe: launches only on the caller-provided current HIP stream, with +// no allocation, compilation, autotuning, packing, host synchronization, +// device synchronization, or default-stream use. Only the caller-provided +// `out` and `workspace` are touched; the split-K specialization uses 6 +// int32 planes (6 * 16 * N * 4 bytes) of the workspace as partial +// accumulators plus a 32-int monotonic arrival-counter tail. Every launch +// overwrites every partial tile, and the counters are zeroed once per +// workspace by the first launch (repair 2: static-pointer-guarded async +// hipMemsetAsync on the caller stream), after which they are monotonic and +// need no reset between Graph replays. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + + // ------------------------------------------------------------------------- + // Exact-shape specialization for glm_tp8_shared_gate_up_proj_m16: + // (M=16, N=512, K=6144). The guard requires all three dimensions so the + // paired M=2 API shape with the same (N, K) and every other (m, n, k) still + // reach the scalar fallback below. For this shape `b` is the iteration-11 + // fragment-slot packed weight b2[(tile*(k/32) + step)*512 + row*32 + kk] + // produced once by launch_pack_w8a8_weight outside the timed region, and + // the scalar fallback decodes the same pack for (n, k) == (512, 6144). The workspace + // contract for this shape guarantees 16 int32 planes (512 KiB); the + // split-K=6 path needs 6 planes (192 KiB) plus a 128-byte counter tail, + // and falls back to the scalar kernel defensively if the caller ever + // passes a smaller buffer. Single dispatch per replay: the fused + // last-arrival combine tail (iteration 3, repair 2) produces the final + // scaled bf16 output inside the partial kernel, and the counter tail is + // zeroed once per workspace (repair 2) without adding a dispatch to the + // timed replay. + // ------------------------------------------------------------------------- + if (m == 16 && n == 512 && k == 6144) { + const int64_t partial_bytes = + static_cast(kDummaSplitK) * kDummaTileM * n * + static_cast(sizeof(int32_t)); + const int64_t counter_bytes = + static_cast(n / kDummaTileN) * sizeof(int32_t); + if (workspace_bytes >= partial_bytes + counter_bytes) { + int32_t* partials = reinterpret_cast(workspace); + int32_t* counters = reinterpret_cast( + reinterpret_cast(partials) + partial_bytes); + // Repair 2: the fused last-arrival tail requires every per-tile + // arrival counter to start at a multiple of kDummaSplitK (0), so the + // modulo-6 fire lands only on the true 6th arrival of each tile per + // replay. The caller workspace is torch.empty (not guaranteed + // zeroed), so zero the 128-byte counter tail once per workspace with + // the validated workspace-pointer-static guard: the first launch with + // this workspace (the pre-capture eager warmup in the normal flow) + // issues an async hipMemsetAsync on the caller stream, and the static + // prevents any later launch from repeating it, so the timed replay + // keeps exactly one dispatch. Once initialized the counters are + // monotonic (6 increments per tile per replay) and never need a reset + // between Graph replays. + static const void* s_zeroed_counters_ws = nullptr; + if (s_zeroed_counters_ws != workspace) { + hipMemsetAsync(counters, 0, counter_bytes, stream); + s_zeroed_counters_ws = workspace; + } + const int grid = (n / kDummaTileN) * kDummaSplitK; + hipLaunchKernelGGL( + w8a8_dumma_m16_sk8_partial_kernel, + dim3(grid), + dim3(kDummaBlockThreads), + 0, + stream, + a, + b, + partials, + x_scale, + weight_scale, + reinterpret_cast<__hip_bfloat16*>(out), + counters, + n, + k); + return; + } + } + + // ------------------------------------------------------------------------- + // Exact-shape specialization for glm_tp8_shared_down_proj_m16: + // (M=16, N=6144, K=256) -- down_proj iteration 8 (multi-N-tile + // amortization round): TWO 16x16 N tiles per block (N=32/block, grid = + // N/32 = 192 blocks), 256 threads = four wavefronts; wave w computes one + // tile t = w>>1 with in-block split-K=2 (k-half h = w&1), A staged once + // cooperatively for both tiles (one 64-k quarter per wave + one top + // barrier), B staged per tile from the [N,K] n-major pack, one combine + // barrier covers both tiles' partial publish/read pairs (odd waves run + // the fused scale/bf16 epilogue). 768 waves total = 1.6 waves/SIMD, 12 + // waves/CU co-resident -- the accepted parallelism class. The guard + // requires all three dimensions so the paired M=2 API shape with the + // same (N, K) and every other (m, n, k) still reach the scalar fallback + // below. For this shape `b` is the [N,K] n-major packed weight + // P[n*K+kk] = raw[kk*n+n] produced once by launch_pack_w8a8_weight + // outside the timed region; the scalar fallback decodes the same pack for + // (n, k) == (6144, 256). No workspace use (in-block LDS combine, no + // split-K planes), single Graph-safe dispatch per replay. + // ------------------------------------------------------------------------- + if (m == 16 && n == 6144 && k == 256) { + hipLaunchKernelGGL( + w8a8_dumma_m16_n32_k256_sk2_kernel, + dim3(static_cast(n / (2 * kDummaTileN))), + dim3(kDummaSk2N32BlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast<__hip_bfloat16*>(out), + n, + k); + return; + } + + const int64_t total = static_cast(m) * n; + const int grid = static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + dim3(grid), + dim3(kScalarBlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast<__hip_bfloat16*>(out), + m, + n, + k); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_gate_up_proj.hip new file mode 100644 index 00000000..3455553e --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_gate_up_proj.hip @@ -0,0 +1,982 @@ +// @@variant shape=glm_tp8_shared_gate_up_proj_m16 commit=fa6ea50630e35036b22141edc48f72975fca041d added=2026-08-31 +// median_us=10.82 p90_us=10.87 +// source=glm5-2-dsh-tp8-m16-test1-cb4d262c +// W8A8 INT8 GEMM for Hygon K500SM_AI / gfx928 (worker_3). +// +// Exact-shape split-K DUMMA specialization for +// glm_tp8_shared_gate_up_proj_m16 (M=16, N=512, K=6144) established before +// the generic scalar fallback: +// +// glm_tp8_shared_gate_up_proj_m16 : (M=16, N=512, K=6144) +// glm_tp8_shared_down_proj_m16 : (M=16, N=6144, K=256) +// +// Iteration 1 measured the minimal gfx928 INT8 m16n16k32 DUMMA tile (one +// 64-thread wavefront per 16x16 output tile, grid = N/16 = 32 blocks) at +// 173.6 us median: PMC showed grbm_count == grbm_gui_active with one +// wavefront per CU and 88 of 120 CUs idle, i.e. a latency-bound launch +// geometry where the 192-step K loop exposes its full load->vmcnt->mmac +// chain with no co-resident wave to hide it. +// +// Iteration 2 (architecture round) split K=6144 into 8 uniform 768-element +// slices (24 m16n16k32 steps each) so grid = 32 N-tiles x 8 splits = 256 +// independent one-wave zero-barrier blocks = 2.13 blocks/CU, reaching the +// two-blocks-per-CU latency-hiding target with all 120 CUs covered. Each +// block kept the minimal zero-barrier DUMMA K loop over its own slice and +// stored its int32 partial plane to the caller workspace (8 x 16 x 512 x 4 = +// 256 KiB, within the 16-plane contract budget); a separate 128-thread +// combine kernel (second dispatch) summed the 8 planes in ascending split +// order (bit-exact int32 accumulation order, identical to the unsplit +// kernel) with the fused scale/bf16 epilogue: 76.9 us median (1.29x over the +// 99.3 us Triton Graph baseline). +// +// Iteration 3 (architecture round, this source): the fresh PMC showed the +// two-launch wall is dominated by the combine dispatch, not the GEMM -- +// w8a8_dumma_m16_sk8_combine_kernel (grid = 1 block, 128 threads on one CU, +// only 256 KiB of L2-hot planes) profiles at 47.2 us vs 39.4 us for all 256 +// partial blocks (86.6 us profiled aggregate vs 76.9 us unprofiled +// operator): a one-CU serial dependent-load-chain cliff worth roughly half +// the wall. The decision: keep the accepted split-K=8 one-wave zero-barrier +// grid and remove the one-CU combine dispatch by fusing the combine into the +// partial kernel as a validated qkv-style last-arrival tail. Every block +// ends with __threadfence() + atomicAdd on a monotonic per-N-tile counter in +// the workspace tail; the 8th arriver of each tile (32 of 256 blocks) sums +// the tile's 8 planes in ascending split order (bit-exact int32, identical +// to the unsplit kernel) and applies the fused scale/bf16 epilogue. +// Counters are monotonic ((arrived & 7) == 7 fires exactly once per tile per +// replay), so no reset is needed between Graph replays; one dispatch per +// replay, still Graph-safe. +// +// Repair 1 (in-round, iteration 3) suspected uninitialized counters and added +// a hipStreamIsCapturing-gated one-time hipMemsetAsync; the exact check still +// failed with the identical 7168/8192 mismatch signature, so that diagnosis +// was wrong. +// +// Repair 2 (in-round, iteration 3, this source): the real defect is that the +// fused tail let ALL 64 lanes of every block execute the arrival atomicAdd +// (64 increments per block = 512 per tile per replay), so (arrived & 7) == 7 +// fired ~64 times per tile -- the first fires inside the first-arriving +// block, before its 7 siblings had stored their planes, and the premature +// combiner summed unwritten zero planes: exact check failed with +// first_mismatch (0,0) actual=0.0 and 28/32 tiles (7168/8192 elements) wrong, +// the identical signature already documented by the validated hy3 TP8 gate_up +// lineage. Fix (mirroring that lineage): exactly ONE arrival per block -- +// lane 0 fences + atomicAdds + fences and publishes the result through an LDS +// flag with barriers; only the last arriver of each tile runs the combine. +// The one-time counter zero is now a workspace-pointer-static-guarded +// hipMemsetAsync (the validated qkv/gate_up pattern; no capture-status +// query), issued on the caller stream by the first launch with the workspace +// (the pre-capture eager warmup), so the timed replay keeps one dispatch and +// the monotonic counters need no reset between Graph replays. +// +// Iteration 4 (architecture round, this source): A-only staging. The fresh +// iteration-4 PMC/ISA of the accepted iteration-3 kernel shows the K-loop +// body is the documented decode poison: 16 global_load_ubyte per lane per +// step (8 A + 8 B bytes), a 16-deep s_waitcnt vmcnt cascade, ~20 byte-OR +// reassembly VALU, and ~30 address-carry VALU -- ~64 instructions per step +// with the whole A side of the load chain inside the mmac critical path. +// The DTK du_mma.hpp implementation confirms the m16n16k32 row_major +// matrix_a fragment is 8 CONSECUTIVE bytes at (row = lane&15, k = +// (lane>>4)*8 + {0..7}) -- vectorizable -- while the matrix_b fragment's 8 +// bytes are strided by the leading dimension (not vectorizable without +// packing, out of scope). Decision: keep the accepted split-K=8 geometry +// (grid = 256 blocks = 2.13 blocks/CU, all 120 CUs covered, fused +// last-arrival combine tail) and stage each block's 16x768 A slice into LDS +// once with vectorized 16-B loads (12 global_load_dwordx4 + 12 +// ds_write_b128 per lane, one __syncthreads), then read every k-step's A +// fragment from LDS as one 8-B ds_read_b64 (x[0..7] = A[row][k0 + +// (lane>>4)*8 + {0..7}], bit-identical to du_load_matrix_sync). B keeps +// its direct global path. A bytes reused: the 12 KiB slice is loaded from +// global exactly once per block and each byte is read 24x from LDS (once +// per k-step); cross-block A/B reuse via L2 is unchanged. No workspace, +// plane, counter, or combine-tail change; exact int32 accumulation order +// (k-ascending within a split, ascending split sum) is preserved because +// the mmac consumes identical fragment bytes. +// +// Iteration 5 (packed-weight round, this source): the fresh iteration-5 +// PMC of the accepted iteration-4 code object (digest 7d5673b9..., vmem_read +// 53152, valu_instructions 255872) shows the K-loop body is still half +// decode poison: per lane per k-step B contributes 8 global_load_ubyte, ~15 +// of the 16-deep s_waitcnt vmcnt cascade, ~20 byte-OR reassembly VALU and +// ~30 v_add/v_addc address-carry VALU -- the whole B side of the load chain +// inside the mmac critical path, because the m16n16k32 row_major matrix_b +// fragment's 8 bytes are strided by the leading dimension N and cannot be +// vectorized in the raw [K][N] layout (the iteration-4 round staged A only, +// per its mandate). Decision (mandated option: packed layout): pack B once, +// outside the timed region and outside Graph capture, into the n-major +// transpose packed[n*K + kk] = raw[kk*n + n] for the exact (k, n) == +// (6144, 512) pair (byte count unchanged, so captured addresses stay +// valid). In the packed layout the B fragment is 8 CONTIGUOUS bytes at +// packed[(n0 + lane&15)*k + k0 + (lane>>4)*8 + {0..7}] -- the exact byte +// pattern du_load_matrix_sync(b_frag, b + n0*k + k0, k) reads -- so each +// k-step's B fragment becomes ONE aligned global_load_dwordx2 per lane with +// no byte-OR reassembly and no per-byte waitcnt cascade; the A-side LDS +// stage (iteration 4) is unchanged. This is the exact load-path pattern +// validated on the K=6144 m16n16k32 gate_up lineage (MiniMax TP8 gate_up, +// accepted iterations 5-16: n-major pack + one 8-byte dwordx2 load per +// fragment). The generic scalar fallback decodes the same n-major pack +// whenever (n, k) == (512, 6144), so the paired M=2 API shape with the same +// (N, K) stays correct through the packed buffer; every other (K, N) keeps +// the identity pack and the raw row-major decode, and the exact int32 +// accumulation order (k-ascending within a split, ascending split sum) is +// unchanged because the mmac consumes identical fragment bytes. +// +// Iteration 7 (pipeline round): double-buffered B transport. The fresh +// iteration-5 ISA (digest d2ffb4f6...) shows the single-buffered K-loop +// body is exactly global_load_dwordx2 (B) -> ds_read_b64 (A) -> +// s_waitcnt vmcnt(0) lgkmcnt(0) -> v_mmac per step, with no unroll, so the +// next B load cannot issue until the branch returns to the loop head: every +// per-block B load (a zero-reuse HBM stream -- each (tile, split) block +// reads a disjoint 12 KiB B slice, l2_misses 55733 ~ the compulsory 3.57 +// MiB traffic) exposes its full ~590-cycle HBM round trip on the critical +// path. Mandated this round: compare single buffering with double +// buffering (K=6144 >= 1024, L2 hit 55.11% < 70%, doubled LDS budget +// 25,096 B < 48 KiB all hold). Chosen double buffering: depth-1 register +// prefetch of the packed B fragment (the o_proj-validated pattern -- next +// step's loads issued before the current mmac, rotated in after, compiler +// wait lands at the next fill): step i+1's aligned global_load_dwordx2 is +// issued BEFORE step i's ds_read_b64 + v_mmac, so its HBM latency overlaps +// the previous step's compute instead of sitting on the critical path; the +// final step is peeled so the prefetch address never leaves the block's K +// slice (the packed buffer ends exactly at 512*6144 B; an unconditional +// k0+32 read on the last step would be out of bounds). No second LDS +// buffer: the LDS-fed B path regressed twice (MiniMax TP8 gate_up 25.9 us; +// our iteration 6 at 19.624 us vs 16.778 us), so packed-global dwordx2 B +// remains the transport. LDS per block stays 12,544 + 4 B (doubled budget +// 25,096 B < 48 KiB; 2.13 blocks/CU unchanged), barriers per K step stay 0 +// (only the single pre-loop staging barrier and the two tail barriers, +// unchanged from iteration 5), and the mmac consumes bit-identical fragment +// bytes in k-ascending order, so the output stays bit-identical to the CPU +// reference (0 mismatches, exact check). +// +// Iteration 9 (MLP round, this source): grouped depth-8 register prefetch. +// The iteration-8 occupancy probe (split-K 15 -> 480 one-wave blocks = exactly +// 4 blocks/CU, integer-balanced) REGRESSED to 21.54 us median, so co-resident +// wave count is not the lever: the kernel is DRAM-latency-bound on per-wave +// in-flight B bytes, not wave-starved (l2_misses ~55.7k = the compulsory 3.57 +// MiB A+B traffic; each block still serially waits ~24 ~440-540-cycle HBM +// round trips). This round raises memory-level parallelism INSIDE the wave +// while keeping the exact 256-block split-K=8 grid and its contiguous 12 KiB +// per-block B slices (no DRAM scatter): the 24-step K loop is regrouped into +// 3 x 8-step groups; before each group's 8 mmacs consume the previously +// loaded B fragments, the NEXT group's 8 packed-B dwordx2 loads are issued +// back-to-back (8 independent loads in flight per lane = 4 KiB per wave), so +// ONE vmcnt fill covers 8 steps and each group's loads have the whole +// preceding 8-mmac burst (~8x the iteration-7 single-step window) to expire +// off the critical path. Grid, split-K=8, A LDS staging, packed-B transport, +// zero in-loop barriers, workspace, and the fused last-arrival tail are +// untouched; the final group is peeled (no prefetch) so the prefetch address +// never leaves the block's [k_start, k_start + k_per_split) slice, and the +// mmacs consume bit-identical fragment bytes in k-ascending order, so the +// output stays bit-identical to the CPU reference (0 mismatches, exact +// check). +// +// Iteration 11 (layout consolidation, this source): fragment-slot B pack. +// The iteration-10 direct-A probe REGRESSED to 15.06 us median: its ISA shows +// the compiler could not hoist the 48 global A+B dwordx2 loads (register +// pressure) and fell back to a depth-2 rotation with `s_waitcnt vmcnt(8)` +// before EVERY mmac, re-exposing a full global round trip per step -- so the +// A-side LDS staging stays (the 12 KiB stage is served by ~30-40-cycle +// ds_read2_b64 with 2-step lookahead, not by global loads). The accepted +// iteration-9 kernel's remaining B-side inefficiency is LAYOUT, not MLP: +// the 24 packed-B dwordx2 loads are fully hoisted (all in flight from the +// prologue, per-step vmcnt waits mostly satisfied), but the iteration-5 +// n-major pack packed[n*K+kk] places a block's 16 rows x 768 B B-slice in 16 +// chunks 6,144 B apart, so each 64-lane load instruction touches 16 strided +// 64-B lines (8 B used each) and the 256 concurrent blocks stream 4,096 +// strided DRAM row streams -- the kernel delivers only ~250-270 GB/s +// (counter-derived) where the o_proj-validated fragment-slot pack +// (packed[n_tile][k_step][lane][8]) carried ~461 GB/s on this part. This +// round re-packs the exact (k, n) == (6144, 512) pair once (outside the +// timed region, same buffer bytes) into b2[(tile*(k/32) + step)*512 + row*32 +// + kk]: one 512-B slot per m16n16k32 step, so each lane still reads its 8 +// CONSECUTIVE fragment bytes as ONE aligned global_load_dwordx2 (the exact +// du_load_matrix_sync fragment bytes, k-ascending order, bit-exact int32), +// a full 64-lane load instruction covers 512 CONTIGUOUS bytes (8 full 64-B +// lines, 100% line utilization), and each block's 24-slot slice is one +// contiguous 12 KiB stream -- the 256 blocks now stream the 3 MiB pack in +// 256 sequential DRAM row runs. Grid (256), split-K=8, A LDS staging, +// zero in-loop barriers, the b8 grouped rotation, the workspace, and the +// fused last-arrival tail are untouched; the scalar fallback decodes the +// same fragment-slot pack for (n, k) == (512, 6144) so the paired M=2 API +// shape stays correct through the same packed buffer; l2_misses stay the +// compulsory ~55.6k (same bytes, no over-fetch) while the B second-half +// line L2 hits disappear (l2_hits drop), and the output stays bit-identical +// to the CPU reference (0 mismatches, exact check). +// +// Iteration 20 (stream-shape probe, this source): split-K 8 -> 6. The +// measured split-K/stream lattice for this shape is {blocks x per-wave KiB +// slice -> median}: 64 x 48 -> 13.64 us (iter 15), 128 x 24 -> 12.124 us +// (iter 13), 256 x 12 -> 11.736 us (accepted, iter 11), 480 x 4 -> 21.54 us +// (iter 8). Two readings survive every probe so far: (a) the counter-derived +// delivery rate rises with the number of concurrent per-wave streams +// (260 -> 287 -> 302 GB/s at 64 -> 128 -> 256 streams), with diminishing +// elasticity; (b) there is a steep DRAM-efficiency cliff below ~12-KiB +// per-stream length (4-KiB slices collapse to ~165 GB/s even at 480 streams +// and 4 blocks/CU). Those two readings leave exactly ONE untested lattice +// point that keeps every stream >= 12 KiB with an integer uniform split: +// 192 blocks x 16-KiB slices = split-K 6 (6144/6 = 1024 = 32 k32 steps per +// slice; 32 N-tiles x 6 = 192 blocks = 1.6 blocks/CU, fully resident at +// every plausible occupancy >= 2/CU, so no queued tail at any occupancy). +// 6 is deliberately OUTSIDE the round's trusted set {2,3,4,7,8,10,11,15}: +// the set's untested members are either length-cliffed (10/11 would need +// 8.7-9.6-KiB slices), underfilling (3 -> 96 blocks < 120 CUs), or require +// ragged non-constexpr splits (7/10/11: 192 is not divisible by 7/10/11), +// while 6 divides 192 evenly and keeps the exact 8-step grouped structure +// (32 = 4 x 8). The decode skill's non-power-of-two sweep family includes +// 6 explicitly, and the validated wqkv_a winner used split-K 6. Mechanism: +// the 3.55-MiB compulsory traffic, the 3-MiB A L2 re-reads (192 x 16 KiB == +// 256 x 12 KiB), the 512-B per-instruction B request granularity, the +// fragment-slot pack, the grouped depth-8 rotation, the zero-barrier +// single-wave loop, the plane store, and the fused last-arrival tail are +// all unchanged -- only the per-wave B stream grows 12 -> 16 KiB and the +// concurrent stream count falls 256 -> 192. Falsifiable: median stays +// 11.736 +/- 0.3 if ~302 GB/s is a count- and length-independent DRAM-mix +// floor for this 3.55-MiB traffic (the session's second plateau-eligible +// flat candidate after iteration 19), drops toward ~11.5-11.65 if the +// longer 16-KiB rows lift delivery to ~305-310 GB/s, or regresses to +// ~12.0-12.1 if count-dominance dominates (falsifying the +// length-compensation model). Outputs stay bit-identical to the CPU +// reference: the pack is split-agnostic (pure (tile, k-step) slots), each +// split's 32 steps are consumed in k-ascending order, and the fused tail +// sums the 6 planes in ascending split order (int32 addition commutative, +// every partial sum within int32 range as before), firing on +// (arrived % 6) == 5 -- the monotonic modulo fire lands exactly once per +// tile per replay on the true 6th arrival. +// +// The scalar kernel stays the fallback for every unmatched (m, n, k), +// including the paired M=2 API shape with the same (N, K) and the +// down_proj shape. +// +// Math contract: +// out[m, n] = bf16( int32_dot(a[m, :], b[:, n]) * x_scale[m] * weight_scale[n] ) +// The int8 dot product is accumulated exactly in int32 (the largest assigned +// K=6144 keeps every partial sum within int32 range: +// 6144 * 127 * 127 = 99,058,176 < 2^31), so the result is bit-identical to +// the CPU int64 reference regardless of accumulation order. + +// Include the installed headers in this known-good order: the DTK's +// du_mma.h is not self-contained when included before the HIP runtime +// headers. +#include +#include +#include + +#include + +namespace { + +constexpr int kScalarBlockThreads = 128; +constexpr int kPackBlockThreads = 256; + +// DUMMA tile constants for the exact-shape specialization. gfx928 INT8 +// tensor cores execute m16n16k32 (signed char x signed char -> int32) with +// wavefront=64; each 16x16 output tile is owned by exactly one wavefront. +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; +constexpr int kDummaBlockThreads = 64; // one 64-thread wavefront per block + +// Split-K geometry for the exact-shape specialization (iteration 2; +// iteration 20: 8 -> 6). K=6144 / 32 = 192 m16n16k32 steps, split into 6 +// uniform 32-step slices; grid = (N/16) * kDummaSplitK = 32 * 6 = 192 blocks +// = 1.6 blocks/CU on the 120-CU device (fully resident at every occupancy +// >= 2/CU, no queued tail), each slice a contiguous 16-KiB fragment-slot +// run. 6 divides 192 evenly, so every K boundary stays aligned to the +// DUMMA k=32 unit with zero wasted work. Probe rationale (iteration 20): +// the measured stream lattice (64 x 48 KiB -> 13.64 us, 128 x 24 KiB -> +// 12.124 us, 256 x 12 KiB -> 11.736 us, 480 x 4 KiB -> 21.54 us) reads +// monotone in concurrent-stream count (260 -> 287 -> 302 GB/s) with a steep +// length cliff below ~12 KiB; 192 x 16 KiB is the only untested lattice +// point that keeps every stream >= 12 KiB with an integer uniform split. +constexpr int kDummaSplitK = 6; +constexpr int kDummaCombineThreads = 128; + +// Iteration-4 A-only staging constants (iteration 20: k_per_split 768 -> +// 1024 with split-K 6). The exact-shape guard pins k=6144, so each split +// slice is k_per_split = 6144/6 = 1024 int8 columns. The A stage is the +// 16x1024 slice padded to a 1040-column row stride: 1040 is 16-aligned +// (staging stores compile to ds_write_b128) and 8-aligned (the per-k-step +// A-fragment read compiles to one ds_read_b64); the +16 skew bounds LDS +// bank conflicts to 2-way. 16 * 1040 = 16,640 B LDS per block; with ~2-3 +// blocks/CU that is ~33-50 KiB of the 64 KiB/CU LDS, so the accepted +// occupancy class is unchanged (2/CU = 33,280 B < 64 KiB still holds). +constexpr int kPerSplitK = kDummaTileK * 32; // 1024 = 6144 / 6 +constexpr int kAStageStride = kPerSplitK + 16; // 1040 + +// --------------------------------------------------------------------------- +// Generic scalar W8A8 GEMM fallback. +// +// One thread computes exactly one output element out[row, col]. Threads are +// mapped linear = row * n + col, so adjacent lanes sit on adjacent addresses +// in the fastest-changing N dimension: B[kk, col] loads coalesce across a +// wavefront and A[row, kk] broadcasts within a row. blockDim is a multiple +// of the gfx928 wavefront size 64. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear % n); + + const int8_t* __restrict__ a_row = a + static_cast(row) * k; + // The exact gate_up (n, k) == (512, 6144) uses the iteration-11 + // fragment-slot packed weight b2[(tile*(k/32) + step)*512 + row*32 + kk] + // (produced once by launch_pack_w8a8_weight outside the timed region); all + // other shapes keep the identity [K, N] row-major layout. Decode both so + // the generic fallback -- including the paired M=2 API shape with the same + // (N, K) -- stays correct through the same packed pointer. + const bool fragslot_packed = (n == 512 && k == 6144); + const int8_t* __restrict__ b_col = b + col; + int64_t b_stride = static_cast(n); // identity layout per-kk stride + if (fragslot_packed) { + // b2[(col>>4)*(k>>5)*512 + (col&15)*32 + (kk>>5)*512 + (kk&31)] + b_col = b + static_cast(col >> 4) * ((k >> 5) * 512) + + static_cast(col & 15) * 32; + b_stride = 0; // addressed via the per-kk slot formula below + } + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + const int64_t b_off = fragslot_packed + ? (static_cast(kk >> 5) * 512 + (kk & 31)) + : kk * b_stride; + acc += static_cast(a_row[kk]) * + static_cast(b_col[b_off]); + } + + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Exact-shape split-K DUMMA specialization for glm_tp8_shared_gate_up_proj_m16: +// (M=16, N=512, K=6144). Guarded by an exact (m, n, k) match in +// launch_w8a8_gemm so paired M=2 API shapes and every other shape keep the +// scalar fallback. +// +// Architecture round (iteration 2): the iteration-1 grid of 32 one-wave +// blocks left 88 of 120 CUs idle and each 192-step K loop ran with no +// co-resident wave (PMC: grbm_count == grbm_gui_active, 173.6 us). Splitting +// K into 8 uniform 768-element slices turns the launch into 256 independent +// zero-barrier one-wave blocks (2.13 blocks/CU), so while one block's +// fragment loads sit on vmcnt, the co-resident block issues its mmac chain. +// Each block keeps the minimal gfx928 INT8 m16n16k32 DUMMA tile with explicit +// int32 accumulation over its own k-slice (split-major block order keeps the +// first 32 blocks on one 384 KiB B slice: L2-hot), then stores the int32 +// partial plane to the workspace using the verified gfx928 accumulator lane +// ownership (row = lane & 15, col_mod4 = lane >> 4, acc_frag.x[i] -> columns +// col_mod4 + 4*i). +// +// Iteration 3 fused combine (this source): the standalone 128-thread combine +// dispatch profiled at 47.2 us (grid = 1 block on one CU) -- the largest +// single component of the 76.9 us wall. The combine is now a last-arrival +// tail inside this kernel: after the plane store each block publishes one +// monotonic arrival on counters[tile] (tile = the block's N-tile); the block +// whose arrival makes (arrived & 7) == 7 is the last of the tile's 8 splits, +// so it sums the 8 planes in ascending split order (bit-exact int32 +// accumulation order, identical to the unsplit kernel) and applies the fused +// scale/bf16 epilogue. Repair 2: exactly ONE arrival per block (lane 0 fence +// + atomicAdd + LDS flag + barriers; the round-3 draft's all-64-lane atomicAdd +// fired the combine prematurely, before sibling planes were stored -- see the +// tail comment). Only 32 of 256 blocks run the combine; counters live in the +// workspace tail, are zeroed once per workspace by launch_w8a8_gemm, and +// never need a reset (monotonic across Graph replays). Single Graph-safe +// dispatch per replay. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDummaBlockThreads) void +w8a8_dumma_m16_sk8_partial_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int32_t* __restrict__ partials, // [kDummaSplitK][M][N] int32 planes + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + __hip_bfloat16* __restrict__ out, // [M][N] + int32_t* __restrict__ counters, // [N/16] monotonic arrival counters + int n, + int k) { + const int lane = static_cast(threadIdx.x); + const int n_tiles = n / kDummaTileN; + const int split = static_cast(blockIdx.x) / n_tiles; + const int tile = static_cast(blockIdx.x) % n_tiles; + const int n0 = tile * kDummaTileN; + + const int k_per_split = k / kDummaSplitK; // 1024, a multiple of 32 + const int k_start = split * k_per_split; + + // Iteration 4 (A-only staging): stage this block's 16 x k_per_split A + // slice into LDS once with vectorized 16-B loads, then serve every + // k-step's A fragment from LDS as one 8-B read. The DTK du_mma.hpp + // row_major matrix_a loader maps lane l to 8 CONSECUTIVE bytes + // A[lane&15][k0 + ((lane>>4)<<3) + {0..7}], so the staged fragment bytes + // are bit-identical to du_load_matrix_sync while the accepted kernel's 8 + // global ubyte loads + waitcnt cascade + byte-OR reassembly for A leave + // the per-step critical path (see the header comment for the iteration-4 + // ISA evidence). B (iteration 5) is the n-major packed weight + // (packed[n*K+kk] = raw[kk*n+n] for (k, n) == (6144, 512), produced once + // by launch_pack_w8a8_weight outside the timed region), whose fragment is + // 8 CONTIGUOUS bytes at packed[(n0 + lane&15)*k + k0 + (lane>>4)*8 + + // {0..7}], loaded below as one aligned global_load_dwordx2 per lane + // (bit-identical to du_load_matrix_sync on the packed layout). + __shared__ int8_t s_a[kDummaTileM * kAStageStride]; + { + // 16 x 1024 B = 16,384 B = 1024 x 16-B groups; 64 lanes x 16 = 1024. + constexpr int kGroupsPerRow = kPerSplitK / 16; // 64 + constexpr int kStageVecs = + (kDummaTileM * kGroupsPerRow) / kDummaBlockThreads; // 16 + static_assert(kDummaTileM * kGroupsPerRow == + kDummaBlockThreads * kStageVecs, + "A slice must divide evenly across the staging lanes"); +#pragma unroll + for (int t = 0; t < kStageVecs; ++t) { + const int g = lane + kDummaBlockThreads * t; + const int row = g / kGroupsPerRow; + const int c16 = (g % kGroupsPerRow) * 16; + // 16-aligned: a base, k_start (multiple of 1024), row*k (k=6144) and + // c16 are all multiples of 16, so this is one global_load_dwordx4. + const int4 v = + *reinterpret_cast(a + k_start + row * k + c16); + // 16-aligned in LDS: kAStageStride = 1040 is a multiple of 16, so this + // is one ds_write_b128. + *reinterpret_cast(&s_a[row * kAStageStride + c16]) = v; + } + } + __syncthreads(); // every lane reads rows staged by other lanes below + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + const int a_row = lane & (kDummaTileM - 1); + const int a_kq = (lane >> 4) << 3; // 8 consecutive k bytes per lane + // Iteration 5 (packed weight): the same lane bits select the B fragment + // of the n-major packed weight -- n = lane & 15, k offset (lane>>4)*8 -- + // so b_row and b_kq coincide with a_row / a_kq. + const int b_row = lane & (kDummaTileM - 1); + const int b_kq = (lane >> 4) << 3; + // Iteration 11 (layout consolidation, this source): the per-step packed-B + // address is loop-invariant except the step index, so precompute its base + // here. The packed weight is now the o_proj-validated fragment-slot + // layout b2[(tile*(k/32) + step)*512 + row*32 + kk] (one 512-B slot per + // m16n16k32 step): lane l still owns the 8 CONSECUTIVE bytes + // B[n0 + lane&15][k0 + (lane>>4)*8 + {0..7}] -- the exact + // du_load_matrix_sync fragment -- but a full 64-lane load instruction now + // covers 512 CONTIGUOUS bytes (8 full 64-B lines, 100% line utilization) + // instead of 16 rows x 8 B strided 6,144 B apart under the iteration-5 + // n-major pack. Each block's 24-slot slice is one contiguous 12 KiB + // stream (the contiguity the iteration-9 hypothesis assumed), so the 256 + // blocks stream the 3 MiB pack in 256 sequential DRAM row runs instead of + // 4,096 strided row streams -- the same fragment-slot transport that + // carried the validated o_proj kernel at ~461 GB/s effective read + // bandwidth on this part (vs ~250-270 GB/s counter-derived for this + // kernel's strided n-major pack). Current and prefetch addresses stay + // 8-byte aligned (b_base and kBFragStep are multiples of 8). + constexpr int kBFragStep = kDummaTileN * kDummaTileK; // 512 B per k32 slot + const int64_t b_base = + (static_cast(tile) * (k >> 5) + + split * (k_per_split / kDummaTileK)) * + kBFragStep + + static_cast(b_row) * kDummaTileK + b_kq; + + // Iteration 7 (pipeline round): double-buffered B transport. The + // iteration-5 K-loop body (one global_load_dwordx2 (B) + one ds_read_b64 + // (A) + s_waitcnt vmcnt(0) lgkmcnt(0) + one v_mmac per step, no unroll) + // exposes the full ~590-cycle HBM round trip of every B load on the + // critical path because the next load cannot issue until the branch + // returns to the loop head, and B is a zero-reuse HBM stream (disjoint + // 12 KiB slice per (tile, split) block). This round is the mandated + // single-vs-double buffering comparison. Double buffering here = one + // register-resident B fragment in flight (depth-1 software prefetch, the + // o_proj-validated pattern): step i+1's packed dwordx2 is issued BEFORE + // step i's ds_read + mmac and rotated in after, so its HBM latency + // overlaps the previous step's compute and the compiler places the vmcnt + // wait at the next fill point. LDS stays 12,544 + 4 B per block (no + // second LDS buffer: the LDS-fed B path regressed twice -- MiniMax TP8 + // gate_up 25.9 us and our iteration 6 at 19.624 us -- while packed-global + // dwordx2 B is the validated transport), so the doubled LDS budget + // (25,096 B) stays below 48 KiB and the 2.13 blocks/CU geometry is + // unchanged. Barriers per K step: 0 -- the only __syncthreads are the + // single pre-loop A-staging barrier and the two tail barriers, unchanged + // from iteration 5. The final step is peeled so the prefetch address + // never leaves the block's [k_start, k_start + k_per_split) slice (the + // packed buffer ends exactly at 512*6144 B; an unconditional k0+32 read + // on the last step would read past its end). The mmac consumes + // bit-identical fragment bytes in k-ascending order, so the output stays + // bit-identical to the CPU reference (0 mismatches, exact check). + // Iteration 9 (MLP round, this source): grouped depth-8 register prefetch. + // The iteration-8 occupancy probe (split-K 15 -> 480 one-wave blocks = + // exactly 4 blocks/CU, integer-balanced, no tail wave) REGRESSED to 21.54 + // us median, so co-resident wave count is NOT the lever: this kernel is + // DRAM-latency-bound on per-wave in-flight B bytes, not wave-starved + // (l2_misses ~55.7k = the compulsory 3.57 MiB A+B traffic; every block + // still serially waits ~24 ~440-540-cycle HBM round trips). This round + // raises memory-level parallelism INSIDE the wave while keeping the exact + // 256-block split-K=8 grid and its contiguous 12 KiB per-block B slices + // (no DRAM scatter): the 24-step K loop is regrouped into 3 x 8-step + // groups; before each group's 8 mmacs consume the previously loaded B + // fragments, the NEXT group's 8 packed-B dwordx2 loads are issued + // back-to-back (8 independent loads in flight per lane = 4 KiB per wave), + // so ONE vmcnt fill covers 8 steps and each group's loads have the whole + // preceding 8-mmac burst (~8x the iteration-7 single-step window) to expire + // off the critical path. Grid, split-K=8, A LDS staging, packed-B + // transport, zero in-loop barriers, workspace, and the fused last-arrival + // tail are untouched; the final group is peeled (no prefetch) so the + // prefetch address never leaves the block's [k_start, k_start + + // k_per_split) slice (the packed buffer ends exactly at 512*6144 B), and + // the mmacs consume bit-identical fragment bytes in k-ascending order, so + // the output stays bit-identical to the CPU reference (0 mismatches, + // exact check). + constexpr int kUnroll = 8; + constexpr int kGroupK = kUnroll * kDummaTileK; // 256 = 8 k32 steps + static_assert(kPerSplitK % kGroupK == 0, + "32 steps must split evenly into 8-step groups"); + uint64_t b8[kUnroll]; + // b_first is the base of this block's 32-slot (16 KiB) contiguous + // fragment-slot slice; each slot is kBFragStep = 512 B. Declared at + // function scope (repair 2: the draft scoped it inside the prologue block, + // so the grouped loop below could not see it) so the prologue, the grouped + // loop and the peeled tail all reuse the same base. + const int64_t b_first = b_base; + { + // Prologue: issue group 0 (the split slice's first 8 fragment slots) so + // its round trip is in flight before the first group fill. +#pragma unroll + for (int i = 0; i < kUnroll; ++i) { + b8[i] = *reinterpret_cast(b + b_first + i * kBFragStep); + } + } + int k0 = k_start; + int64_t b_off = 0; // B byte offset of the current 8-step group's first slot + // Three full groups: issue the NEXT group's 8 B loads before this group's + // 8 mmacs (each next-group load is in flight during the whole 8-mmac + // burst) and rotate the registers in after; the compiler places the vmcnt + // fill at the group head (the iteration-7 wait placement on this shape). + for (int g = 0; g < kPerSplitK / kGroupK - 1; ++g) { + uint64_t b8_next[kUnroll]; +#pragma unroll + for (int i = 0; i < kUnroll; ++i) { + // 8-aligned: b_first, b_off, kUnroll*kBFragStep, i*kBFragStep and b_kq + // are all multiples of 8, so each of these is one global_load_dwordx2 + // over 8 consecutive bytes of one 512-B fragment slot (8 full 64-B + // lines per full 64-lane load instruction). + b8_next[i] = *reinterpret_cast( + b + b_first + b_off + kUnroll * kBFragStep + i * kBFragStep); + } +#pragma unroll + for (int i = 0; i < kUnroll; ++i) { + // A fragment from LDS: one 8-B read reproduces the exact + // du_load_matrix_sync bytes (x[0..7] = A[row][k0 + a_kq + {0..7}]); + // the offset is 8-aligned (kAStageStride and a_kq are multiples of 8 + // and (k0 - k_start) + i*kDummaTileK is a multiple of 32), so this is + // one ds_read_b64. + const uint64_t a8 = *reinterpret_cast( + &s_a[a_row * kAStageStride + (k0 - k_start) + i * kDummaTileK + + a_kq]); + __builtin_memcpy(a_frag.x, &a8, sizeof(a8)); + // B fragment from the iteration-11 fragment-slot pack: 8 contiguous + // bytes at b2[(tile*(k/32) + 24*split + s)*512 + b_row*32 + b_kq] for + // step s of the block's split slice (i = 0..7) -- the exact bytes + // du_load_matrix_sync(b_frag, b + n0*k + k0, k) reads from the packed + // layout, delivered as ONE aligned global_load_dwordx2 per lane + // instead of 8 strided global_load_ubyte + waitcnt cascade + byte-OR + // reassembly. + __builtin_memcpy(b_frag.x, &b8[i], sizeof(b8[i])); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } +#pragma unroll + for (int i = 0; i < kUnroll; ++i) { + b8[i] = b8_next[i]; + } + k0 += kGroupK; + b_off += kUnroll * kBFragStep; // next 8-step group's first slot + } + { + // Peeled final group (no prefetch): consumes the last group loaded by the + // loop's final issue (b8 holds the split slice's slots 24 .. 31). +#pragma unroll + for (int i = 0; i < kUnroll; ++i) { + const uint64_t a8 = *reinterpret_cast( + &s_a[a_row * kAStageStride + (k0 - k_start) + i * kDummaTileK + + a_kq]); + __builtin_memcpy(a_frag.x, &a8, sizeof(a8)); + __builtin_memcpy(b_frag.x, &b8[i], sizeof(b8[i])); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + } + + // Register-only partial store with the verified gfx928 accumulator + // ownership: lane l owns row = l & 15 and columns col_mod4 + 4*i with + // col_mod4 = l >> 4, so acc_frag.x[i] belongs to output column + // n0 + col_mod4 + 4*i of row (l & 15). + const int plane = kDummaTileM * n; // stride between split planes + int32_t* plane_p = partials + split * plane; + const int row = lane & (kDummaTileM - 1); + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = n0 + col_mod4 + 4 * i; + plane_p[row * n + col] = acc_frag.x[i]; + } + + // Fused last-arrival combine tail (iteration 3; repair 2: exactly ONE + // arrival per block). The iteration-3 draft let all 64 lanes of every + // block execute the arrival atomicAdd (64 increments per block = 512 per + // tile per replay), so (arrived & 7) == 7 fired ~64 times per tile -- the + // first fires inside the first-arriving block, before its 7 siblings had + // stored their planes, and the premature combiner summed unwritten zero + // planes (exact check: 7168/8192 mismatched, first_mismatch (0,0) + // actual=0.0 -- the identical signature documented by the validated hy3 + // TP8 gate_up lineage). Fix: exactly one arrival per block -- lane 0 + // fences (release), atomicAdds, fences (acquire) and publishes the result + // through LDS; the rest of the block waits at a barrier, and only the last + // arriver of each tile (32 of 192 blocks) runs the combine. The monotonic + // modulo fire (arrived % kDummaSplitK) == kDummaSplitK-1 then lands + // exactly once per tile per replay, on the true 6th arrival (iteration 20: + // split-K 6 is non-power-of-two, so the power-of-two mask of iteration 3 + // becomes the generic modulo), after all 6 planes of the tile are visible; + // the counters are zeroed once per + // workspace by launch_w8a8_gemm (repair 2: static-pointer-guarded + // hipMemsetAsync, the validated qkv/gate_up pattern) and stay monotonic + // across Graph replays. + __syncthreads(); + __shared__ int s_is_last; + if (lane == 0) { + __threadfence(); // release: this block's plane stores visible to the + // observer of this block's arrival atomic + const int arrived = atomicAdd(&counters[tile], 1); + __threadfence(); // acquire: reads below see every sibling plane store + // that was released before its arrival atomic + // Iteration 20: kDummaSplitK = 6 is non-power-of-two, so the modulo-6 + // fire replaces the power-of-two mask of iteration 3; the monotonic + // counters still land the fire exactly once per tile per replay on the + // true 6th arrival (counters zeroed once per workspace, monotonic + // across Graph replays). + s_is_last = + ((arrived % kDummaSplitK) == (kDummaSplitK - 1)) ? 1 : 0; + } + __syncthreads(); + if (s_is_last == 0) { + return; + } + { + // Last arriver of this tile: sum the kDummaSplitK planes in ascending + // split order -- bit-exact int32 accumulation order, identical to the + // unsplit kernel -- and fuse the scale/bf16 epilogue. + const int e_local = lane * 4; // 64 lanes x 4 = 256 tile elements + const int tile_row = e_local >> 4; + const int col0 = n0 + (e_local & 15); + int32_t v[kDummaSplitK][4]; +#pragma unroll + for (int s = 0; s < kDummaSplitK; ++s) { + const int32_t* p = partials + s * plane + tile_row * n + col0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + v[s][i] = p[i]; + } + } + const float xs = x_scale[tile_row]; + __hip_bfloat16* orow = out + tile_row * n + col0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kDummaSplitK; ++s) { + acc += v[s][i]; + } + const float scaled = + static_cast(acc) * xs * weight_scale[col0 + i]; + orow[i] = __float2bfloat16(scaled); + } + } +} + +// Superseded by the iteration-3 fused last-arrival tail inside +// w8a8_dumma_m16_sk8_partial_kernel (single dispatch per replay); retained +// only for reference. The iteration-2 version of this kernel (grid = 1 +// block, 128 threads on one CU, serial per-element 8-plane dependent load +// chain) profiled at 47.2 us -- the largest single component of the 76.9 us +// operator wall. +__global__ __launch_bounds__(kDummaCombineThreads) void +w8a8_dumma_m16_sk8_combine_kernel( + const int32_t* __restrict__ partials, // [kDummaSplitK][M][N] int32 + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + __hip_bfloat16* __restrict__ out, // [M][N] + int n) { + const int total = kDummaTileM * n; // 8192 output elements + const int plane = kDummaTileM * n; // stride between split planes + for (int e = static_cast(threadIdx.x); e < total; + e += kDummaCombineThreads) { + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kDummaSplitK; ++s) { + acc += partials[s * plane + e]; + } + const int row = e / n; + const int col = e - row * n; + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[e] = __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// Device-to-device packing kernels (optional pack_weight op, outside the +// timed region). Iteration 5 introduced the packed weight for the exact +// gate_up (k, n) == (6144, 512) pair so every m16n16k32 B fragment is 8 +// contiguous bytes; iteration 11 re-packs that pair into the o_proj-validated +// fragment-slot layout (w8a8_pack_fragslot_i8_kernel) so each full 64-lane +// load instruction reads 512 contiguous bytes (see the kernel comment); every +// other (K, N) keeps the identity copy of raw_weight [K, N] int8 and +// weight_scale [N, 1] fp32, the generic fallback layout. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kPackBlockThreads) void +w8a8_pack_identity_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { + dst[i] = src[i]; + } +} + +__global__ __launch_bounds__(kPackBlockThreads) void +w8a8_pack_identity_scale_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { + dst[i] = src[i]; + } +} + +// Fragment-slot pack (iteration 11) for the exact gate_up (k, n) == +// (6144, 512) pair: dst[(tile*(k/32) + step)*512 + row*32 + kk] = +// src[kk*n + tile*16 + row], with tile = n_idx/16, row = n_idx%16, +// step = k_idx/32, kk = k_idx%32. Runs once, outside the timed region and +// outside Graph capture, into the same caller-owned packed_weight buffer +// (byte count unchanged, so captured addresses stay valid). Each +// m16n16k32 B-fragment lane still owns 8 CONTIGUOUS bytes at +// b2[(tile*(k/32) + 24*split + s)*512 + (lane&15)*32 + (lane>>4)*8 + {0..7}] +// (one 512-B slot per k32 step, see w8a8_dumma_m16_sk8_partial_kernel), so +// each fragment stays ONE aligned global_load_dwordx2 per lane -- while a +// full 64-lane load instruction now covers 512 contiguous bytes (8 full +// 64-B lines, 100% line utilization) and each block's 24-slot slice is one +// contiguous 12 KiB stream instead of 16 rows x 768 B strided 6,144 B +// apart under the iteration-5 n-major pack. This is the o_proj-validated +// fragment-slot transport that carried ~461 GB/s effective read bandwidth on +// this part. The identity copy remains the fallback for every other (K, N). +__global__ __launch_bounds__(kPackBlockThreads) void +w8a8_pack_fragslot_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count, + int n, + int k) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { + const int n_idx = static_cast(i / k); + const int k_idx = static_cast(i - static_cast(n_idx) * k); + const int tile = n_idx >> 4; // n_idx / kDummaTileN (16) + const int row = n_idx & 15; // n_idx % kDummaTileN + const int step = k_idx >> 5; // k_idx / kDummaTileK (32) + const int kk = k_idx & 31; // k_idx % kDummaTileK + const int64_t dst_idx = + (static_cast(tile) * (k >> 5) + step) * 512 + + static_cast(row) * 32 + kk; + dst[dst_idx] = src[static_cast(k_idx) * n + n_idx]; + } +} + +} // namespace + +// Optional out-of-timed-region weight pack (stable host symbol used by +// csrc/bindings.cpp). Bootstrap: identity device-to-device copy. +// Iteration 5 packed the exact gate_up (k, n) == (6144, 512) pair into the +// n-major transpose packed[n*K+kk] = raw[kk*n+n]; iteration 11 re-packs that +// pair into the o_proj-validated fragment-slot layout +// b2[(tile*(k/32) + step)*512 + row*32 + kk] (one contiguous 512-B slot per +// m16n16k32 step; byte count unchanged); every other (K, N) keeps the +// identity copy. The matching GEMM interpretation is selected by the exact +// (m, n, k) guard in launch_w8a8_gemm and by the scalar fallback's +// (n, k) == (512, 6144) decode. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_count = static_cast(k) * n; + const int weight_grid = static_cast( + (weight_count + kPackBlockThreads - 1) / kPackBlockThreads); + if (k == 6144 && n == 512) { + hipLaunchKernelGGL( + w8a8_pack_fragslot_i8_kernel, + dim3(weight_grid), + dim3(kPackBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count, + n, + k); + } else { + hipLaunchKernelGGL( + w8a8_pack_identity_kernel, + dim3(weight_grid), + dim3(kPackBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count); + } + + const int64_t scale_count = static_cast(n); + const int scale_grid = static_cast( + (scale_count + kPackBlockThreads - 1) / kPackBlockThreads); + hipLaunchKernelGGL( + w8a8_pack_identity_scale_kernel, + dim3(scale_grid), + dim3(kPackBlockThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + scale_count); +} + +// Timed W8A8 GEMM entry point (stable host symbol used by csrc/bindings.cpp). +// Graph-safe: launches only on the caller-provided current HIP stream, with +// no allocation, compilation, autotuning, packing, host synchronization, +// device synchronization, or default-stream use. Only the caller-provided +// `out` and `workspace` are touched; the split-K specialization uses 6 +// int32 planes (6 * 16 * N * 4 bytes) of the workspace as partial +// accumulators plus a 32-int monotonic arrival-counter tail. Every launch +// overwrites every partial tile, and the counters are zeroed once per +// workspace by the first launch (repair 2: static-pointer-guarded async +// hipMemsetAsync on the caller stream), after which they are monotonic and +// need no reset between Graph replays. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + + // ------------------------------------------------------------------------- + // Exact-shape specialization for glm_tp8_shared_gate_up_proj_m16: + // (M=16, N=512, K=6144). The guard requires all three dimensions so the + // paired M=2 API shape with the same (N, K) and every other (m, n, k) still + // reach the scalar fallback below. For this shape `b` is the iteration-11 + // fragment-slot packed weight b2[(tile*(k/32) + step)*512 + row*32 + kk] + // produced once by launch_pack_w8a8_weight outside the timed region, and + // the scalar fallback decodes the same pack for (n, k) == (512, 6144). The workspace + // contract for this shape guarantees 16 int32 planes (512 KiB); the + // split-K=6 path needs 6 planes (192 KiB) plus a 128-byte counter tail, + // and falls back to the scalar kernel defensively if the caller ever + // passes a smaller buffer. Single dispatch per replay: the fused + // last-arrival combine tail (iteration 3, repair 2) produces the final + // scaled bf16 output inside the partial kernel, and the counter tail is + // zeroed once per workspace (repair 2) without adding a dispatch to the + // timed replay. + // ------------------------------------------------------------------------- + if (m == 16 && n == 512 && k == 6144) { + const int64_t partial_bytes = + static_cast(kDummaSplitK) * kDummaTileM * n * + static_cast(sizeof(int32_t)); + const int64_t counter_bytes = + static_cast(n / kDummaTileN) * sizeof(int32_t); + if (workspace_bytes >= partial_bytes + counter_bytes) { + int32_t* partials = reinterpret_cast(workspace); + int32_t* counters = reinterpret_cast( + reinterpret_cast(partials) + partial_bytes); + // Repair 2: the fused last-arrival tail requires every per-tile + // arrival counter to start at a multiple of kDummaSplitK (0), so the + // modulo-6 fire lands only on the true 6th arrival of each tile per + // replay. The caller workspace is torch.empty (not guaranteed + // zeroed), so zero the 128-byte counter tail once per workspace with + // the validated workspace-pointer-static guard: the first launch with + // this workspace (the pre-capture eager warmup in the normal flow) + // issues an async hipMemsetAsync on the caller stream, and the static + // prevents any later launch from repeating it, so the timed replay + // keeps exactly one dispatch. Once initialized the counters are + // monotonic (6 increments per tile per replay) and never need a reset + // between Graph replays. + static const void* s_zeroed_counters_ws = nullptr; + if (s_zeroed_counters_ws != workspace) { + hipMemsetAsync(counters, 0, counter_bytes, stream); + s_zeroed_counters_ws = workspace; + } + const int grid = (n / kDummaTileN) * kDummaSplitK; + hipLaunchKernelGGL( + w8a8_dumma_m16_sk8_partial_kernel, + dim3(grid), + dim3(kDummaBlockThreads), + 0, + stream, + a, + b, + partials, + x_scale, + weight_scale, + reinterpret_cast<__hip_bfloat16*>(out), + counters, + n, + k); + return; + } + } + + const int64_t total = static_cast(m) * n; + const int grid = static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + dim3(grid), + dim3(kScalarBlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast<__hip_bfloat16*>(out), + m, + n, + k); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/fused_qkv_a_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/fused_qkv_a_proj.hip new file mode 100644 index 00000000..aaa4a2a8 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/fused_qkv_a_proj.hip @@ -0,0 +1,1396 @@ +// @@variant shape=glm_tp8_fused_qkv_a_proj_m4096 commit=935c285f78b3c36ce8f49ef13a68b0cf4e8ff4d5 added=2026-08-31 +// median_us=937.6 p90_us=939.5 speedup=73.23 baseline_us=6.866e+04 +// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 +// @@variant shape=glm_tp8_fused_qkv_a_proj_m4096 +// MetaInfer W8A8 INT8 GEMM for Hygon K500SM_AI / gfx928 (worker_0, GPU 0). +// +// Assigned shape: +// glm_tp8_fused_qkv_a_proj_m4096 : M=4096, N=2624, K=6144 +// +// Operator contract (fixed by the control plane): +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// +// Iteration 1 (correctness-first DUMMA bootstrap): +// * Packing bootstrap: launch_pack_w8a8_weight is the mandated identity +// device-to-device copy for every (K, N) (raw [K, N] row-major weight is +// copied unchanged and the fp32 scales are copied unchanged). The GEMM +// therefore interprets packed_weight as the logical [K, N] row-major +// layout everywhere. Later Parallel explore rounds may change the pack +// kernel and the matching GEMM interpretation together. +// * Iteration 1 dispatched tile <64, 64, 128> (41 x 64 = 2624 blocks, +// 16,384 B LDS/block, 4 blocks/CU, official median 6391.92 us = 20.66 +// logical TOPS). Its exact-source ISA shows the stage loop is +// LDS-issue-bound, not MMA-bound: per stage per wavefront 136 ds_read +// (mostly du_load_matrix_sync byte-reassembly ds_read_u8) + 107 waitcnt +// vs only 32 v_mmac and 2 barriers; 2624 blocks x 96 barriers = +// 251,904 barriers and 2.06 GB of staged global traffic per replay. +// * The epilogue is the coalesced fragment store: the verified gfx928 +// accumulator mapping (row = lane & 15, col_mod4 = lane >> 4, +// frag.x[i] -> column col_mod4 + 4*i) is transposed in registers (8 +// shfl_xor + 8 v_cndmask) so each lane owns four contiguous bf16 columns +// and writes ONE 8-byte store (100% store sector efficiency). The +// per-element float multiply order (float(dot) * x_scale[m] * +// weight_scale[n]) and the bf16 rounding are the reference order, so the +// stored bits are identical to the scalar path. +// * Generic arms: the <64,64,128> tiled kernel also serves every +// large-M shape with N % 64 == 0 and K % 128 == 0 (identity layout); a +// scalar int8/int32 grid-stride kernel is the fallback for every other +// (m, n, k), including the paired M=2 and M=16 API shapes with the same +// (N, K), which must never reach the tiled path (m >= 128 guard). +// * The timed operator (launch_w8a8_gemm) performs no allocation, +// compilation, autotuning, packing, host/device synchronization, or +// default-stream launch: it only dispatches kernels on the +// caller-provided HIP stream and ignores the workspace (split-K capacity +// is zero for this shape: m*n*4 = 42,990,080 B > 16 MiB budget). +// +// Iteration 1 (mandated 2-D macro-tile baseline: benchmark 64x64 / 64x128 / +// 128x64, pick with code-object resources + occupancy): +// * Tile comparison (all three instantiations already compiled by the +// iteration-1 build; VGPR/SGPR/LDS/spills are exact code-object facts +// from profiles/.../iteration1/current-best-isa/metadata.txt; grid and +// staged traffic are arithmetic for M=4096, N=2624, K=6144): +// tile VGPR SGPR LDS B blocks/CU waves/CU grid staged GB +// <64, 64,128> 57 36 16384 4 16 2624 2.06 +// <64,128,128> 92 36 24576 2 8 1344 1.56* +// <128,64,128> 88 34 24576 2 8 1312 1.55 +// (*) 64x128 needs an N tail: N=2624 = 41*64 is not a multiple of +// 128, so the n0=2560 block would zero-fill B columns >= 2624 and an +// epilogue column guard, wasting ~2.4% of all MMAs. No split-K in any +// arm: 1312 blocks dwarf the 120 CUs. +// * Iteration 1 dispatched the exact-fit 2-D macro-tile <128, 64, 128>: +// M=4096 = 32*128 and N=2624 = 41*64 are both exact, so no tail +// predication is introduced anywhere (the generic A row guard and the +// epilogue row guard are retained for the generic path). Grid = (N/64) x +// (M/128) = 41 x 32 = 1312 blocks; 256 threads = 4 wavefronts, each wave +// owns a 64x32 quadrant = eight m16n16k32 int32 accumulator fragments; +// single-buffered 128-K LDS stage (A[128,128] + B[128,64] = 24,576 +// B/block) with TWO __syncthreads per stage, K=6144 = 48 stages. Versus +// the iteration-1 64x64 tile this halves B re-reads (B staged by 32 +// M-blocks instead of 64; staged global traffic 2.06 -> 1.55 GB per +// replay) and cuts total barriers 251,904 -> 125,952, at the cost of +// occupancy 4 -> 2 blocks/CU (16 -> 8 resident waves/CU). The int32 +// accumulation order (k0-outer over 128-K stages, kk-inner ascending, +// same m16n16k32 element-to-slot mapping) is bit-identical to the +// iteration-1 kernel and the scalar reference. +// * The <64,128,128> instantiation stays compiled as the next benchmark +// candidate (needs the N-tail handling described above before dispatch); +// the generic <64,64,128> arm and the scalar fallback are untouched. +// +// Iteration 2 (mandated operand-reuse round: cooperative A+B LDS staging +// vs direct loads, quantified A/B reuse, vectorized coalesced global loads, +// bank-safe LDS layout): +// * Operand-reuse accounting for the <128,64,128> macro-tile (M=4096, +// N=2624, K=6144): per block A = 128 x 6144 = 786,432 B and B = 6144 x +// 64 = 393,216 B of unique bytes; the 41 N-blocks re-read every A +// row-slab (A global traffic M*K*41 = 1.03 GB/replay, 41x amplification +// of the 25.2 MB unique A) and the 32 M-blocks re-read every B +// column-slab (K*N*32 = 0.52 GB/replay, 32x of the 16.1 MB unique B); +// total staged traffic 1.55 GB/replay. Cooperative staging additionally +// buys 4x wave-level reuse inside each block (one 16-B vectorized global +// load per byte, four waves re-read it from LDS), so it is retained: a +// direct per-wave global fragment path would issue 4x more vmem (4 waves +// x 48 x 8-B fragment loads per stage = ~9,216 vs ~288 staged +// 16-B loads per wave) and every 8-B fragment load would be a +// row-strided 16-sector access (vs 4 sectors for the coalesced 16-B +// staging) -- both directions strictly worse. +// * The exact-source ISA of the iteration-1 <128,64,128> kernel shows the +// per-stage per-wavefront consume issues 80 ds_read (64 ds_read_u8 + +// 16 ds_read2_b32) + 99 waitcnt + ~96 byte-reassembly VALU +// (v_and/v_or3/v_lshlrev) per stage against only 32 v_mmac: the +// row_major matrix_b loader of du_load_matrix_sync lowers each B +// fragment to 8 ds_read_u8 at 16-B spacing plus a byte-reassembly chain, +// and the LDS row strides 128 (A) / 64 (B) alias every row onto one +// bank phase (32r and 16r mod 32 == 0), matching the PMC +// lds_bank_conflicts = 274,071,552 (~12.5 per LDS instruction). The +// kernel is LDS-issue/conflict-bound, not MMA- or HBM-bound. +// * Focused change (one mechanism, all inside csrc/w8a8_gemm_hip.hip): +// (1) launch_pack_w8a8_weight now packs the exact (k,n) == (6144,2624) +// weight once, outside the timed region and out of Graph capture, into +// the n-major layout packed[n*K + kk] = raw[kk*N + n] (same byte count +// and buffer, captured addresses unchanged; every other (K, N) keeps +// the identity copy); +// (2) the exact-shape guard dispatches a new +// w8a8_dumma_prefill_packedb_kernel<128,64,128> (grid 41x32 = 1312 +// blocks, 256 threads = 4 waves, same single-buffered 128-K stage with +// two __syncthreads, same k0-outer/kk-inner int32 order, same epilogue) +// that stages A row-major and B n-major into bank-safe LDS rows of 136 B +// (StageK + 8: 8-byte-aligned so every fragment read is one ds_read_b64, +// 34 dwords so the 16 rows of each 16-lane ds_read_b64 phase land on 16 +// distinct bank pairs -> 4-phase conflict-free floor) and replaces both +// library loaders with the lineage-validated explicit load_frag8 (one +// aligned 8-B load per lane per fragment: x[0..7] = 8 consecutive bytes +// of row (lane&15) at k-offset ((lane>>4)*8) for row_major A and the +// same 8 k-bytes of column n0+(lane&15) for col_major B over the +// packed [N, K] tile -- identical fragment values, so the m16n16k32 +// v_mmac inputs and the exact int32 accumulation order are bit-identical +// to the iteration-1 kernel and the scalar reference); +// (3) the scalar fallback decodes the n-major layout when +// (k,n) == (6144,2624) so the paired M=2/M=16 API shapes with the same +// (N, K) stay byte-exact; the generic <64,64,128> identity arm and the +// <64,128,128>/<128,64,128> sibling instantiations are untouched. +// * Expected ISA/PMC deltas: ds_read per stage per wavefront 80 -> 24 +// (16 A + 8 B ds_read_b64), waitcnt 99 -> ~24-32, reassembly VALU ~0, +// lds_bank_conflicts 274M -> conflict-free floor (~1-3M), LDS/block +// 24,576 -> 26,112 B (2 blocks/CU unchanged; no spills). +// +// Iteration 3 (mandated pipeline round: single vs double buffering across K +// tiles; retain double buffering only on ISA/PMC evidence of reduced VMEM +// stalls without harmful LDS or occupancy growth): +// * Bottleneck evidence from the accepted iteration-2 source (fresh PMC +// digest 20f3bc9a..., profiles/.../iteration3/pmc.json: profiled +// 2041.095 us, vmem_read 1,574,400, lds_instructions 7,892,992 with +// lds_wait 4,571,574, lds_bank_conflicts 0, 26,112 B LDS, 2 blocks/CU): +// the exact-source ISA of w8a8_dumma_prefill_packedb_kernel<128,64,128> +// shows the stage loop is [global_load_dwordx4 ... -> s_waitcnt +// vmcnt(0) -> ds_write2_b64 -> s_barrier -> 24x ds_read_b64 -> 32x +// v_mmac -> s_barrier] -- the FULL global-load round trip (issue, vmcnt +// wait, LDS flush) is serialized in front of every consume with ZERO +// compute overlap, and the 14 vmcnt + 57 lgkmcnt waits in the symbol sit +// on that critical path. The K=6144 grid (1312 blocks) and 2-block +// co-residency do not hide this per-stage latency. +// * Focused change (one mechanism, all inside csrc/w8a8_gemm_hip.hip): +// the exact-shape guard now dispatches a new +// w8a8_dumma_prefill_packedb_db_kernel<128,64,64> -- SAME macro-tile +// (grid 41x32 = 1312 blocks, 256 threads = 4 waves, 64x32 quadrant per +// wave, eight m16n16k32 int32 accumulators), SAME load_frag8 fragment +// reads, SAME coalesced epilogue -- with the K stage changed from a +// single-buffered 128-K stage (2 __syncthreads, full VMEM latency +// exposed) to a DOUBLE-buffered 64-K stage with ONE __syncthreads per +// stage: the stage-(s+1) global loads are issued into 3 int4 staging +// VGPR per thread BEFORE the stage-s consume, and the prefetched vectors +// are flushed into the idle LDS buffer after the consume (the validated +// M=3072 prefetch order global_load_dwordx4 -> v_mmac -> vmcnt(0) -> +// ds_write2_b64, here in pure HIP: the vmcnt wait lands on the data +// dependency before the ds_write, i.e. AFTER the MMAC burst). +// * Resources: LDS = 2 x (A[128,72] + B[64,72]) = 2 x 13,824 = 27,648 +// B/block (26,112 -> 27,648 B, +5.9%); 2 x 27,648 = 55,296 <= 65,536 B +// keeps TWO resident blocks/CU = 8 waves/CU, occupancy UNCHANGED (the +// StageK=128 double-buffer alternative needs 52,224 B/block -> 1 +// block/CU = 4 waves/CU and is rejected up front as harmful occupancy +// loss). kLdsStride = 72 = 64 + 8: 8-aligned with an odd 8-byte count +// (72/8 = 9) so the 16 rows of each 16-lane ds_read_b64 phase land on 16 +// distinct bank pairs (18r mod 32, r = 0..15) -- the same conflict-free +// 4-phase floor as iteration 2 -- and 72 % 16 == 8 keeps the two-half +// ds_write2_b64 staging form. VGPR grows by the 12 staging registers +// (~72 -> ~84), far below the 128/thread budget for 2 blocks/CU; 0 +// spills expected. Barriers per block: 1 (prologue) + 95 (loop) = 96, +// the same as iteration 2's 2 x 48 (the final stage is consumed without +// a trailing barrier). +// * Correctness: k0-outer over 96 ASCENDING 64-K stages with kk-inner +// ascending over two m16n16k32 steps gives the identical (k0+kk) +// k-chunk sequence 0, 32, 64, 96, ... as iteration 2's 48 x 128-K +// stages, and the element-to-slot fragment mapping is unchanged, so +// every int32 accumulator is bit-identical (mismatch 0 / max_abs_error +// 0.0 expected). +// * Falsifiable retention: double buffering is retained only if the exact +// shape shows median_us < 2276.930923461914 AND p90_us < +// 2540.756378173828 with mismatch_count == 0, max_abs_error == 0.0, +// graph_capture_passed == true, the exact code object shows 0 spills / +// LDS == 27,648 B (2 blocks/CU still fits: 55,296 <= 65,536) and the +// stage ISA shows the prefetch loads issued BEFORE the ds_read_b64 / +// v_mmac burst with the vmcnt wait after the MMACs (load -> mmac -> +// vmcnt(0) -> ds_write2_b64), i.e. reduced exposed VMEM stalls; any +// median/p90 >= current best, any mismatch, any spill, any occupancy +// loss, or any ISA that still serializes the vmcnt wait in front of the +// consume falsifies the round and the control plane restores the +// iteration-2 source (single buffering retained). +// +// Iteration 5 (mandated packing round: test one weight packing/swizzle that +// makes each DUMMA B tile vector-loadable and LDS-bank-safe): +// * Baseline: the accepted iteration-3 kernel +// w8a8_dumma_prefill_packedb_db_kernel<128,64,64> (official median +// 1070.23 us = 123.4 logical TOPS; fresh PMC digest 47542b38..., +// profiles/.../iteration4/pmc.json of the current source: profiled +// 1101.96 us, grid 1312, 96 arch_vgpr / 27,648 B LDS, lds_instructions +// 4,901,632 = 9.7 per wave-stage (6 merged ds_read2_b64 + 3 +// ds_write2_b64), lds_wait 3,701,906, lds_bank_conflicts 12,091,392 = +// 6 x 4 x 4 x 96 x 1312 EXACTLY -- one conflict per 16-lane phase of +// every merged ds_read2_b64, the inherent 2-cycle alias of the +16-row +// (16 x 72 = 1152 B = 9 x 128 B) second-fragment offset, NOT a +// layout defect (the iteration-2 unmerged ds_read_b64 control had 0 +// conflicts at the same 72/136-B strides), vmem_read 1,574,400, L2 hit +// 79.3% with 5,155,404 misses, TCC read 269.3 MB/replay). The kernel is +// LDS-latency-bound, not HBM-bound; the global B operand is the +// remaining layout lever. +// * Mechanism: replace the iteration-2 n-major pack +// packed[n*K + kk] (each 64-K B tile = 64 column chunks of 64 B at +// 6144-B stride -> 64 x 128-B lines touched per tile, 50% line +// utilization, scattered L2 footprint) with a B-PANEL pack +// packed[(n>>6)*(k*64) + (kk>>4)*1024 + (n&63)*16 + (kk&15)] = +// raw[kk*n + n] (41 panels of [K, 64] int8, 16-byte k-runs per column, +// byte-exact bijection verified by simulation): every [64,64] B tile of +// a 64-K stage is then FOUR CONTIGUOUS 1-KiB slabs (one per 16-k-run), +// each staging vector is a 16-byte-aligned contiguous k-run, and a wave +// reads 4 x 256 B of one panel region per stage -- the tile is +// vector-loadable as one coalesced stream with 8x fewer 128-B lines +// touched per tile (32 vs 256) and full line utilization, so L2 tag +// pressure and B HBM fetch traffic drop without any change in LDS +// layout or instruction count. +// * Kernel-side changes are ONLY the two B global-load sites (prologue + +// steady-state prefetch) of w8a8_dumma_prefill_packedb_db_kernel +// (dispatched solely by the exact-shape guard; static_assert BN == 64): +// vector (n_row, k16) is loaded from byte +// (n0>>6)*(k*64) + ((k0>>4)+k16)*1024 + n_row*16 instead of +// (n0+n_row)*k + k0 + k16*16. LDS staging layout (b_tile[64][72], +// two-half ds_write2_b64), the col_major load_frag8 consumes, the +// k0-outer/kk-inner int32 accumulation order, the coalesced epilogue, +// the grid (41x32 = 1312 blocks), threads (256 = 4 waves), occupancy +// (2 blocks/CU, LDS 27,648 B) and VGPR budget are ALL unchanged, so the +// round is a pure packing/swizzle A/B. +// * Correctness: the same weight elements reach the same fragment slots +// in the same ascending k-chunk sequence, so every int32 accumulator is +// bit-identical (mismatch 0 / max_abs_error 0.0 expected); the scalar +// fallback decodes the panel layout when (k,n) == (6144,2624) (ascending +// 16-byte k-runs keep the exact int32 order for the paired M=2/M=16 API +// shapes); the generic identity arms and sibling kernels are untouched. +// * Packing stays outside the timed region and out of Graph capture +// (launch_pack_w8a8_weight, same byte count k*n = 16,121,856 and same +// buffer -> graph-stable captured addresses unchanged; pack relation +// verified by a byte-exact index simulation before dispatch). +// * Predicted deltas: vmem_read_instructions unchanged (~1,574,400), LDS +// counters unchanged (conflicts stay at the 12,091,392 merged-read alias +// floor), L2 misses < 5,155,404 with hit rate UP, TCC read bytes < +// 269,329,728 (B line fetches halve), median ~1000-1070 us. Falsifiable: +// retained only if median_us < 1070.2328491210938 AND p90_us < +// 1081.9308471679688 with mismatch_count == 0, max_abs_error == 0.0, +// graph_capture_passed == true, the exact code object shows 0 spills / +// LDS == 27,648 B / VGPR <= 128 (2 blocks/CU), and PMC confirms L2 +// misses and TCC read bytes strictly below the iteration-4 values with +// vmem_read/lds/conflicts not growing; any median/p90 >= current best, +// any mismatch, any spill, any occupancy loss, or any L2/HBM regression +// falsifies the pack and the control plane restores the iteration-3 +// source. +// +// Include order is fixed by the control plane: hip_runtime, hip_bfloat16, +// then du_mma (this DTK's du_mma.h is not self-contained otherwise). + +#include +#include +#include + +#include + +namespace { + +constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; // INT8 DUMMA unit: m16n16k32, int32 accumulation +// 4 wavefronts; must remain a multiple of the gfx928 wavefront size (64). +constexpr int kThreadsPerBlock = 256; + +// Exact assigned shape (K, N) after TP=8 partitioning. +constexpr int kExactN = 2624; +constexpr int kExactK = 6144; + +using namespace du::dumma; + +// Coalesced accumulator epilogue for gfx928 INT8 m16n16k32 (verified against +// du_store_matrix_sync in the accepted worker-29 lineage): the accumulator +// lane mapping is row = lane & 15, col_mod4 = lane >> 4, frag.x[i] maps to +// columns col_mod4 + 4*i. A register 4x4 transpose (two 2x2 steps with +// shfl_xor 16 then 32, one v_cndmask per element per step) re-routes the +// int32 values so lane (r, c4) owns the four CONTIGUOUS columns +// base_col + 4*c4 .. +3; they are scaled (float(dot) * x_scale[row] * +// weight_scale[col], the exact reference multiply order), converted to bf16, +// packed and stored with ONE 8-byte store per lane. The row >= m guard is +// wavefront-uniform (lane & 15 cycles the same 16 rows in every 16-lane +// group), so the shuffles never mix active and inactive lanes. base_col is a +// multiple of 64 for the dispatched tiles, so the float4 weight_scale load +// (col0 % 4 == 0) and the 8-byte store (n even -> out offset a multiple of +// 8) are aligned. +template +__device__ __forceinline__ void store_prefill_fragment_coalesced( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 64, n*2 is a multiple of 8). + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + static_cast(row) * n + col0) = + packed; +} + +// Simple native INT8 DUMMA m16n16k32 tiled prefill kernel (identity [K, N] +// weight layout; launch_pack_w8a8_weight is the identity bootstrap, so +// packed_weight == raw logical weight). Template parameters: BM x BN output +// tile per block, StageK K rows per single-buffered LDS stage. 256 threads = +// 4 wavefronts; each wave owns a (BM/2) x (BN/2) quadrant of m16n16k32 +// int32 accumulator fragments. Grid = (N/BN) x (ceil(M/BM)). +// +// K is staged cooperatively in a SINGLE LDS buffer at StageK granularity +// with TWO __syncthreads per stage (one after the staging stores, one after +// all waves consumed the buffer): +// * A[BM, StageK] is staged row-major with flat 16-byte vectors (each +// thread strides over kAVectors = BM*StageK/16 vectors); rows past M are +// zero-filled and masked in the epilogue. +// * B[StageK, BN] is staged row-major from the logical [K, N] weight with +// flat 16-byte vectors (each thread strides over kBVectors = +// StageK*BN/16 vectors). The dispatch guarantees n % BN == 0 and +// k % StageK == 0, so every global vector is aligned and in-bounds. +// * Fragment loads use the library du_load_matrix_sync (row_major for both +// operands) and accumulate with du_mma_sync in ascending k order -- +// k0-outer over stages, kk-inner over kTileK steps -- preserving the +// exact int32 accumulation order of the scalar reference. +// * The epilogue is the coalesced fragment store. +// +// Bootstrap keeps this direct/single-buffered by design (no double buffering, +// no split-K, no raw asm); later rounds tune tile, stage, and packing. +template +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_tiled_kernel( + const int8_t* __restrict__ x_q, // [M, K] row-major + const int8_t* __restrict__ weight, // [K, N] row-major (identity pack) + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [M, N] + int m, + int n, + int k) { + constexpr int kWaveM16 = BM / 32; + constexpr int kWaveN16 = BN / 32; + static_assert(BM % 32 == 0 && BN % 32 == 0, + "block tile must be a multiple of the wave quadrant"); + static_assert(kThreadsPerBlock % kWaveSize == 0, + "blockDim must be a multiple of the gfx928 wavefront size"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * BM; + const int n0 = static_cast(blockIdx.x) * BN; + + __shared__ __align__(16) int8_t a_tile[BM * StageK]; + __shared__ __align__(16) int8_t b_tile[StageK * BN]; + + DUFragment + a_frag[kWaveM16]; + DUFragment + b_frag[kWaveN16]; + DUFragment + acc[kWaveM16][kWaveN16]; +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_fill_fragment(acc[j][i], 0); + } + } + + // Flat 16-byte staging vectors per stage (compile-time; the dispatch + // guarantees k % StageK == 0 and n % BN == 0). + constexpr int kAVectors = BM * StageK / static_cast(sizeof(int4)); + constexpr int kBVectors = StageK * BN / static_cast(sizeof(int4)); + + for (int k0 = 0; k0 < k; k0 += StageK) { + // Stage A[BM, StageK]: thread `vec` owns the 16-byte chunk at flat byte + // offset vec*16 (row = offset / StageK, k-run = offset % StageK). + for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { + const int byte_off = vec * static_cast(sizeof(int4)); + const int local_row = byte_off / StageK; + const int kk = byte_off - local_row * StageK; + const int g_row = m0 + local_row; + reinterpret_cast(a_tile)[vec] = + (g_row < m) + ? *reinterpret_cast( + x_q + static_cast(g_row) * k + k0 + kk) + : int4{0, 0, 0, 0}; + } + // Stage B[StageK, BN] from the logical [K, N] weight: thread `vec` owns + // the 16-byte chunk at (kk = offset / BN, col = offset % BN). + for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { + const int byte_off = vec * static_cast(sizeof(int4)); + const int kk = byte_off / BN; + const int col = byte_off - kk * BN; + reinterpret_cast(b_tile)[vec] = + *reinterpret_cast( + weight + static_cast(k0 + kk) * n + n0 + col); + } + // Make the staged tiles visible to every wavefront. + __syncthreads(); + + // Consume the stage: each wave accumulates its (BM/2) x (BN/2) quadrant + // with kk-ascending m16n16k32 MMAs. + const int local_row = wave_row * (BM / 2); + const int local_col = wave_col * (BN / 2); +#pragma unroll + for (int kk = 0; kk < StageK; kk += kTileK) { +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { + du_load_matrix_sync( + a_frag[j], a_tile + (local_row + j * kTileM) * StageK + kk, + StageK); + } +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_load_matrix_sync( + b_frag[i], b_tile + kk * BN + local_col + i * kTileN, BN); + } +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_mma_sync(acc[j][i], a_frag[j], b_frag[i], acc[j][i]); + } + } + } + // Single buffer: every wavefront must finish reading LDS before the next + // stage's stores overwrite it. + __syncthreads(); + } + + const int base_row = m0 + wave_row * (BM / 2); + const int base_col = n0 + wave_col * (BN / 2); +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + store_prefill_fragment_coalesced( + acc[j][i], x_scale, weight_scale, out, m, n, + base_row + j * kTileM, base_col + i * kTileN, lane); + } + } +} + +// Explicit 8-byte int8 fragment loader (lineage-validated load_frag8): for +// both the row_major matrix_a and the col_major matrix_b m16n16k32 +// fragments, lane l -> row (l & 15), k-quarter (l >> 4) holds the EIGHT +// CONTIGUOUS bytes at p[row*ldm + (l>>4)*8 .. +7]. Loading the 8 bytes as +// one int64 forces the ds_read_b64 form (16-lane phases, distinct bank +// pairs per group -> 4-phase conflict-free floor with an odd 8-byte LDS row +// stride), identical fragment values to the library loader, so the v_mmac +// inputs and the exact int32 accumulation order are unchanged. +template +__device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, + unsigned ldm) { + const unsigned row = __lane_id() & 0xf; + const unsigned kq = __lane_id() >> 4; + const int64_t v = + *reinterpret_cast(p + row * ldm + (kq << 3)); + reinterpret_cast(f.x)[0] = v; +} + +// Iteration-2 exact-shape packed-B prefill kernel (dispatched only for +// (m >= 128, n == 2624, k == 6144), where launch_pack_w8a8_weight produced +// the n-major packed_weight[n*K + kk] = raw[kk*N + n] layout). Same +// geometry as w8a8_dumma_prefill_tiled_kernel<128,64,128> (grid (N/64) x +// (M/128), 256 threads = 4 wavefronts, 64x32 quadrant per wave, eight +// m16n16k32 int32 accumulators, single-buffered 128-K LDS stage with TWO +// __syncthreads per stage, same k0-outer/kk-inner int32 accumulation order, +// same coalesced fragment epilogue). The two differences, both inside the +// operand path: +// * B is staged n-major (b_tile[n][kk] from packed [N, K]) so col_major +// B fragments read 8 contiguous K bytes per lane; +// * A and B LDS row strides are 136 B (StageK + 8): 8-byte aligned for +// the ds_read_b64 fragment loads (odd 8-byte stride count -> the 16 +// rows of each 16-lane phase land on 16 distinct bank pairs, 34r mod 32 +// = 2r mod 32, r = 0..15 -> all 16 even dword starts, conflict-free +// 4-phase floor) and 136 % 16 == 8 so staging stores are the two-half +// int64 ds_write2_b64 form (single ds_write_b128 would be misaligned on +// odd rows). +// Both fragment loaders are the explicit load_frag8 above; the staged bytes +// are the same matrix elements in the same fragment slots as the library +// loaders produced from the identity layout, so every v_mmac input and the +// exact int32 accumulation (k0-outer over 48 stages, kk-inner ascending +// over four m16n16k32 steps) are bit-identical to the accepted iteration-1 +// kernel and to the scalar reference. +template +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_packedb_kernel( + const int8_t* __restrict__ x_q, // [M, K] row-major + const int8_t* __restrict__ packed_w, // [N, K] n-major (exact shape) + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [M, N] + int m, + int n, + int k) { + constexpr int kWaveM16 = BM / 32; + constexpr int kWaveN16 = BN / 32; + static_assert(BM % 32 == 0 && BN % 32 == 0, + "block tile must be a multiple of the wave quadrant"); + static_assert(kThreadsPerBlock % kWaveSize == 0, + "blockDim must be a multiple of the gfx928 wavefront size"); + // Bank-safe LDS row stride: StageK + 8 bytes = 34 dwords. 8-aligned with + // an odd 8-byte count -> every lane's 8-byte fragment is ds_read_b64 at + // the 4-phase conflict-free floor (see the kernel comment above). + constexpr int kLdsStride = StageK + 8; + static_assert(kLdsStride % 8 == 0 && (kLdsStride / 8) % 2 == 1, + "LDS row stride must be 8-aligned with an odd 8-byte count"); + static_assert(kLdsStride % 16 == 8, + "LDS row stride must be 8 mod 16 for the two-half staging"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * BM; + const int n0 = static_cast(blockIdx.x) * BN; + + __shared__ __align__(16) int8_t a_tile[BM * kLdsStride]; + __shared__ __align__(16) int8_t b_tile[BN * kLdsStride]; + + DUFragment + a_frag[kWaveM16]; + DUFragment + b_frag[kWaveN16]; + DUFragment + acc[kWaveM16][kWaveN16]; +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_fill_fragment(acc[j][i], 0); + } + } + + // Flat 16-byte staging vectors per stage (dispatch guarantees + // k % StageK == 0 and n % BN == 0). 8 vectors per 128-K row. + constexpr int kAVectors = BM * StageK / static_cast(sizeof(int4)); + constexpr int kBVectors = StageK * BN / static_cast(sizeof(int4)); + constexpr int kVecsPerRow = StageK / static_cast(sizeof(int4)); + + for (int k0 = 0; k0 < k; k0 += StageK) { + // Stage A[BM, StageK] row-major with row stride kLdsStride: thread + // `vec` owns the 16-byte chunk at (local_row, k16); rows past M are + // zero-filled and masked in the epilogue. + for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { + const int local_row = vec / kVecsPerRow; + const int k16 = vec - local_row * kVecsPerRow; + const int g_row = m0 + local_row; + int4 v{0, 0, 0, 0}; + if (g_row < m) { + v = *reinterpret_cast( + x_q + static_cast(g_row) * k + k0 + k16 * 16); + } + // Two 8-byte halves (stride % 16 == 8 -> odd rows are not 16-byte + // aligned, so a single int4 store is illegal; the int64 pair lowers + // to one ds_write2_b64 per thread). + const int64_t* src64 = reinterpret_cast(&v); + int64_t* dst64 = reinterpret_cast( + a_tile + local_row * kLdsStride + k16 * 16); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + // Stage B[BN, StageK] n-major from the packed [N, K] weight: thread + // `vec` owns the 16-byte chunk at (n_row, k16); the K axis is always + // in range. + for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { + const int n_row = vec / kVecsPerRow; + const int k16 = vec - n_row * kVecsPerRow; + const int4 v = *reinterpret_cast( + packed_w + static_cast(n0 + n_row) * k + k0 + k16 * 16); + const int64_t* src64 = reinterpret_cast(&v); + int64_t* dst64 = reinterpret_cast( + b_tile + n_row * kLdsStride + k16 * 16); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + // Make the staged tiles visible to every wavefront. + __syncthreads(); + + // Consume the stage: each wave accumulates its (BM/2) x (BN/2) quadrant + // with kk-ascending m16n16k32 MMAs; one ds_read_b64 per fragment. + const int local_row = wave_row * (BM / 2); + const int local_col = wave_col * (BN / 2); +#pragma unroll + for (int kk = 0; kk < StageK; kk += kTileK) { +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { + load_frag8( + a_frag[j], + a_tile + (local_row + j * kTileM) * kLdsStride + kk, + kLdsStride); + } +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + load_frag8( + b_frag[i], + b_tile + (local_col + i * kTileN) * kLdsStride + kk, + kLdsStride); + } +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_mma_sync(acc[j][i], a_frag[j], b_frag[i], acc[j][i]); + } + } + } + // Single buffer: every wavefront must finish reading LDS before the + // next stage's stores overwrite it. + __syncthreads(); + } + + const int base_row = m0 + wave_row * (BM / 2); + const int base_col = n0 + wave_col * (BN / 2); +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + store_prefill_fragment_coalesced( + acc[j][i], x_scale, weight_scale, out, m, n, + base_row + j * kTileM, base_col + i * kTileN, lane); + } + } +} + +// Iteration-3 double-buffered packed-B prefill kernel (dispatched only for +// (m >= 128, n == 2624, k == 6144)): SAME macro-tile as the accepted +// iteration-2 w8a8_dumma_prefill_packedb_kernel<128,64,128> (grid (N/64) x +// (M/128) = 41x32 = 1312 blocks, 256 threads = 4 wavefronts, 64x32 quadrant +// per wave, eight m16n16k32 int32 accumulators, same load_frag8 fragment +// reads, same bank-safe LDS row strides, same coalesced epilogue). The ONLY +// change is the K pipeline (see the file-header iteration-3 section); +// iteration 5 additionally reads the packed B operand from the B-panel +// layout (see the file-header iteration-5 section and +// w8a8_pack_panel_i8_kernel) -- the LDS staging layout, fragment consumes, +// and accumulation order are untouched: +// * The single-buffered 128-K stage (two __syncthreads; the exact-source +// ISA serializes global_load_dwordx4 -> s_waitcnt vmcnt(0) -> +// ds_write2_b64 -> s_barrier -> consume with no compute overlapping the +// VMEM round trip) becomes a 64-K DOUBLE-buffered stage with ONE +// __syncthreads per stage: the stage-(s+1) global loads are issued into +// 3 int4 staging VGPR per thread BEFORE the stage-s consume, and the +// prefetched vectors are flushed into the idle buffer after the consume, +// so the VMEM latency of stage s+1 overlaps the LDS/MMAC work of stage s +// (validated M=3072 prefetch order, in pure HIP: the vmcnt wait lands on +// the register data dependency right before the ds_write2_b64, i.e. +// after the v_mmac burst). +// * LDS = 2 x (A[128,72] + B[64,72]) = 2 x 13,824 = 27,648 B/block +// (26,112 -> 27,648 B, +5.9%); 2 x 27,648 = 55,296 <= 65,536 B keeps +// TWO resident blocks/CU = 8 waves/CU, occupancy unchanged. kLdsStride = +// 72 = StageK + 8 (8-aligned, odd 8-byte count 9 -> the 16 rows of each +// 16-lane ds_read_b64 phase land on 16 distinct bank pairs, 18r mod 32, +// r = 0..15 -> the same 4-phase conflict-free floor; 72 % 16 == 8 keeps +// the two-half ds_write2_b64 staging form). 12 extra VGPR for the +// staging registers (~72 -> ~84, still <= 128 for 2 blocks/CU). +// * k0-outer over 96 ASCENDING 64-K stages with kk-inner ascending over +// two m16n16k32 steps: the (k0+kk) k-chunk sequence 0, 32, 64, 96, ... +// is identical to iteration 2 (48 x 128-K stages), so every int32 +// accumulator is bit-identical (mismatch 0 / max_abs_error 0.0). +// * Barriers per block: 1 (prologue) + 95 (loop) = 96, the same as +// iteration 2's 2 x 48; the final stage is consumed without a trailing +// barrier. +template +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_packedb_db_kernel( + const int8_t* __restrict__ x_q, // [M, K] row-major + const int8_t* __restrict__ packed_w, // [N/64][K][64] B panels (exact shape, iter 5) + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [M, N] + int m, + int n, + int k) { + constexpr int kWaveM16 = BM / 32; + constexpr int kWaveN16 = BN / 32; + static_assert(BM % 32 == 0 && BN % 32 == 0, + "block tile must be a multiple of the wave quadrant"); + static_assert(kThreadsPerBlock % kWaveSize == 0, + "blockDim must be a multiple of the gfx928 wavefront size"); + // Bank-safe LDS row stride: StageK + 8 bytes = 18 dwords for StageK = 64. + constexpr int kLdsStride = StageK + 8; + static_assert(kLdsStride % 8 == 0 && (kLdsStride / 8) % 2 == 1, + "LDS row stride must be 8-aligned with an odd 8-byte count"); + static_assert(kLdsStride % 16 == 8, + "LDS row stride must be 8 mod 16 for the two-half staging"); + // Iteration 5: the packed B weight is the panel layout (see + // w8a8_pack_panel_i8_kernel), which is defined for BN == 64 columns per + // panel and 16-byte k-runs (StageK % 16 == 0). Only <128,64,64> is ever + // dispatched by the exact-shape guard, so the panel index math below is + // exact for every dispatch. + static_assert(BN == 64 && StageK % 16 == 0, + "panel B pack requires BN == 64 and StageK % 16 == 0"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * BM; + const int n0 = static_cast(blockIdx.x) * BN; + const int kStages = k / StageK; // dispatch guarantees k % StageK == 0 + + __shared__ __align__(16) int8_t a_tile[2][BM * kLdsStride]; + __shared__ __align__(16) int8_t b_tile[2][BN * kLdsStride]; + + DUFragment + a_frag[kWaveM16]; + DUFragment + b_frag[kWaveN16]; + DUFragment + acc[kWaveM16][kWaveN16]; +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_fill_fragment(acc[j][i], 0); + } + } + + // Flat 16-byte staging vectors per stage. For the dispatched <128,64,64> + // instantiation the 512 A vectors and 256 B vectors divide evenly over the + // 256 threads (2 + 1 per thread), so the prefetch payload is 3 int4 + // staging registers held live across the consume. + constexpr int kVecsPerRow = StageK / static_cast(sizeof(int4)); + constexpr int kAVectors = BM * StageK / static_cast(sizeof(int4)); + constexpr int kBVectors = StageK * BN / static_cast(sizeof(int4)); + constexpr int kAVecPerThread = kAVectors / kThreadsPerBlock; + constexpr int kBVecPerThread = kBVectors / kThreadsPerBlock; + static_assert(kAVectors % kThreadsPerBlock == 0 && + kBVectors % kThreadsPerBlock == 0, + "staging vectors must divide evenly across the block"); + + int4 a_pre[kAVecPerThread]; + int4 b_pre[kBVecPerThread]; + + // ---- Prologue: issue the stage-0 global loads, flush them into buffer 0 + // and make the tile visible to every wavefront. + { + const int k0 = 0; +#pragma unroll + for (int v = 0; v < kAVecPerThread; ++v) { + const int vec = tid + v * kThreadsPerBlock; + const int local_row = vec / kVecsPerRow; + const int k16 = vec - local_row * kVecsPerRow; + const int g_row = m0 + local_row; + a_pre[v] = (g_row < m) + ? *reinterpret_cast( + x_q + static_cast(g_row) * k + k0 + + k16 * 16) + : int4{0, 0, 0, 0}; + } +#pragma unroll + for (int v = 0; v < kBVecPerThread; ++v) { + const int vec = tid + v * kThreadsPerBlock; + const int n_row = vec / kVecsPerRow; + const int k16 = vec - n_row * kVecsPerRow; + // Iteration 5: B is read from the panel-packed layout; the [64, 64] + // tile is 4 contiguous 1-KiB slabs (one per 16-k-run) and vector + // (n_row, k16) sits at byte (n0>>6)*(k*64) + ((k0>>4)+k16)*1024 + + // n_row*16 (16-byte aligned, contiguous k-run of column n0+n_row). + const int64_t panel_off = + static_cast(n0 >> 6) * (static_cast(k) * 64) + + (static_cast(k0 >> 4) + k16) * 1024 + n_row * 16; + b_pre[v] = *reinterpret_cast(packed_w + panel_off); + } +#pragma unroll + for (int v = 0; v < kAVecPerThread; ++v) { + const int vec = tid + v * kThreadsPerBlock; + const int local_row = vec / kVecsPerRow; + const int k16 = vec - local_row * kVecsPerRow; + const int64_t* src64 = reinterpret_cast(&a_pre[v]); + int64_t* dst64 = reinterpret_cast( + a_tile[0] + local_row * kLdsStride + k16 * 16); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } +#pragma unroll + for (int v = 0; v < kBVecPerThread; ++v) { + const int vec = tid + v * kThreadsPerBlock; + const int n_row = vec / kVecsPerRow; + const int k16 = vec - n_row * kVecsPerRow; + const int64_t* src64 = reinterpret_cast(&b_pre[v]); + int64_t* dst64 = reinterpret_cast( + b_tile[0] + n_row * kLdsStride + k16 * 16); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + } + __syncthreads(); + + // ---- Steady state: for s = 0..kStages-2 issue the stage-(s+1) global + // loads into the staging registers BEFORE consuming stage s from buffer + // s&1, then flush the prefetched vectors into the idle buffer (s+1)&1 and + // close the stage with ONE __syncthreads (the barrier's vmcnt wait lands + // after the consume, so the VMEM latency of stage s+1 overlaps the + // LDS/MMAC work of stage s). The last stage is consumed after the loop + // without a trailing barrier. + for (int s = 0; s < kStages - 1; ++s) { + // Prefetch stage s+1: issue the global loads (kept in VGPR while the + // stage-s consume runs). + const int k0 = (s + 1) * StageK; +#pragma unroll + for (int v = 0; v < kAVecPerThread; ++v) { + const int vec = tid + v * kThreadsPerBlock; + const int local_row = vec / kVecsPerRow; + const int k16 = vec - local_row * kVecsPerRow; + const int g_row = m0 + local_row; + a_pre[v] = (g_row < m) + ? *reinterpret_cast( + x_q + static_cast(g_row) * k + k0 + + k16 * 16) + : int4{0, 0, 0, 0}; + } +#pragma unroll + for (int v = 0; v < kBVecPerThread; ++v) { + const int vec = tid + v * kThreadsPerBlock; + const int n_row = vec / kVecsPerRow; + const int k16 = vec - n_row * kVecsPerRow; + // Iteration 5: panel-packed B tile read (see the prologue comment). + const int64_t panel_off = + static_cast(n0 >> 6) * (static_cast(k) * 64) + + (static_cast(k0 >> 4) + k16) * 1024 + n_row * 16; + b_pre[v] = *reinterpret_cast(packed_w + panel_off); + } + + // Consume stage s from buffer s&1: each wave accumulates its (BM/2) x + // (BN/2) quadrant with kk-ascending m16n16k32 MMAs; one ds_read_b64 per + // fragment. + const int local_row = wave_row * (BM / 2); + const int local_col = wave_col * (BN / 2); +#pragma unroll + for (int kk = 0; kk < StageK; kk += kTileK) { +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { + load_frag8( + a_frag[j], + a_tile[s & 1] + (local_row + j * kTileM) * kLdsStride + kk, + kLdsStride); + } +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + load_frag8( + b_frag[i], + b_tile[s & 1] + (local_col + i * kTileN) * kLdsStride + kk, + kLdsStride); + } +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_mma_sync(acc[j][i], a_frag[j], b_frag[i], acc[j][i]); + } + } + } + + // Flush the prefetched stage-(s+1) vectors into the idle buffer. + { + const int buf = (s + 1) & 1; +#pragma unroll + for (int v = 0; v < kAVecPerThread; ++v) { + const int vec = tid + v * kThreadsPerBlock; + const int local_row = vec / kVecsPerRow; + const int k16 = vec - local_row * kVecsPerRow; + const int64_t* src64 = reinterpret_cast(&a_pre[v]); + int64_t* dst64 = reinterpret_cast( + a_tile[buf] + local_row * kLdsStride + k16 * 16); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } +#pragma unroll + for (int v = 0; v < kBVecPerThread; ++v) { + const int vec = tid + v * kThreadsPerBlock; + const int n_row = vec / kVecsPerRow; + const int k16 = vec - n_row * kVecsPerRow; + const int64_t* src64 = reinterpret_cast(&b_pre[v]); + int64_t* dst64 = reinterpret_cast( + b_tile[buf] + n_row * kLdsStride + k16 * 16); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + } + __syncthreads(); + } + + // ---- Tail: consume the last stage (its data was flushed before the last + // barrier; no trailing barrier needed). + { + const int s = kStages - 1; + const int local_row = wave_row * (BM / 2); + const int local_col = wave_col * (BN / 2); +#pragma unroll + for (int kk = 0; kk < StageK; kk += kTileK) { +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { + load_frag8( + a_frag[j], + a_tile[s & 1] + (local_row + j * kTileM) * kLdsStride + kk, + kLdsStride); + } +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + load_frag8( + b_frag[i], + b_tile[s & 1] + (local_col + i * kTileN) * kLdsStride + kk, + kLdsStride); + } +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_mma_sync(acc[j][i], a_frag[j], b_frag[i], acc[j][i]); + } + } + } + } + + const int base_row = m0 + wave_row * (BM / 2); + const int base_col = n0 + wave_col * (BN / 2); +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + store_prefill_fragment_coalesced( + acc[j][i], x_scale, weight_scale, out, m, n, + base_row + j * kTileM, base_col + i * kTileN, lane); + } + } +} + +// Generic scalar int8/int32 fallback for unmatched shapes and small-M API +// cases (including the paired M=2 and M=16 shapes with the same (N, K) -- +// they never reach the m >= 128 tiled path). One output element per +// grid-stride step; exact int32 accumulation (max assigned K = 6144 keeps +// the int8 dot well within int32 range: |dot| <= 6144*127*127 < 2^31), then +// the fused float scale (dot * x_scale[m] * weight_scale[n]) and bf16 store. +// The bootstrap pack is the identity copy, so packed_weight is always the +// logical [K, N] row-major layout here. +__global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, // [M, K] row-major + const int8_t* __restrict__ b, // [K, N] row-major (identity pack) + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [M, N] + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; + idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = a + static_cast(row) * k; + // Iteration 5: for the exact (k, n) == (6144, 2624) the packed_weight + // buffer holds the B-panel layout packed[(col>>6)*(k*64) + (kk>>4)*1024 + // + (col&63)*16 + (kk&15)] = raw[kk*n + col] (41 panels of [K, 64] with + // 16-byte k-runs per column, produced once out of the timed region), so + // the fallback decodes that layout there in ascending-kk 16-byte runs + // (keeps the paired M=2/M=16 API shapes with the same (N, K) + // byte-exact); every other (k, n) keeps the identity [K, N] row-major + // copy. The int32 accumulation stays kk-ascending 0..k-1 in both arms. + const bool packed_panel = (k == kExactK && n == kExactN); + int32_t acc = 0; + if (packed_panel) { + const int8_t* panel = + b + static_cast(col >> 6) * (static_cast(k) * 64) + + (col & 63) * 16; + for (int kkb = 0; kkb < k; kkb += 16) { + const int8_t* p = panel + (kkb >> 4) * 1024; + for (int bb = 0; bb < 16; ++bb) { + acc += static_cast(a_row[kkb + bb]) * + static_cast(p[bb]); + } + } + } else { + const int8_t* b_ptr = b + col; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_ptr[0]); + b_ptr += n; + } + } + const float scaled = static_cast(acc) * + x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// Identity device-to-device byte copy (pack_weight bootstrap, valid for any +// (K, N)). +__global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +// Exact-shape n-major pack (iteration 2): for (k, n) == (6144, 2624) the +// weight is transposed once, outside the timed region and out of Graph +// capture, into packed[n * k + kk] = raw[kk * n + n_row] so every col_major +// B fragment is 8 contiguous K bytes in the packed buffer. Same byte count +// and buffer as the identity pack, so captured addresses are unchanged. One +// thread per output byte; runs once during weight prep. (Kept compiled as +// the iteration-2/3 layout control; iteration 5 dispatches the panel pack +// below for the exact shape.) +__global__ __launch_bounds__(256) void w8a8_pack_nmajor_i8_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t total = static_cast(k) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; + idx += stride) { + const int col = static_cast(idx / k); + const int kk = static_cast(idx - static_cast(col) * k); + packed[idx] = raw[static_cast(kk) * n + col]; + } +} + +// Exact-shape B-panel pack (iteration 5): for (k, n) == (6144, 2624) the +// weight is packed once, outside the timed region and out of Graph capture, +// into 41 panels of [K, 64] int8 with 16-byte k-runs per column: +// packed[(n >> 6) * (k * 64) + (kk >> 4) * 1024 + (n & 63) * 16 + +// (kk & 15)] = raw[kk * n + n] +// (the byte-exact bijection 41 x 384 x 64 x 16 = k*n is verified by +// simulation; same byte count and buffer, so captured addresses are +// unchanged). Every [64, 64] B tile of a 64-K stage is then 4 CONTIGUOUS +// 1-KiB slabs (one per 16-k-run) and every 16-byte staging vector is a +// 16-byte-aligned contiguous k-run of one column, so each DUMMA B tile is +// vector-loadable as one coalesced stream with full 128-B line utilization +// (vs 64 scattered 64-B chunks at 6144-B stride in the n-major layout). +// One thread per output byte; runs once during weight prep. +__global__ __launch_bounds__(256) void w8a8_pack_panel_i8_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t total = static_cast(k) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; + idx += stride) { + const int col = static_cast(idx / k); + const int kk = static_cast(idx - static_cast(col) * k); + // (col >> 6) < n/64 <= 41, (kk >> 4) < k/16 <= 384, (col & 63) < 64, + // (kk & 15) < 16 -> the mixed-radix index stays < k*n in int. + packed[(col >> 6) * (k * 64) + (kk >> 4) * 1024 + (col & 63) * 16 + + (kk & 15)] = raw[static_cast(kk) * n + col]; + } +} + +// Identity device-to-device float copy for the weight scales. +__global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +} // namespace + +// Stable host launch symbol consumed by csrc/bindings.cpp. Dispatches on the +// caller-provided stream only; never allocates, synchronizes, packs, or +// touches the default stream. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + auto* out_ptr = static_cast(out); + + // Exact assigned shape guard: glm_tp8_fused_qkv_a_proj_m4096 + // (M=4096, N=2624, K=6144), m >= 128 so M=2/M=16 paired API shapes with + // the same (N, K) keep reaching the scalar fallback below. Iteration 5: + // dispatch the exact-fit 2-D macro-tile <128, 64, 128> geometry with the + // B-PANEL packed weight (launch_pack_w8a8_weight packs (6144, 2624) once, + // out of the timed region, into packed[(n>>6)*(k*64) + (kk>>4)*1024 + + // (n&63)*16 + (kk&15)] = raw[kk*N + n] so every 64-K B tile is 4 + // contiguous 1-KiB slabs), the bank-safe LDS row strides, the explicit + // 8-byte load_frag8 fragment loads, and the double-buffered 64-K K + // pipeline (one __syncthreads per stage, stage-(s+1) global loads + // prefetched into staging VGPR before the stage-s consume) -- the + // mandated packing/swizzle round at unchanged tile/occupancy (2 + // blocks/CU, 27,648 B LDS/block). The single-buffered iteration-2 + // control kernel stays compiled (see the m < 0 forcing block below) for + // the A/B and for rollback. + if (m >= 128 && n == kExactN && k == kExactK) { + const dim3 grid( + static_cast(n / 64), + static_cast((m + 127) / 128)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64>), + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Force the sibling tile instantiations to compile (<64,128,128> is the + // next benchmark candidate; this branch can never run because m <= 0 + // already returned above). The single-buffered iteration-2 packed-B kernel + // is kept here so the control arm of the iteration-3 A/B stays compiled + // (and rollback only needs the guard repointed). + if (m < 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 128, 128>), + dim3(1), + dim3(static_cast(kThreadsPerBlock)), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<128, 64, 128>), + dim3(1), + dim3(static_cast(kThreadsPerBlock)), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_kernel<128, 64, 128>), + dim3(1), + dim3(static_cast(kThreadsPerBlock)), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Generic native INT8 DUMMA prefill path for every large-M shape with + // N % 64 == 0 and K % 128 == 0 (identity [K, N] layout). The exact shape + // above already returned; this covers compatible sibling shapes. + if (m >= 128 && (n % 64) == 0 && (k % 128) == 0) { + const dim3 grid( + static_cast(n / 64), + static_cast((m + 63) / 64)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 64, 128>), + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Scalar int8/int32 fallback for every unmatched (m, n, k) and all + // small-M API cases. + const int64_t total = static_cast(m) * n; + constexpr int kFallbackThreads = 256; + int64_t blocks = (total + kFallbackThreads - 1) / kFallbackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + const dim3 grid(static_cast(blocks)); + const dim3 block(kFallbackThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); +} + +// Stable host launch symbol consumed by csrc/bindings.cpp. Iteration 5: the +// exact (k, n) == (6144, 2624) int8 weight is packed once (outside the +// timed region and out of Graph capture) into the B-panel layout +// packed[(n>>6)*(k*64) + (kk>>4)*1024 + (n&63)*16 + (kk&15)] = raw[kk*N + n] +// consumed by the packed-B prefill kernel's panel tile reads and by the +// scalar fallback's panel decode; every other (K, N) keeps the identity copy +// (the logical [K, N] row-major weight). The [N] fp32 scales are always +// copied unchanged. Same byte count (k*n and n), same allocation, +// graph-stable addresses. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kPackThreads = 256; + const int64_t weight_count = static_cast(k) * n; + // Iteration 5: the exact (k, n) == (6144, 2624) weight is packed once, + // outside the timed region and out of Graph capture, into the B-panel + // layout (same byte count and buffer addresses, so captured pointers stay + // valid; the packed-B prefill kernel and the scalar fallback for this + // (k, n) both decode it). Every other (K, N) keeps the identity copy + // (raw row-major [K, N]). + if (k == kExactK && n == kExactN) { + int64_t blocks = (weight_count + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + hipLaunchKernelGGL( + w8a8_pack_panel_i8_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else if (weight_count > 0) { + int64_t blocks = (weight_count + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + hipLaunchKernelGGL( + w8a8_pack_identity_i8_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count); + } + if (n > 0) { + int64_t blocks = (static_cast(n) + kPackThreads - 1) / + kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + hipLaunchKernelGGL( + w8a8_pack_identity_f32_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip new file mode 100644 index 00000000..7de05137 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip @@ -0,0 +1,1871 @@ +// @@variant shape=glm_tp8_kv_b_proj_m4096 commit=bb4aec7f2d04a07b6c3c5cf4c19e1b9fc97fc8a3 added=2026-08-31 +// median_us=137.4 p90_us=137.9 speedup=50.51 baseline_us=6940 +// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 +// INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. +// +// Worker: worker_1 (physical GPU 1), assigned shapes: +// glm_tp8_q_b_proj_m4096 : M=4096, N=2048, K=2048 +// glm_tp8_kv_b_proj_m4096 : M=4096, N=3584, K=512 +// +// Logical operation (exact contract): +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// +// Iteration 3 (first valid HIP experiment; iterations 1-2 were killed in the +// agent infrastructure before any proposal). Accepted-best source digest +// a64b6250... with one bounded mechanism: n-major packed B + col_major B +// fragments for the exact glm_tp8_q_b_proj_m4096 (k,n) == (2048, 2048). +// * Large-prefill path (m >= 128 with exact 128x64x64 geometry): native +// INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 output tile +// per block; four wavefronts (256 threads); each wave owns a 64x32 +// quadrant built from eight m16n16k32 int32 accumulator fragments; the +// block cooperatively vector-loads A[128,64] and B[64,64] into one +// single-buffered LDS stage (15,360 B total: 128*80 + 64*80 with 16 B +// padding per row for bank skew); two __syncthreads per stage; fused +// dot * x_scale[m] * weight_scale[n] epilogue stored directly as bf16 +// from the accumulator fragments using the verified gfx928 int8 +// m16n16k32 lane mapping (row = lane & 15, col = (lane >> 4) + 4*i). +// A is read with the library du_load_matrix_sync row-major loader. +// B is layout-templated: +// - kNMajorB == true (exact (k,n) == (2048, 2048), q_b_proj): weight +// is packed once to n-major packed[n*K + k] == raw[k*N + n] outside +// the timed region; the stage stores B n-major (n row, k contiguous, +// 80 B row stride); each B fragment is one contiguous 8-byte LDS run +// per lane, loaded with the lineage-validated load_fragment8 (one +// ds_read2_b64) - kills the 32 ds_read_u8 byte-gathers + mask/OR +// reassembly VALU seen in the accepted-best ISA. +// - kNMajorB == false (any other shape, e.g. kv_b_proj (512, 3584)): +// raw [K, N] row-major identity layout, row-major B loader (the +// accepted-best path, byte-identical). +// Grid dim3(N/64, M/128) = 1024 blocks (q_b_proj) / 1792 blocks +// (kv_b_proj) dwarfs the 120 CUs, so no split-K is needed. +// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), +// including all small-M API cases (M=2, M=16), M tails, and M in +// (0, 128) with the same (K, N); it decodes the n-major packed layout +// for the exact (2048, 2048) and the identity [K, N] layout otherwise. +// * launch_pack_w8a8_weight: n-major device-to-device permutation for +// (k, n) == (2048, 2048), identity copy for every other (k, n) (both the +// int8 weight and the fp32 scale). Packing never happens inside the +// timed GEMM and keeps the same byte count (graph-stable addresses). +// +// Iteration 4 (tile-aspect round): the exact (k, n) == (2048, 2048) q_b_proj +// arm switches from the 128x64 tile to a 128x128 tile +// (w8a8_dumma_prefill_128x128_kernel). With the n-major packed B + 8-byte +// load_fragment8 B reads in place (iteration 3), the remaining LDS-read cost +// is dominated by the A-side row-major fragment loads (4 x ds_read2_b32 per +// kk per wave). A 128x128 tile with four 64x64 quadrants keeps the same +// 4-wave x 8-accumulator-fragment count per wave (8 -> 16 m16n16k32 MMACs +// per kk from 8 LDS fragment reads, i.e. per-MMAC LDS reads drop 6/8 -> 8/16 +// = -33%), and the A-side global re-read halves (A reuse 32 -> 16, A traffic +// 268 -> 134 MiB) while B reuse stays 32 (134 MiB), balancing the per-byte +// A:B traffic at 1:1 and cutting total tile traffic 402 -> 268 MiB (-33%). +// LDS grows 15,360 -> 20,480 B/block (still 2 blocks/CU = 40,960 B <= 64 +// KiB); the grid becomes (N/128, M/128) = 512 blocks (~4.3/CU). The 16 +// accumulator fragments cost ~+32 VGPR; arch VGPR is expected <= 128 so the +// accepted kernel's 2-blocks/CU co-residency is preserved (falsified if the +// exact code object shows > 128 VGPR -> 1 block/CU). kv_b_proj keeps the +// byte-identical 128x64 kNMajorB=false arm and the scalar fallback is +// untouched. +// +// Iteration 5 (packing round): the exact (k, n) == (2048, 2048) q_b_proj +// weight is re-packed once (outside timing/Graph, same byte count, same +// graph-stable buffer) into a swizzled 64-k-stage-major layout +// packed[((k0*8 + kc)*n + col)*8 + b] == raw[kk*n + col], +// kk = k0*64 + kc*8 + b (k0 = 64-k stage, kc = 8-byte sub-chunk, b = byte), +// and the 128x128 kernel stages each B tile as 16-byte plane chunks into an +// LDS plane layout b_tile[kc][n][8] (plane stride 1024 B, n stride 8 B). +// The exact code object of the accepted kernel shows every B fragment read +// is an 8-byte LDS access at (lane&15)*80 + (lane>>4)*8 (n-major stride-80 +// rows), and PMC counts 6,291,456 bank conflicts over 1,048,576 LDS slots +// (6.0/slot) with 4,078,494 LDS waits (3.9/slot): the loop is LDS-latency +// bound and the stride-80 rows alias bank phases every 8 rows (80*8 == 640 +// == 0 mod 128), so every fragment read conflicts ~4-6-way. In the plane +// layout a lane's 8 bytes sit at (lane>>4)*1024 + (lane&15)*8 relative to +// the fragment origin: 16 lanes per 128-B phase cover all 32 banks exactly +// once, i.e. B fragment reads become lane-linear and zero-conflict (4-cycle +// minimum for 512 B). Staging reads stay coalesced (each 16-B chunk is two +// n rows of one sub-chunk -> one aligned int4) and LDS writes are 2-way +// (the 16-B optimum). Same fragment operand bytes, same k0-outer/kk-inner +// int32 accumulation -> bit-identical results. The scalar fallback decodes +// the new pack for (2048,2048) only; kv_b_proj keeps the byte-identical +// 128x64 kNMajorB=false identity path. +// +// Iteration 6 (epilogue round): the fused dot * x_scale[m] * weight_scale[n] +// -> bf16 epilogue has been in-kernel since iteration 1 and the workspace is +// unused ((void)workspace; no split-K, no combine pass anywhere in the call +// chain), so the remaining epilogue inefficiency of the exact (2048, 2048) +// 128x128 arm is the STORE pattern: store_prefill_fragment issues one +// 2-byte bf16 store per lane per element - 64 scattered stores per wave per +// fragment set, each wavefront store touching 16 rows x 4 columns so every +// 32-B sector is only 25% utilized (PMC: 131,072 vmem_write_instructions = +// 512 blocks x 4 waves x 64). The new epilogue transposes each 16x16 +// fragment's 4-element groups inside the 4-lane column group (lanes r, +// r+16, r+32, r+48; two 2x2 steps with __shfl_xor 16 then 32 - the +// lineage-validated 4x4 register transpose accepted on the sibling TP8 +// workers 0/2, same DTK 26.04/gfx928), so lane (r, c4) owns the four +// CONTIGUOUS columns 4*c4 .. +3 and writes ONE 8-byte store per lane (100% +// store sector efficiency; vmem_write 131,072 -> 32,768). Only int32 +// values move between lanes: the per-element multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are +// unchanged, so stored bits are identical. The per-row x_scale (4 rows per +// lane) and per-column weight_scale float4s (4 per lane) are additionally +// register-batched on the last K stage (before the final protective +// __syncthreads), so the epilogue is pure compute + 16 coalesced stores with +// no interleaved vmem loads (vmem_read ~413,696 -> ~278k). kv_b_proj keeps +// the byte-identical 128x64 kNMajorB=false arm (old per-element epilogue) +// and the generic scalar fallback is untouched. +// +// Iteration 7 (compute-pipeline round): the exact (k, n) == (2048, 2048) +// 128x128 arm loads its four A fragments per kk with the lineage-validated +// direct load_fragment8 (one ds_read2_b64 straight into the v_mmac operand) +// instead of the library du_load_matrix_sync row_major loader. The exact +// code object shows the library loader lowers to 8 x ds_read2_b32 plus a +// redundant byte-reassembly chain per fragment (~7 VALU: v_and 0xff00 / +// 0xff0000 / 0xff000000 + v_or_b32_sdwa + v_or3, ~50-56 VALU per stage per +// wave) between the LDS read and the first MMAC, and PMC counts 3,009,148 +// LDS waits (2.3/slot) against a latency-bound loop running at 1 block/CU +// (160 VGPR). du_mma.hpp defines matrix_a row_major int8 as x[i] = +// p[(lane&15)*ldm + (lane>>4)*8 + i] (8 consecutive bytes, memory order) +// and du_mma_sync passes reinterpret(x) unchanged to v_mmac, so +// load_fragment8 produces byte-identical operand values and the int32 +// accumulation is bit-identical; only the redundant VALU reassembly (and its +// lgkmcnt wait states) is removed from the LDS->MMAC critical path, letting +// the 32-MMAC burst issue back-to-back. Expected PMC: valu_instructions +// 8.66M -> ~4.8-5.6M (-35..-45%), lds_instructions 1.31M -> ~1.05M (-20%, +// 8 A ds_read2_b32 + 4 B ds_read2_b64 -> 8 ds_read2_b64 per stage), +// lds_wait_instructions 3.01M -> ~2.2-2.5M, lds_bank_conflicts +// approximately unchanged (A rows keep the 2-way stride-80 aliasing), +// vmem_read/vmem_write and the 32 v_mmac/stage unchanged. kv_b_proj keeps +// the byte-identical 128x64 kNMajorB=false arm (unchanged lowering) and the +// generic scalar fallback is untouched. +// +// Iteration 1 (kv_b_proj tile-baseline round): the exact (k, n) == (512, +// 3584) kv_b_proj arm is still the iteration-2 bootstrap pipeline (library +// fragment loaders, identity [K, N] B, per-element scattered epilogue) and +// runs at 65 TOPS vs 156 TOPS on the fully-modernized q_b_proj arm (same +// GPU, same DTK, measured 220.14 us on K=2048/N=2048). This round ports the +// proven pipeline (swizzled plane-pack B, direct 8-byte A/B fragment loads, +// register-batched coalesced epilogue, fully unrolled kk) to kv_b_proj as +// one templated 2-D macro-tile family (w8a8_dumma_prefill_tile_kernel) over +// the three mandated block tiles plus the iteration-4 aspect flip, all +// compiled, active selected by kActiveKvTile (0 = 64x64, 1 = 64x128, +// 2 = 128x64, 3 = 64x128 with 32x32 quadrants): +// config tile waves threads LDS B grid (blocks) est. VGPR +// 0 64x64 1 64 9,216 (56, 64) = 3584 ~128-140 +// 1 64x128 2 128 13,312 (28, 64) = 1792 ~128-140 +// 2 128x64 4 256 14,336 (56, 32) = 1792 ~88-100 +// 3 64x128 8 512 13,312 (28, 64) = 1792 ~55-60 +// Per-wave LDS fragment-read bytes per MMAC: 256 B for configs 0/1 (4 A + 4 +// B fragments feed 16 MMACs per kk) vs 384 B for config 2 (4 A + 2 B feed 8 +// MMACs) and vs 384 B for the old 128x64 arm; B reads are zero-conflict in +// the plane layout for every plane stride, and the A stride-80 2-way +// aliasing matches the accepted q_b_proj arm. Occupancy: 64x64 -> 7 +// blocks/CU = 7 waves/CU (LDS- and VGPR-limited), 64x128 -> 4 blocks/CU = 8 +// waves/CU (LDS-limited; VGPR <= 128 keeps 4 blocks), 128x64 -> 2 blocks/CU +// = 8 waves/CU (VGPR-limited) - vs 2 blocks/CU = 8 waves/CU today but with +// ~2.9x fewer LDS instructions (direct 8-byte reads vs library byte-gathers) +// and 4x fewer store instructions (coalesced epilogue, 229,376 -> 57,344 +// vmem_write). Active config for the measurement: 64x128 (config 1): +// deepest per-kk MMAC burst of the three alongside the lowest per-MMAC LDS +// read volume, half the A-side global re-read of 128x64, and 8 waves/CU at +// 4 blocks/CU. +// * Iteration-1 measurement: config 1 landed at arch_vgpr 136 (code object +// 135) -> 3 blocks/CU = 6 waves/CU (1.5/SIMD), median 209.4 us = 71.8 +// TOPS. The kernel is latency-bound (per-stage wall ~2,454 cycles vs +// ~250 cycles of pure issue), so 6 waves/CU leaves the barrier + LDS + +// global-latency chain exposed on the SIMDs that hold only one wave. +// * Iteration-2 (operand-reuse) round selects config 2 (128x64): same +// cooperative A+B staging machinery and swizzled-pack B, same +// bit-identical int32 order; its +50% A-side LDS fragment-read bytes per +// MMAC (384 vs 256 B; A rows are re-read by the two n-half waves in both +// configs, but config 2's narrower 64x32 quadrant halves the MMACs per +// A-frag load) is absorbed by the LDS pipe (~10% of per-stage cycles) +// while the compiled 78 VGPR / 14,336 B gives 3 blocks x 4 waves = 12 +// waves/CU (3/SIMD) - a 2x occupancy increase over the measured config-1 +// residency. +// The swizzled pack is extended to (512, 3584) in +// launch_pack_w8a8_weight (same byte count, outside the timed region) and +// the scalar fallback decodes it for both swizzled pairs, so the paired +// M=2/(3584,512) API case stays exact. All other arms (q_b_proj 128x128, +// generic 128x64, identity packs) are byte-identical. +// +// Iteration 5 (kv_b_proj packing round): the exact (k, n) == (512, 3584) +// weight is re-packed once (outside timing/Graph, same byte count, same +// graph-stable buffer) from the k-stage-major swizzle to a TILE-CONTIGUOUS +// stage-major layout +// packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] == raw[kk*n + col] +// kk = k0*64 + kc*8 + b, nt = col/128, row = col%128, +// so every (64-k stage, 128-n tile) B tile is ONE contiguous 8192-B global +// region. The tile-family kernel (w8a8_dumma_prefill_tile_kernel, all four +// configs) stages it as 8 consecutive 1024-B wavefront streams instead of 8 +// streams 28,672 B apart in the k-stage-major layout; the LDS plane tile +// [kc][n][8], the zero-conflict lane-linear 8-byte B fragment reads +// (load_fragment8_plane_stride), the 16-B staging writes and the +// k0-outer/kk-inner int32 accumulation are byte-identical, so mismatch 0 / +// max_abs_error 0 is expected without tolerance debate. q_b_proj (2048, +// 2048) keeps the old k-stage-major swizzle byte-identical (its 128x128 +// kernel, pack kernel and fallback decode are untouched); the scalar +// fallback decodes the new tile-contiguous layout for (512,3584) and the +// old swizzle for (2048,2048). +// +// Iteration 6 (kv_b_proj epilogue round): the mandate (per-row x_scale, +// per-column weight_scale, bf16 conversion and the final coalesced store +// fused into the compute kernel; no workspace/combine pass) has been in +// place since iteration 1 - the launcher is a single kernel with +// (void)workspace and the epilogue is register-batched (xs scalar + ws +// float4 prefetched on the last K stage) with one 8-byte bf16 store per lane +// per fragment (57,344 vmem_write). The remaining epilogue inefficiency of +// the exact (512, 3584) tile family is the bf16 CONVERSION lowering: the +// exact code object shows hip's __float2bfloat16 emits, per element, an +// exec-masked inf/NaN fixup (v_and exp + v_cmp_ne + s_and_saveexec + s_xor + +// v_mov 0 + v_or 0x10000 + v_cmp_eq_u32_sdwa + v_cndmask + s_or_b64) around +// the 3-instruction RNE rounding (v_bfe + v_add3 + shift-in-pack) - ~11 +// instructions and three scalar exec manipulations that serialize the +// block-tail epilogue. This round replaces that path, ONLY in the +// w8a8_dumma_prefill_tile_kernel epilogue, with a plain RNE conversion +// (bf16_rne_u16: u += 0x7fff + ((u >> 16) & 1); u >> 16) behind a new +// kRneOnly template flag on store_prefill_fragment_coalesced_scaled +// (default false, so the q_b_proj 128x128 / 128x64 arms and the scalar +// fallback stay byte-identical). The RNE formula is bit-identical to hip's +// own non-masked branch for every finite float - the only values this GEMM +// can produce (int32 accumulators x finite scales) - so mismatch 0 / +// max_abs_error 0 is expected on the unchanged pre-seeded reference data. +// PMC prediction: valu_instructions 7.31M -> ~5.9-6.1M (-15..-20%), +// lds_instructions 1,089,536 unchanged, lds_bank_conflicts 2,293,760 +// unchanged, vmem_read 229,376 / vmem_write 57,344 unchanged, arch_vgpr <= +// 64 / 0 scratch / 13,312 B LDS / 2 blocks/CU = 16 waves/CU residency +// unchanged; the exact code object must show no s_and_saveexec/s_xor/ +// s_andn2_saveexec around the conversion (v_bfe + v_add3 straight into the +// 8-byte pack). + +#include +#include +#include + +#include + +namespace { + +using namespace du::dumma; + +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 +constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kBlockM = 128; +constexpr int kBlockN = 64; +constexpr int kStageK = 64; +constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) +constexpr int kBStride = kBlockN + kBPad; +constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride +constexpr int kBlockThreads = 4 * kWaveSize; + +// Exact (k, n) pairs whose weight buffer is packed with the iteration-5 +// swizzled 64-k-stage layout (packed[((k0*8+kc)*n+col)*8+b] == raw[kk*n+col]): +// the q_b_proj pair (2048, 2048) and - since this kv_b_proj tile round - the +// kv_b_proj pair (512, 3584). Every other (k, n) keeps the identity layout. +constexpr int kPackSwizzleK1 = 2048; +constexpr int kPackSwizzleN1 = 2048; +constexpr int kPackSwizzleK2 = 512; +constexpr int kPackSwizzleN2 = 3584; + +inline bool is_swizzled_pack(int k, int n) { + return (k == kPackSwizzleK1 && n == kPackSwizzleN1) || + (k == kPackSwizzleK2 && n == kPackSwizzleN2); +} + +// Iteration 5 (kv_b_proj packing round): the (512, 3584) pair switched from +// the k-stage-major swizzle to the tile-contiguous layout +// packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b]; (2048, 2048) keeps +// the k-stage-major swizzle byte-identical. +inline bool is_tile_major_pack(int k, int n) { + return k == kPackSwizzleK2 && n == kPackSwizzleN2; +} + +// Active 2-D macro-tile for the exact glm_tp8_kv_b_proj_m4096 shape among the +// three mandated candidates plus the iteration-4 aspect flip (0 = 64x64, +// 1 = 64x128, 2 = 128x64, 3 = 64x128 with 32x32 quadrants). All four +// instantiations are compiled and dispatched behind the exact (k, n) guard; +// flipping this constant re-routes the kv_b_proj arm without source surgery. +constexpr int kActiveKvTile = 3; // 3 = 64x128 with 32x32 quadrants + // (iteration-4 tile-shape round): the + // down_proj-validated 16-waves/CU recipe + // (512 threads, 8 wavefronts/block, 4 + // accs/wave, arch_vgpr <= 64 -> 2 blocks/CU + // = 16 waves/CU = 4/SIMD) applied as the + // M-tile 128 -> 64 / N-tile 64 -> 128 aspect + // flip; config 2 (128x64, 3 blocks x 4 waves + // = 12 waves/CU, accepted best 169.257 us) + // stays compiled and is the retention + // fallback if config 3 regresses (arch_vgpr + // > 64 -> 1 block/CU = 8 waves -> reject). +constexpr int kKvTileM = 64; +constexpr int kKvTileN = 128; +constexpr int kKvQuadN = 64; + +using bf16_t = hip_bfloat16; + +// --------------------------------------------------------------------------- +// Direct fragment epilogue for one m16n16k32 accumulator fragment. +// Verified gfx928 int8 m16n16k32 accumulator ownership (matches +// du_store_matrix_sync): lane & 15 selects the row, lane >> 4 selects +// col % 4, and x[i] maps to columns col%4 + 4*i. The scale multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are +// identical to the harness reference ((dot.float() * x_scale) * ws.T then +// .to(bfloat16)), so stored bf16 bits match exactly. +// --------------------------------------------------------------------------- +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int base_row, + int base_col, + int m, + int n, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M masking: padded rows never write + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// Iteration 6 (epilogue round): register-batched COALESCED direct-fragment +// epilogue for the exact (2048, 2048) 128x128 arm. The per-element variant +// (store_prefill_fragment) above issues one 2-byte bf16 store per lane per +// element: 64 scattered stores per wave per fragment set, each touching 16 +// rows x 4 columns, i.e. every 32-B store sector only 25% utilized. This +// epilogue transposes the 4-element groups within each 4-lane column group +// (lanes r, r+16, r+32, r+48 - a 4x4 transpose, two 2x2 steps with +// shfl_xor 16 then 32; the lineage-validated pattern accepted on the sibling +// TP8 workers 0/2 on this DTK, which lowers __shfl_xor to ds_bpermute at the +// block tail where the LDS pipe is idle), so lane (r, c4) ends up holding +// the four CONTIGUOUS columns 4*c4 .. 4*c4+3, converts them to bf16, packs +// 4 bf16 (8 B) and writes ONE 8-byte store per lane (100% store sector +// efficiency). Only the int32 values are re-routed between lanes; the +// per-element scale multiply order (float(dot) * x_scale[row] * +// weight_scale[col]) and the __float2bfloat16 rounding are unchanged, so +// the stored bits are identical to the per-element store. xs/ws come from +// caller registers (batched per wave on the last K stage), so the epilogue +// issues no vmem loads. The row >= m guard is wavefront-uniform (all 64 +// lanes of a wave share the same 16-row window), so the shuffles never mix +// active and inactive lanes; base_col is a multiple of 64 and n*2 a +// multiple of 8, so the 8-byte store is aligned. +// --------------------------------------------------------------------------- +// Iteration 6 (kv_b_proj epilogue round): plain round-to-nearest-even bf16 +// conversion without the inf/NaN fixup path. hip's __float2bfloat16 lowers +// (verified on this exact DTK 26.04 / gfx928 toolchain with a micro compile) +// to an exec-masked sequence per element: v_and (exp mask) + v_cmp_ne + +// s_and_saveexec + s_xor + v_bfe + v_add3 + v_mov 0 + v_or 0x10000 + +// v_cmp_eq_u32_sdwa + v_cndmask + s_or_b64 (~11 instructions including three +// scalar exec manipulations that serialize the wavefront). The masking only +// changes the result for inf/NaN inputs; for every finite float - the only +// values this GEMM can produce from int32 accumulators and finite scales - +// the RNE rounding is exactly u += 0x7fff + ((u >> 16) & 1); return u >> 16, +// bit-identical to hip's own non-masked branch (same v_bfe + v_add3 + +// shift). The kernel epilogue is the only place that converts all M*N = +// 14.68M elements; dropping the mask removes ~6 VALU + 5 SALU per element +// (~1.38M VALU + 1.15M SALU per launch, ~19% of the measured 7.31M VALU +// stream) and - the falsifiable axis - removes the +// s_and_saveexec / s_xor / s_andn2_saveexec / s_or_b64 pairs that serialize +// the block-tail epilogue behind scalar exec manipulation. +// --------------------------------------------------------------------------- +__device__ __forceinline__ unsigned short bf16_rne_u16(float x) { + union { + float f; + uint32_t u; + } c; + c.f = x; + const uint32_t u = c.u + 0x7fff + ((c.u >> 16) & 1); + return static_cast(u >> 16); +} + +template +__device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( + const AccFragment& frag, + float xs, + const float4& ws, + bf16_t* __restrict__ out, + int base_row, + int base_col, + int m, + int n, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M masking: padded rows never write (wavefront-uniform) + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 16, n*2 is a multiple of 8); ws is + // the float4 at weight_scale + base_col + 4*c4 preloaded by the caller. + const int col0 = base_col + 4 * c4; + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + uint64_t packed; + if constexpr (kRneOnly) { + // Iteration 6 (kv_b_proj epilogue round): plain RNE (see bf16_rne_u16) - + // bit-identical to __float2bfloat16 for every finite input, without the + // exec-masked inf/NaN fixup path. + packed = static_cast(bf16_rne_u16(v0)) | + (static_cast(bf16_rne_u16(v1)) << 16) | + (static_cast(bf16_rne_u16(v2)) << 32) | + (static_cast(bf16_rne_u16(v3)) << 48); + } else { + packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + } + *reinterpret_cast(out + static_cast(row) * n + col0) = + packed; +} + +// --------------------------------------------------------------------------- +// Direct 8-byte LDS fragment load (lineage-validated on gfx928/DTK 26.04 for +// m16n16k32 int8 fragments). du_load_matrix_sync's int8 loaders assign +// x[0..7] = 8 consecutive bytes at (lane & 15) * ldm + ((lane >> 4) << 3) for +// matrix_b col_major, but the compiler lowers that to per-byte ds_read_u8 + +// mask/OR reassembly VALU. Writing the same 8 bytes directly into the +// fragment storage keeps the operand bit pattern identical (exact int32 +// accumulation unchanged) and lets the compiler emit one ds_read2_b64 per +// fragment straight into the v_mmac operand. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void load_fragment8( + DUFragmentBase& frag, + const signed char* __restrict__ base, + int stride, + int lane) { + const int off = (lane & 15) * stride + ((lane >> 4) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// --------------------------------------------------------------------------- +// Iteration 5 (packing round): 8-byte fragment load for the swizzled plane +// LDS layout b_tile[kc][n][8] (kc = 8-byte k sub-chunk, n row stride 8 B, +// plane stride 1024 B). Lane l reads the same 8 consecutive k bytes as +// load_fragment8 (n row = lane&15, k chunk = lane>>4) but from +// (lane >> 4) * 1024 + (lane & 15) * 8 +// relative to the fragment origin, so the 64 lanes hit 16 distinct bank +// pairs per 128-B phase (all 32 banks once) -> zero bank conflicts, the +// 4-cycle minimum for a 512-B fragment. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void load_fragment8_plane( + DUFragmentBase& frag, + const signed char* __restrict__ base, + int lane) { + const int off = ((lane >> 4) << 10) + ((lane & 15) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// --------------------------------------------------------------------------- +// Plane-layout 8-byte fragment load with a parameterized plane stride +// (plane_stride = kBlockN * 8 bytes). Identical to load_fragment8_plane for +// kBlockN == 128 (plane stride 1024 B); for narrower tiles the plane is +// smaller but the bank pattern is unchanged: lane l reads the 8 bytes at +// (lane >> 4) * plane_stride + (lane & 15) * 8, so each 16-lane group of a +// 128-B phase covers all 32 banks exactly once (zero bank conflicts) for any +// plane stride. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void load_fragment8_plane_stride( + DUFragmentBase& frag, + const signed char* __restrict__ base, + int plane_stride, + int lane) { + const int off = ((lane >> 4) * plane_stride) + ((lane & 15) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// --------------------------------------------------------------------------- +// Compile-time B-fragment layout selector. row_major / col_major are tag +// types in du::dumma, so the layout template argument must be selected as a +// type (a conditional expression over type names is not a valid template +// argument). kNMajorB == true (n-major packed B, exact (2048, 2048)) -> +// col_major fragments (8 contiguous k bytes per lane, load_fragment8); +// kNMajorB == false (raw [K, N] row-major B) -> row_major fragments (the +// accepted-best loader). +// --------------------------------------------------------------------------- +template +struct b_frag_layout { + using type = row_major; +}; +template <> +struct b_frag_layout { + using type = col_major; +}; + +// --------------------------------------------------------------------------- +// Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. +// Each wave owns a 64x32 quadrant (eight m16n16k32 int32 accumulators); the +// block cooperatively stages A[128,64] from x_q (row-major, stride k) and +// B[64,64] into a single-buffered LDS stage. Two barriers per stage: one +// after the cooperative load, one before the next stage overwrites LDS. +// A fragments use the library du_load_matrix_sync row-major loader. B is +// templated on its staged layout: +// * kNMajorB == true: B is packed n-major (packed[n*K + k]) and staged +// n-major (n row, 80 B stride, k contiguous); each B fragment is 8 +// contiguous k bytes per lane, loaded by load_fragment8 (one 8-byte LDS +// read per fragment), eliminating the ds_read_u8 byte-gather + VALU +// reassembly of the accepted-best kernel. +// * kNMajorB == false: B stays raw [K, N] row-major, staged k-major with +// row-major library fragments (accepted-best path, unchanged). +// Dispatch guarantees m % 128 == 0, n % 64 == 0, k % 64 == 0, so every +// global load/store is in-bounds and 16-byte aligned. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row * 64; + const int local_col = wave_col * 32; + + // Single-buffered stage: A[128, 80] + B[64, 80] = 15,360 B/block + // (4 blocks/CU fit the 64 KiB LDS budget; the padded 80-byte strides are + // 16-byte-aligned and break the 64-byte LDS bank periodicity). + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment::type> + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + + // Cooperative staging: A[128,64] is 512 int4s (two per thread), B[64,64] + // is 256 int4s (one per thread). + // A: thread tid owns row (tid*16)/64 and 16-byte column group + // (tid*16)%64 for rows [0,64) and [64,128). + // B (kNMajorB == false): raw [K, N] row-major; thread tid owns K row + // tid>>2 and the 16-byte N-column group (tid&3)*16 (16 consecutive N + // values at a fixed K row -> one aligned int4 in global and in LDS). + // B (kNMajorB == true): packed n-major; thread tid owns N row tid>>2 and + // the 16-byte K-column group (tid&3)*16 (16 consecutive K values at a + // fixed N row -> one aligned int4 in global and in the n-major LDS). + const int vector_byte_offset = tid * static_cast(sizeof(int4)); + const int stage_row = vector_byte_offset / kStageK; // 0..63 + const int stage_col = vector_byte_offset - stage_row * kStageK; + const int b_k = tid >> 2; // K row (false) / N row (true) + const int b_nc = (tid & 3) * 16; // N group (false) / K group (true) + + for (int k0 = 0; k0 < k; k0 += kStageK) { + *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = + *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); + *reinterpret_cast(a_tile + + (stage_row + kBlockM / 2) * kAStride + stage_col) = + *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + + k0 + stage_col); + if constexpr (kNMajorB) { + // n-major packed B: 16 consecutive k bytes of the (n0 + b_k) row. + *reinterpret_cast(b_tile + b_k * kBStride + b_nc) = + *reinterpret_cast( + weight + static_cast(n0 + b_k) * k + k0 + b_nc); + } else { + // raw [K, N] row-major B: 16 consecutive n bytes of the (k0 + b_k) row. + *reinterpret_cast(b_tile + b_k * kBStride + b_nc) = + *reinterpret_cast( + weight + static_cast(k0 + b_k) * n + n0 + b_nc); + } + __syncthreads(); + + // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. + // Accumulation order: k0-outer over 64-K stages, kk-inner (kk=0 then + // kk=32), matching the reference int32 accumulation. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + if constexpr (kNMajorB) { + // n-major LDS tile: n rows stride kBStride, k contiguous within a + // row; col_major fragment slots are 8 consecutive k bytes at + // (lane & 15) * kBStride + ((lane >> 4) << 3) relative to the + // fragment origin -> one ds_read2_b64 per fragment per lane. + load_fragment8(b_frag0, b_tile + local_col * kBStride + kk, kBStride, + lane); + load_fragment8(b_frag1, + b_tile + (local_col + kTileN) * kBStride + kk, + kBStride, lane); + } else { + du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, + kBStride); + du_load_matrix_sync(b_frag1, + b_tile + kk * kBStride + local_col + kTileN, + kBStride); + } + du_load_matrix_sync(a_frag0, a_tile + local_row * kAStride + kk, + kAStride); + du_load_matrix_sync(a_frag1, + a_tile + (local_row + kTileM) * kAStride + kk, + kAStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_load_matrix_sync(a_frag0, + a_tile + (local_row + 2 * kTileM) * kAStride + kk, + kAStride); + du_load_matrix_sync(a_frag1, + a_tile + (local_row + 3 * kTileM) * kAStride + kk, + kAStride); + du_mma_sync(acc20, a_frag0, b_frag0, acc20); + du_mma_sync(acc21, a_frag0, b_frag1, acc21); + du_mma_sync(acc30, a_frag1, b_frag0, acc30); + du_mma_sync(acc31, a_frag1, b_frag1, acc31); + } + + // Protect the LDS buffer from the next stage's cooperative overwrite. + __syncthreads(); + } + + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + store_prefill_fragment(acc00, x_scale, weight_scale, out, base_row, + base_col, m, n, lane); + store_prefill_fragment(acc01, x_scale, weight_scale, out, base_row, + base_col + kTileN, m, n, lane); + store_prefill_fragment(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_prefill_fragment(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, lane); + store_prefill_fragment(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, m, n, + lane); + store_prefill_fragment(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, lane); + store_prefill_fragment(acc31, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + kTileN, m, n, + lane); +} + +// --------------------------------------------------------------------------- +// Iteration 4 (tile-aspect round) + iteration 5 (packing round): 128x128 +// tile for the exact q_b_proj (k, n) == (2048, 2048) swizzled-pack B arm. +// Four wavefronts of 64 lanes; each wave owns a 64x64 quadrant (four A x +// four B m16n16k32 fragments = 16 MMACs per kk, 16 int32 accumulator +// fragments). The block cooperatively stages A[128,64] (2 int4/thread, +// 80-byte-strided rows) and B[128,64] from the iteration-5 swizzled pack (2 +// int4/thread, 16-byte plane chunks) into a single-buffered LDS stage: +// A[128,80] (10,240 B) + B plane tile [8][128][8] (8,192 B) = 18,432 +// B/block (2 blocks/CU = 36,864 B <= 64 KiB); two __syncthreads per stage. +// A fragments use the library du_load_matrix_sync row-major loader +// (4 x ds_read2_b32 per kk per wave), B fragments use load_fragment8_plane +// (lane-linear 8-byte reads, zero LDS bank conflicts). A-side global +// re-read halves (A reuse 32 -> 16, 268 -> 134 MiB) while B reuse stays 32 +// (134 MiB): per-byte A:B tile traffic balances 1:1, total 402 -> 268 MiB +// (-33%). Grid dim3(N/128, M/128) = 512 blocks (~4.3/CU). Dispatch +// guarantees m % 128 == 0, n == 2048, k % 64 == 0, so every global +// load/store is in-bounds and 16-byte aligned. Accumulation order stays +// k0-outer / kk-inner with the same element-to-slot fragment mapping, so the +// int32 accumulation is bit-identical to the accepted kernel. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + constexpr int kBlockM128 = 128; + constexpr int kBlockN128 = 128; + constexpr int kTileStride = kStageK + kBPad; // 64 -> 80-byte row stride + constexpr int kTileBytes = kBlockM128 * kTileStride; // 10,240 B per tile + // Iteration 5 (packing round): B tile is the swizzled plane layout + // [kc][n][8] (8 planes of 128 n rows x 8 B, plane stride 1024 B) so each + // lane's 8-byte fragment chunk is lane-linear and bank-conflict-free. + constexpr int kBTileBytes = 8 * kBlockN128 * 8; // 8,192 B/block + constexpr int kBTilePlane = 1024; // 128 n rows x 8 B + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM128; + const int n0 = static_cast(blockIdx.x) * kBlockN128; + const int local_row = wave_row * (kBlockM128 / 2); // 0 or 64 + const int local_col = wave_col * (kBlockN128 / 2); // 0 or 64 + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + + // Iteration 6 (epilogue round): per-lane register-batched scales. Each + // lane owns 4 rows (base_row + 16*i + (lane & 15), i = 0..3) and 4 + // weight_scale float4s (base_col + 16*j + 4*(lane >> 4), j = 0..3) across + // its 16 fragments; they are loaded exactly once on the last K stage + // before the final protective barrier, so the coalesced epilogue below is + // pure compute + 16 eight-byte stores with no interleaved vmem loads. + float xs_m[4]; + float4 ws_m[4]; + + // Single-buffered stage: A[128,80] (10,240 B) + B plane tile [8][128][8] + // (8,192 B) = 18,432 B/block (2 blocks/CU = 36,864 B <= 64 KiB; the A rows + // keep the padded 80-byte stride, the B tile is plane-swizzled so fragment + // reads are lane-linear with zero bank conflicts). + __shared__ __align__(16) int8_t a_tile[kTileBytes]; + __shared__ __align__(16) int8_t b_tile[kBTileBytes]; + + DUFragment + a_frag0, a_frag1, a_frag2, a_frag3; + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; + DUFragment + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13, + acc20, acc21, acc22, acc23, acc30, acc31, acc32, acc33; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc22, 0); + du_fill_fragment(acc23, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + du_fill_fragment(acc32, 0); + du_fill_fragment(acc33, 0); + + // Cooperative staging: A[128,64] is 512 int4s (two per thread) and + // B[128,64] from the iteration-5 swizzled pack is 512 int4s (two per + // thread). Thread tid owns linear int4 slots tid and tid + 256; for A a + // slot maps to row linear >> 2 and the 16-byte column group + // (linear & 3) * 16 (an M row of x_q, one aligned int4 in global and in + // LDS); for B a slot maps to plane linear >> 6 and n-pair linear & 63 + // (two n rows of one 8-byte k sub-chunk, one aligned int4 in the pack and + // in the plane LDS tile). + for (int k0 = 0; k0 < k; k0 += kStageK) { + const int kstage = k0 >> 6; // 64-k group index inside the swizzled pack +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int linear = tid + i * kBlockThreads; + const int row = linear >> 2; // 0..127 + const int col = (linear & 3) * 16; // 0/16/32/48 + *reinterpret_cast(a_tile + row * kTileStride + col) = + *reinterpret_cast( + x_q + static_cast(m0 + row) * k + k0 + col); + // Swizzled pack: slot linear -> (kc = linear >> 6, j = linear & 63); + // the 16-byte chunk is two n rows (n0 + 2j, n0 + 2j + 1) of the 8-byte + // k sub-chunk kc of stage kstage: one aligned int4 in global (two + // consecutive n rows of one plane) and one aligned int4 in the plane + // LDS tile (plane kc, rows 2j..2j+1 -> contiguous 16 B). + const int b_kc = linear >> 6; // 0..7 + const int b_j = linear & 63; // 0..63 (n-pair within the plane) + *reinterpret_cast(b_tile + b_kc * kBTilePlane + b_j * 16) = + *reinterpret_cast( + weight + (static_cast(kstage * 8 + b_kc) * n + + (n0 + 2 * b_j)) * 8); + } + __syncthreads(); + + // Each wave consumes its 64x64 quadrant: sixteen m16n16k32 MMACs per kk. + // Accumulation order: k0-outer over 64-K stages, kk-inner (kk=0 then + // kk=32), matching the reference int32 accumulation. All eight fragment + // loads are issued before the sixteen MMACs (load-all-then-MMAC-all). + // + // Iteration 7 (compute-pipeline round): the A fragments are loaded with + // the same direct 8-byte LDS read as B (load_fragment8) instead of the + // library du_load_matrix_sync row_major loader. The library's int8 + // matrix_a row_major loader assigns x[i] = p[(lane&15)*ldm + + // (lane>>4)*8 + i] (8 consecutive bytes, memory order) and du_mma_sync + // feeds reinterpret(x) straight into v_mmac, but on this DTK the + // loader lowers to 8 x ds_read2_b32 + a redundant per-dword byte + // reassembly chain (~7 VALU: v_and 0xff00/0xff0000/0xff000000 + + // v_or_b32_sdwa + v_or3 per second dword) sitting between the LDS read + // and the MMAC issue (exact code object, both the 128x128 and the 128x64 + // symbols). load_fragment8 fills the same x[0..7] with one 64-bit + // little-endian write, so the operand bytes are bit-identical and the + // compiler emits one ds_read2_b64 straight into the v_mmac operand (the + // same lineage as the B side since iteration 3/5): the LDS->MMAC + // critical path shortens by ~50 VALU + their lgkmcnt wait states per + // stage, and the 32-MMAC burst can issue back-to-back after the barrier. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + load_fragment8(a_frag0, a_tile + local_row * kTileStride + kk, + kTileStride, lane); + load_fragment8(a_frag1, + a_tile + (local_row + kTileM) * kTileStride + kk, + kTileStride, lane); + load_fragment8(a_frag2, + a_tile + (local_row + 2 * kTileM) * kTileStride + kk, + kTileStride, lane); + load_fragment8(a_frag3, + a_tile + (local_row + 3 * kTileM) * kTileStride + kk, + kTileStride, lane); + // B fragments in the plane layout: fragment (f, kk) origin is + // b_tile + (kk >> 3) * 1024 + (local_col + 16f) * 8 and lane l reads + // the 8 bytes at (lane >> 4) * 1024 + (lane & 15) * 8 (lane-linear, + // zero bank conflicts); operand bytes are identical to the accepted + // n-major load_fragment8, so the int32 accumulation is bit-identical. + const int b_kc0 = kk >> 3; // 0 (kk == 0) or 4 (kk == 32) + load_fragment8_plane(b_frag0, + b_tile + b_kc0 * kBTilePlane + local_col * 8, + lane); + load_fragment8_plane(b_frag1, + b_tile + b_kc0 * kBTilePlane + + (local_col + kTileN) * 8, + lane); + load_fragment8_plane(b_frag2, + b_tile + b_kc0 * kBTilePlane + + (local_col + 2 * kTileN) * 8, + lane); + load_fragment8_plane(b_frag3, + b_tile + b_kc0 * kBTilePlane + + (local_col + 3 * kTileN) * 8, + lane); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + du_mma_sync(acc20, a_frag2, b_frag0, acc20); + du_mma_sync(acc21, a_frag2, b_frag1, acc21); + du_mma_sync(acc22, a_frag2, b_frag2, acc22); + du_mma_sync(acc23, a_frag2, b_frag3, acc23); + du_mma_sync(acc30, a_frag3, b_frag0, acc30); + du_mma_sync(acc31, a_frag3, b_frag1, acc31); + du_mma_sync(acc32, a_frag3, b_frag2, acc32); + du_mma_sync(acc33, a_frag3, b_frag3, acc33); + } + + // Last stage only: prefetch the epilogue's per-row x_scale and + // per-column weight_scale values into registers (uniform branch: k0 is + // block-uniform, so the barrier below is reached by every thread; the + // row < m guard keeps the x_scale reads in bounds for any m tail). The + // vmem latency overlaps the barrier below (s_barrier waits on + // lgkmcnt/arrival, not vmcnt) and the dead a_frag/b_frag VGPR slots are + // reused, so the epilogue is pure compute + coalesced stores. + if (k0 + kStageK >= k) { + const int r = lane & 15; + const int c4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int row = base_row + i * kTileM + r; + xs_m[i] = (row < m) ? x_scale[row] : 0.0f; + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + ws_m[j] = *reinterpret_cast( + weight_scale + base_col + j * kTileN + 4 * c4); + } + } + // Protect the LDS buffers from the next stage's cooperative overwrite. + __syncthreads(); + } + + // Iteration 6 (epilogue round): coalesced direct-fragment epilogue. Each + // lane owns four contiguous bf16 columns of its fragment row and writes + // ONE 8-byte store per fragment (16 stores per wave vs 64 scattered 2-byte + // stores before); the scales come from the registers batched on the last + // K stage, so no vmem loads are interleaved with the stores. The int32 + // accumulation is untouched, and the multiply order / bf16 rounding are + // identical to the old per-element store, so output bits are unchanged. + store_prefill_fragment_coalesced_scaled(acc00, xs_m[0], ws_m[0], out, + base_row, base_col, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc01, xs_m[0], ws_m[1], out, + base_row, base_col + kTileN, m, n, + lane); + store_prefill_fragment_coalesced_scaled(acc02, xs_m[0], ws_m[2], out, + base_row, base_col + 2 * kTileN, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc03, xs_m[0], ws_m[3], out, + base_row, base_col + 3 * kTileN, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc10, xs_m[1], ws_m[0], out, + base_row + kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced_scaled(acc11, xs_m[1], ws_m[1], out, + base_row + kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc12, xs_m[1], ws_m[2], out, + base_row + kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc13, xs_m[1], ws_m[3], out, + base_row + kTileM, + base_col + 3 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc20, xs_m[2], ws_m[0], out, + base_row + 2 * kTileM, base_col, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc21, xs_m[2], ws_m[1], out, + base_row + 2 * kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc22, xs_m[2], ws_m[2], out, + base_row + 2 * kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc23, xs_m[2], ws_m[3], out, + base_row + 2 * kTileM, + base_col + 3 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc30, xs_m[3], ws_m[0], out, + base_row + 3 * kTileM, base_col, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc31, xs_m[3], ws_m[1], out, + base_row + 3 * kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc32, xs_m[3], ws_m[2], out, + base_row + 3 * kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc33, xs_m[3], ws_m[3], out, + base_row + 3 * kTileM, + base_col + 3 * kTileN, m, n, lane); +} + +// --------------------------------------------------------------------------- +// Iteration 1 (kv_b_proj tile-baseline round): 2-D macro-tile DUMMA family +// for the exact glm_tp8_kv_b_proj_m4096 (k, n) == (512, 3584). The prior +// kv_b_proj arm (w8a8_dumma_prefill_128x64_kernel) is the iteration-2 +// bootstrap pipeline: library du_load_matrix_sync fragment loaders (per-byte +// ds_read_u8 byte-gather + mask/OR reassembly VALU, ~2.47M LDS instructions +// and 6.88M bank conflicts measured), raw [K, N] B identity layout, and a +// per-element 2-byte scattered epilogue (229,376 vmem_write instructions = +// 25% store sector efficiency) - i.e. 65 TOPS vs 156 TOPS on the sibling +// q_b_proj arm that already carries the full modern pipeline (plane-swizzled +// B, direct 8-byte fragment loads, coalesced epilogue). This round ports +// that proven pipeline to kv_b_proj as one templated family over the three +// mandated macro-tiles: +// config 0: 64x64 (1 wavefront, 64 threads; wave owns the whole tile) +// config 1: 64x128 (2 wavefronts, 128 threads; each wave owns 64x64) +// config 2: 128x64 (4 wavefronts, 256 threads; each wave owns 64x32) +// config 3: 64x128 (8 wavefronts, 512 threads; each wave owns 32x32; +// iteration-4 tile-shape round) +// (kActiveKvTile selects the active config; all four are compiled.) +// +// Pipeline (byte-identical to the validated q_b_proj arm): +// * B is staged from the iteration-5 swizzled 64-k-stage pack, now also +// produced for (512, 3584) by launch_pack_w8a8_weight, into the plane +// LDS tile b_tile[kc][n][8] (plane stride kBlockN*8 B, n stride 8 B); +// every B fragment read is one 8-byte lane-linear ds_read (zero bank +// conflicts) via load_fragment8_plane_stride. +// * A is staged [kBlockM, 80] (16 B padded row stride) and each A fragment +// is one direct 8-byte LDS read (load_fragment8), no VALU reassembly. +// * int32 accumulation stays resident across the full K loop (k0-outer over +// 64-K stages, kk-inner over two m16n16k32 DUMMA steps), fully unrolled +// per stage like the q_b_proj winner (deep 16-MMAC bursts per kk). +// * The fused dot * x_scale * weight_scale -> bf16 epilogue is the +// register-batched coalesced variant (one 8-byte store per lane per +// fragment; scales prefetched on the last K stage). +// * Single-buffered LDS, two __syncthreads per stage, grid dim3(N/BN, +// M/BM) = 1792 blocks (64x128 / 128x64) or 3584 blocks (64x64) - no +// split-K (the MxN grid already over-subscribes the 120 CUs). +// +// Dispatch guarantees m % kBlockM == 0, n % kBlockN == 0, k % 64 == 0 for the +// exact (512, 3584) pair (m == 4096 in the assigned shape), so every global +// load/store is in bounds and 16-byte aligned; the row < m guards stay for +// generic safety (wavefront-uniform, so the coalesced shuffles never mix +// active and inactive lanes). +// +// Iteration 16 (kv_b_proj load-placement round): the exact (k, n) == (512, +// 3584) active config-3 kernel (w8a8_dumma_prefill_tile_kernel<64,128,32,32, +// true>, median 139.567 us = 107.7 TOPS, 49.72x vs fixed Triton 6939.939 us; +// code object 61 VGPR / 28 SGPR / 13,312 B LDS / 0 scratch / 2 blocks/CU = +// 16 waves/CU; grid 1792, workgroup 512) is LDS-latency/issue-bound (lineage +// rule 5) and its stage critical path still exposes, per 64-K stage, one +// grouped A+B staging L2 round trip at the stage top (iteration-7 grouping +// gained +2.24% by merging the two serialized round trips of iteration 6, +// so the round trip sits on the barrier-to-barrier wall). The double- +// buffered one-barrier pipeline is arithmetic-infeasible at 2 blocks/CU +// (iters 9/10/12 compiled to 72/72/68 VGPR -> 1 block/CU = 8 waves/CU, the +// measured regression regime; every full double-buffer costs >= 7 VGPR on +// this 64-VGPR-constrained 512-thread config) and the 128-K stage-depth +// rework was attempted in the killed iteration-15 candidate (NOT replayed). +// This round keeps the SINGLE-buffered two-barrier stage byte-identical and +// moves ONLY the global-load ISSUE point: the grouped A+B staging loads for +// stage s+1 are issued at the END of stage s (after the MMAC burst, before +// the protective trailing __syncthreads) instead of at the top of stage s+1. +// Global loads never touch LDS, so the early issue is hazard-free; the +// vmcnt wait is inserted by the compiler before the first LDS write of stage +// s+1 that consumes the chunk, and the in-flight round trip overlaps the +// trailing barrier + loop back-edge of stage s (partial hiding of each of +// the 8 exposed staging round trips). Crucially NO register payload is +// carried across the burst: a_chunk/b_chunk are dead during the burst and +// are re-issued after it, so peak VGPR stays <= 64 -> 2 blocks/CU = 16 +// waves/CU (this is NOT the iteration-9/10/12 register-carrying placement, +// and NOT the iteration-12 double-buffer: LDS stays 13,312 B/block, 16 +// s_barriers/block). Bit-identity: identical global bytes -> identical LDS +// writes at identical (row, kk) slots -> identical operand bytes in the +// same k0-outer/kk-inner int32 accumulation -> identical register-batched +// RNE epilogue (only the load issue point moves), so mismatch 0 / +// max_abs_error 0 is expected without tolerance debate. Scoping: a new +// defaulted template parameter kPrefetchNext = false keeps configs 0-2 on +// the exact iteration-7 issue-at-stage-top placement (compiled identically +// via the if constexpr retention branch); the q_b_proj 128x128 / 128x64 +// arms, the generic scalar fallback and the pack kernels are untouched. +// Iteration 17 (kv_b_proj consolidation round): the two independently +// measured single-buffered wins are consolidated into ONE active config-3 +// instantiation (w8a8_dumma_prefill_tile_kernel<64,128,32,32,true,24>): +// (a) iteration 16's next-stage load-issue prefetch (kPrefetchNext = true, +// the current shadow candidate, median 138.229 us / p90 138.781, +// +0.968% vs the accepted best): the grouped A+B staging loads for +// stage s+1 are issued after the MMAC burst of stage s, so each of the +// 8 staging L2 round trips overlaps the trailing barrier + loop +// back-edge instead of stalling the stage-top LDS writes; +// (b) iteration 13's A-side LDS bank skew (kALdsPad = 24 -> 88-byte A row +// stride, standalone median 138.783 us / p90 139.185, +0.565% vs the +// accepted best): stride-80 rows are 8-periodic in the 128-B bank +// phase, so the 16 fragment rows (lane & 15) collapse onto 8 phases +// (r and r+8 collide, 2-way); 88 = 64+24 is 8-byte aligned and 88r mod +// 128 is 16-periodic, so the 16 fragment rows hit 16 DISTINCT phases +// -> zero-conflict A fragment reads (lineage rule 3, the qkv stride-68 +// / down_proj stride-72 pattern). Because stride % 16 == 8, the +// 16-byte A staging chunk is written as two int64 halves +// (ds_write2_b64, lineage rule 3), with the SAME 16 bytes at the SAME +// logical (row, kk) slots. +// The two mechanisms are orthogonal: (a) removes the staging L2 round trip +// from the stage-top critical path, (b) removes the 2,293,760-cycle LDS +// bank-conflict term from the LDS pipe (the conflict model: A fragment reads +// 32 x 2-way 8-B reads/stage -> 8 vs 4 cycles, + A staging writes 4 x 2-way +// 16-B writes/stage -> 16 vs 8 cycles, x 8 stages = 1,280 extra cycles/block +// x 1792 blocks). Both keep the single-buffered two-barrier stage (16 +// s_barriers/block), 2 blocks/CU = 16 waves/CU (vgpr <= 64: the write split +// reuses the same 4-VGPR a_chunk and NO payload crosses the burst), and LDS +// 13,824 B/block (2 blocks = 27,648 B <= 64 KiB). Bit-identity: identical +// global bytes -> identical logical A tile content at every (row, kk) (only +// inter-row padding changes) -> identical 8-byte m16n16k32 operand bytes per +// lane -> identical k0-outer/kk-inner int32 accumulation and register-batched +// RNE epilogue -> mismatch 0 / max_abs_error 0 expected without tolerance +// debate. Scoping: kALdsPad defaults to kBPad = 16, so configs 0-2 keep the +// exact iteration-7 stride-80 int4 write path (token-identical); the +// q_b_proj 128x128 / 128x64 arms, the generic scalar fallback and the pack +// kernels are untouched; the exact-shape guard (k == 512 && n == 3584 && +// m >= 64 && m % 64 == 0) is preserved. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__( + ((kBlockM / kQuadM) * (kBlockN / kQuadN)) * kWaveSize) +void w8a8_dumma_prefill_tile_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + // kQuadM = rows owned by each wavefront (64 in configs 0-2; 32 in the + // iteration-4 config 3 so a 64x128 block hosts 8 wavefronts). + constexpr int kThreads = + (kBlockM / kQuadM) * (kBlockN / kQuadN) * kWaveSize; + // Iteration 17 (kv_b_proj consolidation round): A row stride is + // kStageK + kALdsPad. The default pad 16 keeps stride 80 (16-byte + // aligned, 2-way bank aliasing) for the retention configs 0-2; the active + // config-3 instantiation passes kALdsPad = 24 -> stride 88 (8-byte + // aligned, 16 distinct 128-B phases for the 16 fragment rows -> zero + // conflict A fragment reads; lineage rule: stride % 16 == 8 stages with + // two int64 halves, see the A staging write below). + constexpr int kTileStride = kStageK + kALdsPad; // 64 -> 80/88 B A stride + constexpr int kBPlaneStride = kBlockN * 8; // n rows x 8 B per plane + constexpr int kATileBytes = kBlockM * kTileStride; + constexpr int kBTileBytes = 8 * kBlockN * 8; // 8 k-planes of n rows x 8 B + constexpr int kASlots = (kBlockM * kStageK) / 16; // int4 A chunks + constexpr int kBSlots = kBTileBytes / 16; // int4 B chunks + constexpr int kALoads = (kASlots + kThreads - 1) / kThreads; + constexpr int kBLoads = (kBSlots + kThreads - 1) / kThreads; + constexpr int kAccM = kQuadM / kTileM; // A fragments per wave (4 / 2) + constexpr int kAccN = kQuadN / kTileN; // B fragments per wave + constexpr int kWavesN = kBlockN / kQuadN; + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWavesN; + const int wave_col = wave % kWavesN; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row * kQuadM; + const int local_col = wave_col * kQuadN; + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + + // Single-buffered stage: A[kBlockM, 80] + B plane tile [8][kBlockN][8]. + __shared__ __align__(16) int8_t a_tile[kATileBytes]; + __shared__ __align__(16) int8_t b_tile[kBTileBytes]; + + DUFragment + a_frag[kAccM]; + DUFragment + b_frag[kAccN]; + DUFragment acc[kAccM][kAccN]; +#pragma unroll + for (int i = 0; i < kAccM; ++i) { +#pragma unroll + for (int j = 0; j < kAccN; ++j) { + du_fill_fragment(acc[i][j], 0); + } + } + + // Register-batched epilogue scales (prefetched on the last K stage). + float xs_m[kAccM]; + float4 ws_m[kAccN]; + + if constexpr (kPrefetchNext) { + // ----------------------------------------------------------------------- + // Iteration 16 (kv_b_proj load-placement round) - ACTIVE config-3 path: + // single-buffered two-barrier stage with the grouped A+B staging loads + // for stage s+1 issued at the END of stage s (after the MMAC burst, + // before the protective trailing barrier). a_chunk/b_chunk live across + // the trailing barrier + the loop back-edge only (never across a burst), + // so peak VGPR stays <= 64 -> 2 blocks/CU = 16 waves/CU unchanged; LDS + // 13,312 B/block and 16 s_barriers/block unchanged. The vmcnt wait for + // the in-flight loads is inserted by the compiler before the first LDS + // write of the next stage, i.e. each staging L2 round trip overlaps the + // trailing barrier + back-edge instead of stalling the LDS writes. + // ----------------------------------------------------------------------- + int4 a_chunk[kALoads]; + bool a_ok[kALoads]; + int4 b_chunk[kBLoads]; + { + // Prologue: issue the stage-0 A+B loads (the same grouped pair and the + // same addresses as the iteration-7 stage-top placement) so the first + // iteration starts directly with the LDS writes. + const int k0 = 0; + const int kstage = 0; // 64-k stage index inside the swizzled pack +#pragma unroll + for (int i = 0; i < kALoads; ++i) { + const int slot = tid + i * kThreads; + a_ok[i] = slot < kASlots; + if (a_ok[i]) { + const int row = slot >> 2; // 0..kBlockM-1 + const int col = (slot & 3) * 16; // 0/16/32/48 + a_chunk[i] = *reinterpret_cast( + x_q + static_cast(m0 + row) * k + k0 + col); + } + } +#pragma unroll + for (int i = 0; i < kBLoads; ++i) { + const int slot = tid + i * kThreads; + if (slot < kBSlots) { + const int b_kc = slot / (kBlockN / 2); // plane 0..7 + const int b_j = slot - b_kc * (kBlockN / 2); // n-pair 0..BN/2-1 + const int b_row0 = (n0 + 2 * b_j) & 127; + const int b_nt = (n0 + 2 * b_j) >> 7; + b_chunk[i] = *reinterpret_cast( + weight + (((static_cast(kstage) * (n >> 7) + b_nt) * 8 + + b_kc) * 128 + b_row0) * 8); + } + } + } + for (int k0 = 0; k0 < k; k0 += kStageK) { + const int kstage = k0 >> 6; // 64-k stage index inside the swizzled pack +#pragma unroll + for (int i = 0; i < kALoads; ++i) { + if (a_ok[i]) { + const int slot = tid + i * kThreads; + const int row = slot >> 2; + const int col = (slot & 3) * 16; + if constexpr (kTileStride % 16 == 8) { + // Iteration 17 (kv_b_proj consolidation round): 88-byte A rows + // are only 8-byte aligned, so stage the 16-byte chunk as two + // int64 halves (ds_write2_b64, not ds_write_b128 - lineage rule + // 3). Same 16 staging bytes at the same logical (row, kk) + // slots, so bit-identity holds. + int8_t* dst = a_tile + row * kTileStride + col; + const int64_t* src = + reinterpret_cast(&a_chunk[i]); + *reinterpret_cast(dst) = src[0]; + *reinterpret_cast(dst + 8) = src[1]; + } else { + *reinterpret_cast(a_tile + row * kTileStride + col) = + a_chunk[i]; + } + } + } +#pragma unroll + for (int i = 0; i < kBLoads; ++i) { + const int slot = tid + i * kThreads; + if (slot < kBSlots) { + const int b_kc = slot / (kBlockN / 2); + const int b_j = slot - b_kc * (kBlockN / 2); + *reinterpret_cast(b_tile + b_kc * kBPlaneStride + b_j * 16) = + b_chunk[i]; + } + } + __syncthreads(); + + // Each wave consumes its 64 x kQuadN quadrant: kAccM * kAccN + // m16n16k32 MMACs per kk. Accumulation order: k0-outer over 64-K + // stages, kk-inner (kk = 0 then 32), matching the reference int32 + // accumulation. All fragment loads are issued before the MMAC burst + // (load-all-then-MMAC-all, fully unrolled like the q_b_proj winner). +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { +#pragma unroll + for (int i = 0; i < kAccM; ++i) { + load_fragment8(a_frag[i], + a_tile + (local_row + i * kTileM) * kTileStride + kk, + kTileStride, lane); + } +#pragma unroll + for (int j = 0; j < kAccN; ++j) { + load_fragment8_plane_stride( + b_frag[j], + b_tile + (kk >> 3) * kBPlaneStride + + (local_col + j * kTileN) * 8, + kBPlaneStride, lane); + } +#pragma unroll + for (int j = 0; j < kAccN; ++j) { +#pragma unroll + for (int i = 0; i < kAccM; ++i) { + du_mma_sync(acc[i][j], a_frag[i], b_frag[j], acc[i][j]); + } + } + } + + // Last stage only: prefetch the epilogue's per-row x_scale and + // per-column weight_scale values into registers (uniform branch: k0 is + // block-uniform, so the barrier below is reached by every thread; the + // row < m guard keeps the x_scale reads in bounds for any m tail). + if (k0 + kStageK >= k) { + const int r = lane & 15; + const int c4 = lane >> 4; +#pragma unroll + for (int i = 0; i < kAccM; ++i) { + const int row = base_row + i * kTileM + r; + xs_m[i] = (row < m) ? x_scale[row] : 0.0f; + } +#pragma unroll + for (int j = 0; j < kAccN; ++j) { + ws_m[j] = *reinterpret_cast( + weight_scale + base_col + j * kTileN + 4 * c4); + } + } + // Iteration 16 (kv_b_proj load-placement round): issue the grouped A+B + // staging loads for the NEXT stage here - after the burst and the + // epilogue scale prefetch, before the protective trailing barrier - so + // their L2 round trip overlaps the barrier + loop back-edge instead of + // stalling the next stage's LDS writes. Uniform branch (k0 is + // block-uniform); only GLOBAL memory is touched, so the barrier + // semantics are unchanged. On the last stage no loads are issued and + // the placement degenerates to the accepted schedule. + if (k0 + kStageK < k) { + const int k0n = k0 + kStageK; + const int kstagen = k0n >> 6; // 64-k stage index of the next stage +#pragma unroll + for (int i = 0; i < kALoads; ++i) { + const int slot = tid + i * kThreads; + a_ok[i] = slot < kASlots; + if (a_ok[i]) { + const int row = slot >> 2; // 0..kBlockM-1 + const int col = (slot & 3) * 16; // 0/16/32/48 + a_chunk[i] = *reinterpret_cast( + x_q + static_cast(m0 + row) * k + k0n + col); + } + } +#pragma unroll + for (int i = 0; i < kBLoads; ++i) { + const int slot = tid + i * kThreads; + if (slot < kBSlots) { + const int b_kc = slot / (kBlockN / 2); // plane 0..7 + const int b_j = slot - b_kc * (kBlockN / 2); // n-pair 0..BN/2-1 + const int b_row0 = (n0 + 2 * b_j) & 127; + const int b_nt = (n0 + 2 * b_j) >> 7; + b_chunk[i] = *reinterpret_cast( + weight + + (((static_cast(kstagen) * (n >> 7) + b_nt) * 8 + + b_kc) * 128 + b_row0) * 8); + } + } + } + // Protect the LDS buffers from the next stage's cooperative overwrite. + __syncthreads(); + } + } else { + // ----------------------------------------------------------------------- + // Retention path (configs 0-2): EXACT iteration-7 loop, byte-identical. + // ----------------------------------------------------------------------- + for (int k0 = 0; k0 < k; k0 += kStageK) { + const int kstage = k0 >> 6; // 64-k stage index inside the swizzled pack + // Iteration 7 (compute-pipeline round): issue the A and B staging global + // loads for the WHOLE stage before either vmcnt wait (issue grouping / + // prefetch distance = one VMEM round trip), then one wait and both LDS + // writes. The iteration-6 exact code object serializes the staging as + // global_load_dwordx4 A -> s_waitcnt vmcnt(0) -> ds_write_b128 A + // global_load_dwordx4 B -> s_waitcnt vmcnt(0) -> ds_write_b128 B + // (zero instructions between each load and its wait, and the compiler + // reuses one 4-VGPR slot for both loads), so each 64-K stage exposes ~2 + // L2 round trips back-to-back on the barrier-to-barrier critical path. + // The A and B loads are independent (A: x_q rows; B: swizzled-pack tile), + // so grouping them as one issued pair overlaps the two latencies into ~1 + // round trip. Staging content, LDS addresses, the int32 accumulation + // order (k0-outer / kk-inner) and the epilogue are byte-identical; the + // loads simply complete concurrently instead of serially. + int4 a_chunk[kALoads]; + bool a_ok[kALoads]; +#pragma unroll + for (int i = 0; i < kALoads; ++i) { + const int slot = tid + i * kThreads; + a_ok[i] = slot < kASlots; + if (a_ok[i]) { + const int row = slot >> 2; // 0..kBlockM-1 + const int col = (slot & 3) * 16; // 0/16/32/48 + a_chunk[i] = *reinterpret_cast( + x_q + static_cast(m0 + row) * k + k0 + col); + } + } + int4 b_chunk[kBLoads]; +#pragma unroll + for (int i = 0; i < kBLoads; ++i) { + const int slot = tid + i * kThreads; + if (slot < kBSlots) { + const int b_kc = slot / (kBlockN / 2); // plane 0..7 + const int b_j = slot - b_kc * (kBlockN / 2); // n-pair 0..BN/2-1 + // Iteration 5 (packing round): tile-contiguous (512, 3584) pack - + // the (kstage, n-tile) B tile is ONE contiguous 8192-B region + // packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b]; the 16-B + // chunk is two consecutive n rows (n0 + 2j, +1) of one 8-byte k + // sub-chunk. It never crosses a 128-row tile boundary (n0 % 64 == 0 + // and a 64/128-row window starting at a multiple of 64 stays inside + // one 128-row pack tile), so the int4 is aligned in global and LDS. + const int b_row0 = (n0 + 2 * b_j) & 127; + const int b_nt = (n0 + 2 * b_j) >> 7; + b_chunk[i] = *reinterpret_cast( + weight + (((static_cast(kstage) * (n >> 7) + b_nt) * 8 + + b_kc) * 128 + b_row0) * 8); + } + } +#pragma unroll + for (int i = 0; i < kALoads; ++i) { + if (a_ok[i]) { + const int slot = tid + i * kThreads; + const int row = slot >> 2; + const int col = (slot & 3) * 16; + if constexpr (kTileStride % 16 == 8) { + // Iteration 17 (kv_b_proj consolidation round): 88-byte A rows + // are only 8-byte aligned, so stage the 16-byte chunk as two + // int64 halves (ds_write2_b64, not ds_write_b128 - lineage rule + // 3). Same 16 staging bytes at the same logical (row, kk) + // slots, so bit-identity holds. + int8_t* dst = a_tile + row * kTileStride + col; + const int64_t* src = + reinterpret_cast(&a_chunk[i]); + *reinterpret_cast(dst) = src[0]; + *reinterpret_cast(dst + 8) = src[1]; + } else { + *reinterpret_cast(a_tile + row * kTileStride + col) = + a_chunk[i]; + } + } + } +#pragma unroll + for (int i = 0; i < kBLoads; ++i) { + const int slot = tid + i * kThreads; + if (slot < kBSlots) { + const int b_kc = slot / (kBlockN / 2); + const int b_j = slot - b_kc * (kBlockN / 2); + *reinterpret_cast(b_tile + b_kc * kBPlaneStride + b_j * 16) = + b_chunk[i]; + } + } + __syncthreads(); + + // Each wave consumes its 64 x kQuadN quadrant: kAccM * kAccN + // m16n16k32 MMACs per kk. Accumulation order: k0-outer over 64-K + // stages, kk-inner (kk = 0 then 32), matching the reference int32 + // accumulation. All fragment loads are issued before the MMAC burst + // (load-all-then-MMAC-all, fully unrolled like the q_b_proj winner). +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { +#pragma unroll + for (int i = 0; i < kAccM; ++i) { + load_fragment8(a_frag[i], + a_tile + (local_row + i * kTileM) * kTileStride + kk, + kTileStride, lane); + } +#pragma unroll + for (int j = 0; j < kAccN; ++j) { + load_fragment8_plane_stride( + b_frag[j], + b_tile + (kk >> 3) * kBPlaneStride + + (local_col + j * kTileN) * 8, + kBPlaneStride, lane); + } +#pragma unroll + for (int j = 0; j < kAccN; ++j) { +#pragma unroll + for (int i = 0; i < kAccM; ++i) { + du_mma_sync(acc[i][j], a_frag[i], b_frag[j], acc[i][j]); + } + } + } + + // Last stage only: prefetch the epilogue's per-row x_scale and + // per-column weight_scale values into registers (uniform branch: k0 is + // block-uniform, so the barrier below is reached by every thread; the + // row < m guard keeps the x_scale reads in bounds for any m tail). + if (k0 + kStageK >= k) { + const int r = lane & 15; + const int c4 = lane >> 4; +#pragma unroll + for (int i = 0; i < kAccM; ++i) { + const int row = base_row + i * kTileM + r; + xs_m[i] = (row < m) ? x_scale[row] : 0.0f; + } +#pragma unroll + for (int j = 0; j < kAccN; ++j) { + ws_m[j] = *reinterpret_cast( + weight_scale + base_col + j * kTileN + 4 * c4); + } + } + // Protect the LDS buffers from the next stage's cooperative overwrite. + __syncthreads(); + } + } + + // Coalesced direct-fragment epilogue: one 8-byte bf16 store per lane per + // fragment (100% store sector efficiency), scales from caller registers. + // Iteration 6 (kv_b_proj epilogue round): kRneOnly = true - plain RNE bf16 + // conversion without the exec-masked inf/NaN fixup (bit-identical for all + // finite inputs; see bf16_rne_u16). +#pragma unroll + for (int j = 0; j < kAccN; ++j) { +#pragma unroll + for (int i = 0; i < kAccM; ++i) { + store_prefill_fragment_coalesced_scaled( + acc[i][j], xs_m[i], ws_m[j], out, base_row + i * kTileM, + base_col + j * kTileN, m, n, lane); + } + } +} + +// --------------------------------------------------------------------------- +// Generic scalar fallback: one thread per output element. Exact int32 dot +// over K, then fused x_scale * weight_scale, then bf16 store. Handles every +// unmatched (m, n, k), including all small-M API cases (M=2, M=16), M tails +// and M < 128 with the same (K, N). kNMajorPack == true decodes the +// iteration-5 swizzled 64-k-stage layout +// packed[((k0*8+kc)*n+col)*8+b] == raw[kk*n+col] (kk = k0*64+kc*8+b) for the +// exact (k, n) == (2048, 2048); otherwise the weight is the raw [K, N] +// row-major identity layout. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + int32_t acc = 0; + const int8_t* a_row = x_q + static_cast(row) * k; + if constexpr (kNMajorPack) { + if (k == kPackSwizzleK2 && n == kPackSwizzleN2) { + // Tile-contiguous pack (iteration-5 kv_b_proj packing round, exact + // (512, 3584)): packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] == + // raw[kk*n + col] with kk = k0*64 + kc*8 + b, nt = col/128, + // row = col%128, so raw[kk][col] = packed[...]. + for (int kk = 0; kk < k; ++kk) { + const int k0 = kk >> 6; + const int kc = (kk >> 3) & 7; + const int b = kk & 7; + const int nt = col >> 7; + const int row = col & 127; + const int8_t* p = + weight + + ((((static_cast(k0) * (n >> 7) + nt) * 8 + kc) * 128 + + row) * 8) + + b; + acc += static_cast(a_row[kk]) * static_cast(*p); + } + } else { + // Swizzled 64-k-stage pack (iteration 5; exact (2048,2048) q_b_proj + // pair): packed[((k0*8 + kc)*n + col)*8 + b] == raw[kk*n + col] with + // kk = k0*64 + kc*8 + b, so raw[kk][col] = + // packed[((k0*8+kc)*n+col)*8+b]. + for (int kk = 0; kk < k; ++kk) { + const int k0 = kk >> 6; + const int kc = (kk >> 3) & 7; + const int b = kk & 7; + const int8_t* p = weight + + (((static_cast(k0) * 8 + kc) * n + col) * 8) + + b; + acc += static_cast(a_row[kk]) * static_cast(*p); + } + } + } else { + // identity [K, N] row-major: column col is strided by n. + const int8_t* b_col = weight + col; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Identity device-to-device weight packing (outside the timed region and +// outside Graph capture) for every (k, n) except the exact swizzled pairs +// (2048, 2048) q_b_proj and (512, 3584) kv_b_proj (which use the +// iteration-5 swizzled pack). The packed buffer keeps the same byte count +// K*N and the same allocated address, so the layout is graph-stable. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int64_t count) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +// Iteration 5 (packing round): swizzled 64-k-stage-major permutation for the +// exact swizzled (k, n) pair q_b_proj (2048, 2048) only (kv_b_proj +// (512, 3584) switched to the tile-contiguous w8a8_pack_tile_major_kernel +// in the kv_b_proj iteration-5 packing round): +// packed[((k0*8 + kc)*n + col)*8 + b] = raw[kk*n + col], +// kk = k0*64 + kc*8 + b (k0 = 64-k stage, kc = 8-byte sub-chunk, b = byte). +// Byte-wise (one thread per byte) so the permutation is trivially correct; +// runs once outside the timed region and outside Graph capture, keeping the +// same byte count K*N and the same graph-stable buffer address. The layout +// lets the 128x128 kernel stage each B tile as 16-byte plane chunks (one +// aligned int4 global read per chunk, perfectly coalesced) into the plane +// LDS tile [kc][n][8] whose fragment loads are lane-linear with zero bank +// conflicts. +__global__ __launch_bounds__(256) void w8a8_pack_swizzle_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(k) * n; + if (linear < total) { + const int kk = static_cast(linear / n); + const int col = static_cast(linear - static_cast(kk) * n); + const int k0 = kk >> 6; + const int kc = (kk >> 3) & 7; + const int b = kk & 7; + packed[(((static_cast(k0) * 8 + kc) * n + col) * 8) + b] = + raw[linear]; + } +} + +// Iteration 5 (kv_b_proj packing round): tile-contiguous 64-k-stage-major +// permutation for the exact (k, n) == (512, 3584) pair only: +// packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] = raw[kk*n + col], +// kk = k0*64 + kc*8 + b, nt = col/128, row = col%128. +// Every (64-k stage, 128-n tile) B tile is ONE contiguous 8192-B global +// region; the tile-family kernel stages it as 8 consecutive 1024-B +// wavefront streams (the k-stage-major swizzle read 8 streams 28,672 B +// apart). Byte-wise (one thread per byte) so the permutation is trivially +// correct; runs once outside the timed region and outside Graph capture, +// keeping the same byte count K*N and the same graph-stable buffer address. +// The (2048, 2048) q_b_proj pair keeps the k-stage-major swizzle +// (w8a8_pack_swizzle_kernel) byte-identical. +__global__ __launch_bounds__(256) void w8a8_pack_tile_major_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(k) * n; + if (linear < total) { + const int kk = static_cast(linear / n); + const int col = static_cast(linear - static_cast(kk) * n); + const int k0 = kk >> 6; + const int kc = (kk >> 3) & 7; + const int b = kk & 7; + const int nt = col >> 7; + const int row = col & 127; + packed[((((static_cast(k0) * (n >> 7) + nt) * 8 + kc) * 128 + + row) * + 8) + + b] = raw[linear]; + } +} + +__global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( + const float* __restrict__ raw, + float* __restrict__ packed, + int count) { + const int linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Stable host launch symbols consumed by csrc/bindings.cpp. +// Both launchers are pure dispatch: no allocation, no packing, no +// synchronization, no default-stream launch; they run on the caller-provided +// PyTorch stream and are CUDA-Graph safe. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; // no split-K: the GEMM does not use the workspace + (void)workspace_bytes; + auto* out_bf16 = reinterpret_cast(out); + + // Explicit dispatch. Large-M shapes with exact tiled geometry (both + // assigned shapes: M=4096, N in {2048, 3584}, K in {2048, 512}) take the + // native INT8 DUMMA tiled path; every other (m, n, k) - including small-M + // API cases (M=2, M=16) and M < 64 - takes the scalar fallback. The + // exact (k, n) == (2048, 2048) q_b_proj pair takes the iteration-4 128x128 + // tile with the iteration-5 swizzled-pack plane B layout; the exact + // (k, n) == (512, 3584) kv_b_proj pair takes the iteration-1 tile family + // (w8a8_dumma_prefill_tile_kernel, active config kActiveKvTile, also + // swizzled-pack B); every other (k, n) uses the byte-identical 128x64 + // kNMajorB=false arm. The scalar fallback decodes the swizzled pack for + // both swizzled (k, n) pairs. + if (k == kPackSwizzleK2 && n == kPackSwizzleN2 && m >= 64 && m % 64 == 0) { + // Exact kv_b_proj (512, 3584): 2-D macro-tile family. All four + // instantiations are referenced (so all are compiled); only + // kActiveKvTile is launched. grid (N/BN, M/BM), block (waves*64): + // 0 -> 64x64, 1 wavefront/block, (56, 64) = 3584 blocks + // 1 -> 64x128, 2 wavefronts/block, (28, 64) = 1792 blocks + // 2 -> 128x64, 4 wavefronts/block, (56, 32) = 1792 blocks + // 3 -> 64x128, 8 wavefronts/block, (28, 64) = 1792 blocks + const dim3 grid64x64(static_cast(n / 64), + static_cast(m / 64)); + const dim3 block64x64(1 * kWaveSize); + const dim3 grid64x128(static_cast(n / 128), + static_cast(m / 64)); + const dim3 block64x128(2 * kWaveSize); + const dim3 block64x128_8w(8 * kWaveSize); + const dim3 grid128x64(static_cast(n / 64), + static_cast(m / 128)); + const dim3 block128x64(4 * kWaveSize); + if (kActiveKvTile == 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tile_kernel<64, 64, 64>), + grid64x64, block64x64, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else if (kActiveKvTile == 1) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_tile_kernel), + grid64x128, block64x128, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else if (kActiveKvTile == 2) { + // config 2 (128x64) additionally needs m % 128 == 0; any (512, 3584) + // m not covered by the active tile falls back to the swizzle decode. + if (m % 128 == 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tile_kernel<128, 64, 32>), + grid128x64, block128x64, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + constexpr int kBlock = 256; + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + } else { + // config 3 (iteration-4 tile-shape round): 64x128 with 32x32 + // quadrants, 512 threads. m % 64 == 0 is guaranteed by the exact-shape + // guard above, so no tail fallback is needed. + // Iteration 16 (kv_b_proj load-placement round): kPrefetchNext = true + // -> the grouped A+B staging loads for stage s+1 are issued after the + // burst of stage s (before the protective trailing barrier), partially + // hiding each of the 8 staging L2 round trips behind the barrier + + // loop back-edge at unchanged 2 blocks/CU = 16 waves/CU (no register + // payload crosses the burst). Iteration 17 (kv_b_proj consolidation + // round): kALdsPad = 24 additionally applies the iteration-13 A-side + // LDS bank skew (88-byte A row stride -> zero-conflict A fragment + // reads; ds_write2_b64 halves at the staging writes) on top of the + // prefetch. LDS 13,312 -> 13,824 B/block (2 blocks = 27,648 B <= 64 + // KiB); 16 s_barriers/block unchanged. Configs 0-2 keep the defaults + // (kPrefetchNext = false, kALdsPad = 16) -> the exact iteration-7 + // issue-at-stage-top / stride-80 placement (byte-identical). + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_tile_kernel<64, 128, 32, 32, true, 24>), + grid64x128, block64x128_8w, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + } else if (m >= kBlockM && m % kBlockM == 0 && n % kBlockN == 0 && + k % kStageK == 0) { + const dim3 grid(static_cast(n / kBlockN), + static_cast(m / kBlockM)); + const dim3 block(kBlockThreads); + if (is_swizzled_pack(k, n)) { + // Exact q_b_proj (2048, 2048): iteration-4 tile-aspect arm, 128x128 + // tile with the iteration-5 swizzled-pack plane B layout; grid + // (N/128, M/128) = 512 blocks for M=4096. n == 2048 is divisible by + // 128 by the exact guard; the outer arm already guarantees + // m % 128 == 0, k % 64 == 0. + const dim3 grid128(static_cast(n / 128), + static_cast(m / 128)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_128x128_kernel), + grid128, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_128x64_kernel), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + } else { + constexpr int kBlock = 256; + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + if (is_swizzled_pack(k, n)) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + } + (void)hipGetLastError(); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kBlock = 256; + const dim3 block(kBlock); + // Pack: iteration-5 kv_b_proj packing round - tile-contiguous permutation + // for (512, 3584) kv_b_proj, k-stage-major swizzle for (2048, 2048) + // q_b_proj, identity device-to-device copy for every other (k, n). + // Packing runs once outside the timed region and keeps the same byte count + // (graph-stable addresses). + const int64_t weight_count = static_cast(k) * n; + const dim3 weight_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + if (is_tile_major_pack(k, n)) { + // Iteration 5 (kv_b_proj packing round): tile-contiguous pack for the + // exact (512, 3584) pair (one contiguous 8192-B region per (stage, + // n-tile) B tile); (2048, 2048) keeps the k-stage-major swizzle below. + hipLaunchKernelGGL(w8a8_pack_tile_major_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else if (is_swizzled_pack(k, n)) { + hipLaunchKernelGGL(w8a8_pack_swizzle_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else { + hipLaunchKernelGGL(w8a8_pack_identity_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, weight_count); + } + + const int64_t scale_count = n; + const dim3 scale_grid( + static_cast((scale_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, + scale_grid, block, 0, stream, + weight_scale, packed_weight_scale, n); + (void)hipGetLastError(); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/o_proj.hip new file mode 100644 index 00000000..49710bb2 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/o_proj.hip @@ -0,0 +1,1347 @@ +// @@variant shape=glm_tp8_o_proj_m4096 commit=486f3ec14b6d88c045aa35a5c85d1d44816c8c2d added=2026-08-31 +// median_us=759.9 p90_us=760.8 speedup=72.08 baseline_us=5.478e+04 +// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 +// @@variant shape=glm_tp8_o_proj_m4096 +// INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. +// +// Worker: worker_2 (physical GPU 2), assigned shape: +// glm_tp8_o_proj_m4096 : M=4096, N=6144, K=2048 +// +// Bootstrap strategy (iteration 1, correctness-first but a usable profiling +// baseline - a large-Prefill scalar K loop would not be): +// * Large-prefill path (exact (m, n, k) == (4096, 6144, 2048)): +// native INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 +// output tile per block; four wavefronts (256 threads); each wave owns +// a 64x32 quadrant built from eight m16n16k32 accumulator fragments; +// the block cooperatively vector-loads A[128,64] and B[64,64] into one +// single-buffered 64-K LDS stage (15,360 B total: 128*80 + 64*80 with +// 16 B padding per row for bank skew); two __syncthreads per stage; +// fused dot * x_scale[m] * weight_scale[n] epilogue stored directly as +// bf16 from the accumulator fragments. Grid dim3(96, 32) = 3072 blocks +// dwarfs the 120 CUs, so no split-K is needed. +// The weight is read in the identity [K, N] int8 layout (bootstrap: +// launch_pack_w8a8_weight is an identity device-to-device copy for +// every (k, n)). A is staged m-major into a_tile[m][k] and its eight +// fragment k-values per lane are contiguous, so A fragments use the +// direct 8-byte LDS fill (load_fragment8, the exact pattern validated +// on the worker-29 TP4 gate_up lineage iteration 13 and the TP8 hy3 +// o_proj lineage iteration 12). B is staged k-major into +// b_tile[k][n] and its fragments use the library du_load_matrix_sync +// row_major loader (8 per-byte LDS reads strided by ldm per lane - +// correct, direct, and simple; the packed n-major layout + 8-byte B +// fragment fill is a later-round optimization once the packed layout +// is introduced). +// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), +// including all small-M API cases (M=2, M=16) and M=3072 with the same +// (K, N) = (2048, 6144); the fallback reads the identity [K, N] layout +// for every (k, n). +// * launch_pack_w8a8_weight: identity device-to-device copy of the raw +// [K, N] weight and the [N] weight_scale for every (k, n) (bootstrap; +// later rounds may change only the pack HIP implementation and the +// matching GEMM interpretation). Packing never happens inside the +// timed GEMM. +// Iteration 5 (packing round): for the exact (k, n) == (2048, 6144) +// launch_pack_w8a8_weight now produces the n-major [N, K] layout +// packed[n*K + kk] = raw[kk*N + n] (same byte count and buffer -> +// graph-stable), and the exact-shape dispatch uses the new +// w8a8_dumma_256x64x64_packedb_kernel, which stages B n-major into +// bank-safe 72-byte LDS rows and fills each col_major B fragment with +// one contiguous 8-byte ds_read_b64 (zero per-byte reads, zero +// reassembly VALU); the scalar fallback decodes the n-major layout +// when (k, n) matches so the paired M=2/M=16 and M=3072 API shapes +// stay byte-exact. Every other (k, n) keeps the identity copy. +// Iteration 6 (epilogue round): the fused scale->bf16->coalesced +// store epilogue of w8a8_dumma_256x64x64_packedb_kernel is +// restructured to register-batch the per-row and per-column scales: +// each lane loads its 8 x_scale rows and its 2 weight_scale float4 +// columns ONCE (issued on the last K stage, before the final +// protective barrier) and the 16 per-fragment coalesced stores reuse +// them; the previous per-call epilogue re-loaded the scales inside +// each store (PMC vmem_read surplus ~153,600 wavefront loads above +// the pure staging floor). Output bits are identical (same transpose, +// same float(dot) * x_scale[row] * weight_scale[col] order, same bf16 +// rounding). No workspace/combine pass exists: the workspace +// argument stays unused and the epilogue is fully fused in-kernel. +// +// Rules honored: headers in the known-good order (hip/hip_runtime.h, +// hip/hip_bfloat16.h, du_mma.h); wavefront = 64 and blockDim a multiple of +// 64; int32 accumulation with float conversion only for +// dot * x_scale[m] * weight_scale[n]; hipLaunchKernelGGL on the +// caller-provided PyTorch stream; no stream-0 / synchronization / temporary +// allocation in gemm_out; no NVIDIA wmma/mma.sync/PTX/warp-32 masks; all +// barriers on paths reached by every thread in the block; explicit dispatch +// for the assigned shape with a scalar generic fallback for everything else. + +#include +#include +#include + +#include + +namespace { + +using namespace du::dumma; + +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 +constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kTargetM = 4096; +constexpr int kTargetN = 6144; +constexpr int kTargetK = 2048; +constexpr int kBlockM = 128; +constexpr int kBlockN = 64; +constexpr int kStageK = 64; +constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) +constexpr int kBStride = kBlockN + kBPad; +constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride +constexpr int kBlockThreads = 4 * kWaveSize; + +// Iteration-4 tile-shape constants: 256x64 output tile (M-major doubling of +// the accepted 128x64 tile) at 256 threads = 4 wavefronts, each wave owning a +// 128x32 quadrant (16 m16n16k32 int32 accumulators). +constexpr int kBlockM256 = 2 * kBlockM; + +// Iteration-5 packed-B constants: for the exact (k, n) == (2048, 6144) the +// weight is packed once (outside the timed region) into the n-major [N, K] +// layout packed[n*K + kk] = raw[kk*N + n]; B is then staged n-major into LDS +// rows of 72 B (64 + 8). 72 is an odd 8-byte count, so the 16 lanes of each +// ds_read_b64 fragment-load phase land on 16 distinct bank pairs (18r mod 32, +// r = 0..15 -> all 16 even dword starts: the conflict-free 4-phase floor), +// and 72 % 16 == 8, so staging stores use two int64 halves (ds_write2_b64) +// instead of one misaligned int4. +constexpr int kBPadPacked = 8; // 64 -> 72-byte n-major B LDS row stride +constexpr int kBStridePacked = kStageK + kBPadPacked; + +// Iteration-9 (A-stride round) constant: packedb A LDS row stride 64 -> 88 +// bytes, replacing the shared 80-byte kAStride for w8a8_dumma_256x64x64_ +// packedb_kernel only (the identity-layout kernels keep kAStride = 80). +// 88 = 22 dwords: 22r mod 32 is a permutation of the 16 even dword starts +// over r = 0..15, so the 16 lanes of each ds_read_b64 fragment-load phase +// land on 16 distinct bank pairs (the same conflict-free 4-phase floor the +// packed-B stride 72 already uses); the current 80-byte (20-dword) stride +// has 20r mod 32 with period 8, so rows r and r+8 alias onto one bank phase +// and EVERY A fragment read pays a 2-way bank conflict (PMC +// lds_bank_conflicts 20,447,232, deterministic across iterations 5..8). +// 88 % 16 == 8, so the 16-byte staging stores go through two int64 halves +// (ds_write2_b64, the exact accepted packed-B staging pattern) - the same +// instruction count as the current ds_write_b128. A[256,88] = 22,528 B + +// B[64,72] = 4,608 B = 27,136 B/block; 2 blocks/CU = 54,272 <= 65,536, so +// co-residency, VGPR, barrier count and the global-load pattern are all +// unchanged; only the physical LDS addresses move (same bytes stored, same +// bytes read, fragment values bit-identical). +constexpr int kAStridePacked = kStageK + 24; // 64 -> 88-byte A LDS row stride + +using bf16_t = hip_bfloat16; + +// Direct 8-byte LDS fragment fill for matrix_a on the m-major A tile. +// The m16n16k32 matrix_a row_major loader assigns x[i] = +// p[row*ldm + col_group*8 + i] (row = lane & 15, col_group = lane >> 4), +// i.e. eight consecutive k-values per lane; the library path makes the +// compiler emit per-byte ds_read_u8 loads plus a mask/OR reassembly chain +// per dword that is an arithmetic identity on the v_mmac A operand. Writing +// the same 8 bytes straight into the fragment storage keeps the operand bit +// pattern and the int32 accumulation identical and lets the compiler feed a +// single 8-byte LDS read to the v_mmac. This is the exact pattern validated +// on the worker-29 TP4 gate_up lineage (accepted iteration 13) and the TP8 +// hy3 o_proj lineage (accepted iteration 12). +template +__device__ __forceinline__ void load_fragment8( + Frag& frag, + const signed char* __restrict__ base, + int stride, + int lane) { + const int off = (lane & 15) * stride + ((lane >> 4) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// Coalesced fragment store. The m16n16k32 accumulator lane mapping +// (row = lane & 15, column group c4 = lane >> 4, frag.x[i] -> column +// c4 + 4*i) gives each lane four elements strided by 4 columns, so a direct +// store would be four 2-byte scalar stores per lane whose wavefront +// addresses touch each 32-B sector at 25% utilization. This epilogue +// transposes the 4-element groups within each 4-lane column group (lanes r, +// r+16, r+32, r+48 -- a 4x4 transpose, two 2x2 steps with shfl_xor 16 then +// 32), so lane (r, c4) ends up holding the four CONTIGUOUS columns +// 4*c4 .. 4*c4+3, converts them to bf16, packs 4 bf16 (8 B), and writes ONE +// 8-byte store per lane (100% store sector efficiency). This DTK lowers +// __shfl_xor to ds_bpermute at the block tail where the LDS pipe is idle. +// Only the int32 values are re-routed between lanes; the per-element scale +// multiply order (float(dot) * x_scale[row] * weight_scale[col]) and the +// bf16 rounding are unchanged, so the stored bits are identical to the +// direct per-element store. The row >= m guard is wavefront-uniform (all +// 64 lanes of a wave share the same 16-row window), so the shuffles never +// mix active and inactive lanes; base_col is a multiple of 16 and n*2 a +// multiple of 8, so the float4 weight_scale load and the 8-byte store are +// aligned. +template +__device__ __forceinline__ void store_prefill_fragment_coalesced( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int base_row, + int base_col, + int m, + int n, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M masking: padded rows never write (wavefront-uniform) + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 16, n*2 is a multiple of 8). + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + static_cast(row) * n + col0) = + packed; +} + +// Register-batched variant of the coalesced fragment store (iteration 6, +// epilogue round). Identical to store_prefill_fragment_coalesced in every +// output bit: same 4x4 shuffle transpose, same per-element multiply order +// float(dot) * x_scale[row] * weight_scale[col], same __float2bfloat16 +// rounding, same single 8-byte coalesced store. The only difference is that +// the per-row x_scale and per-column weight_scale values are passed in from +// registers: the caller loads each lane's 8 x_scale rows and 2 weight_scale +// float4 columns exactly once per wave (the per-call epilogue re-loaded the +// scales inside each of the 16 fragment stores), so the stored bits are +// unchanged while the epilogue vmem read count per lane per wave drops from +// up to 8 x_scale + 16 float4 loads to exactly 8 + 2 and the loads are +// independent of (and issue before) the shuffle/convert/store chains. +template +__device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( + const AccFragment& frag, + float xs, + const float4& ws, + bf16_t* __restrict__ out, + int base_row, + int base_col, + int m, + int n, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M masking: padded rows never write (wavefront-uniform) + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) owns columns base_col + 4*c4 .. +3 (8 B, 8-byte aligned: + // base_col is a multiple of 16, n*2 is a multiple of 8); ws is the float4 + // at weight_scale + base_col + 4*c4 preloaded by the caller. + const int col0 = base_col + 4 * c4; + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + static_cast(row) * n + col0) = + packed; +} + +// --------------------------------------------------------------------------- +// Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. +// Each wave owns a 64x32 quadrant (eight m16n16k32 int32 accumulators); the +// block cooperatively stages A[128,64] from x_q (row-major, stride k) and +// B[64,64] from the identity [K, N] weight into a single-buffered LDS +// buffer. Two barriers per stage: one after the cooperative load, one before +// the next stage overwrites LDS. Direct fragment epilogue (coalesced store). +// +// Stage layout (bootstrap, identity [K, N] weight): +// A[128, 80]: m-major, 512 int4s; thread t loads int4 #0 at +// a_tile[(t/4)*80 + (t%4)*16] from x_q[(m0 + t/4)*k + k0 + +// (t%4)*16] and int4 #1 at row +64 (rows 64..127). +// B[64, 80]: k-major, 256 int4s; thread t loads int4 at +// b_tile[(t/4)*80 + (t%4)*16] from +// weight[(k0 + t/4)*N + n0 + (t%4)*16] (16 consecutive n at +// fixed k; one int4 per thread). B fragments are loaded with +// the library row_major loader (ldm = 80): lane (r, c4) holds +// x[i] = b_tile[(kk + r + i)*80 + local_col + c4], i.e. eight +// k-rows strided by ldm - correct for the identity [K, N] +// layout; later rounds can switch to an n-major packed B with +// the 8-byte fragment fill. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row * 64; + const int local_col = wave_col * 32; + + // Single-buffered 64-K stage: A[128, 80] + B[64, 80] = 15,360 B/block + // (4 blocks/CU fit the 64 KiB LDS budget; the padded 80-byte strides are + // 16-byte-aligned and break the 64-byte LDS bank periodicity). + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + + // Cooperative staging: each thread owns one 16-byte vector in A rows + // [0,64), one in A rows [64,128), and one 16-byte vector of B. + // A[128,64] is 512 int4s; B[64,64] is 256 int4s. + const int vector_byte_offset = tid * static_cast(sizeof(int4)); + const int stage_row = vector_byte_offset / kStageK; // tid / 4, 0..63 + const int stage_col = vector_byte_offset - stage_row * kStageK; // (tid%4)*16 + const int b_row = tid >> 2; // k row of this thread's B int4 + const int b_n16 = (tid & 3) * 16; // n offset (stage-local) of B int4 + + for (int k0 = 0; k0 < k; k0 += kStageK) { + *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = + *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); + *reinterpret_cast(a_tile + + (stage_row + kBlockM / 2) * kAStride + stage_col) = + *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + + k0 + stage_col); + *reinterpret_cast(b_tile + b_row * kBStride + b_n16) = + *reinterpret_cast( + weight + (static_cast(k0 + b_row) * n + n0 + b_n16)); + __syncthreads(); + + // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + // A fragments: direct 8-byte LDS fill (m-major A tile; each lane's + // eight fragment k-values are contiguous). + load_fragment8(a_frag0, a_tile + local_row * kAStride + kk, kAStride, + lane); + load_fragment8(a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, + kAStride, lane); + // B fragments: library row_major loader on the k-major [K, N] LDS + // tile (bootstrap, identity layout; ldm = kBStride). + du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, + kBStride); + du_load_matrix_sync(b_frag1, + b_tile + kk * kBStride + local_col + kTileN, + kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + load_fragment8( + a_frag0, a_tile + (local_row + 2 * kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8( + a_frag1, a_tile + (local_row + 3 * kTileM) * kAStride + kk, + kAStride, lane); + du_mma_sync(acc20, a_frag0, b_frag0, acc20); + du_mma_sync(acc21, a_frag0, b_frag1, acc21); + du_mma_sync(acc30, a_frag1, b_frag0, acc30); + du_mma_sync(acc31, a_frag1, b_frag1, acc31); + } + + // Protect the LDS buffer from the next stage's cooperative overwrite. + __syncthreads(); + } + + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + // Coalesced epilogue store (one 8-byte store per lane per fragment; scales + // + bf16 conversion fused in-kernel, no workspace pass). + store_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, + base_row, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, + base_row, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, + lane); + store_prefill_fragment_coalesced(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc31, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + kTileN, + m, n, lane); +} + +// --------------------------------------------------------------------------- +// Iteration-3 pipeline-control variant: double-buffered K-stage of the same +// 128x64 tile / 64-K stage / fragment loads / epilogue as the accepted +// single-buffered kernel. Stage s+1's three global loads (2 A int4s + 1 B +// int4 per thread) are issued BEFORE the stage s MMAC burst and the loaded +// vectors stay in VGPR; the publish (ds_write into the idle buffer) happens +// AFTER the burst, so the compiler's vmcnt(0) wait sits behind 16 v_mmac + +// ~70 LDS reads instead of directly behind the load issue (the accepted +// single-buffered kernel waits immediately after 3 loads per wave per +// stage). One __syncthreads per stage (32/block vs 64/block: the barrier +// doubles as read-done for stage s+1 and write-done for stage s; the last +// stage needs no fence before the register epilogue). Costs: LDS +// 15,360 -> 30,720 B/block (A[2][128,80] + B[2][64,80]), residency +// 4 -> 2 blocks/CU = 8 waves/CU; VGPR cap at 2 blocks/CU is 128, so the +// ~12-VGPR loop-carried prefetch payload cannot change residency. Retained +// only if ISA/PMC evidence shows the vmcnt wait moved off the load->ds_write +// path (covered by the MMAC burst) and the measured median/p90 beat the +// 1162.15/1164.03 guard. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_db_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row * 64; + const int local_col = wave_col * 32; + + // Double-buffered 64-K stage: A[2][128,80] + B[2][64,80] = 30,720 B/block + // (2 blocks/CU fit the 64 KiB LDS budget; padded 80-byte strides and bank + // phases unchanged from the accepted kernel). + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + + // Cooperative staging indexing is identical to the accepted kernel: each + // thread owns two 16-byte A vectors (rows [0,64) and [64,128)) and one + // 16-byte B vector. + const int vector_byte_offset = tid * static_cast(sizeof(int4)); + const int stage_row = vector_byte_offset / kStageK; // tid / 4, 0..63 + const int stage_col = vector_byte_offset - stage_row * kStageK; // (tid%4)*16 + const int b_row = tid >> 2; // k row of this thread's B int4 + const int b_n16 = (tid & 3) * 16; // n offset (stage-local) of B int4 + + const int k_stages = k / kStageK; + + // Loop-carried prefetch payload (stage s+1 global vectors in VGPR). + int4 pa0, pa1, pb0; + + // Prologue: fetch stage 0 into registers, publish into buffer 0, fence. + { + const int k0 = 0; + pa0 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); + pa1 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + k0 + + stage_col); + pb0 = *reinterpret_cast( + weight + (static_cast(k0 + b_row) * n + n0 + b_n16)); + *reinterpret_cast(a_tile[0] + stage_row * kAStride + stage_col) = + pa0; + *reinterpret_cast(a_tile[0] + + (stage_row + kBlockM / 2) * kAStride + + stage_col) = pa1; + *reinterpret_cast(b_tile[0] + b_row * kBStride + b_n16) = pb0; + __syncthreads(); + } + + for (int s = 0; s < k_stages; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + + // Prefetch stage s+1: issue the global loads BEFORE the stage s MMAC + // burst; the loaded vectors stay in VGPR and the compiler's vmcnt(0) wait + // is inserted only at the publish below (covered by the burst). + if (s + 1 < k_stages) { + const int k0 = (s + 1) * kStageK; + pa0 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); + pa1 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + k0 + + stage_col); + pb0 = *reinterpret_cast( + weight + (static_cast(k0 + b_row) * n + n0 + b_n16)); + } + + // Consume stage s from the current buffer: bit-identical fragment values, + // MMAC sequence, int32 accumulation order and lane mapping to the + // accepted kernel (only the buffer base changes). +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + load_fragment8(a_frag0, a_tile[cur] + local_row * kAStride + kk, + kAStride, lane); + load_fragment8(a_frag1, a_tile[cur] + (local_row + kTileM) * kAStride + + kk, + kAStride, lane); + du_load_matrix_sync(b_frag0, + b_tile[cur] + kk * kBStride + local_col, kBStride); + du_load_matrix_sync(b_frag1, + b_tile[cur] + kk * kBStride + local_col + kTileN, + kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + load_fragment8( + a_frag0, a_tile[cur] + (local_row + 2 * kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8( + a_frag1, a_tile[cur] + (local_row + 3 * kTileM) * kAStride + kk, + kAStride, lane); + du_mma_sync(acc20, a_frag0, b_frag0, acc20); + du_mma_sync(acc21, a_frag0, b_frag1, acc21); + du_mma_sync(acc30, a_frag1, b_frag0, acc30); + du_mma_sync(acc31, a_frag1, b_frag1, acc31); + } + + // Publish stage s+1 into the idle buffer after the burst (vmcnt wait + // covered by the burst), then one fence: read-done for s+1 and + // write-done for s (buffer cur is free for the s+2 publish). + if (s + 1 < k_stages) { + *reinterpret_cast(a_tile[nxt] + stage_row * kAStride + + stage_col) = pa0; + *reinterpret_cast(a_tile[nxt] + + (stage_row + kBlockM / 2) * kAStride + + stage_col) = pa1; + *reinterpret_cast(b_tile[nxt] + b_row * kBStride + b_n16) = pb0; + __syncthreads(); + } + } + + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + store_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, + base_row, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, + base_row, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, + lane); + store_prefill_fragment_coalesced(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc31, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + kTileN, + m, n, lane); +} + +// --------------------------------------------------------------------------- +// Iteration-4 tile-shape variant: 256x64 output tile (M-major doubling of the +// accepted 128x64 tile) at 256 threads = 4 wavefronts, each wave owning a +// 128x32 quadrant (sixteen m16n16k32 int32 accumulators). This is the +// measured tile axis pushed one step further: per-FLOP B-fragment LDS reads +// scale as 1/(quadrant m-height) (the 64x128 flip's quadrants were 32x64 and +// lost +8% with a doubled per-FLOP B cost), so doubling the quadrant height +// to 128 HALVES the per-FLOP B ds_read_u8 + reassembly VALU - the measured +// LDS-read/wait bottleneck (PMC lds_wait ~= lds_instructions) - while keeping +// the same 64-column quadrant width (B library row_major loader ldm = 80). +// B reuse doubles (each staged B element is consumed by all 256 block rows), +// so the grid-level B re-read drops from M/128 = 32x to M/256 = 16x (staged B +// traffic 402.7 -> 201.3 MB/replay; A re-read stays N/64 = 96x = 805.3 MB; +// total 1.208 -> 1.007 GB, -17%). Single-buffered 64-K stage: A[256,80] + +// B[64,80] = 25,600 B/block -> 2 blocks/CU = 8 waves/CU, the SAME +// co-residency structure as the accepted double-buffered 128x64 kernel +// (2 blocks x 4 waves), so the comparison is not residency-confounded. +// Two __syncthreads per stage, direct 8-byte A fragment fills, fused +// coalesced scale->bf16 epilogue, row>=m guard. Grid dim3(96,16) = 1536 +// blocks = 12.8 blocks/CU, so the MxN output-tile grid still dwarfs the +// 120 CUs. Fragment values, MMAC sequence, int32 accumulation order +// (k0-outer 64, kk-inner 32, hardware k-order in m16n16k32), +// element-to-lane mapping, barrier count and the fused epilogue are +// bit-identical to the accepted kernels (only the tile height, the quadrant +// split and the staging vectorization change), so exact bf16 equality +// (mismatch 0 / max_abs_error 0.0) is expected without tolerance debate. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kBlockThreads) + void w8a8_dumma_256x64x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + // 4 waves: wave_row = (wave>>1)*128 covers rows 0..255 in two 128-row bands + // (waves 2w, 2w+1 share the band), wave_col picks the 32-col half. + const int wave_row = (wave >> 1) * (kBlockM256 / 2); + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM256; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row; + const int local_col = wave_col * 32; + + // Single-buffered 64-K stage: A[256, 80] + B[64, 80] = 25,600 B/block + // (2 blocks/CU fit the 64 KiB LDS budget; padded 80-byte strides and bank + // phases unchanged from the accepted kernels). + __shared__ __align__(16) int8_t a_tile[kBlockM256 * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kBStride]; + + DUFragment + a_frag0, a_frag1, a_frag2, a_frag3, a_frag4, a_frag5, a_frag6, a_frag7; + DUFragment + b_frag0, b_frag1; + // Sixteen m16n16k32 int32 accumulators (8 m-groups x 2 n-groups per wave). + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31, + acc40, acc41, acc50, acc51, acc60, acc61, acc70, acc71; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + du_fill_fragment(acc40, 0); + du_fill_fragment(acc41, 0); + du_fill_fragment(acc50, 0); + du_fill_fragment(acc51, 0); + du_fill_fragment(acc60, 0); + du_fill_fragment(acc61, 0); + du_fill_fragment(acc70, 0); + du_fill_fragment(acc71, 0); + + // Cooperative staging: A[256,64] is 1,024 int4s (4 per thread: rows + // [0,64), [64,128), [128,192), [192,256)) and B[64,64] is 256 int4s (one + // per thread), 5 loads per thread per stage. + const int a_row0 = tid >> 2; // 0..63 + const int a_col = (tid & 3) * 16; // 0,16,32,48 + + for (int k0 = 0; k0 < k; k0 += kStageK) { + *reinterpret_cast(a_tile + a_row0 * kAStride + a_col) = + *reinterpret_cast( + x_q + static_cast(m0 + a_row0) * k + k0 + a_col); + *reinterpret_cast(a_tile + + (a_row0 + kBlockM256 / 4) * kAStride + a_col) = + *reinterpret_cast( + x_q + static_cast(m0 + a_row0 + kBlockM256 / 4) * k + + k0 + a_col); + *reinterpret_cast(a_tile + + (a_row0 + kBlockM256 / 2) * kAStride + a_col) = + *reinterpret_cast( + x_q + static_cast(m0 + a_row0 + kBlockM256 / 2) * k + + k0 + a_col); + *reinterpret_cast(a_tile + + (a_row0 + 3 * kBlockM256 / 4) * kAStride + + a_col) = + *reinterpret_cast( + x_q + static_cast(m0 + a_row0 + 3 * kBlockM256 / 4) * k + + k0 + a_col); + const int b_row = tid >> 2; // k row of this thread's B int4, 0..63 + const int b_n16 = (tid & 3) * 16; + *reinterpret_cast(b_tile + b_row * kBStride + b_n16) = + *reinterpret_cast( + weight + (static_cast(k0 + b_row) * n + n0 + b_n16)); + __syncthreads(); + + // Each wave consumes its 128x32 quadrant: eight m-groups x two n-groups, + // 16 m16n16k32 MMACs per kk. The A fragment loads and the MMAC sequence + // keep the exact int32 accumulation order (k0-outer 64, kk-inner 32, + // hardware k-order in m16n16k32) and the element-to-lane mapping of the + // accepted kernels; only the number of m-groups per wave changes. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + load_fragment8(a_frag0, a_tile + local_row * kAStride + kk, kAStride, + lane); + load_fragment8(a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8(a_frag2, a_tile + (local_row + 2 * kTileM) * kAStride + + kk, + kAStride, lane); + load_fragment8(a_frag3, a_tile + (local_row + 3 * kTileM) * kAStride + + kk, + kAStride, lane); + du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, + kBStride); + du_load_matrix_sync(b_frag1, + b_tile + kk * kBStride + local_col + kTileN, + kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc20, a_frag2, b_frag0, acc20); + du_mma_sync(acc21, a_frag2, b_frag1, acc21); + du_mma_sync(acc30, a_frag3, b_frag0, acc30); + du_mma_sync(acc31, a_frag3, b_frag1, acc31); + load_fragment8(a_frag4, a_tile + (local_row + 4 * kTileM) * kAStride + + kk, + kAStride, lane); + load_fragment8(a_frag5, a_tile + (local_row + 5 * kTileM) * kAStride + + kk, + kAStride, lane); + load_fragment8(a_frag6, a_tile + (local_row + 6 * kTileM) * kAStride + + kk, + kAStride, lane); + load_fragment8(a_frag7, a_tile + (local_row + 7 * kTileM) * kAStride + + kk, + kAStride, lane); + du_mma_sync(acc40, a_frag4, b_frag0, acc40); + du_mma_sync(acc41, a_frag4, b_frag1, acc41); + du_mma_sync(acc50, a_frag5, b_frag0, acc50); + du_mma_sync(acc51, a_frag5, b_frag1, acc51); + du_mma_sync(acc60, a_frag6, b_frag0, acc60); + du_mma_sync(acc61, a_frag6, b_frag1, acc61); + du_mma_sync(acc70, a_frag7, b_frag0, acc70); + du_mma_sync(acc71, a_frag7, b_frag1, acc71); + } + + // Protect the LDS buffer from the next stage's cooperative overwrite. + __syncthreads(); + } + + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + store_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, + base_row, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, + base_row, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, + lane); + store_prefill_fragment_coalesced(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc31, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc40, x_scale, weight_scale, out, + base_row + 4 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc41, x_scale, weight_scale, out, + base_row + 4 * kTileM, base_col + kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc50, x_scale, weight_scale, out, + base_row + 5 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc51, x_scale, weight_scale, out, + base_row + 5 * kTileM, base_col + kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc60, x_scale, weight_scale, out, + base_row + 6 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc61, x_scale, weight_scale, out, + base_row + 6 * kTileM, base_col + kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc70, x_scale, weight_scale, out, + base_row + 7 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc71, x_scale, weight_scale, out, + base_row + 7 * kTileM, base_col + kTileN, + m, n, lane); +} + +// --------------------------------------------------------------------------- +// Iteration-5 packed-B variant of the accepted 256x64 kernel: the B operand +// switches to the one-time n-major [N, K] pack (packed[n*K + kk] = +// raw[kk*N + n], produced by launch_pack_w8a8_weight outside the timed +// region and Graph capture; same byte count and same buffer, so captured +// addresses are unchanged) and the B fragment path switches from the +// library row_major loader (8 per-byte ds_read_u8 strided by ldm per lane +// plus mask/OR reassembly VALU per dword - the measured LDS-read axis) to +// the lineage-validated 8-byte fragment fill (load_fragment8, one +// ds_read_b64 per col_major fragment, zero reassembly VALU) on the n-major +// staged tile. The m16n16k32 matrix_b hardware slot for lane +// (r = lane & 15, kq = lane >> 4) is B[k = kk + kq*8 + i][n = local_col + +// i*16 + r] - eight consecutive k bytes at fixed n - so staging B n-major +// (b_tile[n][k]) makes each lane's eight fragment bytes CONTIGUOUS in LDS: +// the fill reads the same elements into the same fragment slots with the +// same v_mmac inputs and the exact int32 accumulation order (k0-outer 64, +// kk-inner 32, hardware k-order), so exact bf16 equality (mismatch 0 / +// max_abs_error 0.0) is expected without tolerance debate. LDS bank +// safety: the n-major B tile uses the 72-byte row stride (odd 8-byte +// count) so the 16 lanes of each ds_read_b64 phase read 18r mod 32 +// (r = 0..15) distinct bank pairs - conflict-free; 72 % 16 == 8 forces the +// two-int64-half staging stores (ds_write2_b64). Iteration 9 gives the A +// tile the same conflict-free floor: the A row stride is 88 bytes (22 +// dwords; 22r mod 32 over r = 0..15 is a permutation of the 16 even dword +// starts), replacing the 80-byte stride's 2-way class (20r mod 32, period +// 8, rows r and r+8 alias), and A staging uses the same two-int64-half +// ds_write2_b64 pattern (88 % 16 == 8). Everything else is bit-identical +// to the accepted kernel: same bytes staged/read (only the physical LDS +// addresses move), same 128x32 quadrant per wave, same MMAC sequence, two +// __syncthreads per stage, fused coalesced scale->bf16 epilogue, row>=m +// guard, grid dim3(96,16) = 1536 blocks. LDS 22,528 (A) + 4,608 (B) = +// 27,136 B/block -> 2 blocks/CU = 8 waves/CU, the same co-residency +// structure as the accepted kernel. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kBlockThreads) + void w8a8_dumma_256x64x64_packedb_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_w, // [N, K] n-major (exact shape) + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + // 4 waves: wave_row = (wave>>1)*128 covers rows 0..255 in two 128-row bands + // (waves 2w, 2w+1 share the band), wave_col picks the 32-col half. + const int wave_row = (wave >> 1) * (kBlockM256 / 2); + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM256; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row; + const int local_col = wave_col * 32; + + // Single-buffered 64-K stage: A[256, 88] (iteration-9 stride) + B[64, 72] + // n-major = 27,136 B/block (2 blocks/CU fit the 64 KiB LDS budget). + __shared__ __align__(16) int8_t a_tile[kBlockM256 * kAStridePacked]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kBStridePacked]; + + DUFragment + a_frag0, a_frag1, a_frag2, a_frag3, a_frag4, a_frag5, a_frag6, a_frag7; + DUFragment + b_frag0, b_frag1; + // Sixteen m16n16k32 int32 accumulators (8 m-groups x 2 n-groups per wave). + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31, + acc40, acc41, acc50, acc51, acc60, acc61, acc70, acc71; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + du_fill_fragment(acc40, 0); + du_fill_fragment(acc41, 0); + du_fill_fragment(acc50, 0); + du_fill_fragment(acc51, 0); + du_fill_fragment(acc60, 0); + du_fill_fragment(acc61, 0); + du_fill_fragment(acc70, 0); + du_fill_fragment(acc71, 0); + + // Cooperative staging: A[256,64] is 1,024 int4s (4 per thread: rows + // [0,64), [64,128), [128,192), [192,256)) and B[64,64] is 256 int4s (one + // per thread), 5 loads per thread per stage. + const int a_row0 = tid >> 2; // 0..63 + const int a_col = (tid & 3) * 16; // 0,16,32,48 + // B staging is n-major: thread owns the 16-byte k-chunk at + // (n_row = tid >> 2, k16 = (tid & 3) * 16). + const int b_n_row = tid >> 2; // 0..63 + const int b_k16 = (tid & 3) * 16; // 0,16,32,48 + + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + // Iteration-6 (epilogue round) register-batched scales: each lane's 16 + // fragment stores cover the 8 m-groups (rows base_row + r + 16*i, + // r = lane & 15) x 2 n-groups (cols base_col + 4*c4 and + // base_col + 16 + 4*c4, c4 = lane >> 4), so each lane needs exactly 8 + // x_scale values and 2 weight_scale float4s. They are loaded ONCE on the + // last K stage, issued BEFORE the final protective barrier (uniform + // branch: k0 is block-uniform, so the barrier below is reached by every + // thread), so the vmem latency overlaps the barrier and the 16 + // shuffle/convert/store chains; the previous per-call epilogue re-loaded + // the scales inside each of the 16 fragment stores (PMC vmem_read + // surplus ~24-25 wavefront loads/wave above the 5/thread/stage staging + // floor matches ~8 x_scale + ~16 float4 re-loads per lane per wave). + float xs_m[8]; + float4 ws0, ws1; + + for (int k0 = 0; k0 < k; k0 += kStageK) { + const int4 av0 = *reinterpret_cast( + x_q + static_cast(m0 + a_row0) * k + k0 + a_col); + const int4 av1 = *reinterpret_cast( + x_q + static_cast(m0 + a_row0 + kBlockM256 / 4) * k + k0 + + a_col); + const int4 av2 = *reinterpret_cast( + x_q + static_cast(m0 + a_row0 + kBlockM256 / 2) * k + k0 + + a_col); + const int4 av3 = *reinterpret_cast( + x_q + static_cast(m0 + a_row0 + 3 * kBlockM256 / 4) * k + + k0 + a_col); + // A rows are 88 B apart (8 mod 16), so each 16-byte staging store is two + // int64 halves (ds_write2_b64, one instruction) - the exact accepted + // packed-B staging pattern below; a single int4 store would be + // misaligned on odd rows. + const int64_t* asrc0 = reinterpret_cast(&av0); + const int64_t* asrc1 = reinterpret_cast(&av1); + const int64_t* asrc2 = reinterpret_cast(&av2); + const int64_t* asrc3 = reinterpret_cast(&av3); + int64_t* adst0 = reinterpret_cast( + a_tile + a_row0 * kAStridePacked + a_col); + int64_t* adst1 = reinterpret_cast( + a_tile + (a_row0 + kBlockM256 / 4) * kAStridePacked + a_col); + int64_t* adst2 = reinterpret_cast( + a_tile + (a_row0 + kBlockM256 / 2) * kAStridePacked + a_col); + int64_t* adst3 = reinterpret_cast( + a_tile + (a_row0 + 3 * kBlockM256 / 4) * kAStridePacked + a_col); + adst0[0] = asrc0[0]; + adst0[1] = asrc0[1]; + adst1[0] = asrc1[0]; + adst1[1] = asrc1[1]; + adst2[0] = asrc2[0]; + adst2[1] = asrc2[1]; + adst3[0] = asrc3[0]; + adst3[1] = asrc3[1]; + // Stage B[64, 64] n-major from the packed [N, K] weight: 16 consecutive + // k bytes per thread at fixed n (one coalesced int4 global load). Row + // stride 72 is 8 mod 16, so the LDS store is two int64 halves + // (ds_write2_b64) - a single int4 store would be misaligned on odd rows. + const int4 bv = *reinterpret_cast( + packed_w + + (static_cast(n0 + b_n_row) * k + k0 + b_k16)); + const int64_t* bsrc64 = reinterpret_cast(&bv); + int64_t* bdst64 = reinterpret_cast( + b_tile + b_n_row * kBStridePacked + b_k16); + bdst64[0] = bsrc64[0]; + bdst64[1] = bsrc64[1]; + __syncthreads(); + + // Each wave consumes its 128x32 quadrant: eight m-groups x two n-groups, + // 16 m16n16k32 MMACs per kk. The A fragment loads and the MMAC sequence + // keep the exact int32 accumulation order (k0-outer 64, kk-inner 32, + // hardware k-order in m16n16k32) and the element-to-lane mapping of the + // accepted kernels; only the B fragment fill (one 8-byte contiguous LDS + // read per col_major fragment on the n-major tile) changes. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + load_fragment8(a_frag0, a_tile + local_row * kAStridePacked + kk, + kAStridePacked, lane); + load_fragment8(a_frag1, + a_tile + (local_row + kTileM) * kAStridePacked + kk, + kAStridePacked, lane); + load_fragment8(a_frag2, + a_tile + (local_row + 2 * kTileM) * kAStridePacked + kk, + kAStridePacked, lane); + load_fragment8(a_frag3, + a_tile + (local_row + 3 * kTileM) * kAStridePacked + kk, + kAStridePacked, lane); + load_fragment8(b_frag0, b_tile + local_col * kBStridePacked + kk, + kBStridePacked, lane); + load_fragment8(b_frag1, + b_tile + (local_col + kTileN) * kBStridePacked + kk, + kBStridePacked, lane); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc20, a_frag2, b_frag0, acc20); + du_mma_sync(acc21, a_frag2, b_frag1, acc21); + du_mma_sync(acc30, a_frag3, b_frag0, acc30); + du_mma_sync(acc31, a_frag3, b_frag1, acc31); + load_fragment8(a_frag4, + a_tile + (local_row + 4 * kTileM) * kAStridePacked + kk, + kAStridePacked, lane); + load_fragment8(a_frag5, + a_tile + (local_row + 5 * kTileM) * kAStridePacked + kk, + kAStridePacked, lane); + load_fragment8(a_frag6, + a_tile + (local_row + 6 * kTileM) * kAStridePacked + kk, + kAStridePacked, lane); + load_fragment8(a_frag7, + a_tile + (local_row + 7 * kTileM) * kAStridePacked + kk, + kAStridePacked, lane); + du_mma_sync(acc40, a_frag4, b_frag0, acc40); + du_mma_sync(acc41, a_frag4, b_frag1, acc41); + du_mma_sync(acc50, a_frag5, b_frag0, acc50); + du_mma_sync(acc51, a_frag5, b_frag1, acc51); + du_mma_sync(acc60, a_frag6, b_frag0, acc60); + du_mma_sync(acc61, a_frag6, b_frag1, acc61); + du_mma_sync(acc70, a_frag7, b_frag0, acc70); + du_mma_sync(acc71, a_frag7, b_frag1, acc71); + } + + // Last stage only: prefetch the epilogue's per-row x_scale and + // per-column weight_scale values into registers (uniform branch: k0 is + // block-uniform, so the barrier below is reached by every thread; the + // row < m guard keeps the x_scale reads in bounds for any m tail). + if (k0 + kStageK >= k) { + const int r = lane & 15; + const int c4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int row = base_row + i * kTileM + r; + xs_m[i] = (row < m) ? x_scale[row] : 0.0f; + } + ws0 = *reinterpret_cast(weight_scale + base_col + + 4 * c4); + ws1 = *reinterpret_cast(weight_scale + base_col + + kTileN + 4 * c4); + } + // Protect the LDS buffer from the next stage's cooperative overwrite. + __syncthreads(); + } + + store_prefill_fragment_coalesced_scaled(acc00, xs_m[0], ws0, out, base_row, + base_col, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc01, xs_m[0], ws1, out, base_row, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc10, xs_m[1], ws0, out, + base_row + kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced_scaled(acc11, xs_m[1], ws1, out, + base_row + kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc20, xs_m[2], ws0, out, + base_row + 2 * kTileM, base_col, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc21, xs_m[2], ws1, out, + base_row + 2 * kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc30, xs_m[3], ws0, out, + base_row + 3 * kTileM, base_col, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc31, xs_m[3], ws1, out, + base_row + 3 * kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc40, xs_m[4], ws0, out, + base_row + 4 * kTileM, base_col, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc41, xs_m[4], ws1, out, + base_row + 4 * kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc50, xs_m[5], ws0, out, + base_row + 5 * kTileM, base_col, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc51, xs_m[5], ws1, out, + base_row + 5 * kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc60, xs_m[6], ws0, out, + base_row + 6 * kTileM, base_col, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc61, xs_m[6], ws1, out, + base_row + 6 * kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc70, xs_m[7], ws0, out, + base_row + 7 * kTileM, base_col, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc71, xs_m[7], ws1, out, + base_row + 7 * kTileM, + base_col + kTileN, m, n, lane); +} + +// --------------------------------------------------------------------------- +// Generic scalar fallback: one thread per output element. Exact int32 dot +// over K, then fused x_scale * weight_scale, then bf16 store. Handles every +// unmatched (m, n, k), including all small-M API cases (M=2, M=16) and any +// M with the same (K, N). For the exact (k, n) == (2048, 6144) the weight is +// the n-major packed [N, K] layout (produced once out of the timed region), +// which the fallback decodes; every other (k, n) keeps the identity [K, N] +// layout. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + int32_t acc = 0; + const int8_t* a_row = x_q + static_cast(row) * k; + // Iteration 5: for the exact (k, n) == (2048, 6144) the packed_weight + // buffer holds the n-major transpose packed[col*k + kk] = raw[kk*n + col] + // (produced once out of the timed region), so the fallback decodes that + // layout there (keeps the paired M=2/M=16 and M=3072 API shapes with the + // same (K, N) byte-exact); every other (k, n) keeps the identity [K, N] + // row-major copy. + const bool packed_nmajor = (k == kTargetK && n == kTargetN); + const int8_t* b_ptr = + packed_nmajor ? weight + static_cast(col) * k : weight + col; + const int b_stride = packed_nmajor ? 1 : n; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * static_cast(b_ptr[0]); + b_ptr += b_stride; + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Identity packing (bootstrap; outside the timed region and outside Graph +// capture; the packed buffer keeps the same byte count K*N and the same +// allocated address, so the layout is graph-stable). launch_pack_w8a8_weight +// is an identity device-to-device copy for every (k, n): the packed weight +// is byte-identical to the raw [K, N] weight and the packed scale is +// byte-identical to the raw [N] scale. Later Parallel explore rounds may +// replace these kernels with a real layout transform and must then update +// the matching GEMM interpretation; the identity D2D copy is the generic +// fallback for unmatched (k, n). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int64_t count) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +// Exact-shape n-major pack (iteration 5): for (k, n) == (2048, 6144) the +// weight is transposed once, outside the timed region and out of Graph +// capture, into packed[col * k + kk] = raw[kk * n + col] so every col_major +// B fragment is 8 contiguous K bytes in the packed buffer (and in the +// staged n-major LDS tile). Same byte count and same buffer as the +// identity pack, so captured addresses are unchanged; correctness is +// checked against the raw logical [K, N] weight. One thread per output +// byte; runs once during weight prep. +__global__ __launch_bounds__(256) void w8a8_pack_nmajor_i8_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t total = static_cast(k) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; + idx += stride) { + const int col = static_cast(idx / k); + const int kk = static_cast(idx - static_cast(col) * k); + packed[idx] = raw[static_cast(kk) * n + col]; + } +} + +__global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( + const float* __restrict__ raw, + float* __restrict__ packed, + int count) { + const int linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Stable host launch symbols consumed by csrc/bindings.cpp. +// Both launchers are pure dispatch: no allocation, no packing, no +// synchronization, no default-stream launch; they run on the caller-provided +// PyTorch stream and are CUDA-Graph safe. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; // no split-K: the GEMM does not use the workspace + (void)workspace_bytes; + auto* out_bf16 = reinterpret_cast(out); + + // Explicit dispatch. The assigned shape (M=4096, N=6144, K=2048) takes the + // iteration-5 packed-B DUMMA 256x64 path (w8a8_dumma_256x64x64_packedb_kernel + // consumes the n-major packed [N, K] weight produced by + // launch_pack_w8a8_weight; the identity-layout kernels stay compiled for + // re-dispatch only together with an identity pack revert); every other + // (m, n, k) - including small-M API cases (M=2, M=16) and M=3072 with the + // same (K, N) - takes the scalar fallback, which decodes the n-major + // packed layout when (k, n) == (2048, 6144) and the identity [K, N] + // layout otherwise. + if (m == kTargetM && n == kTargetN && k == kTargetK) { + const dim3 grid(kTargetN / kBlockN, kTargetM / kBlockM256); + const dim3 block(kBlockThreads); + hipLaunchKernelGGL(w8a8_dumma_256x64x64_packedb_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + constexpr int kBlock = 256; + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + (void)hipGetLastError(); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Iteration 5: the exact (k, n) == (2048, 6144) weight is packed once, + // outside the timed region and out of Graph capture, into the n-major + // [N, K] layout packed[col*k + kk] = raw[kk*n + col] (same byte count and + // same buffer -> graph-stable addresses) so every col_major B fragment is + // 8 contiguous K bytes; every other (k, n) keeps the identity + // device-to-device copy (the scalar fallback decodes the n-major layout + // when (k, n) matches and the identity layout otherwise). + constexpr int kBlock = 256; + const dim3 block(kBlock); + + const int64_t weight_count = static_cast(k) * n; + const dim3 weight_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + if (k == kTargetK && n == kTargetN) { + hipLaunchKernelGGL(w8a8_pack_nmajor_i8_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else { + hipLaunchKernelGGL(w8a8_pack_identity_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, weight_count); + } + + const int64_t scale_count = n; + const dim3 scale_grid( + static_cast((scale_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, + scale_grid, block, 0, stream, + weight_scale, packed_weight_scale, n); + (void)hipGetLastError(); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/q_b_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/q_b_proj.hip new file mode 100644 index 00000000..72932e37 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/q_b_proj.hip @@ -0,0 +1,1019 @@ +// @@variant shape=glm_tp8_q_b_proj_m4096 commit=a9fcfc5e2ac1234af51849905834a850b5ef3448 added=2026-08-31 +// median_us=220.1 p90_us=220.5 speedup=65.35 baseline_us=1.439e+04 +// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 +// INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. +// +// Worker: worker_1 (physical GPU 1), assigned shapes: +// glm_tp8_q_b_proj_m4096 : M=4096, N=2048, K=2048 +// glm_tp8_kv_b_proj_m4096 : M=4096, N=3584, K=512 +// +// Logical operation (exact contract): +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// +// Iteration 3 (first valid HIP experiment; iterations 1-2 were killed in the +// agent infrastructure before any proposal). Accepted-best source digest +// a64b6250... with one bounded mechanism: n-major packed B + col_major B +// fragments for the exact glm_tp8_q_b_proj_m4096 (k,n) == (2048, 2048). +// * Large-prefill path (m >= 128 with exact 128x64x64 geometry): native +// INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 output tile +// per block; four wavefronts (256 threads); each wave owns a 64x32 +// quadrant built from eight m16n16k32 int32 accumulator fragments; the +// block cooperatively vector-loads A[128,64] and B[64,64] into one +// single-buffered LDS stage (15,360 B total: 128*80 + 64*80 with 16 B +// padding per row for bank skew); two __syncthreads per stage; fused +// dot * x_scale[m] * weight_scale[n] epilogue stored directly as bf16 +// from the accumulator fragments using the verified gfx928 int8 +// m16n16k32 lane mapping (row = lane & 15, col = (lane >> 4) + 4*i). +// A is read with the library du_load_matrix_sync row-major loader. +// B is layout-templated: +// - kNMajorB == true (exact (k,n) == (2048, 2048), q_b_proj): weight +// is packed once to n-major packed[n*K + k] == raw[k*N + n] outside +// the timed region; the stage stores B n-major (n row, k contiguous, +// 80 B row stride); each B fragment is one contiguous 8-byte LDS run +// per lane, loaded with the lineage-validated load_fragment8 (one +// ds_read2_b64) - kills the 32 ds_read_u8 byte-gathers + mask/OR +// reassembly VALU seen in the accepted-best ISA. +// - kNMajorB == false (any other shape, e.g. kv_b_proj (512, 3584)): +// raw [K, N] row-major identity layout, row-major B loader (the +// accepted-best path, byte-identical). +// Grid dim3(N/64, M/128) = 1024 blocks (q_b_proj) / 1792 blocks +// (kv_b_proj) dwarfs the 120 CUs, so no split-K is needed. +// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), +// including all small-M API cases (M=2, M=16), M tails, and M in +// (0, 128) with the same (K, N); it decodes the n-major packed layout +// for the exact (2048, 2048) and the identity [K, N] layout otherwise. +// * launch_pack_w8a8_weight: n-major device-to-device permutation for +// (k, n) == (2048, 2048), identity copy for every other (k, n) (both the +// int8 weight and the fp32 scale). Packing never happens inside the +// timed GEMM and keeps the same byte count (graph-stable addresses). +// +// Iteration 4 (tile-aspect round): the exact (k, n) == (2048, 2048) q_b_proj +// arm switches from the 128x64 tile to a 128x128 tile +// (w8a8_dumma_prefill_128x128_kernel). With the n-major packed B + 8-byte +// load_fragment8 B reads in place (iteration 3), the remaining LDS-read cost +// is dominated by the A-side row-major fragment loads (4 x ds_read2_b32 per +// kk per wave). A 128x128 tile with four 64x64 quadrants keeps the same +// 4-wave x 8-accumulator-fragment count per wave (8 -> 16 m16n16k32 MMACs +// per kk from 8 LDS fragment reads, i.e. per-MMAC LDS reads drop 6/8 -> 8/16 +// = -33%), and the A-side global re-read halves (A reuse 32 -> 16, A traffic +// 268 -> 134 MiB) while B reuse stays 32 (134 MiB), balancing the per-byte +// A:B traffic at 1:1 and cutting total tile traffic 402 -> 268 MiB (-33%). +// LDS grows 15,360 -> 20,480 B/block (still 2 blocks/CU = 40,960 B <= 64 +// KiB); the grid becomes (N/128, M/128) = 512 blocks (~4.3/CU). The 16 +// accumulator fragments cost ~+32 VGPR; arch VGPR is expected <= 128 so the +// accepted kernel's 2-blocks/CU co-residency is preserved (falsified if the +// exact code object shows > 128 VGPR -> 1 block/CU). kv_b_proj keeps the +// byte-identical 128x64 kNMajorB=false arm and the scalar fallback is +// untouched. +// +// Iteration 5 (packing round): the exact (k, n) == (2048, 2048) q_b_proj +// weight is re-packed once (outside timing/Graph, same byte count, same +// graph-stable buffer) into a swizzled 64-k-stage-major layout +// packed[((k0*8 + kc)*n + col)*8 + b] == raw[kk*n + col], +// kk = k0*64 + kc*8 + b (k0 = 64-k stage, kc = 8-byte sub-chunk, b = byte), +// and the 128x128 kernel stages each B tile as 16-byte plane chunks into an +// LDS plane layout b_tile[kc][n][8] (plane stride 1024 B, n stride 8 B). +// The exact code object of the accepted kernel shows every B fragment read +// is an 8-byte LDS access at (lane&15)*80 + (lane>>4)*8 (n-major stride-80 +// rows), and PMC counts 6,291,456 bank conflicts over 1,048,576 LDS slots +// (6.0/slot) with 4,078,494 LDS waits (3.9/slot): the loop is LDS-latency +// bound and the stride-80 rows alias bank phases every 8 rows (80*8 == 640 +// == 0 mod 128), so every fragment read conflicts ~4-6-way. In the plane +// layout a lane's 8 bytes sit at (lane>>4)*1024 + (lane&15)*8 relative to +// the fragment origin: 16 lanes per 128-B phase cover all 32 banks exactly +// once, i.e. B fragment reads become lane-linear and zero-conflict (4-cycle +// minimum for 512 B). Staging reads stay coalesced (each 16-B chunk is two +// n rows of one sub-chunk -> one aligned int4) and LDS writes are 2-way +// (the 16-B optimum). Same fragment operand bytes, same k0-outer/kk-inner +// int32 accumulation -> bit-identical results. The scalar fallback decodes +// the new pack for (2048,2048) only; kv_b_proj keeps the byte-identical +// 128x64 kNMajorB=false identity path. +// +// Iteration 6 (epilogue round): the fused dot * x_scale[m] * weight_scale[n] +// -> bf16 epilogue has been in-kernel since iteration 1 and the workspace is +// unused ((void)workspace; no split-K, no combine pass anywhere in the call +// chain), so the remaining epilogue inefficiency of the exact (2048, 2048) +// 128x128 arm is the STORE pattern: store_prefill_fragment issues one +// 2-byte bf16 store per lane per element - 64 scattered stores per wave per +// fragment set, each wavefront store touching 16 rows x 4 columns so every +// 32-B sector is only 25% utilized (PMC: 131,072 vmem_write_instructions = +// 512 blocks x 4 waves x 64). The new epilogue transposes each 16x16 +// fragment's 4-element groups inside the 4-lane column group (lanes r, +// r+16, r+32, r+48; two 2x2 steps with __shfl_xor 16 then 32 - the +// lineage-validated 4x4 register transpose accepted on the sibling TP8 +// workers 0/2, same DTK 26.04/gfx928), so lane (r, c4) owns the four +// CONTIGUOUS columns 4*c4 .. +3 and writes ONE 8-byte store per lane (100% +// store sector efficiency; vmem_write 131,072 -> 32,768). Only int32 +// values move between lanes: the per-element multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are +// unchanged, so stored bits are identical. The per-row x_scale (4 rows per +// lane) and per-column weight_scale float4s (4 per lane) are additionally +// register-batched on the last K stage (before the final protective +// __syncthreads), so the epilogue is pure compute + 16 coalesced stores with +// no interleaved vmem loads (vmem_read ~413,696 -> ~278k). kv_b_proj keeps +// the byte-identical 128x64 kNMajorB=false arm (old per-element epilogue) +// and the generic scalar fallback is untouched. +// +// Iteration 7 (compute-pipeline round): the exact (k, n) == (2048, 2048) +// 128x128 arm loads its four A fragments per kk with the lineage-validated +// direct load_fragment8 (one ds_read2_b64 straight into the v_mmac operand) +// instead of the library du_load_matrix_sync row_major loader. The exact +// code object shows the library loader lowers to 8 x ds_read2_b32 plus a +// redundant byte-reassembly chain per fragment (~7 VALU: v_and 0xff00 / +// 0xff0000 / 0xff000000 + v_or_b32_sdwa + v_or3, ~50-56 VALU per stage per +// wave) between the LDS read and the first MMAC, and PMC counts 3,009,148 +// LDS waits (2.3/slot) against a latency-bound loop running at 1 block/CU +// (160 VGPR). du_mma.hpp defines matrix_a row_major int8 as x[i] = +// p[(lane&15)*ldm + (lane>>4)*8 + i] (8 consecutive bytes, memory order) +// and du_mma_sync passes reinterpret(x) unchanged to v_mmac, so +// load_fragment8 produces byte-identical operand values and the int32 +// accumulation is bit-identical; only the redundant VALU reassembly (and its +// lgkmcnt wait states) is removed from the LDS->MMAC critical path, letting +// the 32-MMAC burst issue back-to-back. Expected PMC: valu_instructions +// 8.66M -> ~4.8-5.6M (-35..-45%), lds_instructions 1.31M -> ~1.05M (-20%, +// 8 A ds_read2_b32 + 4 B ds_read2_b64 -> 8 ds_read2_b64 per stage), +// lds_wait_instructions 3.01M -> ~2.2-2.5M, lds_bank_conflicts +// approximately unchanged (A rows keep the 2-way stride-80 aliasing), +// vmem_read/vmem_write and the 32 v_mmac/stage unchanged. kv_b_proj keeps +// the byte-identical 128x64 kNMajorB=false arm (unchanged lowering) and the +// generic scalar fallback is untouched. + +#include +#include +#include + +#include + +namespace { + +using namespace du::dumma; + +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 +constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kBlockM = 128; +constexpr int kBlockN = 64; +constexpr int kStageK = 64; +constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) +constexpr int kBStride = kBlockN + kBPad; +constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride +constexpr int kBlockThreads = 4 * kWaveSize; + +// Exact (k, n) pair whose weight buffer is packed n-major +// (packed[n*K + k] == raw[k*N + n]): glm_tp8_q_b_proj_m4096 only. +constexpr int kPackNMajorK = 2048; +constexpr int kPackNMajorN = 2048; + +using bf16_t = hip_bfloat16; + +// --------------------------------------------------------------------------- +// Direct fragment epilogue for one m16n16k32 accumulator fragment. +// Verified gfx928 int8 m16n16k32 accumulator ownership (matches +// du_store_matrix_sync): lane & 15 selects the row, lane >> 4 selects +// col % 4, and x[i] maps to columns col%4 + 4*i. The scale multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are +// identical to the harness reference ((dot.float() * x_scale) * ws.T then +// .to(bfloat16)), so stored bf16 bits match exactly. +// --------------------------------------------------------------------------- +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int base_row, + int base_col, + int m, + int n, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M masking: padded rows never write + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// Iteration 6 (epilogue round): register-batched COALESCED direct-fragment +// epilogue for the exact (2048, 2048) 128x128 arm. The per-element variant +// (store_prefill_fragment) above issues one 2-byte bf16 store per lane per +// element: 64 scattered stores per wave per fragment set, each touching 16 +// rows x 4 columns, i.e. every 32-B store sector only 25% utilized. This +// epilogue transposes the 4-element groups within each 4-lane column group +// (lanes r, r+16, r+32, r+48 - a 4x4 transpose, two 2x2 steps with +// shfl_xor 16 then 32; the lineage-validated pattern accepted on the sibling +// TP8 workers 0/2 on this DTK, which lowers __shfl_xor to ds_bpermute at the +// block tail where the LDS pipe is idle), so lane (r, c4) ends up holding +// the four CONTIGUOUS columns 4*c4 .. 4*c4+3, converts them to bf16, packs +// 4 bf16 (8 B) and writes ONE 8-byte store per lane (100% store sector +// efficiency). Only the int32 values are re-routed between lanes; the +// per-element scale multiply order (float(dot) * x_scale[row] * +// weight_scale[col]) and the __float2bfloat16 rounding are unchanged, so +// the stored bits are identical to the per-element store. xs/ws come from +// caller registers (batched per wave on the last K stage), so the epilogue +// issues no vmem loads. The row >= m guard is wavefront-uniform (all 64 +// lanes of a wave share the same 16-row window), so the shuffles never mix +// active and inactive lanes; base_col is a multiple of 64 and n*2 a +// multiple of 8, so the 8-byte store is aligned. +// --------------------------------------------------------------------------- +template +__device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( + const AccFragment& frag, + float xs, + const float4& ws, + bf16_t* __restrict__ out, + int base_row, + int base_col, + int m, + int n, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M masking: padded rows never write (wavefront-uniform) + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 16, n*2 is a multiple of 8); ws is + // the float4 at weight_scale + base_col + 4*c4 preloaded by the caller. + const int col0 = base_col + 4 * c4; + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + static_cast(row) * n + col0) = + packed; +} + +// --------------------------------------------------------------------------- +// Direct 8-byte LDS fragment load (lineage-validated on gfx928/DTK 26.04 for +// m16n16k32 int8 fragments). du_load_matrix_sync's int8 loaders assign +// x[0..7] = 8 consecutive bytes at (lane & 15) * ldm + ((lane >> 4) << 3) for +// matrix_b col_major, but the compiler lowers that to per-byte ds_read_u8 + +// mask/OR reassembly VALU. Writing the same 8 bytes directly into the +// fragment storage keeps the operand bit pattern identical (exact int32 +// accumulation unchanged) and lets the compiler emit one ds_read2_b64 per +// fragment straight into the v_mmac operand. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void load_fragment8( + DUFragmentBase& frag, + const signed char* __restrict__ base, + int stride, + int lane) { + const int off = (lane & 15) * stride + ((lane >> 4) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// --------------------------------------------------------------------------- +// Iteration 5 (packing round): 8-byte fragment load for the swizzled plane +// LDS layout b_tile[kc][n][8] (kc = 8-byte k sub-chunk, n row stride 8 B, +// plane stride 1024 B). Lane l reads the same 8 consecutive k bytes as +// load_fragment8 (n row = lane&15, k chunk = lane>>4) but from +// (lane >> 4) * 1024 + (lane & 15) * 8 +// relative to the fragment origin, so the 64 lanes hit 16 distinct bank +// pairs per 128-B phase (all 32 banks once) -> zero bank conflicts, the +// 4-cycle minimum for a 512-B fragment. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void load_fragment8_plane( + DUFragmentBase& frag, + const signed char* __restrict__ base, + int lane) { + const int off = ((lane >> 4) << 10) + ((lane & 15) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// --------------------------------------------------------------------------- +// Compile-time B-fragment layout selector. row_major / col_major are tag +// types in du::dumma, so the layout template argument must be selected as a +// type (a conditional expression over type names is not a valid template +// argument). kNMajorB == true (n-major packed B, exact (2048, 2048)) -> +// col_major fragments (8 contiguous k bytes per lane, load_fragment8); +// kNMajorB == false (raw [K, N] row-major B) -> row_major fragments (the +// accepted-best loader). +// --------------------------------------------------------------------------- +template +struct b_frag_layout { + using type = row_major; +}; +template <> +struct b_frag_layout { + using type = col_major; +}; + +// --------------------------------------------------------------------------- +// Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. +// Each wave owns a 64x32 quadrant (eight m16n16k32 int32 accumulators); the +// block cooperatively stages A[128,64] from x_q (row-major, stride k) and +// B[64,64] into a single-buffered LDS stage. Two barriers per stage: one +// after the cooperative load, one before the next stage overwrites LDS. +// A fragments use the library du_load_matrix_sync row-major loader. B is +// templated on its staged layout: +// * kNMajorB == true: B is packed n-major (packed[n*K + k]) and staged +// n-major (n row, 80 B stride, k contiguous); each B fragment is 8 +// contiguous k bytes per lane, loaded by load_fragment8 (one 8-byte LDS +// read per fragment), eliminating the ds_read_u8 byte-gather + VALU +// reassembly of the accepted-best kernel. +// * kNMajorB == false: B stays raw [K, N] row-major, staged k-major with +// row-major library fragments (accepted-best path, unchanged). +// Dispatch guarantees m % 128 == 0, n % 64 == 0, k % 64 == 0, so every +// global load/store is in-bounds and 16-byte aligned. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row * 64; + const int local_col = wave_col * 32; + + // Single-buffered stage: A[128, 80] + B[64, 80] = 15,360 B/block + // (4 blocks/CU fit the 64 KiB LDS budget; the padded 80-byte strides are + // 16-byte-aligned and break the 64-byte LDS bank periodicity). + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment::type> + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + + // Cooperative staging: A[128,64] is 512 int4s (two per thread), B[64,64] + // is 256 int4s (one per thread). + // A: thread tid owns row (tid*16)/64 and 16-byte column group + // (tid*16)%64 for rows [0,64) and [64,128). + // B (kNMajorB == false): raw [K, N] row-major; thread tid owns K row + // tid>>2 and the 16-byte N-column group (tid&3)*16 (16 consecutive N + // values at a fixed K row -> one aligned int4 in global and in LDS). + // B (kNMajorB == true): packed n-major; thread tid owns N row tid>>2 and + // the 16-byte K-column group (tid&3)*16 (16 consecutive K values at a + // fixed N row -> one aligned int4 in global and in the n-major LDS). + const int vector_byte_offset = tid * static_cast(sizeof(int4)); + const int stage_row = vector_byte_offset / kStageK; // 0..63 + const int stage_col = vector_byte_offset - stage_row * kStageK; + const int b_k = tid >> 2; // K row (false) / N row (true) + const int b_nc = (tid & 3) * 16; // N group (false) / K group (true) + + for (int k0 = 0; k0 < k; k0 += kStageK) { + *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = + *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); + *reinterpret_cast(a_tile + + (stage_row + kBlockM / 2) * kAStride + stage_col) = + *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + + k0 + stage_col); + if constexpr (kNMajorB) { + // n-major packed B: 16 consecutive k bytes of the (n0 + b_k) row. + *reinterpret_cast(b_tile + b_k * kBStride + b_nc) = + *reinterpret_cast( + weight + static_cast(n0 + b_k) * k + k0 + b_nc); + } else { + // raw [K, N] row-major B: 16 consecutive n bytes of the (k0 + b_k) row. + *reinterpret_cast(b_tile + b_k * kBStride + b_nc) = + *reinterpret_cast( + weight + static_cast(k0 + b_k) * n + n0 + b_nc); + } + __syncthreads(); + + // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. + // Accumulation order: k0-outer over 64-K stages, kk-inner (kk=0 then + // kk=32), matching the reference int32 accumulation. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + if constexpr (kNMajorB) { + // n-major LDS tile: n rows stride kBStride, k contiguous within a + // row; col_major fragment slots are 8 consecutive k bytes at + // (lane & 15) * kBStride + ((lane >> 4) << 3) relative to the + // fragment origin -> one ds_read2_b64 per fragment per lane. + load_fragment8(b_frag0, b_tile + local_col * kBStride + kk, kBStride, + lane); + load_fragment8(b_frag1, + b_tile + (local_col + kTileN) * kBStride + kk, + kBStride, lane); + } else { + du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, + kBStride); + du_load_matrix_sync(b_frag1, + b_tile + kk * kBStride + local_col + kTileN, + kBStride); + } + du_load_matrix_sync(a_frag0, a_tile + local_row * kAStride + kk, + kAStride); + du_load_matrix_sync(a_frag1, + a_tile + (local_row + kTileM) * kAStride + kk, + kAStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_load_matrix_sync(a_frag0, + a_tile + (local_row + 2 * kTileM) * kAStride + kk, + kAStride); + du_load_matrix_sync(a_frag1, + a_tile + (local_row + 3 * kTileM) * kAStride + kk, + kAStride); + du_mma_sync(acc20, a_frag0, b_frag0, acc20); + du_mma_sync(acc21, a_frag0, b_frag1, acc21); + du_mma_sync(acc30, a_frag1, b_frag0, acc30); + du_mma_sync(acc31, a_frag1, b_frag1, acc31); + } + + // Protect the LDS buffer from the next stage's cooperative overwrite. + __syncthreads(); + } + + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + store_prefill_fragment(acc00, x_scale, weight_scale, out, base_row, + base_col, m, n, lane); + store_prefill_fragment(acc01, x_scale, weight_scale, out, base_row, + base_col + kTileN, m, n, lane); + store_prefill_fragment(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_prefill_fragment(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, lane); + store_prefill_fragment(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, m, n, + lane); + store_prefill_fragment(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, lane); + store_prefill_fragment(acc31, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + kTileN, m, n, + lane); +} + +// --------------------------------------------------------------------------- +// Iteration 4 (tile-aspect round) + iteration 5 (packing round): 128x128 +// tile for the exact q_b_proj (k, n) == (2048, 2048) swizzled-pack B arm. +// Four wavefronts of 64 lanes; each wave owns a 64x64 quadrant (four A x +// four B m16n16k32 fragments = 16 MMACs per kk, 16 int32 accumulator +// fragments). The block cooperatively stages A[128,64] (2 int4/thread, +// 80-byte-strided rows) and B[128,64] from the iteration-5 swizzled pack (2 +// int4/thread, 16-byte plane chunks) into a single-buffered LDS stage: +// A[128,80] (10,240 B) + B plane tile [8][128][8] (8,192 B) = 18,432 +// B/block (2 blocks/CU = 36,864 B <= 64 KiB); two __syncthreads per stage. +// A fragments use the library du_load_matrix_sync row-major loader +// (4 x ds_read2_b32 per kk per wave), B fragments use load_fragment8_plane +// (lane-linear 8-byte reads, zero LDS bank conflicts). A-side global +// re-read halves (A reuse 32 -> 16, 268 -> 134 MiB) while B reuse stays 32 +// (134 MiB): per-byte A:B tile traffic balances 1:1, total 402 -> 268 MiB +// (-33%). Grid dim3(N/128, M/128) = 512 blocks (~4.3/CU). Dispatch +// guarantees m % 128 == 0, n == 2048, k % 64 == 0, so every global +// load/store is in-bounds and 16-byte aligned. Accumulation order stays +// k0-outer / kk-inner with the same element-to-slot fragment mapping, so the +// int32 accumulation is bit-identical to the accepted kernel. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + constexpr int kBlockM128 = 128; + constexpr int kBlockN128 = 128; + constexpr int kTileStride = kStageK + kBPad; // 64 -> 80-byte row stride + constexpr int kTileBytes = kBlockM128 * kTileStride; // 10,240 B per tile + // Iteration 5 (packing round): B tile is the swizzled plane layout + // [kc][n][8] (8 planes of 128 n rows x 8 B, plane stride 1024 B) so each + // lane's 8-byte fragment chunk is lane-linear and bank-conflict-free. + constexpr int kBTileBytes = 8 * kBlockN128 * 8; // 8,192 B/block + constexpr int kBTilePlane = 1024; // 128 n rows x 8 B + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM128; + const int n0 = static_cast(blockIdx.x) * kBlockN128; + const int local_row = wave_row * (kBlockM128 / 2); // 0 or 64 + const int local_col = wave_col * (kBlockN128 / 2); // 0 or 64 + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + + // Iteration 6 (epilogue round): per-lane register-batched scales. Each + // lane owns 4 rows (base_row + 16*i + (lane & 15), i = 0..3) and 4 + // weight_scale float4s (base_col + 16*j + 4*(lane >> 4), j = 0..3) across + // its 16 fragments; they are loaded exactly once on the last K stage + // before the final protective barrier, so the coalesced epilogue below is + // pure compute + 16 eight-byte stores with no interleaved vmem loads. + float xs_m[4]; + float4 ws_m[4]; + + // Single-buffered stage: A[128,80] (10,240 B) + B plane tile [8][128][8] + // (8,192 B) = 18,432 B/block (2 blocks/CU = 36,864 B <= 64 KiB; the A rows + // keep the padded 80-byte stride, the B tile is plane-swizzled so fragment + // reads are lane-linear with zero bank conflicts). + __shared__ __align__(16) int8_t a_tile[kTileBytes]; + __shared__ __align__(16) int8_t b_tile[kBTileBytes]; + + DUFragment + a_frag0, a_frag1, a_frag2, a_frag3; + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; + DUFragment + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13, + acc20, acc21, acc22, acc23, acc30, acc31, acc32, acc33; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc22, 0); + du_fill_fragment(acc23, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + du_fill_fragment(acc32, 0); + du_fill_fragment(acc33, 0); + + // Cooperative staging: A[128,64] is 512 int4s (two per thread) and + // B[128,64] from the iteration-5 swizzled pack is 512 int4s (two per + // thread). Thread tid owns linear int4 slots tid and tid + 256; for A a + // slot maps to row linear >> 2 and the 16-byte column group + // (linear & 3) * 16 (an M row of x_q, one aligned int4 in global and in + // LDS); for B a slot maps to plane linear >> 6 and n-pair linear & 63 + // (two n rows of one 8-byte k sub-chunk, one aligned int4 in the pack and + // in the plane LDS tile). + for (int k0 = 0; k0 < k; k0 += kStageK) { + const int kstage = k0 >> 6; // 64-k group index inside the swizzled pack +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int linear = tid + i * kBlockThreads; + const int row = linear >> 2; // 0..127 + const int col = (linear & 3) * 16; // 0/16/32/48 + *reinterpret_cast(a_tile + row * kTileStride + col) = + *reinterpret_cast( + x_q + static_cast(m0 + row) * k + k0 + col); + // Swizzled pack: slot linear -> (kc = linear >> 6, j = linear & 63); + // the 16-byte chunk is two n rows (n0 + 2j, n0 + 2j + 1) of the 8-byte + // k sub-chunk kc of stage kstage: one aligned int4 in global (two + // consecutive n rows of one plane) and one aligned int4 in the plane + // LDS tile (plane kc, rows 2j..2j+1 -> contiguous 16 B). + const int b_kc = linear >> 6; // 0..7 + const int b_j = linear & 63; // 0..63 (n-pair within the plane) + *reinterpret_cast(b_tile + b_kc * kBTilePlane + b_j * 16) = + *reinterpret_cast( + weight + (static_cast(kstage * 8 + b_kc) * n + + (n0 + 2 * b_j)) * 8); + } + __syncthreads(); + + // Each wave consumes its 64x64 quadrant: sixteen m16n16k32 MMACs per kk. + // Accumulation order: k0-outer over 64-K stages, kk-inner (kk=0 then + // kk=32), matching the reference int32 accumulation. All eight fragment + // loads are issued before the sixteen MMACs (load-all-then-MMAC-all). + // + // Iteration 7 (compute-pipeline round): the A fragments are loaded with + // the same direct 8-byte LDS read as B (load_fragment8) instead of the + // library du_load_matrix_sync row_major loader. The library's int8 + // matrix_a row_major loader assigns x[i] = p[(lane&15)*ldm + + // (lane>>4)*8 + i] (8 consecutive bytes, memory order) and du_mma_sync + // feeds reinterpret(x) straight into v_mmac, but on this DTK the + // loader lowers to 8 x ds_read2_b32 + a redundant per-dword byte + // reassembly chain (~7 VALU: v_and 0xff00/0xff0000/0xff000000 + + // v_or_b32_sdwa + v_or3 per second dword) sitting between the LDS read + // and the MMAC issue (exact code object, both the 128x128 and the 128x64 + // symbols). load_fragment8 fills the same x[0..7] with one 64-bit + // little-endian write, so the operand bytes are bit-identical and the + // compiler emits one ds_read2_b64 straight into the v_mmac operand (the + // same lineage as the B side since iteration 3/5): the LDS->MMAC + // critical path shortens by ~50 VALU + their lgkmcnt wait states per + // stage, and the 32-MMAC burst can issue back-to-back after the barrier. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + load_fragment8(a_frag0, a_tile + local_row * kTileStride + kk, + kTileStride, lane); + load_fragment8(a_frag1, + a_tile + (local_row + kTileM) * kTileStride + kk, + kTileStride, lane); + load_fragment8(a_frag2, + a_tile + (local_row + 2 * kTileM) * kTileStride + kk, + kTileStride, lane); + load_fragment8(a_frag3, + a_tile + (local_row + 3 * kTileM) * kTileStride + kk, + kTileStride, lane); + // B fragments in the plane layout: fragment (f, kk) origin is + // b_tile + (kk >> 3) * 1024 + (local_col + 16f) * 8 and lane l reads + // the 8 bytes at (lane >> 4) * 1024 + (lane & 15) * 8 (lane-linear, + // zero bank conflicts); operand bytes are identical to the accepted + // n-major load_fragment8, so the int32 accumulation is bit-identical. + const int b_kc0 = kk >> 3; // 0 (kk == 0) or 4 (kk == 32) + load_fragment8_plane(b_frag0, + b_tile + b_kc0 * kBTilePlane + local_col * 8, + lane); + load_fragment8_plane(b_frag1, + b_tile + b_kc0 * kBTilePlane + + (local_col + kTileN) * 8, + lane); + load_fragment8_plane(b_frag2, + b_tile + b_kc0 * kBTilePlane + + (local_col + 2 * kTileN) * 8, + lane); + load_fragment8_plane(b_frag3, + b_tile + b_kc0 * kBTilePlane + + (local_col + 3 * kTileN) * 8, + lane); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + du_mma_sync(acc20, a_frag2, b_frag0, acc20); + du_mma_sync(acc21, a_frag2, b_frag1, acc21); + du_mma_sync(acc22, a_frag2, b_frag2, acc22); + du_mma_sync(acc23, a_frag2, b_frag3, acc23); + du_mma_sync(acc30, a_frag3, b_frag0, acc30); + du_mma_sync(acc31, a_frag3, b_frag1, acc31); + du_mma_sync(acc32, a_frag3, b_frag2, acc32); + du_mma_sync(acc33, a_frag3, b_frag3, acc33); + } + + // Last stage only: prefetch the epilogue's per-row x_scale and + // per-column weight_scale values into registers (uniform branch: k0 is + // block-uniform, so the barrier below is reached by every thread; the + // row < m guard keeps the x_scale reads in bounds for any m tail). The + // vmem latency overlaps the barrier below (s_barrier waits on + // lgkmcnt/arrival, not vmcnt) and the dead a_frag/b_frag VGPR slots are + // reused, so the epilogue is pure compute + coalesced stores. + if (k0 + kStageK >= k) { + const int r = lane & 15; + const int c4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int row = base_row + i * kTileM + r; + xs_m[i] = (row < m) ? x_scale[row] : 0.0f; + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + ws_m[j] = *reinterpret_cast( + weight_scale + base_col + j * kTileN + 4 * c4); + } + } + // Protect the LDS buffers from the next stage's cooperative overwrite. + __syncthreads(); + } + + // Iteration 6 (epilogue round): coalesced direct-fragment epilogue. Each + // lane owns four contiguous bf16 columns of its fragment row and writes + // ONE 8-byte store per fragment (16 stores per wave vs 64 scattered 2-byte + // stores before); the scales come from the registers batched on the last + // K stage, so no vmem loads are interleaved with the stores. The int32 + // accumulation is untouched, and the multiply order / bf16 rounding are + // identical to the old per-element store, so output bits are unchanged. + store_prefill_fragment_coalesced_scaled(acc00, xs_m[0], ws_m[0], out, + base_row, base_col, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc01, xs_m[0], ws_m[1], out, + base_row, base_col + kTileN, m, n, + lane); + store_prefill_fragment_coalesced_scaled(acc02, xs_m[0], ws_m[2], out, + base_row, base_col + 2 * kTileN, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc03, xs_m[0], ws_m[3], out, + base_row, base_col + 3 * kTileN, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc10, xs_m[1], ws_m[0], out, + base_row + kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced_scaled(acc11, xs_m[1], ws_m[1], out, + base_row + kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc12, xs_m[1], ws_m[2], out, + base_row + kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc13, xs_m[1], ws_m[3], out, + base_row + kTileM, + base_col + 3 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc20, xs_m[2], ws_m[0], out, + base_row + 2 * kTileM, base_col, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc21, xs_m[2], ws_m[1], out, + base_row + 2 * kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc22, xs_m[2], ws_m[2], out, + base_row + 2 * kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc23, xs_m[2], ws_m[3], out, + base_row + 2 * kTileM, + base_col + 3 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc30, xs_m[3], ws_m[0], out, + base_row + 3 * kTileM, base_col, m, + n, lane); + store_prefill_fragment_coalesced_scaled(acc31, xs_m[3], ws_m[1], out, + base_row + 3 * kTileM, + base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc32, xs_m[3], ws_m[2], out, + base_row + 3 * kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc33, xs_m[3], ws_m[3], out, + base_row + 3 * kTileM, + base_col + 3 * kTileN, m, n, lane); +} + +// --------------------------------------------------------------------------- +// Generic scalar fallback: one thread per output element. Exact int32 dot +// over K, then fused x_scale * weight_scale, then bf16 store. Handles every +// unmatched (m, n, k), including all small-M API cases (M=2, M=16), M tails +// and M < 128 with the same (K, N). kNMajorPack == true decodes the +// iteration-5 swizzled 64-k-stage layout +// packed[((k0*8+kc)*n+col)*8+b] == raw[kk*n+col] (kk = k0*64+kc*8+b) for the +// exact (k, n) == (2048, 2048); otherwise the weight is the raw [K, N] +// row-major identity layout. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + int32_t acc = 0; + const int8_t* a_row = x_q + static_cast(row) * k; + if constexpr (kNMajorPack) { + // Swizzled 64-k-stage pack (iteration 5, exact (2048,2048) only): + // packed[((k0*8 + kc)*n + col)*8 + b] == raw[kk*n + col] with + // kk = k0*64 + kc*8 + b, so raw[kk][col] = packed[((k0*8+kc)*n+col)*8+b]. + for (int kk = 0; kk < k; ++kk) { + const int k0 = kk >> 6; + const int kc = (kk >> 3) & 7; + const int b = kk & 7; + const int8_t* p = weight + + (((static_cast(k0) * 8 + kc) * n + col) * 8) + + b; + acc += static_cast(a_row[kk]) * static_cast(*p); + } + } else { + // identity [K, N] row-major: column col is strided by n. + const int8_t* b_col = weight + col; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Identity device-to-device weight packing (outside the timed region and +// outside Graph capture) for every (k, n) except the exact (2048, 2048) +// q_b_proj pair (which uses the iteration-5 swizzled pack). The packed +// buffer keeps the same byte count K*N and the same allocated address, so +// the layout is graph-stable. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int64_t count) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +// Iteration 5 (packing round): swizzled 64-k-stage-major permutation for the +// exact (k, n) == (2048, 2048) q_b_proj weight: +// packed[((k0*8 + kc)*n + col)*8 + b] = raw[kk*n + col], +// kk = k0*64 + kc*8 + b (k0 = 64-k stage, kc = 8-byte sub-chunk, b = byte). +// Byte-wise (one thread per byte) so the permutation is trivially correct; +// runs once outside the timed region and outside Graph capture, keeping the +// same byte count K*N and the same graph-stable buffer address. The layout +// lets the 128x128 kernel stage each B tile as 16-byte plane chunks (one +// aligned int4 global read per chunk, perfectly coalesced) into the plane +// LDS tile [kc][n][8] whose fragment loads are lane-linear with zero bank +// conflicts. +__global__ __launch_bounds__(256) void w8a8_pack_swizzle_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(k) * n; + if (linear < total) { + const int kk = static_cast(linear / n); + const int col = static_cast(linear - static_cast(kk) * n); + const int k0 = kk >> 6; + const int kc = (kk >> 3) & 7; + const int b = kk & 7; + packed[(((static_cast(k0) * 8 + kc) * n + col) * 8) + b] = + raw[linear]; + } +} + +__global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( + const float* __restrict__ raw, + float* __restrict__ packed, + int count) { + const int linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Stable host launch symbols consumed by csrc/bindings.cpp. +// Both launchers are pure dispatch: no allocation, no packing, no +// synchronization, no default-stream launch; they run on the caller-provided +// PyTorch stream and are CUDA-Graph safe. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; // no split-K: the GEMM does not use the workspace + (void)workspace_bytes; + auto* out_bf16 = reinterpret_cast(out); + + // Explicit dispatch. Large-M shapes with exact tiled geometry (both + // assigned shapes: M=4096, N in {2048, 3584}, K in {2048, 512}) take the + // native INT8 DUMMA tiled path; every other (m, n, k) - including small-M + // API cases (M=2, M=16) and M < 128 - takes the scalar fallback. The + // exact (k, n) == (2048, 2048) q_b_proj pair takes the iteration-4 128x128 + // tile with the iteration-5 swizzled-pack B layout; every other (k, n) - + // including kv_b_proj (512, 3584) - uses the byte-identical 128x64 + // kNMajorB=false arm (kNMajorPack == true for the scalar fallback's + // (2048,2048) decode). + if (m >= kBlockM && m % kBlockM == 0 && n % kBlockN == 0 && + k % kStageK == 0) { + const dim3 grid(static_cast(n / kBlockN), + static_cast(m / kBlockM)); + const dim3 block(kBlockThreads); + if (k == kPackNMajorK && n == kPackNMajorN) { + // Exact q_b_proj (2048, 2048): iteration-4 tile-aspect arm, 128x128 + // tile with the iteration-5 swizzled-pack plane B layout; grid + // (N/128, M/128) = 512 blocks for M=4096. n == 2048 is divisible by + // 128 by the exact guard; the outer arm already guarantees + // m % 128 == 0, k % 64 == 0. + const dim3 grid128(static_cast(n / 128), + static_cast(m / 128)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_128x128_kernel), + grid128, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_128x64_kernel), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + } else { + constexpr int kBlock = 256; + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + if (k == kPackNMajorK && n == kPackNMajorN) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + } + (void)hipGetLastError(); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kBlock = 256; + const dim3 block(kBlock); + // Bootstrap pack: iteration-5 swizzled 64-k-stage permutation for the exact + // (2048, 2048) q_b_proj weight; identity device-to-device copy for every + // other (k, n). + const int64_t weight_count = static_cast(k) * n; + const dim3 weight_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + if (k == kPackNMajorK && n == kPackNMajorN) { + hipLaunchKernelGGL(w8a8_pack_swizzle_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else { + hipLaunchKernelGGL(w8a8_pack_identity_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, weight_count); + } + + const int64_t scale_count = n; + const dim3 scale_grid( + static_cast((scale_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, + scale_grid, block, 0, stream, + weight_scale, packed_weight_scale, n); + (void)hipGetLastError(); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_down_proj.hip new file mode 100644 index 00000000..a1b49be9 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_down_proj.hip @@ -0,0 +1,1371 @@ +// @@variant shape=glm_tp8_shared_down_proj_m4096 commit=9b56135e1fec933006c477066840f9c84f497cef added=2026-08-31 +// median_us=278.4 p90_us=279.2 speedup=63.57 baseline_us=1.77e+04 +// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 +// MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). +// worker_3 (iteration 3), assigned shapes: +// glm_tp8_shared_gate_up_proj_m4096 : (M, N, K) = (4096, 512, 6144) +// glm_tp8_shared_down_proj_m4096 : (M, N, K) = (4096, 6144, 256) +// +// This file provides the two stable host launch symbols consumed by +// csrc/bindings.cpp: +// * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's +// HIP stream; no allocation, packing, +// autotuning, or synchronization. Only the +// caller-provided out and workspace are +// touched (workspace is unused this round). +// * launch_pack_w8a8_weight(...) - out-of-timed-region packing. Identity +// device-to-device copy for every (k, n) +// pair EXCEPT the assigned gate_up weight +// (k, n) == (6144, 512), which is packed +// once as a [K, N] -> [N, K] transpose +// (n-major rows, k contiguous) so the timed +// GEMM consumes B fragments as +// k-contiguous bytes. The pack stays +// out-of-timed, out-of-Graph, same byte +// count, same graph-stable buffer. +// +// Bootstrap strategy (correctness-first, usable profiling baseline): +// * Native INT8 DUMMA m16n16k32 tiled kernel with int32 accumulation for +// every assigned shape with M >= 128 (both assigned shapes qualify). +// * 2-D macro-tile family (64x64, 64x128, 128x64) with one 32x32 quadrant +// per wavefront (four m16n16k32 int32 accumulators kept resident over the +// whole K loop). Single-buffered K loop (K stage 128, two barriers per +// stage): no split-K, no raw asm, no speculative double-buffer pipeline. +// +// Round 1 (baseline): the 64x128 tile measured 517.4 us median / 49.81 TOPS; +// routing the exact gate_up shape to the 64x64 tile measured 432.5 us +// (accepted; 19.6% improvement). +// +// Round 2 (operand-reuse, mandated decision): cooperative A+B LDS staging is +// kept and direct per-wave global loads are rejected (A is reused 64x per +// macro-tile via kBlockN and B 64x via kBlockM; direct loads would refetch +// both operands once per wavefront, 4x the 28.3 MB staged global traffic -> +// ~1.73M vmem_read instructions vs 432K, and the TP4 gate_up direct-A arm +// measured 7.9x slower). The focused change makes the staged B layout +// bank-safe for the fragment consumers: the gate_up weight is packed once +// [K,N] -> [N,K] (n-major, k contiguous) and B fragments are loaded +// col_major from a 144-B (128+16) row-strided LDS tile, replacing the +// byte-scattered row_major B reads on the identity [K,N] layout that are the +// 17.3M PMC bank-conflict source. Global loads stay 16-byte vectorized and +// coalesced on both operands. LDS stays 19,456 B/block -> 3 blocks/CU +// (58,368 B, unchanged; round-1's "4 blocks/CU" estimate was impossible on +// both LDS (4x19,456 = 77,824 B > 64 KiB) and VGPR (4x256x72 = 73,728 > +// 65,536)). down_proj keeps the identity 64x128 arm; every unmatched shape +// keeps the generic arms and the scalar fallback (which decodes the packed +// [N,K] layout when (k, n) == (6144, 512) so the paired M=2/M=16 gate_up +// tails stay correct). +// * A[kBlockM, kStageK] and B staged in bank-skewed LDS (row strides +// 144/80/144 = 16-byte-aligned non-power-of-two), consumed with the +// library du_load_matrix_sync loaders: A row_major (unchanged, load- +// bearing on every shape), B row_major for the identity [K, N] layout +// (down_proj + unmatched shapes) and B col_major for the gate_up +// [N, K] packed layout (round 2). +// * Fused direct fragment -> x_scale -> weight_scale -> bf16 epilogue using +// the verified gfx928 accumulator ownership (row = lane & 15, +// col_mod4 = lane >> 4, x[i] -> columns col_mod4 + 4*i). The int32 dot is +// exact for every assigned shape (max |dot| = 6144*127*127 << 2^31), the +// int32 -> float conversion is exact, and the single fp32 multiply +// followed by one bfloat16 rounding matches the reference bit-for-bit. +// * Tail M rows are zero-filled on load and masked on store, so any +// M >= 128 is supported by the tiled path. +// * All other (m, n, k) - including every M < 128 API case and the paired +// M=2/M=16 shapes with the same (N, K) - go to a scalar int8/int32 +// fallback that decodes the [N, K] packed layout when (k, n) == +// (6144, 512) and the identity [K, N] layout otherwise. +// +// Round 3 (pipeline, mandated decision): compare single and double buffering +// across K tiles. The exact (4096, 512, 6144) guard now routes to +// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true>: a software- +// pipelined 64-K double-buffered K loop (prologue + one barrier per stage, +// 97 barriers for 96 stages) that issues stage s+1's global loads into the +// idle LDS buffer before the stage-s MMAC burst, overlapping the +// global-load latency that the round-2 single buffer serializes in front of +// every stage (ISA-verified order: global_load_dwordx4 -> ds_write_b128 -> +// s_waitcnt -> s_barrier -> ds_read2 x16 -> v_mmac x16 -> s_barrier, with +// the next stage's loads gated behind the previous stage's barrier). The +// 128-K double-buffered variant is arithmetically impossible at 3 blocks/CU +// (2 x 19,456 = 38,912 B/block x 3 = 116,736 > 64 KiB -> 1 block/CU, a +// harmful occupancy loss), so the pipelined comparison uses the 64-K +// double buffer: LDS 20,480 B/block x 3 = 61,440 <= 64 KiB and +// VGPR 64-72 x 256 x 3 <= 65,536, i.e. 3 blocks/CU preserved (no occupancy +// confound). Everything else is byte-identical: same int4 staging pattern +// (same 28.3 MB global traffic and ~432K vmem_read instructions), same +// k0-outer/kk-inner int32 order (stage size changes only regroup exact +// integer sums, so results stay bit-identical), same col_major B fragments +// (stride 80 = 64+16 bank skew), same fused epilogue, same pack, same +// generic single-buffered arms and scalar fallback. Decision rule (mandate): +// retain double buffering only if PMC shows reduced VMEM/LDS stall evidence +// (lds_wait_instructions drops below ~1.2M, lds_instructions stays ~1.97M, +// occupancy stays 3 blocks/CU) AND the median improves >2%; otherwise the +// round's evidence retains the single-buffered round-2 structure. +// +// Round 10 (fragment-load-form round, restarted from the accepted round-3 +// source after the round-9 infra kill): replace the library +// du_load_matrix_sync fragment loads of the exact-shape gate_up kernel with +// the lineage-validated explicit load_frag8 8-byte ds_read_b64 loader. The +// exact (4096, 512, 6144) guard now routes to +// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true, true> (new +// kDirectFrag template arm). The round-3 ISA shows every fragment load as +// ds_read2_b32 + per-byte v_and/v_or_b32_sdwa/v_or3 mask-OR "identity" +// reassembly (~9.4M of the 15.43M VALU instructions per replay); loading the +// same 8 bytes per lane as one int64 (ds_read_b64) removes that VALU and +// halves the fragment bank-conflict factor (32-lane dword phases at stride +// 80 -> 16-lane 8-byte phases). Element-to-slot fragment mapping, v_mmac +// operands, exact k0-outer/kk-inner int32 accumulation, staging/global +// traffic, barriers, tile, grid, LDS bytes, occupancy (3 blocks/CU) and the +// fused epilogue are unchanged; results stay bit-identical. The generic arms +// (<64,128,128>, <128,64,128>, <64,64,128>, all kBNMajor=false) keep the +// library loaders and byte-identical codegen/resources (if constexpr). +// +// Round 11 (LDS bank-skew round, HIP-only; ISA policy keeps raw asm +// disallowed because the HIP plateau rule - 8 valid HIP rounds within +/-2% +// of the then-current best - is unmet). The accepted round-10 kernel still +// shows PMC lds_bank_conflicts = 9,437,184, EXACTLY unchanged from the +// round-3 code object (the -9.07% round-10 win came from VALU removal, not +// from conflicts). 9,437,184 = 192 x 512 blocks x 96 stages: every +// block-stage pays 192 conflict cycles = 128 from the 16 ds_read2_b64 +// fragment reads + 64 from the 8 ds_write_b128 staging stores (the 16-B/lane +// store floor). The read half is a stride defect: at stride 80 = 20 dwords, +// 20*r mod 32 has period 8, so every 16-lane read phase groups rows r and +// r+8 onto the same bank pair - a 2-way conflict on ALL 4 sub-phases of +// every fragment read. The exact (4096, 512, 6144) guard now routes to +// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true, true, true> (new +// kSkew8 template arm) which stages both operands at the down_proj-validated +// 8-mod-16 row strides A = 88 (64+24) and B = 72 (64+8) instead of 80: +// 22*r and 18*r mod 32 are distinct over r = 0..15 and gcd(22,32) = +// gcd(18,32) = 2 does not divide 1, so every 16-lane read phase lands on 16 +// DISTINCT bank pairs - the fragment reads become conflict-free (128 of the +// 192 cycles removed; only the 2-way 16-B/lane staging-store floor of 64 +// cycles remains -> lds_bank_conflicts expected ~3.15M). The 8-mod-16 row +// starts forbid 16-byte staging stores, so each 16-byte int4 is staged as +// two 8-byte halves (ds_write2_b64), the validated down_proj staging form +// (its B tile is exactly stride 72; its A tile exactly stride 88). Global +// loads (int4, same 432,128 vmem_read), pack layout, barrier structure +// (97/block), tile/grid (512 blocks), LDS bytes (A[64,88] + B[64,72] x 2 = +// 20,480 B/block -> 3 blocks/CU unchanged), VGPR, the k0-outer/kk-inner +// int32 order, the element-to-slot fragment mapping and the fused epilogue +// are untouched; only the LDS row strides and the staging store granularity +// change, so results stay bit-identical. Generic arms (<64,128,128>, +// <128,64,128>, <64,64,128>) keep kSkew8 = false and byte-identical +// codegen/resources (if constexpr). +// +// Round 16 (lead-2 loop-carried payload round, HIP-only; raw asm stays +// DISALLOWED - plateau=false, the recent-valid-improvement window +// [-34.44%, +2.82%, +9.97%] does not meet three valid HIP rounds within +// [-2%, +2%) of the then-current best, and no prior ISA-guided round +// recorded a compiler limitation with target instructions). The accepted +// round-11 exact code object (digest 9889820e, median 293.30 us / p90 +// 293.545) shows the software-pipeline intent is DEFEATED by the compiler's +// wait placement: per stage the steady-state loop is +// global_load_dwordx4 (A prefetch s+1) -> global_load_dwordx4 (B prefetch +// s+1) -> s_waitcnt vmcnt(1) -> s_waitcnt vmcnt(0) <- waits IMMEDIATELY +// -> 2 x ds_write2_b64 (staging stores) -> 4 x ds_read2_b64 (fragment +// reads, stage s) -> 8 x v_mmac CONTIGUOUS (stage s) -> s_barrier -> loop. +// i.e. the stage-(s+1) global loads are waited ~12-16 instructions after +// issue, BEFORE the stage-s MMAC burst, so the full global-load latency +// (~600-1000+ cycles under 120-CU L2 contention) sits on the per-stage +// critical path with zero overlap (the round-12 analysis' +// "prefetch global load -> ds_write2_b64 -> s_barrier -> ds_read2_b64 -> +// v_mmac" chain; PMC: 6.0M VALU + 1.18M LDS + 465K VMEM wave-instructions +// over 293 us x 120 CUs, i.e. ~95-99% issue-stalled). Iteration history +// rules out the other levers on this exact shape: tile aspect (iter 4 +// 32x64: 749.9 us; iter 8 64x128 w8: 395.4 us), 2-wave quadrants (iter 6: +// 629.0 us - wave count per CU is load-bearing), LDS-side fragment +// prefetch (iter 7: 337.97 us, flat/worse), stage 32 / 6 blocks per CU +// (iter 12: 447.4 us - per-stage fixed cost, not occupancy, is the lever). +// The one untested mechanism that directly removes the exposed vmcnt wait: +// a LEAD-2 LOOP-CARRIED REGISTER PAYLOAD (the TP4 o_proj lineage's +// validated "stage-top publish" form). The exact (4096, 512, 6144) guard +// now routes to w8a8_dumma_prefill_tiled_kernel<64,64,64,true,true,true, +// true,true> (new kPayload template arm): each thread keeps the one A int4 +// + one B int4 of stage s+2 in VGPR across the stage-(s+1) barrier, issues +// those global loads at the TOP of iteration s (no wait), and PUBLISHES +// the payload carried from iteration s-1 into the idle LDS buffer at the +// TOP of iteration s (before the stage-s MMAC burst). The first use of the +// payload registers is the next iteration's publish, and the s_barrier +// fences outstanding vmem, so the compiler can no longer place +// s_waitcnt vmcnt(0) immediately after the loads: the wait lands after the +// stage-s fragment reads + MMAC burst (immediately before the barrier), and +// the global-load latency overlaps the whole compute segment instead of +// serializing in front of it. One barrier per 64-K stage is preserved +// (prologue + 96 = 97/block); the publish targets the idle buffer (WAR +// retired by the previous barrier, RAW retired by the stage barrier), so no +// second barrier is introduced. Everything else is byte-identical: 64x64 +// tile / 256 threads / 4 waves / 32x32 quadrant per wave, grid (8,64) = +// 512 blocks, kSkew8 strides 88/72 with ds_write2_b64 staging (8 per +// block-stage), same 16 ds_read2_b64 fragment reads per block-stage (same +// lds_instructions ~1,179,648), same 20,480 B LDS/block -> 3 blocks/CU +// (payload costs 8 VGPRs: ~54 total, 54 x 256 x 3 = 41,472 <= 65,536; no +// spill risk at 3 blocks/CU), same int4 global loads (same ~432,128 +// vmem_read / 32,768 vmem_write), same k0-outer/kk-inner int32 order and +// the same element-to-slot fragment mapping - the payload changes only WHEN +// global bytes land in LDS, never which bytes or the accumulation order, so +// results stay bit-identical (mismatch 0 / max_abs_error 0.0 expected). +// Generic arms (<64,128,128>, <128,64,128>, <64,64,128>, all kBNMajor= +// false) keep kPayload = false and byte-identical codegen/resources +// (if constexpr). +// +// Round 2 (this session, glm_tp8_shared_down_proj_m4096 exact-shape arm): +// the generic <64,128,128> identity-layout arm that served down_proj +// (4096, 6144, 256) measured 314.72 us median with PMC +// lds_bank_conflicts = 7,864,320, ~600 LDS instructions and heavy per-byte +// fragment-reassembly VALU per block-stage (library row_major loaders on +// the identity [K, N] B layout). The exact (4096, 6144, 256) guard now +// routes to w8a8_dumma_prefill_tiled_kernel<64, 128, 64, true, true, true, +// true> - the file's own lineage-validated gate_up machinery transplanted +// to down_proj: B packed once [K,N] -> [N,K] n-major (kBNMajor, same +// transpose kernel as gate_up, same byte count / graph-stable buffer, +// out-of-timed/out-of-Graph), 64-K double-buffered K loop (kDoubleBuffer, +// prologue + 4 stages for K=256, one barrier per stage), explicit 8-byte +// load_frag8 fragment reads (kDirectFrag, one ds_read_b64 per fragment, +// zero reassembly VALU) and the 8-mod-16 bank-skewed LDS strides A=88 / +// B=72 (kSkew8, down_proj-validated, conflict-free 16-lane fragment +// phases). LDS A[64,88] + B[128,72] x 2 = 29,696 B/block -> 2 blocks/CU +// (occupancy unchanged vs the 27,648-B single buffer); grid (48, 64) = +// 3072 blocks; the stage size 128->64 only regroups exact integer sums +// (k0-outer / kk-inner int32 order over k = 0,32,64,96,128,160,192,224 is +// unchanged) and load_frag8 fills the same x[0..7] slots with the same +// little-endian bytes, so results stay bit-identical (mismatch 0 / +// max_abs_error 0.0 expected). +// +// Round 3 (this session, glm_tp8_shared_down_proj_m4096 pipeline round): +// mandated single vs double buffering comparison across K tiles. The +// accepted round-2 arm (<64,128,64,true,true,true,true>, 286.53 us median / +// 286.64 us p90, digest cd9791d0) is the 64-K DOUBLE-buffered pipeline; its +// exact code object shows the software pipeline already works on this shape +// (steady-state loop: ds_write2_b64 publish s+1 -> global_load_dwordx4 x2 +// for s+2 -> ds_read2_b64 x4 (fragment reads s+1) -> v_mmac x8 -> s_waitcnt +// vmcnt(0) lgkmcnt(0) -> s_barrier: the stage-(s+2) loads stay in flight +// across the fragment reads + MMAC burst and are waited only after it, i.e. +// reduced VMEM stalls), and PMC shows no harmful LDS or occupancy growth vs +// the pre-round-2 single buffer (29,696 vs 27,648 B/block, 2 blocks/CU +// pinned; lds_wait 162,664 / lds_instructions 540,672 / conflicts +// 1,179,648). The (4096, 6144, 256) guard now routes to the SINGLE-buffered +// form <64,128,64,true,false,true,true>: identical stage size 64, identical +// kBNMajor/kDirectFrag/kSkew8/tile/grid/epilogue, kDoubleBuffer = false, so +// each stage serializes its global-load round trip (load -> wait -> +// ds_write2_b64) in front of the MMAC burst with two barriers per stage +// (8/block vs 5/block) and LDS halves to 14,848 B/block. The single-buffer +// branch now threads kSkew8 and kDirectFrag through its staging/MMAC calls +// (previously hardcoded to the library-loader defaults; generic arms keep +// byte-identical codegen via if constexpr). Falsifiable gate: if the +// single-buffered median is NOT strictly below 286.53368949890137 (p90 +// guard < 286.6385078430176), double buffering is retained with the ISA/PMC +// evidence above; if it IS below, the single buffer wins on this shape +// (LDS halved; VGPR may drop <= 42 and open 3 blocks/CU, the trusted +// occupancy-probe set [2, 3, 4]). Results stay bit-identical either way +// (same stage size, same k0-outer/kk-inner int32 order, same bytes). +// +// Round 4 (this session, glm_tp8_shared_down_proj_m4096 tile-shape round): +// the balanced 1:1 aspect probe <64,64,64,true,false,true,true> (256 thr = +// 4 waves, 6144 blocks, LDS 10,240 B/block -> 4 blocks/CU, doubled A +// staging) measured 297.15 us median / 297.97 us p90 and was REJECTED +// (guard restored to the round-3 single-buffered <64,128,64,true,false, +// true,true> winner, 281.65 us / 282.52 us): co-residency was already +// saturated at 2 blocks/CU and the +33% A-side staging cost was not repaid. +// +// Round 5 (this session, glm_tp8_shared_down_proj_m4096 PACKING round): test +// one weight packing/swizzle that makes each DUMMA B tile vector-loadable +// and LDS-bank-safe, with the packing outside timing and the graph-stable +// packed layout validated. The accepted round-3 arm stages its 128x64 B +// tile from the plain n-major [N,K] packed weight: per block-stage the 512 +// int4 loads are 16 rows x 4 int4 per warp at 256-B row stride, i.e. every +// warp instruction touches 16 DISTINCT 64-B segments (50% of each 128-B L2 +// line) - vector loads, but scattered at L2 granularity. The round changes +// the one-time pack for (k,n) == (256,6144) ONLY to the panel layout +// [N/128][K/64][128][64] (n-major 64-k-byte rows per panel; byte-exact index +// relation packed[((nt*(k/64)+kt)*128+nn)*64+kk64] == raw[(nt*128+nn)*k + +// kt*64+kk64], verified by simulation): each block's per-stage B tile is now +// ONE contiguous 8 KiB region, so every warp instruction reads 1 KiB +// contiguous (full sector/line utilization; per block-stage the B staging +// collapses from 128 half-used 128-B lines to 64 fully-used lines). The new +// exact-shape arm w8a8_dumma_prefill_tiled_kernel<64,128,64,true,false, +// true,true,false,true> (new kBPanel template flag, default false) is +// byte-identical in EVERYTHING else: same 64-K single-buffered stage, same +// kBNMajor/kDirectFrag/kSkew8, same 64x128 tile / 512 thr / 8 waves / 32x32 +// quadrant, same grid (48,64) = 3072 blocks, same LDS A[64,88]+B[128,72] = +// 14,848 B/block @ 2 blocks/CU, same ds_write2_b64 staging stores and same +// conflict-free load_frag8 ds_read2_b64 fragment reads (LDS-bank-safe +// preserved - only the GLOBAL byte order of the packed buffer changes, +// never the staged LDS bytes), same k0-outer/kk-inner int32 order, same +// fused epilogue -> mismatch 0 / max_abs_error 0.0 expected. The scalar +// fallback decodes the panel layout for (k,n) == (256,6144) so M<128 tails +// stay correct; gate_up (6144,512) keeps its plain [N,K] transpose pack and +// its own arm; generic arms are untouched (kBPanel=false). Packing runs +// out-of-timed/out-of-Graph in launch_pack_w8a8_weight, same byte count +// (256*6144), same graph-stable buffer. Expected effect: B global staging +// becomes perfectly coalesced 1 KiB-per-warp streams with half the L2 line +// count; PMC vmem_read_instructions unchanged (~614,400 - same 512 int4 +// loads/block-stage), lds_* identical (LDS layout untouched), l2_hits/l2_ +// misses shift slightly in the direction of fewer miss sectors, valu ~flat +// (address math per thread unchanged: one base + offset). FALSIFIABLE GATE: +// median_us < 281.6507148742676 AND p90_us < 282.5235176086426 -> B-global +// coalescing was binding on this shape and the panel pack wins; otherwise +// the plain n-major pack is retained and this round records that B reads +// are L2-resident (90.4% hit) and NOT binding on the LDS-bound critical +// path (consistent with the lineage rule that pack layout stops paying once +// fragment loads/strides are fixed). +// +// Header order is fixed by the control plane: hip_runtime first (du_mma.h +// is not self-contained before it), hip_bfloat16 second, du_mma.h last. + +#include +#include +#include + +#include + +namespace { + +// gfx928 INT8 DUMMA primitive: m16n16k32, int8 x int8 -> int32. +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; + +// gfx928 wavefront is 64 lanes; every blockDim must be a multiple of 64. +constexpr int kWaveSize = 64; + +// Large-prefill dispatch geometry. Each macro-tile config uses a 32x32 +// quadrant per wave, so waves per block = (BlockM/32) * (BlockN/32). +constexpr int kPrefillMinM = 128; +constexpr int kPrefillStageK = 128; +constexpr int kDefaultBlockM = 64; +constexpr int kDefaultBlockN = 64; + +constexpr int kScalarThreads = 256; + +using namespace du::dumma; + +// Verified gfx928 INT8 DUMMA accumulator ownership: +// row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> columns +// col_mod4 + 4*i. Scale by x_scale[row] and weight_scale[col] and store +// bf16 directly from the fragment (no accumulator LDS round trip). +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * xs * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Round 10 (fragment-load-form round): explicit 8-byte fragment loader +// replacing the library du_load_matrix_sync for the exact-shape gate_up +// kernel. Validated on the sibling TP8 down_proj accepted kernel (iter 8, +// -25.9%: 230.91 -> 171.19 us) and the TP4 gate_up lineage (iter 13 +// load_fragment8, -21%: 280.05 -> 220.71 us). For BOTH the library row_major +// matrix_a and col_major matrix_b m16n16k32 int8 loaders, lane l -> row +// (l & 15), k-quarter (l >> 4), and the loader fetches the same eight +// contiguous bytes at p[row*ldm + (l>>4)*8 .. +7] and stores them into the +// same x[0..7] slots in the same little-endian order (the round-3 ISA shows +// the compiler lowering this to ds_read2_b32 followed by a per-byte +// v_and/v_or_b32_sdwa/v_or3 mask-OR "identity" reassembly - ~9.4M dead-ish +// VALU instructions per replay, ~61% of the 15.43M VALU total). Loading the +// same 8 bytes as one int64 forces the 8-byte ds_read_b64 form instead: +// (a) the mask/OR reassembly VALU disappears, and (b) the LDS access +// changes from ds_read2_b32 (32-lane dword phases - at the stride-80 row +// pattern every phase groups rows 0-15 of k-quarter q with rows 0-15 of +// k-quarter q+1 onto the same 16 bank phases, a 4-way conflict per phase) +// to ds_read_b64 (16-lane phases: rows 0-15 land on 16 bank-pairs +// {20*r mod 32} = 8 distinct pairs hit twice - a 2-way conflict per phase, +// half of the current fragment-conflict factor). The element-to-slot +// mapping is unchanged (f.x[0..7] = the same 8 bytes in the same +// little-endian order), the v_mmac operands and the exact k0-outer/kk-inner +// int32 accumulation sequence are untouched, and the staging stores, LDS +// strides, barriers, tile, grid, occupancy and epilogue are byte-identical. +template +__device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, + unsigned ldm) { + const unsigned row = __lane_id() & 0xf; + const unsigned kq = __lane_id() >> 4; + const int64_t v = + *reinterpret_cast(p + row * ldm + (kq << 3)); + reinterpret_cast(f.x)[0] = v; +} + +// Stage one K tile (k0) of A[kBlockM, kStageK] and B into the caller-offset +// LDS buffers (bank-skewed row strides; int4 vectorized coalesced global +// loads, zero-filled tail M rows). kBNMajor == true stages B from the packed +// [N, K] n-major gate_up layout (n-major rows, k contiguous); otherwise from +// the identity [K, N] layout. Shared by the single-buffered and +// double-buffered K loops; the buffer offset (0 or kABytes/kBMaxBytes) is +// applied by the caller so one body serves both pipeline forms. +// +// kSkew8 (round 11, exact-shape arm only): A and B row strides become the +// down_proj-validated 8-mod-16 values 88 (kStageK + 24) and 72 (kStageK + 8) +// so every 16-lane fragment-read phase lands on 16 distinct bank pairs (see +// the round-11 header comment). The 8-mod-16 row starts forbid 16-byte +// stores, so each 16-byte int4 chunk is staged as two 8-byte halves +// (ds_write2_b64); the global int4 loads are byte-identical. +template +__device__ __forceinline__ void stage_prefill_tiles( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int8_t* a_tile, + int8_t* b_tile, + int tid, + int m0, + int n0, + int m, + int n, + int k, + int k0) { + constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + // 16-byte bank skew: keeps int4 staging stores aligned while spreading the + // fragment rows across distinct LDS banks. + constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); + constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows + constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); + static_assert(!kSkew8 || kStageK == 64, + "kSkew8 strides 88/72 assume kStageK == 64"); + constexpr int kAInt4PerRow = kStageK / sizeof(int4); + constexpr int kAInt4 = kBlockM * kAInt4PerRow; + constexpr int kBInt4PerRow = kBlockN / sizeof(int4); + constexpr int kBInt4 = kStageK * kBInt4PerRow; + constexpr int kBNInt4PerRow = kStageK / sizeof(int4); + constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; + // Stage A[kBlockM, kStageK] into LDS (zero-filled tail M rows). + for (int vec = tid; vec < kAInt4; vec += kThreads) { + const int local_row = vec / kAInt4PerRow; + const int v = vec - local_row * kAInt4PerRow; + const int global_row = m0 + local_row; + const int4 val = + global_row < m + ? *reinterpret_cast( + x_q + static_cast(global_row) * k + k0 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + if constexpr (kSkew8) { + // 8-mod-16 row starts: write the 16 bytes as two 8-byte halves so every + // store stays 8-byte aligned (ds_write2_b64 staging; same LDS bytes). + int8_t* dst = + a_tile + local_row * kAStride + v * static_cast(sizeof(int4)); + reinterpret_cast(dst)[0] = + reinterpret_cast(&val)[0]; + reinterpret_cast(dst)[1] = + reinterpret_cast(&val)[1]; + } else { + reinterpret_cast(a_tile + local_row * kAStride)[v] = val; + } + } + if constexpr (kBNMajor) { + // Stage B[kBlockN, kStageK] from the packed [N, K] weight (n-major: + // row = n_local, k contiguous). Same vectorized coalesced int4 pattern + // as the identity arm, one per 16 consecutive k bytes. + // + // kBPanel (round 5, down_proj exact-shape arm only): the one-time pack + // for (k,n) == (256,6144) is the panel layout [N/128][K/64][128][64] - + // each block's per-stage B tile is one contiguous 8 KiB region, so a + // 64-lane warp instruction reads 1 KiB contiguous (full L2 line/sector + // utilization) instead of 16 rows x 64 B scattered at 256-B stride + // (50% of each 128-B line). The per-thread address is still one base + + // offset: packed byte offset = panel_base + local_col*64 + v*16 where + // panel_base = (n0/128)*(k/64)*8192 + (k0/64)*8192. The LDS stores are + // untouched (same stride-72 ds_write2_b64 halves -> LDS-bank-safe + // staging and conflict-free fragment reads preserved); only the global + // byte order of the loaded int4 changes, so the staged LDS bytes are + // identical to the plain n-major arm (bit-identical results). + for (int vec = tid; vec < kBNInt4; vec += kThreads) { + const int local_col = vec / kBNInt4PerRow; + const int v = vec - local_col * kBNInt4PerRow; + int4 val; + if constexpr (kBPanel) { + const int64_t panel_base = + (static_cast(n0) / kBlockN) * (k / kStageK) * + (kBlockN * kStageK) + + (k0 / kStageK) * (kBlockN * kStageK); + val = *reinterpret_cast( + weight + panel_base + + static_cast(local_col) * kStageK + + v * static_cast(sizeof(int4))); + } else { + val = *reinterpret_cast( + weight + static_cast(n0 + local_col) * k + k0 + + v * static_cast(sizeof(int4))); + } + if constexpr (kSkew8) { + int8_t* dst = + b_tile + local_col * kBNStride + + v * static_cast(sizeof(int4)); + reinterpret_cast(dst)[0] = + reinterpret_cast(&val)[0]; + reinterpret_cast(dst)[1] = + reinterpret_cast(&val)[1]; + } else { + reinterpret_cast(b_tile + local_col * kBNStride)[v] = val; + } + } + } else { + // Stage B[kStageK, kBlockN] into LDS from the identity [K, N] weight. + for (int vec = tid; vec < kBInt4; vec += kThreads) { + const int kk = vec / kBInt4PerRow; + const int v = vec - kk * kBInt4PerRow; + reinterpret_cast(b_tile + kk * kBStride)[v] = + *reinterpret_cast( + weight + static_cast(k0 + kk) * n + n0 + + v * static_cast(sizeof(int4))); + } + } +} + +// Round 16 (lead-2 loop-carried payload, exact-shape arm only): the register +// half of the lead-2 software pipeline. Each thread prefetches exactly one A +// int4 and one B int4 of a later K stage into thread-local registers and +// returns WITHOUT waiting (the vmcnt wait is emitted by the compiler at the +// first use, one iteration later). The <64,64,64> arm stages kStageK * 64 +// bytes per operand per stage, i.e. kAInt4 == kBNInt4 == kThreads == 256: +// one int4 per thread per operand, so the whole payload costs 8 VGPRs. The +// (row, v) mapping is the same as stage_prefill_tiles, so +// publish_stage_payload writes the exact same LDS addresses the staging +// loop would (bit-identical staging bytes, just loaded one stage earlier). +template +__device__ __forceinline__ void prefetch_stage_payload( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int tid, + int m0, + int n0, + int m, + int k, + int k0, + int4& a_payload, + int4& b_payload) { + constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + constexpr int kAInt4PerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kAInt4 = kBlockM * kAInt4PerRow; + constexpr int kBNInt4PerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; + static_assert(kAInt4 == kThreads && kBNInt4 == kThreads, + "payload arm requires exactly one int4 per thread per " + "operand (exact-shape <64,64,64> arm)"); + static_assert(kBNMajor, "payload arm stages the packed n-major B layout"); + // Same vec -> (local_row, v) mapping as stage_prefill_tiles, so the + // publish half stores to the same LDS addresses. + const int a_local_row = tid / kAInt4PerRow; + const int a_v = tid - a_local_row * kAInt4PerRow; + const int global_row = m0 + a_local_row; + // Tail M rows are zero-filled exactly as in stage_prefill_tiles. + a_payload = global_row < m + ? *reinterpret_cast( + x_q + static_cast(global_row) * k + k0 + + a_v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + const int b_local_col = tid / kBNInt4PerRow; + const int b_v = tid - b_local_col * kBNInt4PerRow; + b_payload = *reinterpret_cast( + weight + static_cast(n0 + b_local_col) * k + k0 + + b_v * static_cast(sizeof(int4))); +} + +// Round 16: the LDS half of the lead-2 payload pipeline. Writes the payload +// carried in registers (loaded by prefetch_stage_payload one stage earlier) +// into the caller-offset LDS tile with the exact same 8-byte-half +// ds_write2_b64 stores (kSkew8 strides 88/72) that stage_prefill_tiles +// emits, so each thread republishes exactly the bytes it loaded. +template +__device__ __forceinline__ void publish_stage_payload( + int4 a_payload, + int4 b_payload, + int8_t* a_tile, + int8_t* b_tile, + int tid) { + constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); + constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); + constexpr int kAInt4PerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kAInt4 = kBlockM * kAInt4PerRow; + constexpr int kBNInt4PerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; + static_assert(kAInt4 == kThreads && kBNInt4 == kThreads, + "payload arm requires exactly one int4 per thread per " + "operand (exact-shape <64,64,64> arm)"); + static_assert(kBNMajor, "payload arm stages the packed n-major B layout"); + const int a_local_row = tid / kAInt4PerRow; + const int a_v = tid - a_local_row * kAInt4PerRow; + int8_t* a_dst = + a_tile + a_local_row * kAStride + a_v * static_cast(sizeof(int4)); + if constexpr (kSkew8) { + // 8-mod-16 row starts: write the 16 bytes as two 8-byte halves + // (ds_write2_b64 staging, identical to stage_prefill_tiles). + reinterpret_cast(a_dst)[0] = + reinterpret_cast(&a_payload)[0]; + reinterpret_cast(a_dst)[1] = + reinterpret_cast(&a_payload)[1]; + } else { + reinterpret_cast(a_dst)[0] = a_payload; + } + const int b_local_col = tid / kBNInt4PerRow; + const int b_v = tid - b_local_col * kBNInt4PerRow; + int8_t* b_dst = + b_tile + b_local_col * kBNStride + b_v * static_cast(sizeof(int4)); + if constexpr (kSkew8) { + reinterpret_cast(b_dst)[0] = + reinterpret_cast(&b_payload)[0]; + reinterpret_cast(b_dst)[1] = + reinterpret_cast(&b_payload)[1]; + } else { + reinterpret_cast(b_dst)[0] = b_payload; + } +} + +// One wave's 32x32-quadrant MMAC burst for the K tile staged at a_tile/b_tile +// (caller-offset to the target buffer). k0-outer / kk-inner int32 accumulation +// order; stage-size changes only regroup exact integer sums (max |dot| = +// 6144*127*127 << 2^31), so results stay bit-identical. +template +__device__ __forceinline__ void mmac_prefill_stage( + DUFragment& + a_frag0, + DUFragment& + a_frag1, + DUFragment& acc00, + DUFragment& acc01, + DUFragment& acc10, + DUFragment& acc11, + int8_t* a_tile, + int8_t* b_tile, + int local_row, + int local_col) { + // kSkew8 (round 11, exact-shape arm only): A 88 / B 72 row strides; the + // load_frag8 ldm arguments below pick them up unchanged (same 8-byte + // aligned fragment reads, conflict-free 16-lane bank phases). + constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); + constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows + constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); + if constexpr (kBNMajor) { + if constexpr (kDirectFrag) { + // Round 10 exact-shape path: explicit load_frag8 8-byte ds_read_b64 + // fragment reads replacing the library loaders (see load_frag8 + // comment). Both steps of the 64-K stage are unrolled into + // independent fragment register sets so all eight reads of a stage + // issue before the MMACs; the kk = 0 MMAC group runs first, then the + // kk = kTileK group - the exact k0-outer/kk-inner int32 order and the + // element-to-slot fragment mapping of the library path are preserved + // (f.x[0..7] = the same 8 bytes in the same little-endian order), so + // results stay bit-identical. + static_assert(kStageK == 2 * kTileK, + "kDirectFrag two-step unroll assumes kStageK == 2*kTileK"); + DUFragment + a_frag0_1, a_frag1_1; + DUFragment + b_frag0, b_frag1, b_frag0_1, b_frag1_1; + load_frag8(a_frag0, a_tile + local_row * kAStride, kAStride); + load_frag8(a_frag1, a_tile + (local_row + kTileM) * kAStride, kAStride); + load_frag8(b_frag0, b_tile + local_col * kBNStride, kBNStride); + load_frag8(b_frag1, b_tile + (local_col + kTileN) * kBNStride, + kBNStride); + load_frag8(a_frag0_1, a_tile + local_row * kAStride + kTileK, + kAStride); + load_frag8(a_frag1_1, + a_tile + (local_row + kTileM) * kAStride + kTileK, kAStride); + load_frag8(b_frag0_1, b_tile + local_col * kBNStride + kTileK, + kBNStride); + load_frag8(b_frag1_1, + b_tile + (local_col + kTileN) * kBNStride + kTileK, + kBNStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc00, a_frag0_1, b_frag0_1, acc00); + du_mma_sync(acc01, a_frag0_1, b_frag1_1, acc01); + du_mma_sync(acc10, a_frag1_1, b_frag0_1, acc10); + du_mma_sync(acc11, a_frag1_1, b_frag1_1, acc11); + } else { + DUFragment + b_frag0, b_frag1; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + du_load_matrix_sync( + b_frag0, b_tile + local_col * kBNStride + kk, kBNStride); + du_load_matrix_sync( + b_frag1, b_tile + (local_col + kTileN) * kBNStride + kk, + kBNStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + } + } else { + DUFragment + b_frag0, b_frag1; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + } +} + +// Large-prefill path: one block computes a kBlockM x kBlockN output tile. +// (kBlockM/32) x (kBlockN/32) wavefronts each own a 32x32 quadrant (four +// m16n16k32 DUMMA accumulators), while the block cooperatively stages +// A[kBlockM, kStageK] and B in bank-skewed LDS. +// +// kDoubleBuffer == false: single-buffered K loop (two barriers per stage); +// stage s+1's global loads cannot start until the stage-s barrier, so each +// stage serializes one global-load round trip in front of the compute +// (bootstrap structure; co-resident block streams hide it at the CU level). +// +// kDoubleBuffer == true: software-pipelined K loop (round 3). Stage s+1's +// global loads are issued into the idle LDS buffer before the stage-s MMAC +// burst, overlapping the serialized global-load latency with compute. One +// barrier per stage (plus one prologue): it both retires stage-s fragment +// reads (WAR on the buffer just read) and makes stage s+1's ds_writes +// visible (RAW on the prefetched buffer). The gate_up shape uses +// <64, 64, 64, true, true>: two 64-K tiles = 20,480 B/block x 3 blocks/CU = +// 61,440 <= 64 KiB (occupancy preserved; a 128-K double buffer would need +// 38,912 B/block -> 1 block/CU, a harmful occupancy loss). +// +// kBNMajor selects the B operand layout: +// * false (identity): B staged [K, N] row-major, library row_major B +// fragments (byte-scattered LDS reads; correct for the raw [K, N] weight +// and every non-gate_up shape). +// * true (gate_up only): B staged [N, K] n-major with k contiguous (row +// stride 144 = 128+16, or 80 = 64+16 for the 64-K double buffer), +// library col_major B fragments -> each lane's elements are k-contiguous +// bytes instead of byte-scattered columns. This is the bank-safe consumer +// layout for the B side (the 17.3M PMC bank-conflict source in round 1). +// The one-time pack in launch_pack_w8a8_weight provides the [N, K] buffer +// for (k, n) == (6144, 512) only. +template +__global__ __launch_bounds__(kBlockM * kBlockN / 1024 * kWaveSize) void +w8a8_dumma_prefill_tiled_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWavesN = kBlockN / 32; + constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + // 16-byte bank skew: keeps int4 staging stores aligned while spreading the + // fragment rows across distinct LDS banks. The staging loops live in + // stage_prefill_tiles below; the strides remain here as the static-asserted + // alignment contract of every instantiation. kSkew8 (round 11, exact-shape + // arm only) switches A/B to the down_proj-validated 8-mod-16 strides + // 88/72; LDS bytes stay A[64,88] + B[64,72] x 2 = 20,480 B/block so the + // 3-blocks/CU occupancy of the accepted kernel is preserved exactly. + constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); + constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows + constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); + constexpr int kBNInt4 = kBlockN * (kStageK / sizeof(int4)); + // Double buffering uses kBuffers = 2 LDS tile copies; the gate_up shape + // pairs it with kStageK = 64 so 3 blocks/CU is preserved (see kernel + // comment above). + constexpr int kBuffers = kDoubleBuffer ? 2 : 1; + constexpr int kABytes = kBlockM * kAStride; + // One shared B tile sized for the larger of the two layouts so an inactive + // arm never allocates dead LDS (keeps <64,64,128> at 19,456 B total and + // 3 blocks/CU). kSkew8 implies kBNMajor, so only the n-major arm is live. + constexpr int kBMaxBytes = + kSkew8 ? kBlockN * kBNStride + : (kStageK * kBStride > kBlockN * kBNStride + ? kStageK * kBStride + : kBlockN * kBNStride); + static_assert(kSkew8 ? (kAStride % 8 == 0) : (kAStride % sizeof(int4) == 0), + "A LDS row stride must stay int4 aligned (or 8-byte aligned " + "under kSkew8)"); + static_assert(kBStride % static_cast(sizeof(int4)) == 0, + "B LDS row stride must stay int4 aligned"); + static_assert( + kSkew8 ? (kBNStride % 8 == 0) : (kBNStride % sizeof(int4) == 0), + "n-major B LDS row stride must stay int4 aligned (or 8-byte aligned " + "under kSkew8)"); + static_assert(!kSkew8 || kBNMajor, + "kSkew8 is only defined for the packed n-major B layout"); + static_assert(!kSkew8 || kStageK == 64, + "kSkew8 strides 88/72 assume kStageK == 64"); + static_assert(kStageK % kTileK == 0, + "K stage must be a multiple of the DUMMA K unit"); + static_assert(kStageK % static_cast(sizeof(int4)) == 0, + "K stage must keep int4 staging aligned"); + static_assert(kBNInt4 % kThreads == 0, + "n-major B staging must divide evenly across threads"); + static_assert(!kPayload || (kBNMajor && kDoubleBuffer && kDirectFrag && + kSkew8 && kStageK == 64 && kBlockM == 64 && + kBlockN == 64), + "payload arm is specialized to the exact-shape <64,64,64> " + "double-buffered load_frag8/kSkew8 arm (one int4 per thread " + "per operand)"); + static_assert( + !kBPanel || (kBNMajor && kStageK == 64 && kBlockN == 128 && !kPayload), + "kBPanel is specialized to the down_proj panel pack (n-major 128x64 " + "per-stage panels, 64-K stages, non-payload pipeline)"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWavesN; + const int wave_col = wave - wave_row * kWavesN; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + + __shared__ __align__(16) int8_t a_tile[kBuffers * kABytes]; + __shared__ __align__(16) int8_t b_tile[kBuffers * kBMaxBytes]; + + DUFragment + a_frag0, a_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + if constexpr (kDoubleBuffer) { + if constexpr (kPayload) { + // Round 16 lead-2 loop-carried payload pipeline (exact-shape arm + // only): stage s+1's global loads are issued a FULL iteration early + // (top of iteration s-1) into a two-int4-per-thread register payload + // and stay outstanding across the stage-(s-1) fragment reads + MMAC + // burst and the stage barrier; the vmcnt wait lands at the top of + // iteration s (first use of the payload registers), immediately + // before the ds_write2 publish into the idle LDS buffer. The + // s_barrier fences outstanding vmem, so the wait can no longer sit + // right after the load issue: the global-load latency overlaps the + // whole compute segment instead of serializing in front of every + // stage (the round-11 ISA defect). Same 2 buffers, same one barrier + // per stage (prologue + 96 = 97/block), same 16 ds_read2_b64 + + // 8 ds_write2_b64 per block-stage, same staging bytes/addresses + // (bit-identical int32 order). + int4 a_payload, b_payload; + // Prologue: issue the stage-1 global loads first (so they ride behind + // the stage-0 inline staging), stage k0 = 0 into buffer 0 inline + // (same load+wait+write path as the accepted kernel), then the first + // barrier. + prefetch_stage_payload( + x_q, weight, tid, m0, n0, m, k, kStageK, a_payload, b_payload); + stage_prefill_tiles( + x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, 0); + __syncthreads(); + for (int k0 = 0; k0 < k; k0 += kStageK) { + const int buf = (k0 / kStageK) & 1; + if (k0 + kStageK < k) { + // Publish the payload carried from the previous iteration (stage + // k0+kStageK) into the idle buffer: the vmcnt wait for the loads + // issued one iteration earlier is emitted here at the first use, + // after the previous MMAC burst + barrier. + publish_stage_payload( + a_payload, b_payload, a_tile + (buf ^ 1) * kABytes, + b_tile + (buf ^ 1) * kBMaxBytes, tid); + } + if (k0 + 2 * kStageK < k) { + // Issue stage k0+2*kStageK's global loads into the payload + // registers (no wait); they stay in flight across the MMAC burst + // and the barrier below and are waited at the next iteration's + // publish (RAW). + prefetch_stage_payload( + x_q, weight, tid, m0, n0, m, k, k0 + 2 * kStageK, a_payload, + b_payload); + } + // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. + // k0-outer / kk-inner int32 accumulation order. + mmac_prefill_stage( + a_frag0, a_frag1, acc00, acc01, acc10, acc11, + a_tile + buf * kABytes, b_tile + buf * kBMaxBytes, + local_row, local_col); + // One barrier per stage: (a) WAR - all stage-s fragment reads of the + // buffer just read are retired before it becomes the publish target + // two stages later; (b) RAW - every thread's stage-(s+1) ds_writes + // are visible before the next iteration reads them. + __syncthreads(); + } + } else { + // Prologue: stage k0 = 0 into buffer 0, then the pipelined loop. + stage_prefill_tiles( + x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, 0); + __syncthreads(); + for (int k0 = 0; k0 < k; k0 += kStageK) { + const int buf = (k0 / kStageK) & 1; + if (k0 + kStageK < k) { + // Prefetch stage s+1 into the idle buffer: global loads issue here, + // before the stage-s MMAC burst, and complete (RAW) behind the + // single barrier below. + stage_prefill_tiles( + x_q, weight, a_tile + (buf ^ 1) * kABytes, + b_tile + (buf ^ 1) * kBMaxBytes, tid, m0, n0, m, n, k, + k0 + kStageK); + } + // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. + // k0-outer / kk-inner int32 accumulation order. + mmac_prefill_stage( + a_frag0, a_frag1, acc00, acc01, acc10, acc11, + a_tile + buf * kABytes, b_tile + buf * kBMaxBytes, + local_row, local_col); + // One barrier per stage: (a) WAR - all stage-s fragment reads of the + // buffer just read are retired before it becomes the prefetch target + // two stages later; (b) RAW - every thread's stage-(s+1) ds_writes are + // visible before the next iteration reads them. + __syncthreads(); + } + } + } else { + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[kBlockM, kStageK] and B into the single LDS buffer + // (kSkew8 8-mod-16 row strides 88/72 with ds_write2_b64 staging). + stage_prefill_tiles( + x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, k0); + __syncthreads(); + + // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. + // k0-outer / kk-inner int32 accumulation order (kDirectFrag: eight + // load_frag8 ds_read2_b64 fragment reads, zero reassembly VALU). + mmac_prefill_stage( + a_frag0, a_frag1, acc00, acc01, acc10, acc11, + a_tile, b_tile, local_row, local_col); + // Retire this stage's reads before the next stage's staging writes + // overwrite the single LDS buffer (WAR hazard across waves). + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Generic scalar fallback: one output element per grid-stride iteration with +// exact int32 accumulation. Correct for any (m, n, k), including the small-M +// API shapes (M=2/16) and any unmatched geometry. The weight is the identity +// [K, N] row-major layout for every (k, n) EXCEPT the gate_up weight +// (k, n) == (6144, 512), whose buffer is packed [N, K] n-major; this kernel +// decodes that layout so the paired gate_up tails stay correct. +__global__ __launch_bounds__(kScalarThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + // Decode the packed layout when this weight was packed that way: + // gate_up (k, n) == (6144, 512) and, since round 2 of this session, + // down_proj (k, n) == (256, 6144) - since round 5 in the PANEL layout + // [N/128][K/64][128][64] (n-major 64-k-byte rows per panel; byte-exact + // relation packed[((nt*4+kt)*128+nn)*64+kk64] == raw[(nt*128+nn)*k + + // kt*64+kk64] with nt = col/128, nn = col%128, kt = kk/64, kk64 = kk%64); + // identity [K, N] otherwise. + const bool b_panel = (k == 256 && n == 6144); + const bool b_nmajor = (k == 6144 && n == 512) || b_panel; + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + int32_t acc = 0; + const int8_t* a_row = x_q + static_cast(row) * k; + for (int kk = 0; kk < k; ++kk) { + int64_t b_off; + if (b_panel) { + const int nt = col / 128; + const int nn = col - nt * 128; + const int kt = kk / 64; + const int kk64 = kk - kt * 64; + b_off = (static_cast(nt * 4 + kt) * 128 + nn) * 64 + kk64; + } else { + b_off = b_nmajor ? static_cast(col) * k + kk + : static_cast(kk) * n + col; + } + acc += static_cast(a_row[kk]) * + static_cast(weight[b_off]); + } + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// Bootstrap pack: identity device-to-device copy for every (k, n) pair +// (raw [K, N] layout preserved). Runs outside the timed/CUDA-Graph region. +__global__ __launch_bounds__(kScalarThreads) void +w8a8_pack_identity_kernel( + const int8_t* __restrict__ raw_weight, + const float* __restrict__ weight_scale, + int8_t* __restrict__ packed_weight, + float* __restrict__ packed_weight_scale, + int64_t weight_elems, + int n) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < weight_elems) { + packed_weight[idx] = raw_weight[idx]; + } + if (idx < n) { + packed_weight_scale[idx] = weight_scale[idx]; + } +} + +// Shape-specific gate_up pack: transpose raw [K, N] -> packed [N, K] +// (n-major rows, k contiguous) for (k, n) == (6144, 512) only. Byte-exact +// index relation: packed[cc * k + kk] == raw[kk * n + cc]. Each thread +// gathers the 16 bytes of one aligned int4 k-block from 16 raw rows (stride +// n) and stores one aligned int4, keeping the packed write side fully +// coalesced (same gather + aligned int4-store pattern as the validated +// o_proj transpose repair). Runs outside the timed/CUDA-Graph region; same +// byte count, same graph-stable buffer. +__global__ __launch_bounds__(kScalarThreads) void +w8a8_pack_gateup_transpose_kernel( + const int8_t* __restrict__ raw_weight, + const float* __restrict__ weight_scale, + int8_t* __restrict__ packed_weight, + float* __restrict__ packed_weight_scale, + int k, + int n) { + const int k4 = k / static_cast(sizeof(int4)); + const int64_t total4 = static_cast(n) * k4; + const int64_t stride4 = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx4 = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx4 < total4; idx4 += stride4) { + const int cc = static_cast(idx4 / k4); + const int kblk = static_cast(idx4 - static_cast(cc) * k4); + const int kk = kblk * static_cast(sizeof(int4)); + const int8_t* src = raw_weight + static_cast(kk) * n + cc; + int4 v; + char* vp = reinterpret_cast(&v); +#pragma unroll + for (int i = 0; i < static_cast(sizeof(int4)); ++i) { + vp[i] = src[static_cast(i) * n]; + } + reinterpret_cast(packed_weight + static_cast(cc) * k + + kk)[0] = v; + } + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < n; idx += stride) { + packed_weight_scale[idx] = weight_scale[idx]; + } +} + +// Round 5 (this session): shape-specific down_proj panel pack for +// (k, n) == (256, 6144) only. Raw [K, N] -> panel layout +// [N/128][K/64][128][64]: n-major rows of 64 k-bytes per panel, panels +// ordered (n-tile outer, k-tile inner). Byte-exact index relation (verified +// by simulation): +// packed[((nt * (k/64) + kt) * 128 + nn) * 64 + kk64] == +// raw[(nt * 128 + nn) * k + (kt * 64 + kk64)] +// with nt in [0, n/128), kt in [0, k/64), nn in [0, 128), kk64 in [0, 64). +// Each packed int4 (idx4) decodes as v = idx4 % 4, row_p = idx4 / 4, +// nn = row_p % 128, kt = (row_p / 128) % (k/64), nt = (row_p / 128) / (k/64), +// kk = kt * 64 + v * 16, cc = nt * 128 + nn, and +// packed[idx4*16 .. +16] == raw[(kk + i) * n + cc] for i in 0..15 (the same +// 16-row gather at column cc as the gate_up transpose, stored into a +// different packed offset; each thread stores one aligned int4, so the +// packed write side stays fully coalesced). The panel layout makes every +// per-block B stage tile (128 n x 64 k) ONE contiguous 8 KiB region, so the +// timed kernel's B staging reads 1 KiB per warp instruction contiguously. +// Runs outside the timed/CUDA-Graph region; same byte count, same +// graph-stable buffer, byte-exact vs the raw logical [K, N] weight. +__global__ __launch_bounds__(kScalarThreads) void +w8a8_pack_downproj_panel_kernel( + const int8_t* __restrict__ raw_weight, + const float* __restrict__ weight_scale, + int8_t* __restrict__ packed_weight, + float* __restrict__ packed_weight_scale, + int k, + int n) { + constexpr int kBlockN = 128; + constexpr int kStageK = 64; + const int64_t total4 = (static_cast(n) * k) / 16; + const int64_t stride4 = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx4 = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx4 < total4; idx4 += stride4) { + const int v = static_cast(idx4 & 3); + const int64_t row_p = idx4 >> 2; + const int nn = static_cast(row_p % kBlockN); + const int64_t kt_tmp = row_p / kBlockN; // nt * (k/kStageK) + kt + const int kt = static_cast(kt_tmp % (k / kStageK)); + const int nt = static_cast(kt_tmp / (k / kStageK)); + const int kk = kt * kStageK + v * 16; + const int cc = nt * kBlockN + nn; + const int8_t* src = raw_weight + static_cast(kk) * n + cc; + int4 val; + char* vp = reinterpret_cast(&val); +#pragma unroll + for (int i = 0; i < static_cast(sizeof(int4)); ++i) { + vp[i] = src[static_cast(i) * n]; + } + reinterpret_cast(packed_weight + idx4 * 16)[0] = val; + } + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < n; idx += stride) { + packed_weight_scale[idx] = weight_scale[idx]; + } +} + +} // namespace + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // No split-K: the provided workspace is not touched. The timed operator + // performs no allocation, synchronization, packing, or default-stream + // launch; only the caller-provided out is written. + (void)workspace; + (void)workspace_bytes; + auto* out_bf16 = static_cast(out); + + // Native INT8 DUMMA m16n16k32 tiled path for every large-M shape. Both + // assigned shapes qualify: gate_up (4096, 512, 6144) and down_proj + // (4096, 6144, 256) both have M >= 128, N % 128 == 0 and K % 128 == 0. + // Explicit dispatch by launch geometry; unmatched shapes fall through to + // the scalar path (no shape is specialized away). + if (m >= kPrefillMinM && (n % kDefaultBlockN == 0) && + (k % kPrefillStageK == 0)) { + // Exact-shape guard for the assigned gate_up shape (round-3 pipeline + // decision, round-10 load_frag8, round-11 kSkew8, round-16 kPayload): + // cooperative A+B staging with the n-major [N, K] packed B operand + // (kBNMajor) consumed as col_major fragments, software-pipelined across + // K tiles (kDoubleBuffer) with a 64-K double buffer: 64x64 tile, 512 + // blocks, 20,480 B LDS/block, 3 blocks/CU (61,440 B) - occupancy + // preserved vs the accepted round-2 single buffer (19,456 B x 3 = + // 58,368 B). One barrier per 64-K stage + prologue (97 for 96 stages). + // Round 11 stages A/B at the 8-mod-16 strides 88/72 (kSkew8) so the + // fragment reads are bank-conflict-free (LDS bytes and occupancy + // unchanged). Round 16 (kPayload) issues the stage-(s+1) global loads a + // full iteration early into a two-int4-per-thread register payload and + // publishes them into the idle LDS buffer at the top of iteration s, so + // the vmcnt wait lands after the stage-s MMAC burst instead of in front + // of it (LDS bytes, barriers, traffic, occupancy and results unchanged). + if (m == 4096 && n == 512 && k == 6144) { + const dim3 grid(n / 64, (m + 63) / 64); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel< + 64, 64, 64, true, true, true, true, true>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + // Exact-shape guard for the assigned down_proj shape (round 2 of this + // session): the generic <64,128,128> identity-layout arm measured + // 314.72 us median / 315.18 us p90 with PMC lds_bank_conflicts = + // 7,864,320, ~600 LDS instructions and heavy per-byte fragment + // reassembly VALU per block-stage (library row_major loaders on the + // identity [K,N] B layout). The round-2 guard transplanted the file's + // lineage-validated gate_up machinery: n-major packed B (kBNMajor, pack + // in launch_pack_w8a8_weight for (k,n) == (256,6144)), 64-K + // double-buffered K loop (kDoubleBuffer), explicit 8-byte load_frag8 + // fragment reads (kDirectFrag), 8-mod-16 bank-skewed LDS strides + // A=88/B=72 (kSkew8); accepted at 286.53 us median / 286.64 us p90 with + // PMC lds_instructions 540,672, conflicts 1,179,648, lds_wait 162,664. + // + // Round 3 (this session, pipeline round): compare single and double + // buffering across K tiles. The (4096, 6144, 256) guard now routes to + // w8a8_dumma_prefill_tiled_kernel<64, 128, 64, true, false, true, true> + // - the SINGLE-buffered form of the accepted round-2 arm: same 64-K + // stage, same kBNMajor/kDirectFrag/kSkew8, same tile/grid/epilogue, but + // kDoubleBuffer = false, so each stage serializes one global-load round + // trip (load -> wait -> ds_write2_b64) in front of its MMAC burst with + // two barriers per stage (8/block vs the accepted 5/block), and the LDS + // halves from 29,696 to 14,848 B/block (A[64,88] + B[128,72] x 1). The + // accepted kernel's exact code object shows the double-buffered loop + // already overlaps stage-(s+1) loads with the stage-s MMAC burst + // (global_load x2 -> ds_read2_b64 x4 -> v_mmac x8 -> s_waitcnt + // vmcnt(0) lgkmcnt(0) -> s_barrier), so this round measures whether that + // overlap is worth the extra 14,848 B of LDS: if the single-buffered + // median is not strictly below the accepted 286.53 us (with the p90 + // noise guard), double buffering is retained per the pipeline mandate. + // Results stay bit-identical: stage size, k0-outer/kk-inner int32 order + // and the element-to-slot fragment mapping are unchanged (only the + // buffer count and barrier count differ), so mismatch 0 is expected. + // + // Round 5 (this session, PACKING round): the (4096, 6144, 256) guard + // now routes to w8a8_dumma_prefill_tiled_kernel<64, 128, 64, true, + // false, true, true, false, true> (new kBPanel template flag): the + // one-time pack for (k, n) == (256, 6144) becomes the panel layout + // [N/128][K/64][128][64] (n-major 64-k-byte rows per panel, byte-exact + // vs the raw [K,N] weight, same byte count / graph-stable buffer, + // out-of-timed/out-of-Graph) so each block's per-stage B tile is one + // contiguous 8 KiB region and every warp's int4 loads read 1 KiB + // contiguous (full L2 line utilization) instead of 16 rows x 64 B at + // 256-B stride (half-used lines). Everything else is byte-identical to + // the round-3 winner: same single-buffered 64-K stage, same + // kBNMajor/kDirectFrag/kSkew8, same 64x128 tile / 512 threads / 8 waves, + // same grid (48,64) = 3072 blocks, same LDS A[64,88] + B[128,72] = + // 14,848 B/block @ 2 blocks/CU, same ds_write2_b64 staging and + // load_frag8 fragment reads (LDS bytes and bank behavior untouched), + // same k0-outer/kk-inner int32 order and element-to-slot fragment + // mapping -> mismatch 0 / max_abs_error 0.0 expected. + if (m == 4096 && n == 6144 && k == 256) { + const dim3 grid(n / 128, (m + 63) / 64); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel< + 64, 128, 64, true, false, true, true, false, true>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (n % 128 == 0) { + // 64x128 tile, 8 waves (32x32 quadrant per wave). + // down_proj: grid (6144/128, 4096/64) = (48, 64) = 3072 blocks. + const dim3 grid(n / 128, (m + 63) / 64); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 128, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (m % 128 == 0) { + // 128x64 tile, 8 waves. + const dim3 grid(n / 64, m / 128); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<128, 64, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + { + // 64x64 tile, 4 waves (generic default). + const dim3 grid(n / 64, (m + 63) / 64); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 64, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + } + + // Generic scalar fallback for every unmatched (m, n, k) and M < 128 + // (including the paired M=2/M=16 API shapes with the same (N, K), which + // decode the packed [N, K] gate_up layout when (k, n) == (6144, 512)). + const int64_t total = static_cast(m) * n; + const unsigned blocks = static_cast( + (total + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Assigned gate_up weight (K=6144, N=512): one-time [K,N] -> [N,K] + // transpose pack (n-major, k contiguous) so the timed GEMM reads B + // fragments as k-contiguous bytes from the staged LDS tile. Runs outside + // the timed/CUDA-Graph region; same byte count, same graph-stable buffer, + // byte-exact vs the raw logical [K,N] weight. + if (k == 6144 && n == 512) { + const int64_t total4 = static_cast(n) * (k / 16); + const unsigned blocks = static_cast( + (total4 + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_pack_gateup_transpose_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + raw_weight, weight_scale, packed_weight, packed_weight_scale, k, n); + return; + } + // Round 5 (this session): the assigned down_proj weight (K=256, N=6144) + // gets the PANEL pack [N/128][K/64][128][64] (n-major 64-k-byte rows per + // panel; each per-block B stage tile is one contiguous 8 KiB region so + // the timed kernel's B staging reads 1 KiB per warp instruction + // contiguously). One-time, out-of-timed/out-of-Graph, same byte count, + // same graph-stable buffer, byte-exact vs the raw [K,N] weight. + if (k == 256 && n == 6144) { + const int64_t total4 = (static_cast(n) * k) / 16; + const unsigned blocks = static_cast( + (total4 + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_pack_downproj_panel_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + raw_weight, weight_scale, packed_weight, packed_weight_scale, k, n); + return; + } + // Identity device-to-device copy for every other (k, n) pair, including + // unmatched (K, N). Runs outside the timed/CUDA-Graph region. + const int64_t weight_elems = static_cast(k) * n; + const int64_t total = weight_elems > n ? weight_elems : n; + const unsigned blocks = static_cast( + (total + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_pack_identity_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + raw_weight, weight_scale, packed_weight, packed_weight_scale, + weight_elems, n); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_gate_up_proj.hip new file mode 100644 index 00000000..fec33f67 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_gate_up_proj.hip @@ -0,0 +1,1083 @@ +// @@variant shape=glm_tp8_shared_gate_up_proj_m4096 commit=df2c88db09d3836df8ff7a4e99e19584fe6d4163 added=2026-08-31 +// median_us=273.8 p90_us=275.3 speedup=77.62 baseline_us=2.125e+04 +// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 +// MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). +// worker_3 (iteration 3), assigned shapes: +// glm_tp8_shared_gate_up_proj_m4096 : (M, N, K) = (4096, 512, 6144) +// glm_tp8_shared_down_proj_m4096 : (M, N, K) = (4096, 6144, 256) +// +// This file provides the two stable host launch symbols consumed by +// csrc/bindings.cpp: +// * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's +// HIP stream; no allocation, packing, +// autotuning, or synchronization. Only the +// caller-provided out and workspace are +// touched (workspace is unused this round). +// * launch_pack_w8a8_weight(...) - out-of-timed-region packing. Identity +// device-to-device copy for every (k, n) +// pair EXCEPT the assigned gate_up weight +// (k, n) == (6144, 512), which is packed +// once as a [K, N] -> [N, K] transpose +// (n-major rows, k contiguous) so the timed +// GEMM consumes B fragments as +// k-contiguous bytes. The pack stays +// out-of-timed, out-of-Graph, same byte +// count, same graph-stable buffer. +// +// Bootstrap strategy (correctness-first, usable profiling baseline): +// * Native INT8 DUMMA m16n16k32 tiled kernel with int32 accumulation for +// every assigned shape with M >= 128 (both assigned shapes qualify). +// * 2-D macro-tile family (64x64, 64x128, 128x64) with one 32x32 quadrant +// per wavefront (four m16n16k32 int32 accumulators kept resident over the +// whole K loop). Single-buffered K loop (K stage 128, two barriers per +// stage): no split-K, no raw asm, no speculative double-buffer pipeline. +// +// Round 1 (baseline): the 64x128 tile measured 517.4 us median / 49.81 TOPS; +// routing the exact gate_up shape to the 64x64 tile measured 432.5 us +// (accepted; 19.6% improvement). +// +// Round 2 (operand-reuse, mandated decision): cooperative A+B LDS staging is +// kept and direct per-wave global loads are rejected (A is reused 64x per +// macro-tile via kBlockN and B 64x via kBlockM; direct loads would refetch +// both operands once per wavefront, 4x the 28.3 MB staged global traffic -> +// ~1.73M vmem_read instructions vs 432K, and the TP4 gate_up direct-A arm +// measured 7.9x slower). The focused change makes the staged B layout +// bank-safe for the fragment consumers: the gate_up weight is packed once +// [K,N] -> [N,K] (n-major, k contiguous) and B fragments are loaded +// col_major from a 144-B (128+16) row-strided LDS tile, replacing the +// byte-scattered row_major B reads on the identity [K,N] layout that are the +// 17.3M PMC bank-conflict source. Global loads stay 16-byte vectorized and +// coalesced on both operands. LDS stays 19,456 B/block -> 3 blocks/CU +// (58,368 B, unchanged; round-1's "4 blocks/CU" estimate was impossible on +// both LDS (4x19,456 = 77,824 B > 64 KiB) and VGPR (4x256x72 = 73,728 > +// 65,536)). down_proj keeps the identity 64x128 arm; every unmatched shape +// keeps the generic arms and the scalar fallback (which decodes the packed +// [N,K] layout when (k, n) == (6144, 512) so the paired M=2/M=16 gate_up +// tails stay correct). +// * A[kBlockM, kStageK] and B staged in bank-skewed LDS (row strides +// 144/80/144 = 16-byte-aligned non-power-of-two), consumed with the +// library du_load_matrix_sync loaders: A row_major (unchanged, load- +// bearing on every shape), B row_major for the identity [K, N] layout +// (down_proj + unmatched shapes) and B col_major for the gate_up +// [N, K] packed layout (round 2). +// * Fused direct fragment -> x_scale -> weight_scale -> bf16 epilogue using +// the verified gfx928 accumulator ownership (row = lane & 15, +// col_mod4 = lane >> 4, x[i] -> columns col_mod4 + 4*i). The int32 dot is +// exact for every assigned shape (max |dot| = 6144*127*127 << 2^31), the +// int32 -> float conversion is exact, and the single fp32 multiply +// followed by one bfloat16 rounding matches the reference bit-for-bit. +// * Tail M rows are zero-filled on load and masked on store, so any +// M >= 128 is supported by the tiled path. +// * All other (m, n, k) - including every M < 128 API case and the paired +// M=2/M=16 shapes with the same (N, K) - go to a scalar int8/int32 +// fallback that decodes the [N, K] packed layout when (k, n) == +// (6144, 512) and the identity [K, N] layout otherwise. +// +// Round 3 (pipeline, mandated decision): compare single and double buffering +// across K tiles. The exact (4096, 512, 6144) guard now routes to +// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true>: a software- +// pipelined 64-K double-buffered K loop (prologue + one barrier per stage, +// 97 barriers for 96 stages) that issues stage s+1's global loads into the +// idle LDS buffer before the stage-s MMAC burst, overlapping the +// global-load latency that the round-2 single buffer serializes in front of +// every stage (ISA-verified order: global_load_dwordx4 -> ds_write_b128 -> +// s_waitcnt -> s_barrier -> ds_read2 x16 -> v_mmac x16 -> s_barrier, with +// the next stage's loads gated behind the previous stage's barrier). The +// 128-K double-buffered variant is arithmetically impossible at 3 blocks/CU +// (2 x 19,456 = 38,912 B/block x 3 = 116,736 > 64 KiB -> 1 block/CU, a +// harmful occupancy loss), so the pipelined comparison uses the 64-K +// double buffer: LDS 20,480 B/block x 3 = 61,440 <= 64 KiB and +// VGPR 64-72 x 256 x 3 <= 65,536, i.e. 3 blocks/CU preserved (no occupancy +// confound). Everything else is byte-identical: same int4 staging pattern +// (same 28.3 MB global traffic and ~432K vmem_read instructions), same +// k0-outer/kk-inner int32 order (stage size changes only regroup exact +// integer sums, so results stay bit-identical), same col_major B fragments +// (stride 80 = 64+16 bank skew), same fused epilogue, same pack, same +// generic single-buffered arms and scalar fallback. Decision rule (mandate): +// retain double buffering only if PMC shows reduced VMEM/LDS stall evidence +// (lds_wait_instructions drops below ~1.2M, lds_instructions stays ~1.97M, +// occupancy stays 3 blocks/CU) AND the median improves >2%; otherwise the +// round's evidence retains the single-buffered round-2 structure. +// +// Round 10 (fragment-load-form round, restarted from the accepted round-3 +// source after the round-9 infra kill): replace the library +// du_load_matrix_sync fragment loads of the exact-shape gate_up kernel with +// the lineage-validated explicit load_frag8 8-byte ds_read_b64 loader. The +// exact (4096, 512, 6144) guard now routes to +// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true, true> (new +// kDirectFrag template arm). The round-3 ISA shows every fragment load as +// ds_read2_b32 + per-byte v_and/v_or_b32_sdwa/v_or3 mask-OR "identity" +// reassembly (~9.4M of the 15.43M VALU instructions per replay); loading the +// same 8 bytes per lane as one int64 (ds_read_b64) removes that VALU and +// halves the fragment bank-conflict factor (32-lane dword phases at stride +// 80 -> 16-lane 8-byte phases). Element-to-slot fragment mapping, v_mmac +// operands, exact k0-outer/kk-inner int32 accumulation, staging/global +// traffic, barriers, tile, grid, LDS bytes, occupancy (3 blocks/CU) and the +// fused epilogue are unchanged; results stay bit-identical. The generic arms +// (<64,128,128>, <128,64,128>, <64,64,128>, all kBNMajor=false) keep the +// library loaders and byte-identical codegen/resources (if constexpr). +// +// Round 11 (LDS bank-skew round, HIP-only; ISA policy keeps raw asm +// disallowed because the HIP plateau rule - 8 valid HIP rounds within +/-2% +// of the then-current best - is unmet). The accepted round-10 kernel still +// shows PMC lds_bank_conflicts = 9,437,184, EXACTLY unchanged from the +// round-3 code object (the -9.07% round-10 win came from VALU removal, not +// from conflicts). 9,437,184 = 192 x 512 blocks x 96 stages: every +// block-stage pays 192 conflict cycles = 128 from the 16 ds_read2_b64 +// fragment reads + 64 from the 8 ds_write_b128 staging stores (the 16-B/lane +// store floor). The read half is a stride defect: at stride 80 = 20 dwords, +// 20*r mod 32 has period 8, so every 16-lane read phase groups rows r and +// r+8 onto the same bank pair - a 2-way conflict on ALL 4 sub-phases of +// every fragment read. The exact (4096, 512, 6144) guard now routes to +// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true, true, true> (new +// kSkew8 template arm) which stages both operands at the down_proj-validated +// 8-mod-16 row strides A = 88 (64+24) and B = 72 (64+8) instead of 80: +// 22*r and 18*r mod 32 are distinct over r = 0..15 and gcd(22,32) = +// gcd(18,32) = 2 does not divide 1, so every 16-lane read phase lands on 16 +// DISTINCT bank pairs - the fragment reads become conflict-free (128 of the +// 192 cycles removed; only the 2-way 16-B/lane staging-store floor of 64 +// cycles remains -> lds_bank_conflicts expected ~3.15M). The 8-mod-16 row +// starts forbid 16-byte staging stores, so each 16-byte int4 is staged as +// two 8-byte halves (ds_write2_b64), the validated down_proj staging form +// (its B tile is exactly stride 72; its A tile exactly stride 88). Global +// loads (int4, same 432,128 vmem_read), pack layout, barrier structure +// (97/block), tile/grid (512 blocks), LDS bytes (A[64,88] + B[64,72] x 2 = +// 20,480 B/block -> 3 blocks/CU unchanged), VGPR, the k0-outer/kk-inner +// int32 order, the element-to-slot fragment mapping and the fused epilogue +// are untouched; only the LDS row strides and the staging store granularity +// change, so results stay bit-identical. Generic arms (<64,128,128>, +// <128,64,128>, <64,64,128>) keep kSkew8 = false and byte-identical +// codegen/resources (if constexpr). +// +// Round 16 (lead-2 loop-carried payload round, HIP-only; raw asm stays +// DISALLOWED - plateau=false, the recent-valid-improvement window +// [-34.44%, +2.82%, +9.97%] does not meet three valid HIP rounds within +// [-2%, +2%) of the then-current best, and no prior ISA-guided round +// recorded a compiler limitation with target instructions). The accepted +// round-11 exact code object (digest 9889820e, median 293.30 us / p90 +// 293.545) shows the software-pipeline intent is DEFEATED by the compiler's +// wait placement: per stage the steady-state loop is +// global_load_dwordx4 (A prefetch s+1) -> global_load_dwordx4 (B prefetch +// s+1) -> s_waitcnt vmcnt(1) -> s_waitcnt vmcnt(0) <- waits IMMEDIATELY +// -> 2 x ds_write2_b64 (staging stores) -> 4 x ds_read2_b64 (fragment +// reads, stage s) -> 8 x v_mmac CONTIGUOUS (stage s) -> s_barrier -> loop. +// i.e. the stage-(s+1) global loads are waited ~12-16 instructions after +// issue, BEFORE the stage-s MMAC burst, so the full global-load latency +// (~600-1000+ cycles under 120-CU L2 contention) sits on the per-stage +// critical path with zero overlap (the round-12 analysis' +// "prefetch global load -> ds_write2_b64 -> s_barrier -> ds_read2_b64 -> +// v_mmac" chain; PMC: 6.0M VALU + 1.18M LDS + 465K VMEM wave-instructions +// over 293 us x 120 CUs, i.e. ~95-99% issue-stalled). Iteration history +// rules out the other levers on this exact shape: tile aspect (iter 4 +// 32x64: 749.9 us; iter 8 64x128 w8: 395.4 us), 2-wave quadrants (iter 6: +// 629.0 us - wave count per CU is load-bearing), LDS-side fragment +// prefetch (iter 7: 337.97 us, flat/worse), stage 32 / 6 blocks per CU +// (iter 12: 447.4 us - per-stage fixed cost, not occupancy, is the lever). +// The one untested mechanism that directly removes the exposed vmcnt wait: +// a LEAD-2 LOOP-CARRIED REGISTER PAYLOAD (the TP4 o_proj lineage's +// validated "stage-top publish" form). The exact (4096, 512, 6144) guard +// now routes to w8a8_dumma_prefill_tiled_kernel<64,64,64,true,true,true, +// true,true> (new kPayload template arm): each thread keeps the one A int4 +// + one B int4 of stage s+2 in VGPR across the stage-(s+1) barrier, issues +// those global loads at the TOP of iteration s (no wait), and PUBLISHES +// the payload carried from iteration s-1 into the idle LDS buffer at the +// TOP of iteration s (before the stage-s MMAC burst). The first use of the +// payload registers is the next iteration's publish, and the s_barrier +// fences outstanding vmem, so the compiler can no longer place +// s_waitcnt vmcnt(0) immediately after the loads: the wait lands after the +// stage-s fragment reads + MMAC burst (immediately before the barrier), and +// the global-load latency overlaps the whole compute segment instead of +// serializing in front of it. One barrier per 64-K stage is preserved +// (prologue + 96 = 97/block); the publish targets the idle buffer (WAR +// retired by the previous barrier, RAW retired by the stage barrier), so no +// second barrier is introduced. Everything else is byte-identical: 64x64 +// tile / 256 threads / 4 waves / 32x32 quadrant per wave, grid (8,64) = +// 512 blocks, kSkew8 strides 88/72 with ds_write2_b64 staging (8 per +// block-stage), same 16 ds_read2_b64 fragment reads per block-stage (same +// lds_instructions ~1,179,648), same 20,480 B LDS/block -> 3 blocks/CU +// (payload costs 8 VGPRs: ~54 total, 54 x 256 x 3 = 41,472 <= 65,536; no +// spill risk at 3 blocks/CU), same int4 global loads (same ~432,128 +// vmem_read / 32,768 vmem_write), same k0-outer/kk-inner int32 order and +// the same element-to-slot fragment mapping - the payload changes only WHEN +// global bytes land in LDS, never which bytes or the accumulation order, so +// results stay bit-identical (mismatch 0 / max_abs_error 0.0 expected). +// Generic arms (<64,128,128>, <128,64,128>, <64,64,128>, all kBNMajor= +// false) keep kPayload = false and byte-identical codegen/resources +// (if constexpr). +// +// Header order is fixed by the control plane: hip_runtime first (du_mma.h +// is not self-contained before it), hip_bfloat16 second, du_mma.h last. + +#include +#include +#include + +#include + +namespace { + +// gfx928 INT8 DUMMA primitive: m16n16k32, int8 x int8 -> int32. +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; + +// gfx928 wavefront is 64 lanes; every blockDim must be a multiple of 64. +constexpr int kWaveSize = 64; + +// Large-prefill dispatch geometry. Each macro-tile config uses a 32x32 +// quadrant per wave, so waves per block = (BlockM/32) * (BlockN/32). +constexpr int kPrefillMinM = 128; +constexpr int kPrefillStageK = 128; +constexpr int kDefaultBlockM = 64; +constexpr int kDefaultBlockN = 64; + +constexpr int kScalarThreads = 256; + +using namespace du::dumma; + +// Verified gfx928 INT8 DUMMA accumulator ownership: +// row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> columns +// col_mod4 + 4*i. Scale by x_scale[row] and weight_scale[col] and store +// bf16 directly from the fragment (no accumulator LDS round trip). +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * xs * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Round 10 (fragment-load-form round): explicit 8-byte fragment loader +// replacing the library du_load_matrix_sync for the exact-shape gate_up +// kernel. Validated on the sibling TP8 down_proj accepted kernel (iter 8, +// -25.9%: 230.91 -> 171.19 us) and the TP4 gate_up lineage (iter 13 +// load_fragment8, -21%: 280.05 -> 220.71 us). For BOTH the library row_major +// matrix_a and col_major matrix_b m16n16k32 int8 loaders, lane l -> row +// (l & 15), k-quarter (l >> 4), and the loader fetches the same eight +// contiguous bytes at p[row*ldm + (l>>4)*8 .. +7] and stores them into the +// same x[0..7] slots in the same little-endian order (the round-3 ISA shows +// the compiler lowering this to ds_read2_b32 followed by a per-byte +// v_and/v_or_b32_sdwa/v_or3 mask-OR "identity" reassembly - ~9.4M dead-ish +// VALU instructions per replay, ~61% of the 15.43M VALU total). Loading the +// same 8 bytes as one int64 forces the 8-byte ds_read_b64 form instead: +// (a) the mask/OR reassembly VALU disappears, and (b) the LDS access +// changes from ds_read2_b32 (32-lane dword phases - at the stride-80 row +// pattern every phase groups rows 0-15 of k-quarter q with rows 0-15 of +// k-quarter q+1 onto the same 16 bank phases, a 4-way conflict per phase) +// to ds_read_b64 (16-lane phases: rows 0-15 land on 16 bank-pairs +// {20*r mod 32} = 8 distinct pairs hit twice - a 2-way conflict per phase, +// half of the current fragment-conflict factor). The element-to-slot +// mapping is unchanged (f.x[0..7] = the same 8 bytes in the same +// little-endian order), the v_mmac operands and the exact k0-outer/kk-inner +// int32 accumulation sequence are untouched, and the staging stores, LDS +// strides, barriers, tile, grid, occupancy and epilogue are byte-identical. +template +__device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, + unsigned ldm) { + const unsigned row = __lane_id() & 0xf; + const unsigned kq = __lane_id() >> 4; + const int64_t v = + *reinterpret_cast(p + row * ldm + (kq << 3)); + reinterpret_cast(f.x)[0] = v; +} + +// Stage one K tile (k0) of A[kBlockM, kStageK] and B into the caller-offset +// LDS buffers (bank-skewed row strides; int4 vectorized coalesced global +// loads, zero-filled tail M rows). kBNMajor == true stages B from the packed +// [N, K] n-major gate_up layout (n-major rows, k contiguous); otherwise from +// the identity [K, N] layout. Shared by the single-buffered and +// double-buffered K loops; the buffer offset (0 or kABytes/kBMaxBytes) is +// applied by the caller so one body serves both pipeline forms. +// +// kSkew8 (round 11, exact-shape arm only): A and B row strides become the +// down_proj-validated 8-mod-16 values 88 (kStageK + 24) and 72 (kStageK + 8) +// so every 16-lane fragment-read phase lands on 16 distinct bank pairs (see +// the round-11 header comment). The 8-mod-16 row starts forbid 16-byte +// stores, so each 16-byte int4 chunk is staged as two 8-byte halves +// (ds_write2_b64); the global int4 loads are byte-identical. +template +__device__ __forceinline__ void stage_prefill_tiles( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int8_t* a_tile, + int8_t* b_tile, + int tid, + int m0, + int n0, + int m, + int n, + int k, + int k0) { + constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + // 16-byte bank skew: keeps int4 staging stores aligned while spreading the + // fragment rows across distinct LDS banks. + constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); + constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows + constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); + static_assert(!kSkew8 || kStageK == 64, + "kSkew8 strides 88/72 assume kStageK == 64"); + constexpr int kAInt4PerRow = kStageK / sizeof(int4); + constexpr int kAInt4 = kBlockM * kAInt4PerRow; + constexpr int kBInt4PerRow = kBlockN / sizeof(int4); + constexpr int kBInt4 = kStageK * kBInt4PerRow; + constexpr int kBNInt4PerRow = kStageK / sizeof(int4); + constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; + // Stage A[kBlockM, kStageK] into LDS (zero-filled tail M rows). + for (int vec = tid; vec < kAInt4; vec += kThreads) { + const int local_row = vec / kAInt4PerRow; + const int v = vec - local_row * kAInt4PerRow; + const int global_row = m0 + local_row; + const int4 val = + global_row < m + ? *reinterpret_cast( + x_q + static_cast(global_row) * k + k0 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + if constexpr (kSkew8) { + // 8-mod-16 row starts: write the 16 bytes as two 8-byte halves so every + // store stays 8-byte aligned (ds_write2_b64 staging; same LDS bytes). + int8_t* dst = + a_tile + local_row * kAStride + v * static_cast(sizeof(int4)); + reinterpret_cast(dst)[0] = + reinterpret_cast(&val)[0]; + reinterpret_cast(dst)[1] = + reinterpret_cast(&val)[1]; + } else { + reinterpret_cast(a_tile + local_row * kAStride)[v] = val; + } + } + if constexpr (kBNMajor) { + // Stage B[kBlockN, kStageK] from the packed [N, K] weight (n-major: + // row = n_local, k contiguous). Same vectorized coalesced int4 pattern + // as the identity arm, one per 16 consecutive k bytes. + for (int vec = tid; vec < kBNInt4; vec += kThreads) { + const int local_col = vec / kBNInt4PerRow; + const int v = vec - local_col * kBNInt4PerRow; + const int4 val = *reinterpret_cast( + weight + static_cast(n0 + local_col) * k + k0 + + v * static_cast(sizeof(int4))); + if constexpr (kSkew8) { + int8_t* dst = + b_tile + local_col * kBNStride + + v * static_cast(sizeof(int4)); + reinterpret_cast(dst)[0] = + reinterpret_cast(&val)[0]; + reinterpret_cast(dst)[1] = + reinterpret_cast(&val)[1]; + } else { + reinterpret_cast(b_tile + local_col * kBNStride)[v] = val; + } + } + } else { + // Stage B[kStageK, kBlockN] into LDS from the identity [K, N] weight. + for (int vec = tid; vec < kBInt4; vec += kThreads) { + const int kk = vec / kBInt4PerRow; + const int v = vec - kk * kBInt4PerRow; + reinterpret_cast(b_tile + kk * kBStride)[v] = + *reinterpret_cast( + weight + static_cast(k0 + kk) * n + n0 + + v * static_cast(sizeof(int4))); + } + } +} + +// Round 16 (lead-2 loop-carried payload, exact-shape arm only): the register +// half of the lead-2 software pipeline. Each thread prefetches exactly one A +// int4 and one B int4 of a later K stage into thread-local registers and +// returns WITHOUT waiting (the vmcnt wait is emitted by the compiler at the +// first use, one iteration later). The <64,64,64> arm stages kStageK * 64 +// bytes per operand per stage, i.e. kAInt4 == kBNInt4 == kThreads == 256: +// one int4 per thread per operand, so the whole payload costs 8 VGPRs. The +// (row, v) mapping is the same as stage_prefill_tiles, so +// publish_stage_payload writes the exact same LDS addresses the staging +// loop would (bit-identical staging bytes, just loaded one stage earlier). +template +__device__ __forceinline__ void prefetch_stage_payload( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int tid, + int m0, + int n0, + int m, + int k, + int k0, + int4& a_payload, + int4& b_payload) { + constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + constexpr int kAInt4PerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kAInt4 = kBlockM * kAInt4PerRow; + constexpr int kBNInt4PerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; + static_assert(kAInt4 == kThreads && kBNInt4 == kThreads, + "payload arm requires exactly one int4 per thread per " + "operand (exact-shape <64,64,64> arm)"); + static_assert(kBNMajor, "payload arm stages the packed n-major B layout"); + // Same vec -> (local_row, v) mapping as stage_prefill_tiles, so the + // publish half stores to the same LDS addresses. + const int a_local_row = tid / kAInt4PerRow; + const int a_v = tid - a_local_row * kAInt4PerRow; + const int global_row = m0 + a_local_row; + // Tail M rows are zero-filled exactly as in stage_prefill_tiles. + a_payload = global_row < m + ? *reinterpret_cast( + x_q + static_cast(global_row) * k + k0 + + a_v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + const int b_local_col = tid / kBNInt4PerRow; + const int b_v = tid - b_local_col * kBNInt4PerRow; + b_payload = *reinterpret_cast( + weight + static_cast(n0 + b_local_col) * k + k0 + + b_v * static_cast(sizeof(int4))); +} + +// Round 16: the LDS half of the lead-2 payload pipeline. Writes the payload +// carried in registers (loaded by prefetch_stage_payload one stage earlier) +// into the caller-offset LDS tile with the exact same 8-byte-half +// ds_write2_b64 stores (kSkew8 strides 88/72) that stage_prefill_tiles +// emits, so each thread republishes exactly the bytes it loaded. +template +__device__ __forceinline__ void publish_stage_payload( + int4 a_payload, + int4 b_payload, + int8_t* a_tile, + int8_t* b_tile, + int tid) { + constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); + constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); + constexpr int kAInt4PerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kAInt4 = kBlockM * kAInt4PerRow; + constexpr int kBNInt4PerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; + static_assert(kAInt4 == kThreads && kBNInt4 == kThreads, + "payload arm requires exactly one int4 per thread per " + "operand (exact-shape <64,64,64> arm)"); + static_assert(kBNMajor, "payload arm stages the packed n-major B layout"); + const int a_local_row = tid / kAInt4PerRow; + const int a_v = tid - a_local_row * kAInt4PerRow; + int8_t* a_dst = + a_tile + a_local_row * kAStride + a_v * static_cast(sizeof(int4)); + if constexpr (kSkew8) { + // 8-mod-16 row starts: write the 16 bytes as two 8-byte halves + // (ds_write2_b64 staging, identical to stage_prefill_tiles). + reinterpret_cast(a_dst)[0] = + reinterpret_cast(&a_payload)[0]; + reinterpret_cast(a_dst)[1] = + reinterpret_cast(&a_payload)[1]; + } else { + reinterpret_cast(a_dst)[0] = a_payload; + } + const int b_local_col = tid / kBNInt4PerRow; + const int b_v = tid - b_local_col * kBNInt4PerRow; + int8_t* b_dst = + b_tile + b_local_col * kBNStride + b_v * static_cast(sizeof(int4)); + if constexpr (kSkew8) { + reinterpret_cast(b_dst)[0] = + reinterpret_cast(&b_payload)[0]; + reinterpret_cast(b_dst)[1] = + reinterpret_cast(&b_payload)[1]; + } else { + reinterpret_cast(b_dst)[0] = b_payload; + } +} + +// One wave's 32x32-quadrant MMAC burst for the K tile staged at a_tile/b_tile +// (caller-offset to the target buffer). k0-outer / kk-inner int32 accumulation +// order; stage-size changes only regroup exact integer sums (max |dot| = +// 6144*127*127 << 2^31), so results stay bit-identical. +template +__device__ __forceinline__ void mmac_prefill_stage( + DUFragment& + a_frag0, + DUFragment& + a_frag1, + DUFragment& acc00, + DUFragment& acc01, + DUFragment& acc10, + DUFragment& acc11, + int8_t* a_tile, + int8_t* b_tile, + int local_row, + int local_col) { + // kSkew8 (round 11, exact-shape arm only): A 88 / B 72 row strides; the + // load_frag8 ldm arguments below pick them up unchanged (same 8-byte + // aligned fragment reads, conflict-free 16-lane bank phases). + constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); + constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows + constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); + if constexpr (kBNMajor) { + if constexpr (kDirectFrag) { + // Round 10 exact-shape path: explicit load_frag8 8-byte ds_read_b64 + // fragment reads replacing the library loaders (see load_frag8 + // comment). Both steps of the 64-K stage are unrolled into + // independent fragment register sets so all eight reads of a stage + // issue before the MMACs; the kk = 0 MMAC group runs first, then the + // kk = kTileK group - the exact k0-outer/kk-inner int32 order and the + // element-to-slot fragment mapping of the library path are preserved + // (f.x[0..7] = the same 8 bytes in the same little-endian order), so + // results stay bit-identical. + static_assert(kStageK == 2 * kTileK, + "kDirectFrag two-step unroll assumes kStageK == 2*kTileK"); + DUFragment + a_frag0_1, a_frag1_1; + DUFragment + b_frag0, b_frag1, b_frag0_1, b_frag1_1; + load_frag8(a_frag0, a_tile + local_row * kAStride, kAStride); + load_frag8(a_frag1, a_tile + (local_row + kTileM) * kAStride, kAStride); + load_frag8(b_frag0, b_tile + local_col * kBNStride, kBNStride); + load_frag8(b_frag1, b_tile + (local_col + kTileN) * kBNStride, + kBNStride); + load_frag8(a_frag0_1, a_tile + local_row * kAStride + kTileK, + kAStride); + load_frag8(a_frag1_1, + a_tile + (local_row + kTileM) * kAStride + kTileK, kAStride); + load_frag8(b_frag0_1, b_tile + local_col * kBNStride + kTileK, + kBNStride); + load_frag8(b_frag1_1, + b_tile + (local_col + kTileN) * kBNStride + kTileK, + kBNStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc00, a_frag0_1, b_frag0_1, acc00); + du_mma_sync(acc01, a_frag0_1, b_frag1_1, acc01); + du_mma_sync(acc10, a_frag1_1, b_frag0_1, acc10); + du_mma_sync(acc11, a_frag1_1, b_frag1_1, acc11); + } else { + DUFragment + b_frag0, b_frag1; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + du_load_matrix_sync( + b_frag0, b_tile + local_col * kBNStride + kk, kBNStride); + du_load_matrix_sync( + b_frag1, b_tile + (local_col + kTileN) * kBNStride + kk, + kBNStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + } + } else { + DUFragment + b_frag0, b_frag1; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + } +} + +// Large-prefill path: one block computes a kBlockM x kBlockN output tile. +// (kBlockM/32) x (kBlockN/32) wavefronts each own a 32x32 quadrant (four +// m16n16k32 DUMMA accumulators), while the block cooperatively stages +// A[kBlockM, kStageK] and B in bank-skewed LDS. +// +// kDoubleBuffer == false: single-buffered K loop (two barriers per stage); +// stage s+1's global loads cannot start until the stage-s barrier, so each +// stage serializes one global-load round trip in front of the compute +// (bootstrap structure; co-resident block streams hide it at the CU level). +// +// kDoubleBuffer == true: software-pipelined K loop (round 3). Stage s+1's +// global loads are issued into the idle LDS buffer before the stage-s MMAC +// burst, overlapping the serialized global-load latency with compute. One +// barrier per stage (plus one prologue): it both retires stage-s fragment +// reads (WAR on the buffer just read) and makes stage s+1's ds_writes +// visible (RAW on the prefetched buffer). The gate_up shape uses +// <64, 64, 64, true, true>: two 64-K tiles = 20,480 B/block x 3 blocks/CU = +// 61,440 <= 64 KiB (occupancy preserved; a 128-K double buffer would need +// 38,912 B/block -> 1 block/CU, a harmful occupancy loss). +// +// kBNMajor selects the B operand layout: +// * false (identity): B staged [K, N] row-major, library row_major B +// fragments (byte-scattered LDS reads; correct for the raw [K, N] weight +// and every non-gate_up shape). +// * true (gate_up only): B staged [N, K] n-major with k contiguous (row +// stride 144 = 128+16, or 80 = 64+16 for the 64-K double buffer), +// library col_major B fragments -> each lane's elements are k-contiguous +// bytes instead of byte-scattered columns. This is the bank-safe consumer +// layout for the B side (the 17.3M PMC bank-conflict source in round 1). +// The one-time pack in launch_pack_w8a8_weight provides the [N, K] buffer +// for (k, n) == (6144, 512) only. +template +__global__ __launch_bounds__(kBlockM * kBlockN / 1024 * kWaveSize) void +w8a8_dumma_prefill_tiled_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWavesN = kBlockN / 32; + constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + // 16-byte bank skew: keeps int4 staging stores aligned while spreading the + // fragment rows across distinct LDS banks. The staging loops live in + // stage_prefill_tiles below; the strides remain here as the static-asserted + // alignment contract of every instantiation. kSkew8 (round 11, exact-shape + // arm only) switches A/B to the down_proj-validated 8-mod-16 strides + // 88/72; LDS bytes stay A[64,88] + B[64,72] x 2 = 20,480 B/block so the + // 3-blocks/CU occupancy of the accepted kernel is preserved exactly. + constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); + constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows + constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); + constexpr int kBNInt4 = kBlockN * (kStageK / sizeof(int4)); + // Double buffering uses kBuffers = 2 LDS tile copies; the gate_up shape + // pairs it with kStageK = 64 so 3 blocks/CU is preserved (see kernel + // comment above). + constexpr int kBuffers = kDoubleBuffer ? 2 : 1; + constexpr int kABytes = kBlockM * kAStride; + // One shared B tile sized for the larger of the two layouts so an inactive + // arm never allocates dead LDS (keeps <64,64,128> at 19,456 B total and + // 3 blocks/CU). kSkew8 implies kBNMajor, so only the n-major arm is live. + constexpr int kBMaxBytes = + kSkew8 ? kBlockN * kBNStride + : (kStageK * kBStride > kBlockN * kBNStride + ? kStageK * kBStride + : kBlockN * kBNStride); + static_assert(kSkew8 ? (kAStride % 8 == 0) : (kAStride % sizeof(int4) == 0), + "A LDS row stride must stay int4 aligned (or 8-byte aligned " + "under kSkew8)"); + static_assert(kBStride % static_cast(sizeof(int4)) == 0, + "B LDS row stride must stay int4 aligned"); + static_assert( + kSkew8 ? (kBNStride % 8 == 0) : (kBNStride % sizeof(int4) == 0), + "n-major B LDS row stride must stay int4 aligned (or 8-byte aligned " + "under kSkew8)"); + static_assert(!kSkew8 || kBNMajor, + "kSkew8 is only defined for the packed n-major B layout"); + static_assert(!kSkew8 || kStageK == 64, + "kSkew8 strides 88/72 assume kStageK == 64"); + static_assert(kStageK % kTileK == 0, + "K stage must be a multiple of the DUMMA K unit"); + static_assert(kStageK % static_cast(sizeof(int4)) == 0, + "K stage must keep int4 staging aligned"); + static_assert(kBNInt4 % kThreads == 0, + "n-major B staging must divide evenly across threads"); + static_assert(!kPayload || (kBNMajor && kDoubleBuffer && kDirectFrag && + kSkew8 && kStageK == 64 && kBlockM == 64 && + kBlockN == 64), + "payload arm is specialized to the exact-shape <64,64,64> " + "double-buffered load_frag8/kSkew8 arm (one int4 per thread " + "per operand)"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWavesN; + const int wave_col = wave - wave_row * kWavesN; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + + __shared__ __align__(16) int8_t a_tile[kBuffers * kABytes]; + __shared__ __align__(16) int8_t b_tile[kBuffers * kBMaxBytes]; + + DUFragment + a_frag0, a_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + if constexpr (kDoubleBuffer) { + if constexpr (kPayload) { + // Round 16 lead-2 loop-carried payload pipeline (exact-shape arm + // only): stage s+1's global loads are issued a FULL iteration early + // (top of iteration s-1) into a two-int4-per-thread register payload + // and stay outstanding across the stage-(s-1) fragment reads + MMAC + // burst and the stage barrier; the vmcnt wait lands at the top of + // iteration s (first use of the payload registers), immediately + // before the ds_write2 publish into the idle LDS buffer. The + // s_barrier fences outstanding vmem, so the wait can no longer sit + // right after the load issue: the global-load latency overlaps the + // whole compute segment instead of serializing in front of every + // stage (the round-11 ISA defect). Same 2 buffers, same one barrier + // per stage (prologue + 96 = 97/block), same 16 ds_read2_b64 + + // 8 ds_write2_b64 per block-stage, same staging bytes/addresses + // (bit-identical int32 order). + int4 a_payload, b_payload; + // Prologue: issue the stage-1 global loads first (so they ride behind + // the stage-0 inline staging), stage k0 = 0 into buffer 0 inline + // (same load+wait+write path as the accepted kernel), then the first + // barrier. + prefetch_stage_payload( + x_q, weight, tid, m0, n0, m, k, kStageK, a_payload, b_payload); + stage_prefill_tiles( + x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, 0); + __syncthreads(); + for (int k0 = 0; k0 < k; k0 += kStageK) { + const int buf = (k0 / kStageK) & 1; + if (k0 + kStageK < k) { + // Publish the payload carried from the previous iteration (stage + // k0+kStageK) into the idle buffer: the vmcnt wait for the loads + // issued one iteration earlier is emitted here at the first use, + // after the previous MMAC burst + barrier. + publish_stage_payload( + a_payload, b_payload, a_tile + (buf ^ 1) * kABytes, + b_tile + (buf ^ 1) * kBMaxBytes, tid); + } + if (k0 + 2 * kStageK < k) { + // Issue stage k0+2*kStageK's global loads into the payload + // registers (no wait); they stay in flight across the MMAC burst + // and the barrier below and are waited at the next iteration's + // publish (RAW). + prefetch_stage_payload( + x_q, weight, tid, m0, n0, m, k, k0 + 2 * kStageK, a_payload, + b_payload); + } + // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. + // k0-outer / kk-inner int32 accumulation order. + mmac_prefill_stage( + a_frag0, a_frag1, acc00, acc01, acc10, acc11, + a_tile + buf * kABytes, b_tile + buf * kBMaxBytes, + local_row, local_col); + // One barrier per stage: (a) WAR - all stage-s fragment reads of the + // buffer just read are retired before it becomes the publish target + // two stages later; (b) RAW - every thread's stage-(s+1) ds_writes + // are visible before the next iteration reads them. + __syncthreads(); + } + } else { + // Prologue: stage k0 = 0 into buffer 0, then the pipelined loop. + stage_prefill_tiles( + x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, 0); + __syncthreads(); + for (int k0 = 0; k0 < k; k0 += kStageK) { + const int buf = (k0 / kStageK) & 1; + if (k0 + kStageK < k) { + // Prefetch stage s+1 into the idle buffer: global loads issue here, + // before the stage-s MMAC burst, and complete (RAW) behind the + // single barrier below. + stage_prefill_tiles( + x_q, weight, a_tile + (buf ^ 1) * kABytes, + b_tile + (buf ^ 1) * kBMaxBytes, tid, m0, n0, m, n, k, + k0 + kStageK); + } + // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. + // k0-outer / kk-inner int32 accumulation order. + mmac_prefill_stage( + a_frag0, a_frag1, acc00, acc01, acc10, acc11, + a_tile + buf * kABytes, b_tile + buf * kBMaxBytes, + local_row, local_col); + // One barrier per stage: (a) WAR - all stage-s fragment reads of the + // buffer just read are retired before it becomes the prefetch target + // two stages later; (b) RAW - every thread's stage-(s+1) ds_writes are + // visible before the next iteration reads them. + __syncthreads(); + } + } + } else { + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[kBlockM, kStageK] and B into the single LDS buffer. + stage_prefill_tiles( + x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, k0); + __syncthreads(); + + // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. + // k0-outer / kk-inner int32 accumulation order. + mmac_prefill_stage( + a_frag0, a_frag1, acc00, acc01, acc10, acc11, + a_tile, b_tile, local_row, local_col); + // Retire this stage's reads before the next stage's staging writes + // overwrite the single LDS buffer (WAR hazard across waves). + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Generic scalar fallback: one output element per grid-stride iteration with +// exact int32 accumulation. Correct for any (m, n, k), including the small-M +// API shapes (M=2/16) and any unmatched geometry. The weight is the identity +// [K, N] row-major layout for every (k, n) EXCEPT the gate_up weight +// (k, n) == (6144, 512), whose buffer is packed [N, K] n-major; this kernel +// decodes that layout so the paired gate_up tails stay correct. +__global__ __launch_bounds__(kScalarThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const bool b_nmajor = (k == 6144 && n == 512); + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + int32_t acc = 0; + const int8_t* a_row = x_q + static_cast(row) * k; + const int8_t* b_base = + weight + (b_nmajor ? static_cast(col) * k : col); + const int64_t b_stride = b_nmajor ? 1 : n; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_base[kk * b_stride]); + } + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// Bootstrap pack: identity device-to-device copy for every (k, n) pair +// (raw [K, N] layout preserved). Runs outside the timed/CUDA-Graph region. +__global__ __launch_bounds__(kScalarThreads) void +w8a8_pack_identity_kernel( + const int8_t* __restrict__ raw_weight, + const float* __restrict__ weight_scale, + int8_t* __restrict__ packed_weight, + float* __restrict__ packed_weight_scale, + int64_t weight_elems, + int n) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < weight_elems) { + packed_weight[idx] = raw_weight[idx]; + } + if (idx < n) { + packed_weight_scale[idx] = weight_scale[idx]; + } +} + +// Shape-specific gate_up pack: transpose raw [K, N] -> packed [N, K] +// (n-major rows, k contiguous) for (k, n) == (6144, 512) only. Byte-exact +// index relation: packed[cc * k + kk] == raw[kk * n + cc]. Each thread +// gathers the 16 bytes of one aligned int4 k-block from 16 raw rows (stride +// n) and stores one aligned int4, keeping the packed write side fully +// coalesced (same gather + aligned int4-store pattern as the validated +// o_proj transpose repair). Runs outside the timed/CUDA-Graph region; same +// byte count, same graph-stable buffer. +__global__ __launch_bounds__(kScalarThreads) void +w8a8_pack_gateup_transpose_kernel( + const int8_t* __restrict__ raw_weight, + const float* __restrict__ weight_scale, + int8_t* __restrict__ packed_weight, + float* __restrict__ packed_weight_scale, + int k, + int n) { + const int k4 = k / static_cast(sizeof(int4)); + const int64_t total4 = static_cast(n) * k4; + const int64_t stride4 = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx4 = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx4 < total4; idx4 += stride4) { + const int cc = static_cast(idx4 / k4); + const int kblk = static_cast(idx4 - static_cast(cc) * k4); + const int kk = kblk * static_cast(sizeof(int4)); + const int8_t* src = raw_weight + static_cast(kk) * n + cc; + int4 v; + char* vp = reinterpret_cast(&v); +#pragma unroll + for (int i = 0; i < static_cast(sizeof(int4)); ++i) { + vp[i] = src[static_cast(i) * n]; + } + reinterpret_cast(packed_weight + static_cast(cc) * k + + kk)[0] = v; + } + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < n; idx += stride) { + packed_weight_scale[idx] = weight_scale[idx]; + } +} + +} // namespace + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // No split-K: the provided workspace is not touched. The timed operator + // performs no allocation, synchronization, packing, or default-stream + // launch; only the caller-provided out is written. + (void)workspace; + (void)workspace_bytes; + auto* out_bf16 = static_cast(out); + + // Native INT8 DUMMA m16n16k32 tiled path for every large-M shape. Both + // assigned shapes qualify: gate_up (4096, 512, 6144) and down_proj + // (4096, 6144, 256) both have M >= 128, N % 128 == 0 and K % 128 == 0. + // Explicit dispatch by launch geometry; unmatched shapes fall through to + // the scalar path (no shape is specialized away). + if (m >= kPrefillMinM && (n % kDefaultBlockN == 0) && + (k % kPrefillStageK == 0)) { + // Exact-shape guard for the assigned gate_up shape (round-3 pipeline + // decision, round-10 load_frag8, round-11 kSkew8, round-16 kPayload): + // cooperative A+B staging with the n-major [N, K] packed B operand + // (kBNMajor) consumed as col_major fragments, software-pipelined across + // K tiles (kDoubleBuffer) with a 64-K double buffer: 64x64 tile, 512 + // blocks, 20,480 B LDS/block, 3 blocks/CU (61,440 B) - occupancy + // preserved vs the accepted round-2 single buffer (19,456 B x 3 = + // 58,368 B). One barrier per 64-K stage + prologue (97 for 96 stages). + // Round 11 stages A/B at the 8-mod-16 strides 88/72 (kSkew8) so the + // fragment reads are bank-conflict-free (LDS bytes and occupancy + // unchanged). Round 16 (kPayload) issues the stage-(s+1) global loads a + // full iteration early into a two-int4-per-thread register payload and + // publishes them into the idle LDS buffer at the top of iteration s, so + // the vmcnt wait lands after the stage-s MMAC burst instead of in front + // of it (LDS bytes, barriers, traffic, occupancy and results unchanged). + if (m == 4096 && n == 512 && k == 6144) { + const dim3 grid(n / 64, (m + 63) / 64); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel< + 64, 64, 64, true, true, true, true, true>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (n % 128 == 0) { + // 64x128 tile, 8 waves (32x32 quadrant per wave). + // down_proj: grid (6144/128, 4096/64) = (48, 64) = 3072 blocks. + const dim3 grid(n / 128, (m + 63) / 64); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 128, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (m % 128 == 0) { + // 128x64 tile, 8 waves. + const dim3 grid(n / 64, m / 128); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<128, 64, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + { + // 64x64 tile, 4 waves (generic default). + const dim3 grid(n / 64, (m + 63) / 64); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 64, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + } + + // Generic scalar fallback for every unmatched (m, n, k) and M < 128 + // (including the paired M=2/M=16 API shapes with the same (N, K), which + // decode the packed [N, K] gate_up layout when (k, n) == (6144, 512)). + const int64_t total = static_cast(m) * n; + const unsigned blocks = static_cast( + (total + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Assigned gate_up weight (K=6144, N=512): one-time [K,N] -> [N,K] + // transpose pack (n-major, k contiguous) so the timed GEMM reads B + // fragments as k-contiguous bytes from the staged LDS tile. Runs outside + // the timed/CUDA-Graph region; same byte count, same graph-stable buffer, + // byte-exact vs the raw logical [K,N] weight. + if (k == 6144 && n == 512) { + const int64_t total4 = static_cast(n) * (k / 16); + const unsigned blocks = static_cast( + (total4 + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_pack_gateup_transpose_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + raw_weight, weight_scale, packed_weight, packed_weight_scale, k, n); + return; + } + // Identity device-to-device copy for every other (k, n) pair, including + // unmatched (K, N). Runs outside the timed/CUDA-Graph region. + const int64_t weight_elems = static_cast(k) * n; + const int64_t total = weight_elems > n ? weight_elems : n; + const unsigned blocks = static_cast( + (total + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_pack_identity_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + raw_weight, weight_scale, packed_weight, packed_weight_scale, + weight_elems, n); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/o_proj.hip new file mode 100644 index 00000000..fc602161 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/o_proj.hip @@ -0,0 +1,696 @@ +// @@variant shape=hy3_tp4_o_proj_m16 commit=cb7af9e4e0a3adf58bc62ec3cf8be1e1a3f4c010 added=2026-08-26 +// median_us=15.51 p90_us=15.55 +// source=hy3-dsh-tp4-m16-2-368c654c +// MetaInfer W8A8 INT8 GEMM - HIP implementation for K500SM_AI / gfx928. +// +// Worker: worker_1 (physical GPU 1) +// Assigned shape: hy3_tp4_o_proj_m16 (M=16, N=4096, K=2048) +// +// Iteration 1 (accepted) = minimal 16x16x32 DUMMA bootstrap: 256 blocks x 1 +// 64-thread wavefront, one N tile per block, library byte-loads of A/B +// straight from the logical row-major globals, 64 ascending-K int32 steps, +// LDS epilogue plane, zero in-loop barriers. +// -> median 89.90 us / p90 90.10 us, 95.3 GB/s algorithmic, 263,424 +// vmem_read/replay = 16.08 loads per K-step per wave, all fenced +// vmcnt(15)..vmcnt(0) before the single v_mmac. +// Iteration 2 (rejected) = grid-parallelism probe: 128 blocks x 2 waves x 2 +// adjacent N tiles (SAME 256 waves, no split-K, no packing). +// -> median 117.81 us (REGRESSION): per-CU wave count is NOT the lever on +// the byte-load design at 2.13 waves/CU; the named next lever was load +// vectorization (packed-B dwordx2 fragment loads) and/or in-block +// split-K. +// Iteration 3 (rejected, measurement-polluted) = architecture round: 256 +// blocks x 2 waves, in-block split-K=2 (uniform K=1024 slices), packed-B +// [n_tile][k_step][lane][8] fragment-slot loads, register-only K-loop +// transport with depth-1 prefetch, K=128 stages. This is byte-for-byte the +// authoritative lineage's ACCEPTED final kernel for this exact shape +// (w8a8_gemm_m16_dumma_packedb_kernel, 18.221/18.242 us median/p90 there). +// -> min 21.37 us (9 uncontended samples tight at 21.4-22.5 us) but median +// 166.13 / p90 216.62: the 30-sample window was bimodal, with ~19 +// samples at ~190-240 us matching an ~8-10x-slower interleaved workload +// on the GPU (external interference, not kernel behavior: the same +// captured graph cannot run 10x slower under its own resources at 2.13 +// blocks/CU all-resident). The fast tail proves the design; the slow +// head is the measurement window. +// Iteration 4 (accepted, 17.155 us median / 17.185 us p90, 499.5 GB/s +// algorithmic - the fastest accepted W8A8 kernel in every worker ledger on +// this node) = architecture/pipeline round, multi-N-tile reuse: each block +// computes TWO adjacent 16x16 N tiles (T=2) so the A fragments loaded per +// K=128 stage are shared across both tiles in registers; everything else +// frozen from iteration 3 (packed-B layout, in-block split-K=2, K=128 +// stages, depth-1 register prefetch). HBM floor analysis: 8,484,352 B of +// unavoidable traffic (8 MiB B read-once + 32 KiB A + 128 KiB bf16 C) at +// 17.155 us = 499.5 GB/s, and no accepted kernel on this device has ever +// exceeded that figure - the kernel sits at the practical streaming +// ceiling. +// Iteration 5 (THIS round) = HIP-only packed-layout comparison (mandate: +// packed layout / A-only staging / B-only staging): repack B from +// [n_tile][k_step][lane][8] to a TILE-PAIR-INTERLEAVED dwordx4 layout +// [n_tile_pair][k_step][lane][16] in which lane's 16 bytes are tile0's 8 +// fragment bytes followed by tile1's 8 fragment bytes for the same +// (k_step, lane). Each lane then covers BOTH tiles with ONE aligned 16-byte +// load per k_step (a single 1 KiB contiguous region per load, a single +// 4 KiB contiguous region per K=128 stage, and one sequential 64 KiB +// stream per block lifetime - was two 2 KiB regions 32 KiB apart per +// stage); B vmem instructions drop from 8 to 4 per stage per wave at an +// identical 24 VGPR in-flight footprint and identical 8 MiB of HBM bytes. +// Geometry frozen from iteration 4 (128 blocks x 2 waves x T=2, in-block +// split-K=2, K=128 stages, depth-1 register prefetch, partial-planes-only +// LDS, one end-of-K barrier). See the kernel comment below. +// +// Scalar fallback strategy (one thread per output element; still the path for +// every shape OTHER than the exact assigned M=16 shape): +// * blockDim = 256 (multiple of the gfx928 64-lane wavefront); +// * grid = (ceil(N/256), M): each block covers 256 adjacent N columns of a +// single M row, so adjacent lanes touch adjacent addresses in the +// fastest-changing N dimension (coalesced B reads / bf16 stores); +// * the complete K loop accumulates exactly in int32. For the assigned +// K=2048 the max |dot| is 2048*127*127 = 33,032,192 << 2^31, so the +// integer accumulation never overflows; +// * after the K loop the two float scales are applied as +// float(dot) * x_scale[m] * weight_scale[n] and the result is stored as +// bf16 (hip_bfloat16 via __float2bfloat16); +// * for the exact (n, k) == (4096, 2048) pair the weight tensor is the +// packed [n_tile_pair][k_step][lane][16] tile-pair-interleaved layout +// (16 bytes = tile0's 8 fragment bytes ++ tile1's 8 fragment bytes for +// the same k_step/lane), which the fallback decodes elementwise (keeps +// the paired M=2 validation exact with the new pack); +// * workspace is untouched by the timed operator (in-block split-K needs no +// partial planes). +// +// pack_weight: for the exact (k, n) == (2048, 4096) o_proj weight it is a +// one-time device permutation into packed[n_tile_pair][k_step][lane][16] B +// fragment slots (byte count unchanged: 8 MiB; runs outside the timed GEMM +// and outside Graph capture). Every other (K, N) keeps the generic identity +// device-to-device copy; packed_weight_scale[n] = weight_scale[n] always. +// +// Host launch symbols (stable, C linkage, consumed by csrc/bindings.cpp): +// launch_w8a8_gemm(...) +// launch_pack_w8a8_weight(...) +// +// Include order is known-good for this DTK: hip_runtime.h first, then +// hip_bfloat16.h, then du_mma.h (du_mma.h is not self-contained when +// included before the HIP runtime headers). + +#include +#include +#include + +#include + +namespace { + +constexpr int kScalarThreads = 256; + +// Exact assigned shape (M=16, N=4096, K=2048). +constexpr int kTargetM = 16; +constexpr int kTargetN = 4096; +constexpr int kTargetK = 2048; + +// --------------------------------------------------------------------------- +// DUMMA geometry for the exact assigned shape: gfx928 INT8 m16n16k32 tile. +// Iteration 4 multi-N-tile reuse: each block owns TWO adjacent 16x16 N tiles +// (T=2, grid = N/(16*2) = 128 blocks >= 120 device CUs) and 2 x 64-thread +// wavefronts; wave w accumulates the int32 dot over its UNIFORM K=1024 slice +// for BOTH tiles, so the A fragments of each K=128 stage are loaded once and +// reused across the two tiles from registers. +// --------------------------------------------------------------------------- +constexpr int kDummaBlockThreads = 64; // one gfx928 wavefront +constexpr int kDummaWavesPerBlock = 2; // in-block split-K=2 +constexpr int kDummaTilesPerBlock = 2; // multi-N-tile reuse (T=2) +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; + +// Bytes of one (n_tile, k_step) B fragment in the packed layout: +// 16 columns * 32 k rows = 512 B, stored as 64 lanes * 8 B. +constexpr int kPackedFragBytes = kDummaTileN * kDummaTileK; // 512 + +// Iteration-5 tile-pair-interleaved layout: one (k_step) slot holds BOTH +// tiles of a block in 16 B per lane - tile0's 8 fragment bytes (bytes 0..7) +// followed by tile1's 8 fragment bytes (bytes 8..15) - so a single aligned +// 16-byte (dwordx4) load per lane covers one k_step of both tiles: 64 lanes +// * 16 B = 1 KiB contiguous per load. Per n_tile_pair the 64 k_steps are +// stored ascending (64 KiB sequential per block lifetime). +constexpr int kPackedPairBytes = kDummaBlockThreads * 16; // 1024 + +// K=128 stage: 4 DUMMA steps per stage, 8 loads per stage per wave (4 A +// dwordx2 shared by both tiles + 4 B dwordx4, one per k_step covering both +// tiles = 6 KB in flight), 8 v_mmac per stage per wave (4 steps x 2 tiles). +constexpr int kStageK = 128; // K bytes per stage +constexpr int kStageSteps = kStageK / kDummaTileK; // 4 DUMMA steps per stage + +// --------------------------------------------------------------------------- +// Packed-B decode for the generic scalar fallback: returns logical +// weight[kk, n] from the packed [n_tile_pair][k_step][lane][16] +// tile-pair-interleaved layout (see the pack kernel for the exact +// permutation). The 16 bytes of a lane slot are tile0's 8 fragment bytes +// (bytes 0..7) followed by tile1's 8 fragment bytes (bytes 8..15). +// --------------------------------------------------------------------------- +__device__ __forceinline__ int8_t +packed_b_element(const int8_t* __restrict__ packed, int n, int kk) { + const int n_tile = n >> 4; + const int pair = n_tile >> 1; // n_tile_pair (block id) + const int sub = n_tile & 1; // 0 -> bytes 0..7, 1 -> bytes 8..15 + const int nn = n & 15; + const int k_step = kk >> 5; + const int kk8 = kk & 31; + const int lane = (kk8 >> 3) * 16 + nn; + const int i = kk8 & 7; + return packed[(static_cast(pair) * (kTargetK / kDummaTileK) + + k_step) * + kPackedPairBytes + + lane * 16 + sub * 8 + i]; +} + +// One thread computes one output element out[row, col]: +// acc = sum_k a[row,k] * b[k,col] (exact int32, k ascending) +// out = bf16(acc * x_scale[row] * weight_scale[col]) +// Bounds-checked, so it is a valid generic fallback for every (m, n, k). +// For the exact (n, k) == (4096, 2048) pair the weight tensor is the packed +// [n_tile_pair][k_step][lane][16] layout (o_proj M=16 pack), decoded +// elementwise. +__global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, // [m, k] row-major + const int8_t* __restrict__ b, // [k, n] (packed for 4096x2048) + const float* __restrict__ x_scale, // [m] + const float* __restrict__ weight_scale, // [n] + hip_bfloat16* __restrict__ out, // [m, n] + int m, + int n, + int k) { + const int row = static_cast(blockIdx.y); + const int col = + static_cast(blockIdx.x) * kScalarThreads + + static_cast(threadIdx.x); + if (row >= m || col >= n) { + return; + } + + const int8_t* __restrict__ a_row = a + static_cast(row) * k; + const bool packed_b = (n == kTargetN && k == kTargetK); + + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + const int8_t b_val = packed_b ? packed_b_element(b, col, kk) + : b[static_cast(kk) * n + col]; + acc += static_cast(a_row[kk]) * static_cast(b_val); + } + + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Iteration-5 stage prefetch helper (multi-N-tile reuse): issue the eight +// loads for one K=128 stage - the 4 A dwordx2 fragments (shared by both +// tiles) plus 4 B dwordx4 fragment-pair loads (one per k_step, each covering +// BOTH tiles) - into caller-provided u64 registers. +// * A: four 8-byte loads of one activation row at ldm = k, 32 B apart (the +// four DUMMA steps of the stage) - lane ownership is loop-invariant +// (du_mma.hpp matrix_a row_major loader: row = lane & 0xf, +// col = (lane >> 4) << 3, x[i] = p[row*ldm + col + i], 8 consecutive +// bytes; du_mma_sync consumes a.x as one 64-bit value, du_mma.hpp line +// 1198). A stays the logical row-major 32 KiB activation (L2-hot). +// These 4 values feed BOTH accumulator tiles - A is read once per block +// instead of once per N tile. +// * B (iteration-5 tile-pair-interleaved layout): per k_step one aligned +// 16-byte load from the [n_tile_pair][k_step][lane][16] slot whose bytes +// 0..7 are tile0's fragment bytes and bytes 8..15 tile1's fragment bytes +// for the same (k_step, lane) - byte-order identical to the library +// matrix_b row_major loader (x[i] = p[(col + i) * ldm + row]) per tile, +// so the two 8-byte halves feed b_frag0.x / b_frag1.x directly. A stage's +// four k_step loads sit at consecutive 1 KiB slot offsets: one 4 KiB +// contiguous region per stage per wave; each wave's 8 stages walk one +// sequential 32 KiB half of the block's 64 KiB pair region, so every B +// byte is read exactly once per replay. +// * The caller's k0_base K-slice offset makes each split wave walk its own +// packed k_step range. The 16-byte loads are 16 B aligned (slot offset is +// k_step * 1024 + lane * 16 on a 16 B-aligned packed buffer). +// --------------------------------------------------------------------------- +__device__ __forceinline__ void load_stage_fragments_t2( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int lane, int a_row, int a_col, int64_t b_base, int k, + int64_t k0_base, int s, + uint64_t& a0, uint64_t& a1, uint64_t& a2, uint64_t& a3, + uint64_t& b00, uint64_t& b01, uint64_t& b02, uint64_t& b03, + uint64_t& b10, uint64_t& b11, uint64_t& b12, uint64_t& b13) { + const int64_t k0 = k0_base + static_cast(s) * kStageK; + const int64_t a_off = static_cast(a_row) * k + k0 + a_col; + a0 = *reinterpret_cast(x_q + a_off); + a1 = *reinterpret_cast(x_q + a_off + kDummaTileK); + a2 = *reinterpret_cast(x_q + a_off + 2 * kDummaTileK); + a3 = *reinterpret_cast(x_q + a_off + 3 * kDummaTileK); + const int64_t pair_off = (k0 / kDummaTileK) * kPackedPairBytes + + static_cast(lane) * 16; + const int64_t b_off = b_base + pair_off; + // One 16-byte load per k_step covers both tiles: bytes 0..7 = tile0's + // fragment, bytes 8..15 = tile1's fragment (little-endian word order). + const uint4 bv0 = *reinterpret_cast(weight + b_off); + const uint4 bv1 = *reinterpret_cast(weight + b_off + kPackedPairBytes); + const uint4 bv2 = *reinterpret_cast(weight + b_off + 2 * kPackedPairBytes); + const uint4 bv3 = *reinterpret_cast(weight + b_off + 3 * kPackedPairBytes); + b00 = (static_cast(bv0.y) << 32) | bv0.x; // tile0, k_step 0 + b01 = (static_cast(bv1.y) << 32) | bv1.x; // tile0, k_step 1 + b02 = (static_cast(bv2.y) << 32) | bv2.x; // tile0, k_step 2 + b03 = (static_cast(bv3.y) << 32) | bv3.x; // tile0, k_step 3 + b10 = (static_cast(bv0.w) << 32) | bv0.z; // tile1, k_step 0 + b11 = (static_cast(bv1.w) << 32) | bv1.z; // tile1, k_step 1 + b12 = (static_cast(bv2.w) << 32) | bv2.z; // tile1, k_step 2 + b13 = (static_cast(bv3.w) << 32) | bv3.z; // tile1, k_step 3 +} + +// --------------------------------------------------------------------------- +// Iteration 17 (HIP-only tail consolidation): float32 -> bf16 round-to-nearest- +// even WITHOUT the __float2bfloat16 NaN/Inf branch. The exact iteration-13 +// code object shows each of the epilogue's four conversions as a ~10- +// instruction guarded sequence ON the block's post-barrier finishing tail: +// the finite path is v_bfe_u32 (bit 16) + v_add3_u32 (u + bit + 0x7fff) with +// the store's d16_hi keeping the high 16 bits, plus the dead-for-this-operator +// s_and_saveexec / v_or_b32 / v_cmp_eq_sdwa / v_cndmask NaN/Inf handling. +// Every converted value here is finite (the int32 dot is bounded by +// 2048*127*127 << 2^31 and both scales are finite), so the guarded branch is +// pure tail cost sitting between the LDS reads and the stores. This helper +// emits only the finite RNE path: r = u + 0x7fff + ((u >> 16) & 1), high 16 +// bits kept - bit-identical to __float2bfloat16 for every finite input (same +// rounding, same carry into the exponent at the bf16 range edge); the only +// divergence is NaN payload bits, which cannot occur here. +// --------------------------------------------------------------------------- +__device__ __forceinline__ hip_bfloat16 bf16_rne_finite(float f) { + uint32_t u; + __builtin_memcpy(&u, &f, sizeof(u)); + u += 0x7fffu + ((u >> 16) & 1u); + const uint16_t h = static_cast(u >> 16); + hip_bfloat16 b; + __builtin_memcpy(&b, &h, sizeof(h)); + return b; +} + +// --------------------------------------------------------------------------- +// Exact assigned shape (M=16, N=4096, K=2048): gfx928 DUMMA INT8 m16n16k32 +// tile, weight pre-packed into tile-pair-interleaved B fragment slots, +// MULTI-N-TILE REUSE. +// * Geometry: 128 blocks (grid = N/(16*2), one block per pair of adjacent +// 16x16 N tiles; 128 >= 120 device CUs, so every CU gets at least one +// block) x 128 threads = 2 x 64-lane wavefronts. +// * In-block split-K=2: wave w owns the UNIFORM K slice of 8 x K=128 stages +// (stages [floor(w*16/2), floor((w+1)*16/2)) = exactly 8 stages, K=1024) +// and accumulates BOTH tiles over that slice - 2 independent accumulator +// fragments per wave. Per K=128 stage per wave: 4 A dwordx2 + 4 B dwordx4 +// (one 16-byte load per k_step covering both tiles) = 8 loads (6 KB in +// flight, identical byte footprint to iteration 4) covering an 8-v_mmac +// burst (4 steps x 2 tiles): load:compute ratio 1.5 (iteration 4) -> 1.0, +// B vmem instructions per replay drop from 16,384 dwordx2 to 8,192 +// dwordx4, and each block's B lifetime stream is ONE sequential 64 KiB +// region (4 KiB contiguous per stage) instead of two 32 KiB regions 32 +// KiB apart; A global traffic unchanged (L2-hot 32 KiB footprint, read +// once per block), B HBM traffic unchanged (8 MiB, every byte read +// exactly once). +// * REGISTER-ONLY K-loop transport with depth-1 prefetch, identical to +// iteration 4: current stage's 8 loads' fragments live in registers, the +// NEXT stage's 8 loads are issued BEFORE the current 8-MMAC burst, +// and the compiler-inserted vmcnt wait lands at the next iteration's +// first fragment fill (global_load -> v_mmac burst -> wait, no LDS hop, +// no lgkmcnt wait in the K loop). The 4 A dwordx2 + 4 B dwordx4 hold the +// same 24 VGPR as iteration 4's 12 dwordx2, so occupancy is unchanged. +// * LDS holds only the four 1 KiB int32 partial planes (2 waves x 2 tiles, +// 4 KiB/block); the ONLY barrier is the single END-of-K __syncthreads() +// before the fused partial combine + scale -> bf16 epilogue (all 128 +// threads, 4 elements each). Barriers per K step = 0. +// * Correctness: per element the int32 accumulation is ascending-K within +// the wave's slice (4 ascending steps per stage, stages ascending), then +// the two waves' partials are summed in LDS (order-independent int32, no +// overflow: max |dot| = 2048*127*127 << 2^31) - bit-identical to +// iteration 4. The iteration-5 repack changed the packed layout +// COHERENTLY across the pack kernel, the DUMMA kernel's loads and the +// scalar fallback decode (same (tile, k_step, lane, byte) -> logical +// (k, n) mapping, new memory order), so the paired M=2 validation stays +// exact. +// * Graph-safe: static smem only, no allocation/sync/default-stream launch, +// pure dispatch on the caller-provided stream. Workspace untouched. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDummaWavesPerBlock * kDummaBlockThreads) void +w8a8_gemm_m16_dumma_packedb_multint_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, // packed [n_tile_pair][k_step][lane][16] + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kDummaBlockThreads; // K-slice owner, 0..1 + const int lane = tid % kDummaBlockThreads; // wavefront lane 0..63 + // Block b owns the two adjacent N tiles 2b and 2b+1. + const int n_tile0 = static_cast(blockIdx.x) * kDummaTilesPerBlock; + const int n0 = n_tile0 * kDummaTileN; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag0; + du::dumma::DUFragment + b_frag1; + du::dumma::DUFragment acc0; + du::dumma::DUFragment acc1; + du::dumma::du_fill_fragment(acc0, 0); + du::dumma::du_fill_fragment(acc1, 0); + + // Each n_tile_pair owns one contiguous packed region of k_steps * 1024 B + // (64 KiB: both tiles' fragments interleaved at 16 B per lane per k_step). + const int64_t k_steps = k / kDummaTileK; + const int64_t b_base = + static_cast(blockIdx.x) * k_steps * kPackedPairBytes; + + // Per-lane A fragment ownership is loop-invariant (du_mma.hpp lines + // 447-460): row = lane & 0xf, col = (lane >> 4) << 3. + const int a_row = lane & 15; + const int a_col = (lane >> 4) << 3; + + // Iteration 13 (HIP-only consolidation): prefetch the four per-thread scale + // values at the kernel top, BEFORE the stage-0 fragment loads. The fused + // epilogue then has ZERO exposed global latency on the block's finishing + // tail: the same 4 vmem reads (x_scale[row], x_scale[row2], + // weight_scale[n0+col], weight_scale[n0+16+col] - each used twice) move + // from after the END-of-K __syncthreads() into the cold-start window, where + // their L2 round trip overlaps the first B-load DRAM latency. Same 4 loads, + // same values, bit-identical math - only the issue time moves (validated + // pattern: the accepted down_proj kernel prefetches its scales at kernel + // top too). Cost: +4 VGPR held live across the K loop (64 -> ~68, register + // file non-binding); scratch must stay 0. + const int row = tid >> 4; // epilogue rows 0..7 (and row2 = +8) + const int col = tid & 15; // epilogue column within the tile + const int row2 = row + (kDummaTileM >> 1); + const float x_scale_lo = x_scale[row]; + const float x_scale_hi = x_scale[row2]; + const float ws_lo = weight_scale[n0 + col]; + const float ws_hi = weight_scale[n0 + kDummaTileN + col]; + + // Four 1 KiB row-major int32 partial planes: [wave][tile][256]. + __shared__ __align__(16) int32_t acc_planes[kDummaWavesPerBlock * + kDummaTilesPerBlock * + kDummaTileM * kDummaTileN]; + + // 16 stages of K=128 partitioned UNIFORMLY across the 2 waves (8 stages of + // K=1024 per wave); ascending-K order per wave and the byte-exact fragment + // loads (A dwordx2, B dwordx4 tile-pair) are unchanged, the end-of-K int32 + // combine is order-independent, so the result is bit-identical to + // iteration 4. + const int total_stages = k / kStageK; // 16 + const int stage_start = (wave * total_stages) / kDummaWavesPerBlock; + const int stage_end = ((wave + 1) * total_stages) / kDummaWavesPerBlock; + const int n_stages = stage_end - stage_start; // 8 stages of K=128 + const int64_t k0_base = static_cast(stage_start) * kStageK; + + // Prologue: stage-0 fragments land directly in registers (one cold-start + // DRAM latency per wave, amortized over 8 stages and staggered across the + // 2 waves). + uint64_t ca0, ca1, ca2, ca3; // current stage's A + uint64_t cb00, cb01, cb02, cb03; // current stage's B, tile 0 + uint64_t cb10, cb11, cb12, cb13; // current stage's B, tile 1 + load_stage_fragments_t2(x_q, weight, lane, a_row, a_col, b_base, k, k0_base, + 0, ca0, ca1, ca2, ca3, cb00, cb01, cb02, cb03, cb10, + cb11, cb12, cb13); + + for (int s = 0; s < n_stages; ++s) { + // Issue the NEXT stage's eight global loads (4 A dwordx2 + 4 B dwordx4) + // BEFORE this stage's MMAC burst so DRAM/L2 latency overlaps the burst; + // the compiler-inserted s_waitcnt vmcnt(0) for these loads lands at the + // next iteration's first read of the prefetched registers (the fragment + // fill below) - 6 KB of in-flight bytes per wait, no LDS hop. + uint64_t na0, na1, na2, na3; + uint64_t nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13; + if (s + 1 < n_stages) { + load_stage_fragments_t2(x_q, weight, lane, a_row, a_col, b_base, k, + k0_base, s + 1, na0, na1, na2, na3, nb00, nb01, + nb02, nb03, nb10, nb11, nb12, nb13); + } + // Stage-s MMAC burst straight from registers, in the exact a_frag.x / + // b_frag.x byte order du_mma_sync consumes: the same A fragment feeds + // both tiles (multi-N-tile reuse - A loaded once per stage). +#pragma unroll + for (int t = 0; t < kStageSteps; ++t) { + const uint64_t av = + (t == 0) ? ca0 : (t == 1) ? ca1 : (t == 2) ? ca2 : ca3; + const uint64_t b0v = + (t == 0) ? cb00 : (t == 1) ? cb01 : (t == 2) ? cb02 : cb03; + const uint64_t b1v = + (t == 0) ? cb10 : (t == 1) ? cb11 : (t == 2) ? cb12 : cb13; + __builtin_memcpy(a_frag.x, &av, 8); + __builtin_memcpy(b_frag0.x, &b0v, 8); + __builtin_memcpy(b_frag1.x, &b1v, 8); + du::dumma::du_mma_sync(acc0, a_frag, b_frag0, acc0); + du::dumma::du_mma_sync(acc1, a_frag, b_frag1, acc1); + } + // Rotate: the prefetched stage s+1 becomes the current stage. No LDS + // WAR hazard; the vmcnt wait for the prefetch is satisfied before the + // next iteration's first use. + if (s + 1 < n_stages) { + ca0 = na0; + ca1 = na1; + ca2 = na2; + ca3 = na3; + cb00 = nb00; + cb01 = nb01; + cb02 = nb02; + cb03 = nb03; + cb10 = nb10; + cb11 = nb11; + cb12 = nb12; + cb13 = nb13; + } + } + + // Combine the two wavefronts' int32 partials in LDS (exact int32 sums, + // order independent) after a single barrier at the END of K - barrier + // count per K step stays 0. Each wave stores its two tiles' planes. + du::dumma::du_store_matrix_sync( + acc_planes + (wave * kDummaTilesPerBlock) * kDummaTileM * kDummaTileN, + acc0, kDummaTileN, du::dumma::mem_row_major); + du::dumma::du_store_matrix_sync( + acc_planes + (wave * kDummaTilesPerBlock + 1) * kDummaTileM * + kDummaTileN, + acc1, kDummaTileN, du::dumma::mem_row_major); + __syncthreads(); + + // Cooperative epilogue over the two 16x16 tiles: all 128 threads, 4 + // elements each - for tile t: (row = tid >> 4, col = tid & 15) for rows + // 0..7 and (row + 8, col) for rows 8..15; sum the 2 planes per element + // (int32 add is order-independent), then fused dot * x_scale[row] * + // weight_scale[col] -> bf16 store. The four scale values were prefetched + // at the kernel top (iteration 13), so no global load sits on this tail; + // iteration 17 replaces the four __float2bfloat16 calls with the finite + // RNE-only bf16_rne_finite (same bits for every finite input, no NaN/Inf + // branch), shortening each element's dependent tail chain. + int32_t t0_lo = 0, t0_hi = 0, t1_lo = 0, t1_hi = 0; +#pragma unroll + for (int w = 0; w < kDummaWavesPerBlock; ++w) { + const int plane_base = + w * kDummaTilesPerBlock * kDummaTileM * kDummaTileN; + t0_lo += acc_planes[plane_base + tid]; + t0_hi += acc_planes[plane_base + tid + kDummaTileM * kDummaTileN / 2]; + t1_lo += acc_planes[plane_base + kDummaTileM * kDummaTileN + tid]; + t1_hi += acc_planes[plane_base + kDummaTileM * kDummaTileN + tid + + kDummaTileM * kDummaTileN / 2]; + } + const float s0_lo = static_cast(t0_lo) * x_scale_lo * ws_lo; + out[static_cast(row) * n + n0 + col] = bf16_rne_finite(s0_lo); + const float s0_hi = static_cast(t0_hi) * x_scale_hi * ws_lo; + out[static_cast(row2) * n + n0 + col] = bf16_rne_finite(s0_hi); + const float s1_lo = static_cast(t1_lo) * x_scale_lo * ws_hi; + out[static_cast(row) * n + n0 + kDummaTileN + col] = + bf16_rne_finite(s1_lo); + const float s1_hi = static_cast(t1_hi) * x_scale_hi * ws_hi; + out[static_cast(row2) * n + n0 + kDummaTileN + col] = + bf16_rne_finite(s1_hi); +} + +// Generic grid-stride device-to-device copy (int8) - identity pack for every +// (K, N) except the exact (2048, 4096) o_proj shape. +__global__ void w8a8_copy_i8_kernel(const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + i < count; i += stride) { + dst[i] = src[i]; + } +} + +// Generic grid-stride device-to-device copy (float). +__global__ void w8a8_copy_f32_kernel(const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + i < count; i += stride) { + dst[i] = src[i]; + } +} + +constexpr int kMaxCopyBlocks = 65535; + +// --------------------------------------------------------------------------- +// One-time pack for the exact (k, n) == (2048, 4096) o_proj weight: permute +// the logical [K, N] int8 layout into packed[n_tile_pair][k_step][lane][16] +// B fragment slots - tile-pair-interleaved: for the same (k_step, lane), +// bytes 0..7 are the 8 fragment bytes of tile (2*pair) and bytes 8..15 the 8 +// fragment bytes of tile (2*pair+1) (see the file header and the kernel +// comment for the exact du_mma.hpp matrix_b row_major lane mapping: +// row = lane & 0xf, col = (lane >> 4) << 3, x[i] = p[(col + i) * ldm + row]). +// One thread per (pair, k_step, lane): 64 k_steps x 128 pairs x 64 lanes = +// 524,288 threads, each writing one aligned 16-byte slot. Runs once per +// weight, outside the timed GEMM and outside Graph capture; byte count is +// unchanged (8 MiB), so the packed tensor is a same-size permutation. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarThreads) void w8a8_pack_o_proj_m16_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = (static_cast(k) / kDummaTileK) * + (n / (2 * kDummaTileN)) * kDummaBlockThreads; + if (idx >= total) { + return; + } + const int pair = static_cast(idx >> 12); // idx / (64 * 64) + const int rem = static_cast(idx & 4095); // idx % (64 * 64) + const int k_step = rem >> 6; // idx / 64 % 64 + const int lane = rem & 63; // idx % 64 + const int row = lane & 15; + const int col = (lane >> 4) << 3; + const int8_t* src0 = + raw + (static_cast(k_step) * kDummaTileK + col) * n + + (pair * 2) * kDummaTileN + row; + const int8_t* src1 = src0 + kDummaTileN; // tile (2*pair + 1) + uint64_t v0 = 0; // tile 2*pair (bytes 0..7) + uint64_t v1 = 0; // tile 2*pair+1 (bytes 8..15) +#pragma unroll + for (int i = 0; i < 8; ++i) { + v0 |= (static_cast(static_cast( + src0[static_cast(i) * n]))) + << (8 * i); + v1 |= (static_cast(static_cast( + src1[static_cast(i) * n]))) + << (8 * i); + } + uint4 slot; + slot.x = static_cast(v0 & 0xffffffffu); + slot.y = static_cast(v0 >> 32); + slot.z = static_cast(v1 & 0xffffffffu); + slot.w = static_cast(v1 >> 32); + *reinterpret_cast( + packed + (static_cast(pair) * (k / kDummaTileK) + k_step) * + kPackedPairBytes + + lane * 16) = slot; +} + +} // namespace + +// Stable host launch symbol used by the trusted binding for +// torch.ops.zth_w8a8.gemm_out. Runs entirely on the caller-provided stream: +// no allocation, no synchronization, no default-stream launch, no workspace +// writes. Returns nothing; the binding returns the caller-provided out. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; // in-block split-K: no cross-block partials, no workspace + (void)workspace_bytes; + + hip_bfloat16* out_bf16 = reinterpret_cast(out); + + if (m == kTargetM && n == kTargetN && k == kTargetK) { + // Exact assigned shape hy3_tp4_o_proj_m16: packed-B two-wave in-block + // split-K=2 DUMMA m16n16k32 kernel with MULTI-N-TILE REUSE (128 blocks, + // two adjacent 16x16 N tiles per block, 2 x 64-thread wavefronts, K=128 + // stages, depth-1 register prefetch). The scalar fallback below stays + // untouched and still covers every unmatched shape, including the paired + // M=2 shape with the same (N, K) (the guard is shape-exact on m too; the + // fallback decodes the packed weight for (n, k) == (4096, 2048)). + const dim3 grid(static_cast( + kTargetN / (kDummaTileN * kDummaTilesPerBlock))); + const dim3 block(kDummaBlockThreads * kDummaWavesPerBlock); + hipLaunchKernelGGL(w8a8_gemm_m16_dumma_packedb_multint_kernel, grid, + block, 0, stream, a, b, x_scale, weight_scale, + out_bf16, n, k); + return; + } + + // Generic scalar fallback for every unmatched (m, n, k). This also covers + // the paired M=2 API shape that shares the same (N, K): it must never be + // routed through an M=16-only specialization. + const dim3 grid( + static_cast((n + kScalarThreads - 1) / kScalarThreads), + static_cast(m)); + const dim3 block(kScalarThreads); + hipLaunchKernelGGL(w8a8_gemm_scalar_kernel, grid, block, 0, stream, a, b, + x_scale, weight_scale, out_bf16, m, n, k); +} + +// Stable host launch symbol used by the optional out-of-timed-region +// torch.ops.zth_w8a8.pack_weight. For the exact (k, n) == (2048, 4096) o_proj +// weight this permutes into the packed-B fragment-slot layout (one-time, +// untimed, outside Graph capture); every other (K, N) keeps the generic +// identity device-to-device copy. Weight scales are always copied through. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + if (k == kTargetK && n == kTargetN) { + // Exact o_proj M=16 weight: permute into [n_tile_pair][k_step][lane][16] + // B fragment slots (one-time, untimed, outside Graph capture). + const int64_t total = (static_cast(k) / kDummaTileK) * + (n / (2 * kDummaTileN)) * kDummaBlockThreads; + const dim3 grid( + static_cast((total + kScalarThreads - 1) / kScalarThreads)); + const dim3 block(kScalarThreads); + hipLaunchKernelGGL(w8a8_pack_o_proj_m16_kernel, grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else { + const int64_t weight_count = static_cast(k) * n; + const int64_t weight_blocks = + (weight_count + kScalarThreads - 1) / kScalarThreads; + const unsigned weight_grid = + static_cast(weight_blocks > kMaxCopyBlocks + ? kMaxCopyBlocks + : weight_blocks); + hipLaunchKernelGGL(w8a8_copy_i8_kernel, dim3(weight_grid), + dim3(kScalarThreads), 0, stream, raw_weight, + packed_weight, weight_count); + } + + const int64_t scale_count = static_cast(n); + const int64_t scale_blocks = + (scale_count + kScalarThreads - 1) / kScalarThreads; + const unsigned scale_grid = + static_cast(scale_blocks > kMaxCopyBlocks ? kMaxCopyBlocks + : scale_blocks); + hipLaunchKernelGGL(w8a8_copy_f32_kernel, dim3(scale_grid), + dim3(kScalarThreads), 0, stream, weight_scale, + packed_weight_scale, scale_count); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/qkv_proj.hip new file mode 100644 index 00000000..6d9acfa9 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/qkv_proj.hip @@ -0,0 +1,1169 @@ +// @@variant shape=hy3_tp4_qkv_proj_m16 commit=7d6959345400cb268cfd9289338d0891db2297f9 added=2026-08-26 +// median_us=26.05 p90_us=26.09 +// source=hy3-dsh-tp4-m16-1-7f1fb1d1 +// MetaInfer W8A8 INT8 GEMM - worker_0 (iteration 17, fused combine tail). +// +// Assigned shape : hy3_tp4_qkv_proj_m16, M=16, N=2560, K=4096, gfx928 +// (K500SM_AI, wavefront=64, 64 KiB LDS/CU, 120 CUs). +// +// Round history (all timings are full-operator Graph median / p90 vs the +// fixed 80.19us Triton baseline; exact int32 accumulation has been preserved +// bit-for-bit in every round): +// R1-R2 direct per-byte fragment loads (grid/occupancy variants): +// ~133-141us; each K step issued 16 scalar global ubyte loads, each +// followed by its own s_waitcnt vmcnt(N) cascade before the single +// v_mmac_i32_16x16x32_i8. +// R3 workspace split-K pair, one 64-thread wavefront per block, grid +// (160 tiles x SPLIT_K): 124.92/125.13us at split-K=2 (320 blocks = +// 2.67 blocks/CU). Falsified grid parallelism as the lever: +// vmem_read stayed 327,680 (16 per-lane ubyte loads per K step, +// invariant to split-K) and per-step wall time ~1.95us was the +// serialized load->pack chain, not bandwidth or MMA issue. +// R4 double-buffered LDS pipeline (s_a[2][16][32] + s_b[2][32][20], +// 2,304 B/block): one 8-byte vector load per lane per tile per K +// step; 54.21/54.32us (1.48x, current best). PMC: vmem_read +// 327,680 -> 40,960; LDS 225,280 instr / 491,520 bank conflicts +// (the matrix_b fragment's 8 strided ds_read_u8 at ldm 20 are +// 4-way conflicted); ~15-VALU byte-OR reassembly per step; 13 +// waitcnts/step; 76.2% L2 hit rate; per-step wall time still +// ~0.85us. +// R5 n-major packed weight + DIRECT vectorized 8-byte B fragment loads +// (A-only LDS staging): 55.73/55.84us - flat/slightly WORSE than +// R4. Removed the B LDS staging, the 8 conflicted ds_read_u8, the +// ~15-VALU OR chain and 2-3 waitcnts per step and nothing moved: +// the limiter is NOT the fragment read path (LSU/L2 request +// throughput or MMA issue). It is the one-wave geometry: 2 loads + +// 1 MMA per step cannot cover the per-step load->stage->fragment-> +// MMA chain, and only 2.67 waves/CU are resident to absorb it. +// +// Round-6 change (staging family choice from L2/VMEM evidence - chosen: +// A+B LDS): port the validated shape-sibling wqkv_a (M=16, K=4096, N=1536; +// 17.3us median, 570 GB/s logical request bandwidth) split-K=10 / +// 4-wavefront / stage-K=64 / block-N=64 architecture to N=2560: +// * block-N=64: one block computes a 16x64 output tile; four 64-thread +// wavefronts each own one 16x16 n-quadrant and SHARE the A fragments. +// * stage-K=64 double-buffered LDS: s_a[2][16][64] + s_b[2][64][64] = +// 10,240 B/block - the sibling's exact unpadded resource profile. +// * one 16-byte cooperative global load per lane per tile per stage (wave +// 0 loads the 1,024 B A stage at 16 B/lane; all four waves load the +// 4,096 B B stage at 16 B/lane) = 5 wave-level vmem_read per stage vs 2 +// per step in R4; vmem_read 40,960 -> 12,800 (-68.75%); request bytes +// 20.97 MB -> 13.11 MB (A is re-read once per 64-N tile = 40x instead of +// 160x; B streamed once, 10.49 MB). +// * stage i+1 loads issue at the TOP of stage i (16 B, 16-B-aligned) into +// VGPRs; the two v_mmac_i32_16x16x32_i8 per wave on the current stage +// cover their latency; the compiler-ordered s_waitcnt vmcnt(0) lands the +// bytes right before the ds_write_b128 pair that publishes the alternate +// slot; one __syncthreads() per stage (cross-wave barrier, no raw asm). +// * split-K default 6: 40 tiles x 6 = 240 blocks = exactly 2 blocks/CU x 4 +// wavefronts = 8 resident waves/CU (the sibling's validated 240-block / +// 2-blocks-per-CU occupancy; K=4096 = 64 stage-64 units split as 4x11 + +// 2x10, all stage-aligned, tiling [0, K) exactly once ascending). +// ZTH_W8A8_QKV_SPLIT_K still selects any trusted candidate +// {2,3,4,5,6,8,9,12} (block counts 80..480) without source edits. +// * combine kernel updated to 16x64 planes (40 blocks x 256 threads x 4 +// elements): exact int32 sum of the SPLIT_K partials in ascending slice +// order + x_scale/weight_scale + bf16 (RN-even), same stream, in Graph. +// * Exact int32 accumulation preserved bit-for-bit: same +// v_mmac_i32_16x16x32_i8 hardware k-order, same row_major fragment lane +// mapping as rounds 1-5, slices tile [0, K) once ascending, combine sums +// in int32 in ascending slice order - every per-tile dot is identical to +// rounds 1-5 and the scalar fallback. +// +// Round-6 result (built, ran, FAILED exact correctness): 72.32us median / +// 72.32us p90 (1.109x vs the 80.19us Triton baseline) but 30,697 / 40,960 +// mismatches, max_abs_error 126.0, first mismatch at flat_index 16 (m=0, +// n=16) - the first element of wave 1's quadrant, with wave 0's whole +// quadrant (n = 0..15) correct. Root cause (fixed this round): the B-stage +// cooperative load's lane role `b_nchunk = lane >> 6` is 0 for EVERY lane +// (lane = tid & 63 ranges 0..63), so all four waves redundantly staged the +// SAME 16 columns [n0, n0+16) of each 64x64 B stage and LDS columns 16..63 +// were never written (garbage). Each wave's fragment read starts at LDS +// column 16*wave, so waves 1..3 (3/4 of the outputs) consumed unstaged +// garbage - exactly matching the first mismatch at n=16 and the ~3/4 wrong +// outputs. Round-7 (this file) fixes the single lane role: `b_nchunk = wave` +// (== tid >> 6), so wave w stages B[kk + krow][n0 + 16w .. +16w + 16) - the +// same 16-column quadrant its fragments read. No instruction count changes +// (each lane still issues exactly one 16-B load per tile per stage; only the +// addresses change), so the R6 raw speed, the 12,800 vmem_read / 13.11 MB +// request profile and the 240-block / 8-waves-per-CU occupancy are preserved +// unchanged. +// +// Expected outcome (falsifiable): the repaired kernel must pass exact +// correctness (0 mismatches) at the R6 raw speed (~72us median, 1.109x vs +// the 80.19us Triton baseline). R6's sibling-scaled ~29-35us target was NOT +// reached (72.32us measured), so the accepted R4 54.21us official best may +// still lead; this round is the mandated smallest-correctness-fix of the +// faster-but-incorrect candidate, not a redesign. +// +// Round-11 change (combine-parallelism round, partial kernel BYTE-IDENTICAL +// to the accepted iteration-7 source): the operator aggregate profile of the +// accepted R7 tree (profiled 26.08us partial + 10.88us combine = 36.96us, +// matching the 36.88us official median, so the split is real, not PMC +// inflation) shows the combine kernel is ~29% of the operator wall despite +// reading only 6 x 40 x 4,096 B = 983,040 B of int32 partials. Root cause: +// the n64 combine launches 40 blocks x 256 threads x 4 elements per thread +// (160 waves on ~40 CUs, 24 strided int32 loads per thread from 6 planes +// 163,840 B apart) - a tiny, latency-bound grid with no wave-level slack. +// Round 11 (this file) splits the 16x64 tile into its four 16x16 quadrants: +// grid 40 -> 160 blocks (exactly 1.33 blocks/CU over 120 CUs), 256 threads +// x 1 element per thread, each element still sums its SPLIT_K planes in +// ascending slice order and applies the identical x_scale/weight_scale bf16 +// (RN-even) store. The element mapping is bit-identical to the R6-R10 +// combine (lin = quad*256 + t == the old t + q*256 set), so every output +// byte is unchanged; only the grid/parallelism changes. +// +// Expected outcome (falsifiable): combine 10.88us -> ~2-4us and operator +// median 36.88us -> ~29-32us (~1.15-1.27x vs the R7 accepted best), p90 +// follows. Null result (combine stays ~10.9us): the combine cost is NOT +// wave-level latency exposure - it is fixed launch/teardown or DRAM +// latency that 4x parallelism cannot hide, and the next round must fuse the +// combine into the partial kernel tail or shrink partial traffic. +// +// Round-12: killed before a valid proposal (no source change; the archived +// iteration-12 source is byte-identical to the accepted iteration-11 tree). +// +// Round-13 (partial-kernel prefetch depth, archived digest 3abfd6bb, NOT +// accepted): issued stage-(i+2)'s loads one full stage earlier via VGPR-held +// a_hold/b_hold. Measured 32.395/32.445us median/p90 - a 2.69% REGRESSION vs +// the accepted 31.522/31.642us best; rejected and the tree reverted to the +// R11 source (this file was byte-identical to iteration 11 before this +// round's edit). Its archived ISA shows why it was a no-op: the compiler +// still orders s_waitcnt vmcnt(0) before the s_barrier, which covers the +// stage-(i+2) fetches issued at the TOP of the same stage, so the load-to- +// wait window was never widened - only +4-8 VGPR and +2 in-flight vmem per +// lane were added. The per-stage wall (~2.44us) is invariant to grid +// parallelism (R3/R9), LDS bank conflicts (R8), stage width (R10) and +// prefetch depth (R13); the operator-aggregate PMC (12.68 MB moved in ~26us +// partial = ~450 GB/s plateau with ~10 KB in flight per CU) means the wall +// tracks REQUEST throughput, not exposed latency. +// +// Round-14 change (this file, B-stage cooperative-load lane coalescing): +// the B stage load is the request hog - each wave's 64 lanes load 16 B at +// stride n = 2560, i.e. 64 DISTINCT 128-B lines per wave-load, so each block +// issues 256 L1->L2 line-requests per stage for only 64 unique lines +// (~26.8 G requests/s kernel-wide, plausibly saturating L2 request slots and +// inflating the queueing latency the per-stage chain exposes). The staging +// lane roles become tid-linear: b_krow = tid >> 2 (each wave stages 16 of +// the 64 stage rows) and b_nchunk = tid & 3 (the four 16-B column chunks of +// that row), so a wave-load's 64 lanes touch only 16 unique lines (4 lanes +// cover one row's 64 B) and the stage publishes the SAME bytes to the SAME +// LDS addresses with 4x fewer L1->L2 requests (256 -> 64 per stage per +// block). Published LDS content, fragment reads, the ascending-K v_mmac +// sequence (bit-identical int32 dots; slices still tile [0, K) exactly +// once), the one __syncthreads()/stage, the 10,240 B/block LDS footprint, +// the 240-block / 2-blocks-per-CU occupancy, the 12,800 vmem_read +// instruction count and the combine kernel are all unchanged. +// +// Expected outcome (falsifiable): if L1/L2 request throughput was the +// per-stage limiter, the partial kernel drops from ~26.08us to ~12-20us and +// the operator median from 31.52us to ~17-25us (~1.26-1.85x vs the R11 +// accepted best), p90 follows. Null result (median flat ~31.5us): L1/L2 +// request count is NOT the limiter (the plateau is DRAM behavior or +// launch/tail), and the next round must attack request BYTES (wider +// n-tiling, block-N=128 with split-K=12 to halve the 40x A re-reads at +// unchanged 2-blocks/CU occupancy) - NOT more prefetch depth and NOT more +// request coalescing. +// +// Round-17 change (this file, fused combine tail): the operator-aggregate +// PMC of the accepted R14 tree (profile iteration15/pmc.json, reused as the +// current-best evidence) splits the 28.483us wall into partial ~22.879us +// profiled and combine 3.68us profiled (160 blocks, 24 VGPR), and the +// combine's L2 hit rate is only 6.3% (l2_hits 1,166 / misses 17,205): the +// 983,040 B of int32 partials were evicted from L2 by the 12.68 MB A/B +// stream of the partial kernel and must be re-read from HBM, on top of the +// second-kernel launch gap. R16's block-N=128 / split-K=12 "more compute +// per stage" attack measured 32.087us - an 11.2% REGRESSION - falsifying +// per-stage compute width as the fix for the partial kernel, exactly as its +// null-result clause predicted ("the wall is launch/tail overhead or DRAM +// latency - the next round must fuse/trim the combine (19% of the wall) or +// the grid tail"). This round FUSES the combine into the partial kernel +// tail, the smallest change that removes the combine kernel launch, its gap +// and its HBM re-read together: +// * The partial kernel's stage loop, fragment mapping, tid-linear B +// staging, K tiling and per-tile int32 dots are BYTE-IDENTICAL to the +// accepted R14 source; only the tail changes. +// * After the (unchanged) private du_store_matrix_sync, each block's +// tid-0 does __threadfence() + atomicAdd(&counters[tile], 1). The block +// whose atomicAdd returns SPLIT_K-1 is the LAST arrival for its tile - +// hardware-ordered by the atomics themselves, so there is NO spin loop +// and NO residency assumption (safe for every trusted split-K, even if +// register pressure ever dropped occupancy). +// * The last block sums the tile's SPLIT_K planes in ascending slice +// order (identical per-element int32 accumulation order as the R11 +// combine, so every output byte is bit-identical) and writes the scaled +// bf16 output with the same x_scale/weight_scale RN-even conversion and +// store addresses. Element mapping lin = tid*4 + e covers the 16x64 +// tile exactly once; per plane each wave loads 64 x 16 B = one +// contiguous 1,024 B chunk (coalesced), so the fused tail's 6 plane +// reads per thread are wide v4 loads. +// * The fused tail reads the sibling planes while they are still L2-hot +// (co-resident blocks write them ~0-2us earlier), so the combine's +// HBM re-read (~2.5us at the measured 6.3% hit rate) largely disappears +// as well. The operator becomes ONE kernel launch on the caller's +// stream; the round-11 combine kernel and its launch gap (~1-2us) are +// gone. +// * The per-tile counters (num_tiles x int32 = 160 B) live in the last +// 160 B of the 16-plane contract workspace (plane 15's tail; max used +// plane index for any trusted split-K <= 12 is 479, far below 600, so +// no overlap with partial data). The combiner self-clears its tile's +// counter (atomicExch 0) after combining, so replays start from zero; +// the very first launch with a given workspace is zeroed by one async +// hipMemsetAsync on the caller's stream, guarded by a workspace-pointer +// static so the captured graph contains only the single kernel launch +// and every replay is memset-free. No allocation, no host/device +// synchronization, no default stream - the timed-region rules are +// unchanged. +// * The two-kernel path (partial kernel with FUSED=false + the round-11 +// quadrant combine kernel) is preserved for workspaces smaller than the +// 16-plane contract, so the fallback behavior is byte-identical to R14. +// +// Expected outcome (falsifiable): if the combine kernel + its launch gap +// (~4-5.5us of the 28.48us wall) were latency/traffic, the operator median +// drops to ~24-26.5us (~1.08-1.19x vs the R14 accepted best, ~3.0-3.35x vs +// the 80.19us Triton baseline), p90 follows, and the profiled kernel count +// goes 2 -> 1. Null result (median flat ~28.5us): the fused tail's own cost +// (last-tile combine ~1-2us on the critical path) cancels the savings - the +// wall is inside the partial stage loop itself, and the next round must +// attack per-stage DRAM latency (wider B lines / different B layout) rather +// than launch structure. +// +// Round-20/21/22: R20 (cold-plane L2 warm, discarded volatile reads of the +// base-stage slices) was killed before measurement; R21 (byte-identical R17 +// re-issue) was killed the same way, proving the kills environmental; R22 +// finally measured the warm on the byte-identical R17 tree: 28.409/28.469us +// median/p90 - a 7.9% REGRESSION (rejected; tree reverted to the R17 source, +// digest 264dbb1f). The warm's 1.92 MB of extra reads cost ~4.3us of wall +// during the bandwidth-saturated main phase and the tail did NOT get faster, +// falsifying the tail-as-DRAM-re-read model: the fused tail is NOT partial +// re-read DRAM traffic (the sibling planes are L2-hot), it is the arrival +// barrier/atomic serialization + last-stage skew + the combiner's serial +// critical path, and the main phase is at the DRAM/request floor (R18's +// packed-B layout - half the L2->DRAM line requests - was flat; R19's 3 +// blocks/CU occupancy was a regression; R14's lane-coalescing and R17's +// fusion are the only real wins). +// +// Round-23 change (this file, fused-tail critical-path trim): with the tail +// established as serial latency, the two smallest arrival-protocol cuts: +// * The trailing `atomicExch(&counters[tile], 0)` self-clear is REMOVED. +// The last-arrival test becomes modulo-SPLIT_K on the monotonic arrival +// count: replay r's arrivals return (r-1)*SPLIT_K + 0..SPLIT_K-1, so the +// SPLIT_K-th arrival of EVERY replay is exactly the one with +// `arrived % SPLIT_K == SPLIT_K - 1` (int32 wraparound would need +// 2^31/6 ~ 358M replays). The combiner no longer pays a dependent +// atomicExch round trip (~0.2-0.4us) after its output stores; the +// one-time host hipMemsetAsync of the counters before the first launch +// is unchanged (and now purely a first-launch precondition), the +// FUSED=false two-kernel fallback never touches counters, and every +// replay stays exactly one kernel launch on the caller's stream. +// * The tail's four per-thread bf16 stores (four `global_store_short_d16_hi` +// with an intervening waitcnt in the archived ISA) are packed into ONE +// 8-B store: the 4 elements (lin = tid*4 + 0..3) are always 4 consecutive +// columns of one row (4 | 64, so a 4-element run never crosses the +// 64-column row boundary) and the byte address is 8-B aligned (n0, col0 +// and row*n are all multiples of 4), so the four RN-even bf16 bit +// patterns pack little-endian into one uint64 - identical output bytes, +// identical store addresses, one store instruction instead of four plus +// their address mads and the intervening wait. +// The stage loop, tid-linear B staging, fragment mapping, K tiling, +// per-tile int32 dots (ascending-K v_mmac + ascending-slice plane sums), +// the arrival atomicAdd/fence order and the plane layout are byte-identical +// to the accepted R17 source, so exact int32 accumulation and every output +// byte are preserved bit-for-bit. +// +// Expected outcome (falsifiable): if the trailing atomic round trip + the +// serialized 2-B store issue are on the last-tile combiner's critical path, +// the operator median drops from 26.331us toward ~25.8-26.0us (~1-2%, +// ~1.02x vs the R17 accepted best, ~3.08-3.11x vs the 80.19us Triton +// baseline), p90 follows, and the profiled fused kernel drops from 25.439us +// toward ~25.0-25.2us with the global_store instruction count down slightly +// in the tail and everything else (vmem_read 14,720, LDS 198,320, grid 240, +// one launch) unchanged. Null result (median flat ~26.3us): the tail wall is +// the arrival atomicAdd + plane-load latency + last-stage skew (the +// self-clear and store issue are off it), and the next round must split the +// combine across the last TWO arrivals (bitmap + scratch + bounded-flag +// handoff; the producer is provably executing when the consumer spins, so +// it is deadlock-free) or attack the per-stage load batch directly. +// +// Timed-region rules honored: launch_w8a8_gemm performs no allocation, +// compilation, autotuning, packing, host/device synchronization or +// default-stream launch; it uses only the caller-provided tensors (the +// workspace holds the SPLIT_K int32 partial planes; round 17 additionally +// uses its last 160 B for the per-tile fused-combine arrival counters) and +// issues the kernel launch(es) on the caller's stream - one fused kernel +// for the contract workspace (round 17), two kernels (partial + combine) +// for the preserved fallback. The one-time counter zero is an async +// hipMemsetAsync on the same stream before the first launch with a given +// workspace, so it is never captured into the replayed Graph. The split-K +// choice is a host-side static dispatch decision (env override or default +// 6) made once per launch call, never inside the Graph replay, so +// capture/replay determinism holds. +// +// Header order is the DTK known-good order: hip_runtime first (du_mma.h is +// not self-contained when included before the HIP runtime headers), then +// the bfloat16 header, then du_mma.h. + +#include +#include +#include + +#include +#include + +namespace { + +// gfx928 native wavefront is 64 lanes; blockDim must be a multiple of 64. +constexpr int kBlockThreads = 256; + +// Minimal DUMMA tile constants (gfx928 INT8 support is m16n16k32). +constexpr int kDummaM = 16; +constexpr int kDummaN = 16; +constexpr int kDummaK = 32; +constexpr int kDummaThreads = 64; // one wavefront (64 lanes) per block +constexpr int kCombineThreads = 256; // 16x16 = 256 output elements per block + +// Round-4 LDS staging (preserved kernel): padded B row stride (bytes). The +// matrix_b row_major fragment reads p[(col + j) * ldm + row]; at ldm = 16 +// all 64 lanes of a ds_read_u8 collide on 4 LDS banks (16-way conflict), +// while ldm = 20 (4 mod 8) spreads each read over 16 banks (4-way). +constexpr int kDummaLdsBRow = kDummaN + 4; // 20 + +// Round-6 geometry constants: block-N=64 (four 16x16 quadrants per block), +// stage-K=64 (two 32-K DUMMA steps per stage), four wavefronts per block. +constexpr int kDummaN64 = 64; // N-tile width per round-6 block +constexpr int kStageK = 64; // K rows staged per double-buffer slot +constexpr int kDummaThreads4 = 256; // 4 wavefronts x 64 lanes per block + +// Trusted occupancy-probe split-K candidates (control plane): all fit the +// contract workspace for this shape (2,621,440 B = 16 partial planes of +// 163,840 B each) and produce 64-aligned K slices. Round-6 block counts = +// 40*S (80..480). +constexpr int kTrustedSplitK[] = {2, 3, 4, 5, 6, 8, 9, 12}; + +// Round-to-nearest-even float -> bfloat16 bit pattern. bf16 tensors store +// uint16, so a manual uint16 store avoids depending on any bf16 type +// conversion API and is bit-identical to __float2bfloat16 in its default +// (RN-even) rounding mode. +__device__ __forceinline__ uint16_t float_to_bf16_bits(float f) { + uint32_t u = 0; + __builtin_memcpy(&u, &f, sizeof(u)); + const uint32_t bias = 0x7FFFu + ((u >> 16) & 1u); + u += bias; + return static_cast(u >> 16); +} + +// 16-byte vector used by the round-6 stage loads/stores (one 16-byte +// cooperative load per lane per tile; both global and LDS addresses are +// 16-B aligned - see the kernel comment). +struct alignas(16) v16_t { + int64_t lo; + int64_t hi; +}; + +// Stage one DUMMA K-step's A (16x32) and B (32x16) int8 tiles from global +// into one LDS double-buffer slot using one 8-byte vector load per lane per +// tile (2 vmem instructions per K step total, vs the 16 scalar +// global_load_ubyte the direct fragment loads issue). All 64 lanes +// participate; both global addresses are 8-byte aligned (kk is a multiple of +// 32, n0 of 16, n of 2560, x_q/weight tensor bases are 256 B aligned): +// A: lane = row(0..15) | chunk(0..3)*16 -> 8 B of A row `row` at k-columns +// kk + chunk*8 .. +8, written to s_a[slot][row][chunk*8 .. +8) +// (LDS offset row*32 + chunk*8, 8 B aligned). +// B: lane = krow(0..31) | half(0..1)*32 -> 8 B of K-row kk+krow at +// n-columns n0 + half*8 .. +8, written to s_b[slot][krow][half*8 .. +8) +// (LDS offset krow*20 + half*8; odd krows are 4 B aligned, so the +// compiler splits those writes into two ds_write_b32 - harmless). +__device__ __forceinline__ void stage_dumma_step( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int kk, + int k, + int n, + int n0, + int lane, + int slot, + int8_t* __restrict__ s_a_slot, + int8_t* __restrict__ s_b_slot) { + const int a_row = lane & 15; + const int a_chunk = lane >> 4; // 0..3 + const int b_krow = lane & 31; + const int b_half = lane >> 5; // 0..1 + const int64_t va = *reinterpret_cast( + x_q + static_cast(a_row) * k + kk + a_chunk * 8); + *reinterpret_cast(s_a_slot + a_row * kDummaK + a_chunk * 8) = va; + const int64_t vb = *reinterpret_cast( + weight + static_cast(kk + b_krow) * n + n0 + b_half * 8); + *reinterpret_cast(s_b_slot + b_krow * kDummaLdsBRow + b_half * 8) = + vb; +} + +// Workspace split-K partial kernel (template over SPLIT_K), M == 16: +// partial[tile, slice][m, n] = int32_dot(x_q[m, k0:k1], weight[k0:k1, n]) +// Layouts (all contiguous): +// x_q [16, K] int8 row-major (ldm = K) +// weight [K, N] int8 row-major (ldm = N, identity pack layout) +// partials [(SPLIT_K * num_tiles) * 256] int32, plane = slice*num_tiles + +// tile, row-major 16x16 tile per plane (workspace) +// One block = one wavefront = one (tile, slice). The K slice is 32-aligned +// (non-uniform: slices differ by at most one 32-K DUMMA step) and covers +// [k0, k0 + steps*32) in ascending order. Round-4: the per-step fragment +// loads are served from a double-buffered LDS pipeline (s_a[2][16][32] + +// s_b[2][32][20], 2,304 B per block) instead of direct 16 scalar byte loads +// from global; step i+1's two staging loads issue at the top of step i and +// complete under the step-i fragment reads + v_mmac, with one __syncthreads +// per step ordering them. Each block still writes exactly one private int32 +// partial plane. +template +__global__ __launch_bounds__(kDummaThreads) void w8a8_dumma_m16_splitk_partial_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int32_t* __restrict__ partials, + int n, + int k) { + const int tile = static_cast(blockIdx.x); + const int slice = static_cast(blockIdx.y); + const int n0 = tile * kDummaN; + const int lane = static_cast(threadIdx.x); // 0..63, one wavefront + + // 32-aligned non-uniform slice bounds over the K dimension. + const int total_steps = k / kDummaK; // 4096 / 32 = 128 + const int base = total_steps / SPLIT_K; + const int rem = total_steps - base * SPLIT_K; + const int steps = base + (slice < rem ? 1 : 0); + const int k0 = (slice * base + (slice < rem ? slice : rem)) * kDummaK; + + // Double-buffered LDS staging slots (round 4). + __shared__ int8_t s_a[2][kDummaM][kDummaK]; + __shared__ int8_t s_b[2][kDummaK][kDummaLdsBRow]; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + const int a_row = lane & 15; + const int a_chunk = lane >> 4; // 0..3 + const int b_krow = lane & 31; + const int b_half = lane >> 5; // 0..1 + + // Prologue: stage step 0 into slot 0, then make it visible. + if (steps > 0) { + stage_dumma_step(x_q, weight, k0, k, n, n0, lane, 0, &s_a[0][0][0], + &s_b[0][0][0]); + } + __syncthreads(); + + for (int i = 0; i < steps; ++i) { + const int cur = i & 1; + const int nxt = cur ^ 1; + + // Prefetch step i+1: issue its two 8-byte global loads FIRST (into + // registers, no wait) so the ~hundreds of cycles of global latency + // overlap the step-i fragment LDS reads and the v_mmac below; the + // trailing ds_writes + barrier land the bytes before the next + // iteration's fragment reads. + int64_t a_nxt = 0; + int64_t b_nxt = 0; + if (i + 1 < steps) { + const int kk_nxt = k0 + (i + 1) * kDummaK; + a_nxt = *reinterpret_cast( + x_q + static_cast(a_row) * k + kk_nxt + a_chunk * 8); + b_nxt = *reinterpret_cast( + weight + static_cast(kk_nxt + b_krow) * n + n0 + + b_half * 8); + } + + du::dumma::du_load_matrix_sync(a_frag, &s_a[cur][0][0], kDummaK); + du::dumma::du_load_matrix_sync(b_frag, &s_b[cur][0][0], kDummaLdsBRow); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + + if (i + 1 < steps) { + *reinterpret_cast(&s_a[nxt][0][0] + a_row * kDummaK + + a_chunk * 8) = a_nxt; + *reinterpret_cast(&s_b[nxt][0][0] + b_krow * kDummaLdsBRow + + b_half * 8) = b_nxt; + } + __syncthreads(); + } + + // Private int32 partial plane: slice-major so a tile's S planes are + // stride-uniform in the combine kernel. + const int num_tiles = static_cast(gridDim.x); + du::dumma::du_store_matrix_sync( + partials + (slice * num_tiles + tile) * (kDummaM * kDummaN), + acc_frag, kDummaN, du::dumma::mem_row_major); +} + +// --------------------------------------------------------------------------- +// Round-6 partial kernel: block-N=64 / 4 wavefronts / stage-K=64 double- +// buffered A+B LDS (the validated sibling wqkv_a split-K=10 architecture, +// ported to N=2560). +// +// One block computes the 16x64 output tile [0..16) x [n0..n0+64). Four +// wavefronts (256 threads) each own one 16x16 n-quadrant (wave w -> columns +// n0 + 16w .. n0 + 16w + 16) and share the A fragments. +// +// Staging (one 16-byte cooperative load per lane per tile per stage): +// A stage 16x64 int8 = 1,024 B: wave 0 loads it at 16 B/lane +// (lane = row(0..15) | chunk(0..3)*16 -> A[row][kk + chunk*16 .. +16), +// 16-B aligned: row*4096 + kk + chunk*16, kk % 64 == 0). +// B stage 64x64 int8 = 4,096 B: all four waves load it at 16 B/lane +// (tid-linear roles: stage row b_krow = tid >> 2, 16-B column chunk +// b_nchunk = tid & 3 -> B[kk + b_krow][n0 + b_nchunk*16 .. +16), +// 16-B aligned: (kk + b_krow)*2560 + n0 + b_nchunk*16). Round-7 +// correctness property preserved: every (row, chunk) of the stage is +// written exactly once (wave w stages rows [16w, 16w+16) x all four +// chunks, so columns [n0, n0+64) are fully covered; the round-6 +// `b_nchunk = lane >> 6` bug that left columns 16..63 unwritten is +// structurally impossible now). Round-14: the tid-linear roles coalesce +// each wave-load onto 16 unique 128-B lines (4 consecutive lanes cover +// one row's 64 B) instead of rounds 6-11's one-lane-per-row mapping +// (64 unique lines per wave-load) - 4x fewer L1->L2 requests, identical +// bytes, identical LDS addresses, identical fragment reads. +// Stage i+1's five wave-level loads issue at the TOP of stage i (into +// VGPRs, no wait); the two v_mmac per wave on the current stage cover +// their latency; the compiler-ordered s_waitcnt vmcnt(0) lands the bytes +// right before the ds_write_b128 pair that publishes the alternate slot; +// one __syncthreads() per stage makes the slot visible to all four waves. +// +// Fragment reads keep the same row_major DUMMA lane mapping as rounds 1-5 +// (so the per-tile int32 dots are bit-identical): a_frag = 8 consecutive K +// bytes of one M row (ldm 64); b_frag = 8 strided ds_read_u8 (ldm 64, 16-way +// LDS bank conflict, accepted: round 5 proved the B fragment read path is +// off the critical path, and the sibling's 10,240 B profile is exactly +// unpadded). +// +// Partial store: each wave du_store_matrix_sync's its 16x16 quadrant into +// the (slice, tile) 16x64 int32 plane (slice-major, same convention as +// rounds 3-5). +template +__global__ __launch_bounds__(kDummaThreads4) +void w8a8_dumma_m16_n64_splitk_partial_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int32_t* __restrict__ partials, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + uint16_t* __restrict__ out, + int32_t* __restrict__ counters, + int n, + int k) { + const int tile = static_cast(blockIdx.x); // 0..(n/64-1) = 0..39 + const int slice = static_cast(blockIdx.y); // 0..SPLIT_K-1 + const int tid = static_cast(threadIdx.x); // 0..255 + const int wave = tid >> 6; // 0..3 + const int lane = tid & 63; + const int n0 = tile * kDummaN64; + + // 64-aligned non-uniform slice bounds over the K dimension (units of + // 64-K stages): slices tile [0, K) exactly once in ascending order. + const int total_stages = k / kStageK; // 4096 / 64 = 64 + const int base = total_stages / SPLIT_K; + const int rem = total_stages - base * SPLIT_K; + const int stages = base + (slice < rem ? 1 : 0); + const int s0 = (slice * base + (slice < rem ? slice : rem)) * kStageK; + + // Double-buffered stage slots (round-6 geometry; 10,240 B/block total). + __shared__ int8_t s_a[2][kDummaM][kStageK]; // 2 x 1,024 B + __shared__ int8_t s_b[2][kStageK][kDummaN64]; // 2 x 4,096 B + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Staging lane roles. + const int a_row = (wave == 0) ? (lane & 15) : 0; + const int a_chunk = (wave == 0) ? (lane >> 4) : 0; // 0..3 + // Round-14 (B-stage load coalescing): tid-linear staging roles - each wave + // stages 16 of the 64 stage rows x all four 16-B column chunks (wave w + // covers rows [16w, 16w+16)). Round-7's correctness property (every stage + // row/column written exactly once, all four waves participate) is + // preserved: b_nchunk = tid & 3 ranges 0..3 inside every wave, and + // b_krow = tid >> 2 ranges 0..63 across the block, so the round-6 + // `lane >> 6` redundant-staging bug cannot recur. + const int b_krow = tid >> 2; // 0..63: 16 stage rows per wave + const int b_nchunk = tid & 3; // 0..3: 16-B column chunk of that row + + // Prologue: stage stage-0 into slot 0, then make it visible. + if (stages > 0) { + v16_t a_v{}; + v16_t b_v; + if (wave == 0) { + a_v = *reinterpret_cast( + x_q + static_cast(a_row) * k + s0 + a_chunk * 16); + } + b_v = *reinterpret_cast( + weight + static_cast(s0 + b_krow) * n + n0 + b_nchunk * 16); + if (wave == 0) { + *reinterpret_cast(&s_a[0][a_row][a_chunk * 16]) = a_v; + } + *reinterpret_cast(&s_b[0][b_krow][b_nchunk * 16]) = b_v; + } + __syncthreads(); + + for (int i = 0; i < stages; ++i) { + const int cur = i & 1; + const int nxt = cur ^ 1; + + // Prefetch stage i+1: issue its 16-byte global loads FIRST (into + // registers, no wait) so their latency overlaps the stage-i fragment LDS + // reads and the two v_mmac below; the compiler orders the s_waitcnt + // vmcnt(0) immediately before the LDS stores that publish the alternate + // slot. + v16_t a_v{}; + v16_t b_v; + if (i + 1 < stages) { + const int kk_nxt = s0 + (i + 1) * kStageK; + if (wave == 0) { + a_v = *reinterpret_cast( + x_q + static_cast(a_row) * k + kk_nxt + a_chunk * 16); + } + b_v = *reinterpret_cast( + weight + static_cast(kk_nxt + b_krow) * n + n0 + + b_nchunk * 16); + } + + // Two 32-K DUMMA steps of the current stage (ascending K). + for (int j = 0; j < 2; ++j) { + du::dumma::du_load_matrix_sync(a_frag, &s_a[cur][0][j * kDummaK], + kStageK); + du::dumma::du_load_matrix_sync(b_frag, + &s_b[cur][j * kDummaK][wave * kDummaN], + kStageK); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + if (i + 1 < stages) { + if (wave == 0) { + *reinterpret_cast(&s_a[nxt][a_row][a_chunk * 16]) = a_v; + } + *reinterpret_cast(&s_b[nxt][b_krow][b_nchunk * 16]) = b_v; + } + __syncthreads(); + } + + // Private int32 partial plane quadrant: slice-major planes of 16x64 so a + // tile's S planes are stride-uniform in the combine kernel. + const int num_tiles = static_cast(gridDim.x); + du::dumma::du_store_matrix_sync( + partials + (slice * num_tiles + tile) * (kDummaM * kDummaN64) + + wave * kDummaN, + acc_frag, kDummaN64, du::dumma::mem_row_major); + + if (FUSED) { + // Round-17 fused combine tail (replaces the separate combine kernel): + // every block signals its arrival for its tile with one tid-0 + // __threadfence() + atomicAdd(&counters[tile], 1); the LAST arrival + // (atomicAdd returns the value that identifies it, hardware-ordered + // after all SPLIT_K sibling stores+fences - no spin loop, no residency + // assumption) sums the tile's SPLIT_K planes in ascending slice order + // and writes the scaled bf16 output with the same x_scale/weight_scale + // RN-even conversion and store addresses as the R11 combine. The + // per-element int32 sums are bit-identical to R11 (ascending slice + // order), so the exact-K-order accumulation contract is unchanged. + // + // Round-23 (arrival-protocol trim): the counters are now MONOTONIC - + // the trailing atomicExch(&counters[tile], 0) self-clear is gone. Each + // replay adds exactly SPLIT_K to its tile's counter, so replay r's + // arrivals return (r-1)*SPLIT_K + 0..SPLIT_K-1 and the SPLIT_K-th + // arrival of EVERY replay is exactly the one with + // (arrived % SPLIT_K) == SPLIT_K - 1 (int32 wraparound would need + // 2^31/6 ~ 358M replays). The combiner no longer pays a dependent + // atomicExch round trip after its output stores; the host still zeroes + // the counters once per workspace before the first launch (see + // launch_w8a8_gemm) and the FUSED=false fallback never touches them. + __syncthreads(); + __shared__ int s_is_last; + if (tid == 0) { + __threadfence(); // release: sibling plane stores visible to the + // observer of this block's arrival atomic + const int arrived = atomicAdd(&counters[tile], 1); + __threadfence(); // acquire: this block's reads (below, after the + // barrier) see every sibling store that was released + // before its arrival atomic + s_is_last = ((arrived % SPLIT_K) == SPLIT_K - 1); + } + __syncthreads(); + if (s_is_last) { + // 1,024 output elements per 16x64 tile; 256 threads x 4 consecutive + // elements (lin = tid*4 + e). Per plane, the 64 lanes of a wave load + // 16 B each = one contiguous 1,024 B plane chunk (perfectly + // coalesced). + const int n0 = tile * kDummaN64; + const int lin0 = tid * 4; + int32_t sums[4] = {0, 0, 0, 0}; + for (int s = 0; s < SPLIT_K; ++s) { + const int32_t* __restrict__ plane = + partials + (s * num_tiles + tile) * (kDummaM * kDummaN64) + lin0; + const int4 p = *reinterpret_cast(plane); + sums[0] += p.x; + sums[1] += p.y; + sums[2] += p.z; + sums[3] += p.w; + } + // Round-23 (store packing): the 4 elements (lin = tid*4 + 0..3) are + // always 4 consecutive columns of one row (4 | 64, so a 4-element run + // never crosses the 64-column row boundary) and the byte address is + // 8-B aligned (n0, col0 and row*n are all multiples of 4), so the + // four 2-B bf16 patterns pack little-endian into ONE 8-B store - + // identical RN-even conversions, identical output bytes and store + // addresses, one store instruction instead of four plus their address + // mads and the intervening waitcnt seen in the R17 ISA. + const int row = lin0 >> 6; + const int col0 = lin0 & 63; + uint64_t packed = 0; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float scaled = static_cast(sums[e]) * x_scale[row] * + weight_scale[n0 + col0 + e]; + packed |= static_cast(float_to_bf16_bits(scaled)) << (16 * e); + } + *reinterpret_cast(out + static_cast(row) * n + n0 + + col0) = packed; + } + } +} + + +// Combine + scale kernel (M == 16): one 16x16 tile per block, one thread per +// output element. Exact int32 sum of the SPLIT_K workspace partials, then +// x_scale/weight_scale scaling and the bf16 (RN-even) store. No barrier is +// needed: every thread reads only its own 16x16 element's S partial values +// (written by the previous kernel on the same stream). +__global__ __launch_bounds__(kCombineThreads) void w8a8_dumma_m16_combine_scale_kernel( + const int32_t* __restrict__ partials, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + uint16_t* __restrict__ out, + int split_k, + int num_tiles, + int n) { + const int tile = static_cast(blockIdx.x); + const int linear = static_cast(threadIdx.x); // 0..255 + const int row = linear >> 4; + const int col = linear & 15; + const int out_col = tile * kDummaN + col; + + int32_t sum = 0; + for (int s = 0; s < split_k; ++s) { + sum += partials[(s * num_tiles + tile) * (kDummaM * kDummaN) + linear]; + } + + const float scaled = + static_cast(sum) * x_scale[row] * weight_scale[out_col]; + out[row * n + out_col] = float_to_bf16_bits(scaled); +} + +// Round-11 combine + scale kernel (M == 16): one 16x16 n-quadrant per block +// (grid = 4 x num_tiles = 160 blocks for the assigned shape), 256 threads x +// 1 element per thread. Exact int32 sum of the SPLIT_K workspace partials in +// ascending slice order, then x_scale/weight_scale scaling and the bf16 +// (RN-even) store. No barrier is needed: every thread reads only its own +// element's S partial values (written by the previous kernel on the same +// stream). Element mapping is bit-identical to the R6-R10 combine: lin = +// quad*256 + t covers exactly the old t + q*256 set, so row/col, the int32 +// sum order and the store addresses are unchanged. Stores are coalesced: +// each block's 256 lanes cover one full 16x16 quadrant (64 B per row-slice +// of 64 lanes x 2 B). +__global__ __launch_bounds__(kCombineThreads) +void w8a8_dumma_m16_n64_combine_scale_kernel( + const int32_t* __restrict__ partials, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + uint16_t* __restrict__ out, + int split_k, + int num_tiles, + int n) { + const int quad = static_cast(blockIdx.x) & 3; // 0..3 + const int tile = static_cast(blockIdx.x) >> 2; // 0..num_tiles-1 + const int t = static_cast(threadIdx.x); // 0..255 + const int lin = quad * kCombineThreads + t; // 0..1023 in the tile + const int row = lin >> 6; + const int col = lin & 63; + const int n0 = tile * kDummaN64; + + int32_t sum = 0; + for (int s = 0; s < split_k; ++s) { + sum += partials[(s * num_tiles + tile) * (kDummaM * kDummaN64) + lin]; + } + const float scaled = + static_cast(sum) * x_scale[row] * weight_scale[n0 + col]; + out[row * n + n0 + col] = float_to_bf16_bits(scaled); +} + +// Scalar W8A8 GEMM (generic fallback for every shape not matched by the +// exact-shape DUMMA launch): +// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// Layouts (all contiguous): +// x_q [M, K] int8 row-major +// weight [K, N] int8 row-major (identity pack layout) +// x_scale [M] fp32 +// weight_scale [N] fp32 +// out [M, N] bf16 (uint16 storage) +// One thread per output element; adjacent lanes map to adjacent N addresses +// (linear index over row-major [M, N] output), keeping stores coalesced. +__global__ __launch_bounds__(kBlockThreads) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + uint16_t* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = static_cast(blockIdx.x) * kBlockThreads + + static_cast(threadIdx.x); + const int64_t total = static_cast(m) * static_cast(n); + if (linear >= total) { + return; + } + const int row = static_cast(linear / static_cast(n)); + const int col = static_cast(linear - static_cast(row) * n); + + const int8_t* __restrict__ a_row = + x_q + static_cast(row) * static_cast(k); + const int8_t* __restrict__ b_col = weight + col; + + // Exact int32 accumulation over the full K dimension. + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = float_to_bf16_bits(scaled); +} + +// Template launch helper for the partial kernel (one wavefront per block, +// grid = (num_tiles, SPLIT_K)). +template +void launch_splitk_partial(const int8_t* a, + const int8_t* b, + int32_t* partials, + int n, + int k, + hipStream_t stream, + int num_tiles) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk_partial_kernel), + dim3(static_cast(num_tiles), static_cast(SPLIT_K)), + dim3(static_cast(kDummaThreads)), + 0, + stream, + a, + b, + partials, + n, + k); +} + +// Template launch helper for the round-6 partial kernel (four wavefronts per +// block, grid = (num_tiles_64, SPLIT_K)). FUSED selects the round-17 fused +// combine tail (single-kernel operator); FUSED=false keeps the round-11 +// two-kernel path for non-contract workspaces (x_scale/weight_scale/out/ +// counters are passed but unused by the kernel when FUSED=false). +template +void launch_splitk_partial_n64(const int8_t* a, + const int8_t* b, + int32_t* partials, + const float* x_scale, + const float* weight_scale, + uint16_t* out, + int32_t* counters, + int n, + int k, + hipStream_t stream, + int num_tiles) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_n64_splitk_partial_kernel), + dim3(static_cast(num_tiles), static_cast(SPLIT_K)), + dim3(static_cast(kDummaThreads4)), + 0, + stream, + a, + b, + partials, + x_scale, + weight_scale, + out, + counters, + n, + k); +} + +// Scalar fallback used by the dispatcher below (defined after it; forward +// declaration keeps the call sites valid). +void launch_scalar_gemm(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + int m, + int n, + int k, + hipStream_t stream); + +// Dispatcher over the trusted split-K set for the n64 partial kernel. +// FUSED=false (the two-kernel path) passes null scale/output/counter +// pointers; the kernel never dereferences them in that instantiation. +template +void launch_splitk_partial_n64_switch(int split_k, + const int8_t* a, + const int8_t* b, + int32_t* partials, + const float* x_scale, + const float* weight_scale, + uint16_t* out, + int32_t* counters, + int m, + int n, + int k, + hipStream_t stream, + int num_tiles) { + switch (split_k) { + case 2: + launch_splitk_partial_n64<2, FUSED>(a, b, partials, x_scale, + weight_scale, out, counters, n, k, + stream, num_tiles); + break; + case 3: + launch_splitk_partial_n64<3, FUSED>(a, b, partials, x_scale, + weight_scale, out, counters, n, k, + stream, num_tiles); + break; + case 4: + launch_splitk_partial_n64<4, FUSED>(a, b, partials, x_scale, + weight_scale, out, counters, n, k, + stream, num_tiles); + break; + case 5: + launch_splitk_partial_n64<5, FUSED>(a, b, partials, x_scale, + weight_scale, out, counters, n, k, + stream, num_tiles); + break; + case 6: + launch_splitk_partial_n64<6, FUSED>(a, b, partials, x_scale, + weight_scale, out, counters, n, k, + stream, num_tiles); + break; + case 8: + launch_splitk_partial_n64<8, FUSED>(a, b, partials, x_scale, + weight_scale, out, counters, n, k, + stream, num_tiles); + break; + case 9: + launch_splitk_partial_n64<9, FUSED>(a, b, partials, x_scale, + weight_scale, out, counters, n, k, + stream, num_tiles); + break; + case 12: + launch_splitk_partial_n64<12, FUSED>(a, b, partials, x_scale, + weight_scale, out, counters, n, k, + stream, num_tiles); + break; + default: + launch_scalar_gemm(a, b, x_scale, weight_scale, out, m, n, k, stream); + break; + } +} + +void launch_scalar_gemm(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + int m, + int n, + int k, + hipStream_t stream) { + const int64_t total = static_cast(m) * static_cast(n); + const int blocks = + static_cast((total + kBlockThreads - 1) / kBlockThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + dim3(static_cast(blocks)), + dim3(static_cast(kBlockThreads)), + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast(out), + m, + n, + k); +} + +} // namespace + +// Host launcher for torch.ops.zth_w8a8.gemm_out (timed / graph-capturable). +// Called on PyTorch's current HIP stream. The exact assigned shape +// (m,n,k) == (16,2560,4096) is dispatched to the workspace split-K DUMMA +// pair (partial kernel + combine/scale kernel, both in the timed Graph); +// every other shape reaches the scalar fallback. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + if (m == 16 && n == 2560 && k == 4096) { + // Round-6 geometry: block-N=64, so 2560 / 64 = 40 tiles. + const int num_tiles = n / kDummaN64; // 2560 / 64 = 40 + const int64_t plane_bytes = + static_cast(num_tiles) * kDummaM * kDummaN64 * + static_cast(sizeof(int32_t)); // 40 * 1024 * 4 = 163,840 + + // Split-K selection: default 6 (mandated primary geometry: 40 x 6 = 240 + // four-wave blocks = exactly 2 blocks/CU x 4 wavefronts = 8 resident + // waves/CU, the validated sibling wqkv_a occupancy). The + // ZTH_W8A8_QKV_SPLIT_K environment variable selects any trusted + // occupancy-probe candidate {2,3,4,5,6,8,9,12} (block counts 80..480) + // so the CU-aligned sweep can be measured without source edits. The read + // happens once per launch call on the host, never during Graph replay, + // so capture/replay determinism is preserved. + int split_k = 6; + if (const char* env = std::getenv("ZTH_W8A8_QKV_SPLIT_K")) { + char* end = nullptr; + const long parsed = std::strtol(env, &end, 10); + if (end != env && *end == '\0') { + for (int t : kTrustedSplitK) { + if (static_cast(parsed) == t) { + split_k = t; + break; + } + } + } + } + + // Workspace fit: clamp to the largest trusted candidate that fits the + // caller workspace (contract guarantees 16 planes = 2,621,440 B, so the + // clamp never fires for the exact shape; kept defensive). + int best_fit = 0; + for (int t : kTrustedSplitK) { + if (static_cast(t) * plane_bytes <= workspace_bytes) { + best_fit = t; + } + } + if (best_fit == 0) { + launch_scalar_gemm(a, b, x_scale, weight_scale, out, m, n, k, stream); + return; + } + if (split_k > best_fit) { + split_k = best_fit; + } + + int32_t* partials = static_cast(workspace); + + // Round-17 fused single-kernel path: the combine runs inside the + // partial kernel (last arrival per tile), so the operator is ONE launch + // and the round-11 combine kernel + its launch gap disappear. The + // per-tile arrival counters (num_tiles x int32) live in the last 160 B + // of the 16-plane contract workspace (plane 15's tail; planes 0..12*40-1 + // are the max used by any trusted split-K, so no overlap). Round-23: the + // counters are monotonic (the combiner no longer self-clears; the + // last-arrival test is modulo-SPLIT_K on the arrival count), so only the + // FIRST launch with a given workspace needs an async zero (guarded by a + // workspace-pointer static: the memset is a stream op on the caller's + // stream - no host sync - and it is NOT part of the captured graph, so + // every replay is just the one kernel). + if (workspace_bytes >= 16 * plane_bytes) { + int32_t* counters = reinterpret_cast( + static_cast(workspace) + 16 * plane_bytes) - + num_tiles; + static const void* s_fused_counters_ws = nullptr; + if (s_fused_counters_ws != workspace) { + hipMemsetAsync(counters, 0, + static_cast(num_tiles) * sizeof(int32_t), + stream); + s_fused_counters_ws = workspace; + } + launch_splitk_partial_n64_switch( + split_k, a, b, partials, x_scale, weight_scale, + static_cast(out), counters, m, n, k, stream, num_tiles); + return; + } + + // Fallback for workspaces smaller than the 16-plane contract: the + // round-11 two-kernel path (partial kernel without the fused tail, + // then the quadrant combine kernel), byte-identical to the accepted + // R14 behavior. The FUSED=false instantiation never dereferences the + // scale/output/counter pointers. + launch_splitk_partial_n64_switch( + split_k, a, b, partials, x_scale, weight_scale, + static_cast(out), static_cast(nullptr), m, n, k, + stream, num_tiles); + hipLaunchKernelGGL( + w8a8_dumma_m16_n64_combine_scale_kernel, + dim3(static_cast(num_tiles * 4)), + dim3(static_cast(kCombineThreads)), + 0, + stream, + partials, + x_scale, + weight_scale, + static_cast(out), + split_k, + num_tiles, + n); + return; + } + launch_scalar_gemm(a, b, x_scale, weight_scale, out, m, n, k, stream); +} + +// Optional out-of-timed-region weight packing. Bootstrap: identity +// device-to-device copy of the logical [K, N] int8 weight and the [N] fp32 +// scales, valid for every (K, N). Later rounds may replace this with a +// packed layout as long as launch_w8a8_gemm interprets it consistently. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + if (k <= 0 || n <= 0) { + return; + } + hipMemcpyAsync( + packed_weight, + raw_weight, + static_cast(k) * static_cast(n) * sizeof(int8_t), + hipMemcpyDeviceToDevice, + stream); + hipMemcpyAsync( + packed_weight_scale, + weight_scale, + static_cast(n) * sizeof(float), + hipMemcpyDeviceToDevice, + stream); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_down_proj.hip new file mode 100644 index 00000000..4a5f7a29 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_down_proj.hip @@ -0,0 +1,321 @@ +// @@variant shape=hy3_tp4_shared_down_proj_m16 commit=8f7319d888b0772736fa86ed3dc42afb6b56aa7b added=2026-08-26 +// median_us=9.174 p90_us=9.754 +// source=hy3-dsh-tp4-m16-2-368c654c +// INT8 W8A8 GEMM - worker_3, round 1 (DUMMA bootstrap). +// +// Assigned shape (worker_3, physical GPU 3): +// hy3_tp4_shared_down_proj_m16 : M=16, N=4096, K=384 +// +// Mandated architecture: minimal native gfx928 DUMMA INT8 m16n16k32 tile, +// ONE 64-lane wavefront per block, one 16x16 output tile per block +// (grid = N/16 = 256 blocks), explicit int32 accumulation, NO cross-wave +// barrier anywhere (single wave per block: no LDS, no __syncthreads). +// +// * A (x_q, logical [16][384] row-major) is loaded with +// du::dumma::du_load_matrix_sync(matrix_a, row_major): each lane's 8 +// fragment bytes A[row][k0 + i] (row = lane&15, k0 = (lane>>4)*8) are +// contiguous in global memory, so one 8-byte load per lane per step. +// * B (weight, logical [384][4096] row-major) is packed ONCE outside the +// timed region by launch_pack_w8a8_weight for the exact (k,n) == +// (384,4096) pair into the [N][K] n-major transpose +// P[n*384 + k] = W[k*4096 + n]. +// The B fragment is loaded with du::dumma::du_load_matrix_sync( +// matrix_b, col_major) over P: lane (row = n column, col = 8-k group) +// reads 8 contiguous bytes P[(n0+row)*384 + k0 + (lane>>4)*8 + i], i.e. +// one 8-byte load per lane per step instead of the strided byte-load +// expansion of a row-major [K][N] B fragment. +// * K=384 = 12 fully-unrolled m16n16k32 steps accumulate into one int32 +// accumulator fragment (du_fill_fragment + du_mma_sync). +// * Fused direct epilogue using the verified gfx928 int8 m16n16k32 +// accumulator ownership (lane%16 = row, lane/16 = col%4, x[i] = columns +// col%4 + 4*i): each lane scales its four int32 accumulators by +// (float(dot) * x_scale[row]) * weight_scale[col] in the same +// left-associative float32 order as the harness reference, then stores +// round-to-nearest-even bf16 via __float2bfloat16. +// +// Exact int32 accumulation order (k-ascending, no split-K) is preserved, so +// the bf16 output is expected bit-identical to the scalar bootstrap. +// +// Dispatch: +// * exact-shape guard (m == 16 && n == 4096 && k == 384) -> DUMMA arm; +// * every other (m, n, k) - including the paired M=2 API shape with the +// same (N, K) - reaches the generic scalar fallback, which decodes the +// n-major pack for (n, k) == (4096, 384) via a b_transposed flag and +// keeps the identity row-major read for every other pair. +// +// Header order is the known-good DTK order: HIP runtime, then bfloat16, +// then du_mma.h (du_mma.h is not self-contained before the HIP runtime +// headers). gfx928 wavefront is 64 lanes; every block size is a multiple +// of 64. + +#include +#include +#include + +#include + +namespace { + +// One block per 256 threads; 256 is a multiple of the gfx928 wavefront (64). +constexpr int kScalarBlockThreads = 256; + +// DUMMA INT8 m16n16k32 tile constants for the M=16 exact-shape arm. +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; +constexpr int kDummaBlockThreads = 64; // one wavefront per block +constexpr int kDummaK = 384; // exact-guard K; 12 m16n16k32 steps + +// Exact pack guard for the assigned (K, N) pair: (384, 4096). +constexpr int kPackK = 384; +constexpr int kPackN = 4096; + +// --------------------------------------------------------------------------- +// Generic scalar fallback: one thread per output element. +// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// x_q is [M, K] int8 row-major (stride K). +// weight is the packed [K, N] int8 buffer (stride N in the identity layout). +// b_transposed selects the exact-shape packed [N][K] layout +// (P[n*K + k] = W[k*N + n], produced by launch_pack_w8a8_weight for +// (k,n) == (384,4096)): the logical column col starts at P[col*K] with unit +// stride; the identity pack keeps the legacy row-major [K][N] layout with +// stride n. The branch is hoisted out of the K loop by the compiler. +// Consecutive threads own consecutive N columns, so adjacent lanes touch +// adjacent addresses in the fastest-changing N dimension. +__global__ __launch_bounds__(kScalarBlockThreads) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k, + int b_transposed) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + + const int8_t* a_row = x_q + static_cast(row) * k; + const int8_t* b_col = + weight + (b_transposed ? static_cast(col) * k : col); + const int64_t b_stride = b_transposed ? 1 : static_cast(n); + + // Exact int32 accumulation over the full K loop. The maximum assigned K + // keeps the exact int8 dot within int32 range (K * 127 * 127 << 2^31). + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * b_stride]); + } + + // Convert to float only for the two scales, then store bf16. + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Exact-shape M=16 DUMMA arm (mandated minimal geometry): blockDim = 64 = +// one 64-lane wavefront, one 16x16 output tile per block, grid = N/16 = 256 +// blocks. Direct global fragment loads (A row_major from logical x_q, B +// col_major from the exact-shape [N][K] n-major pack), explicit int32 +// accumulation over the full K=384 (12 unrolled m16n16k32 steps), fused +// direct scale+bf16 epilogue. No LDS and no barriers anywhere. Only launched +// for (m, n, k) == (16, 4096, 384). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDummaBlockThreads) void +w8a8_dumma_m16n16k32_1wave_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n) { + constexpr int kSteps = kDummaK / kDummaTileK; // 12 + + const int lane = static_cast(threadIdx.x); // 0 .. 63 + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // The col_major B loader adds (lane%16)*ldm internally, so the pointer is + // the tile base only: P[n0*384 + s*32 + (lane%16)*384 + (lane>>4)*8 + i], + // i.e. each lane's 8 fragment bytes are contiguous in P (8-byte aligned). + const int64_t b_tile_base = static_cast(n0) * kDummaK; + +#pragma unroll + for (int s = 0; s < kSteps; ++s) { + du::dumma::du_load_matrix_sync(a_frag, x_q + s * kDummaTileK, kDummaK); + du::dumma::du_load_matrix_sync( + b_frag, packed_b + b_tile_base + s * kDummaTileK, kDummaK); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // gfx928 int8 m16n16k32 accumulator ownership, established against + // du_store_matrix_sync: lane%16 selects the row, lane/16 selects col%4, + // and x[i] holds the columns col%4 + 4*i. + const int row = lane & 15; + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int out_col = n0 + col_mod4 + 4 * i; + const float scaled = static_cast(acc_frag.x[i]) * + x_scale[row] * weight_scale[out_col]; + out[static_cast(row) * n + out_col] = + __float2bfloat16(scaled); + } +} + +// Simple device-to-device copy used by the identity packing (every (K, N) +// pair that is not the exact-shape transpose pair, plus all scales). +template +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_identity_copy_kernel(const T* __restrict__ src, + T* __restrict__ dst, + int64_t numel) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < numel) { + dst[i] = src[i]; + } +} + +// Exact-shape pack: writes the [N][K] transpose +// dst[n*K + k] = src[k*N + n] +// for (k, n) == (384, 4096). One thread per 16-byte k-chunk of one column; +// runs once per weight outside the timed region/Graph, so the strided byte +// loads are acceptable. All other shapes keep the identity copy kernel. +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_pack_b_transpose_kernel(const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int k, + int n) { + const int kchunks = (k + 15) >> 4; + const int64_t total = static_cast(n) * kchunks; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t t = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + t < total; t += stride) { + const int col = static_cast(t / kchunks); + const int k0 = static_cast(t % kchunks) * 16; + alignas(16) int8_t chunk[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) { + chunk[j] = src[(static_cast(k0) + j) * n + col]; + } + *reinterpret_cast(dst + static_cast(col) * k + k0) = + *reinterpret_cast(chunk); + } +} + +} // namespace + +// Optional out-of-timed-region weight packing. For the exact (k, n) == +// (384, 4096) pair it produces the [N][K] n-major transpose +// P[n*384 + k] = W[k*4096 + n] that the M=16 DUMMA arm (col_major B +// fragments) and the scalar fallback (b_transposed flag) both decode; every +// other (K, N) keeps the identity layout, and all scales keep the identity +// copy. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_numel = static_cast(k) * n; + const dim3 block(kScalarBlockThreads); + const dim3 weight_grid(static_cast( + (weight_numel + kScalarBlockThreads - 1) / kScalarBlockThreads)); + if (k == kPackK && n == kPackN) { + // Exact-shape pack: [N][K] transpose P[n*384 + k] = W[k*4096 + n], so + // every B fragment is one contiguous 8-byte load per lane per step. + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_pack_b_transpose_kernel), + weight_grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + weight_grid, block, 0, stream, + raw_weight, packed_weight, weight_numel); + } + + const dim3 scale_grid(static_cast( + (n + kScalarBlockThreads - 1) / kScalarBlockThreads)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + scale_grid, block, 0, stream, + weight_scale, packed_weight_scale, n); +} + +// Timed GEMM operator: no allocation, no compilation, no autotuning, no +// weight packing, no host/device synchronization, no default-stream launch. +// Uses only the caller-provided out and workspace, on the caller's stream. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // The one-wave DUMMA arm performs no split-K, so the workspace is unused. + (void)workspace; + (void)workspace_bytes; + + // The exact (k, n) == (384, 4096) pack is the [N][K] n-major transpose; + // any other (k, n) keeps the identity pack (row-major [K][N]). + const int b_transposed = (k == kPackK && n == kPackN) ? 1 : 0; + + // Exact-shape M=16 DUMMA specialization: one 64-lane wavefront per 16x16 + // output tile -> grid = N/16 = 256 blocks, full K=384 reduction, no + // barriers. Guarded by the full (m, n, k) triple so the paired M=2 API + // shape with the same (N, K) still reaches the generic scalar fallback. + if (m == 16 && n == kPackN && k == kPackK) { + const dim3 grid(static_cast(n / kDummaTileN)); + const dim3 block(kDummaBlockThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16n16k32_1wave_kernel), + grid, block, 0, stream, + a, b, x_scale, weight_scale, + reinterpret_cast<__hip_bfloat16*>(out), n); + return; + } + + // Generic scalar fallback for every other (m, n, k), including the paired + // M=2 shape with the same (N, K): decodes the n-major pack when + // b_transposed, row-major otherwise. + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads)); + const dim3 block(kScalarBlockThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_scalar_gemm_kernel), + grid, block, 0, stream, + a, b, x_scale, weight_scale, + reinterpret_cast<__hip_bfloat16*>(out), m, n, k, b_transposed); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_gate_up_proj.hip new file mode 100644 index 00000000..7d26ecac --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_gate_up_proj.hip @@ -0,0 +1,686 @@ +// @@variant shape=hy3_tp4_shared_gate_up_proj_m16 commit=f52225ed98ffd619fa1318d3eca7e56addd32f82 added=2026-08-26 +// median_us=11.31 p90_us=11.32 +// source=hy3-dsh-tp4-m16-1-7f1fb1d1 +// MetaInfer W8A8 INT8 GEMM for gfx928 (K500SM_AI). +// +// worker_2 — physical GPU 2 +// assigned shape: hy3_tp4_shared_gate_up_proj_m16 (M=16, N=768, K=4096) +// +// Iteration 1 = establish the minimal 16x16x32 DUMMA tile: +// * one 64-lane wavefront per block, one 16x16 output tile per wave, +// explicit int32 accumulation (du_fill_fragment / du_load_matrix_sync / +// du_mma_sync / du_store_matrix_sync from ); +// * A/B fragments are loaded straight from global memory every K step, so +// there is no LDS staging, no double buffering, and no cross-wave +// barrier; the accumulator fragment is drained through LDS in a +// coalesced epilogue with the two float scales and a bf16 store; +// * a scalar generic fallback stays reachable for every unmatched (m,n,k) +// (including the paired M=2 API shape with the same (N,K)); the M=16 +// DUMMA path is guarded by the exact (m,n,k) triple. +// +// Iteration 2 = architecture round on grid parallelism. N=768 limits the +// independent 16x16 output tiles to 48, so no {1,2,4} waves x {1,2,4} adjacent +// N-tile geometry can cover all 120 CUs (max grid_blocks=48, max active CUs +// with one tile per wave = 48). Within that space this round packs 2 +// wavefronts per block, each owning 2 adjacent 16-wide N tiles (a 16x32 +// column span) that share one A fragment: grid = 768/64 = 12 blocks, 128 +// threads, 2 waves per CU on 2 SIMDs, 2 independent accumulator chains per +// SIMD. Both latency-hiding mechanisms are engaged without any split-K: +// * hardware: two co-resident wavefronts on the same CU issue their own +// loads/MMACs while the sibling wave stalls on its fragment loads; +// * compiler: two independent per-wave accumulator chains plus the shared +// A fragment (1 A load + 2 B loads + 2 MMACs per K step) give the +// scheduler real ILP, unlike iteration 1 where a single dependent chain +// degraded to serialized global_load_ubyte -> s_waitcnt vmcnt chain -> +// one v_mmac per step. +// A/B fragments stay direct-from-global (no LDS staging, no cross-wave +// barrier in the K loop); only the accumulator drain reuses the iteration-1 +// LDS epilogue (4 KiB/block, one block-wide __syncthreads), so per-element +// int32 accumulation order (k0 ascending) is bit-identical to iteration 1. +// Measured result: median 215.7us (24 waves total on 12 CUs) vs 250.2us for +// grid=48 one-wave blocks; per-step cost stays ~1.7us/wave at K=4096 (128 +// DUMMA steps), i.e. the kernel remains latency-bound with only 24 waves in +// flight on a 120-CU device, far from the 73.3us Triton Graph baseline. +// +// Iteration 3 = architecture round: split-K with combine. The unsplit grid +// (12 blocks) is below the two-blocks-per-CU latency-hiding target, so the +// K=4096 loop (128 DUMMA steps) is partitioned across S runtime-selectable +// splits: +// * grid = 12*S blocks; blockIdx.x encodes (split = x/12, N-group = x%12); +// every split's K range is an exact multiple of kDummaK via step-based +// boundaries [(s*128)/S, ((s+1)*128)/S), so each block keeps the exact +// iteration-2 inner loop (2 wavefronts, shared A fragment, 2 independent +// accumulator chains, no cross-wave barrier in the K loop) over a shorter +// K slice; +// * each block drains its four int32 accumulator tiles through LDS into +// caller-workspace plane `split` (row-major [16,n] int32, plane stride +// 16*n); accumulation stays k0 ascending inside the slice and summing the +// S planes in split order reproduces the full-K int32 dot exactly (int32 +// addition is exact and |full dot| <= 4096*127*127 ~= 6.6e7 << 2^31), so +// results stay bit-identical to iterations 1-2; +// * a tiny combine kernel (96 blocks x 128 threads) then sums the S planes, +// applies x_scale*weight_scale and stores bf16. Both kernels launch on +// the caller stream inside the timed Graph, so the combine cost is part +// of the measured median/P90. +// Split count is selected at launch by METAINFER_W8A8_SPLIT_K (default 16 as +// of iteration 4, legal 1..16, clamped to workspace planes and K/32 stage +// count). Trusted occupancy-probe candidates [2,3,5,7,8,10,15] give grids +// 24/36/60/84/96/120/180: S=10 is CU-aligned (120 blocks, exactly one per +// CU), S=5 lands on 60 CUs, S=15 oversubscribes to 180. +// +// Iteration 4 = architecture/pipeline round: CU-filling default grid plus +// bounded LDS prefetch. The control plane flags the S=2 default (24 blocks +// on 120 CUs) as far below the two-blocks-per-CU latency-hiding target, and +// the per-step window is load-latency bound (28 global_load_ubyte -> waitcnt +// chain -> 2 v_mmac per step; ~1.7us per step at 2 waves/CU). This round +// * moves the default split to 16 -> grid 192 blocks x 128 threads = 384 +// waves (1.6 blocks/CU, all 120 CUs covered; combine cost of summing 16 +// int32 planes is inside the timed Graph); +// * stages each block's whole K slice in LDS once (A 16x256 = 4 KiB, B +// 256x64 = 16 KiB, 16 B/thread vector loads), then the 8-step K loop +// reads fragments from LDS only (du_load_matrix_sync on __shared__ +// pointers emits ds_read): zero global loads and zero barriers in the +// loop. Per block LDS = 4 KiB drain + 20 KiB staged = 24 KiB, so two +// blocks per CU still fit. The staged copy keeps the raw row-major +// layout with the same ldm semantics, so every fragment holds the exact +// iteration-3 bytes and int32 accumulation order (k0 ascending inside +// the split, splits summed in order by combine) is bit-identical. +// * A bytes are reused 2x per K step (one fragment feeds the wave's two +// 16-wide N-tiles); B bytes are consumed once per step from the staged +// copy. Splits S < 16 keep the iteration-3 direct-global loop untouched. +// Falsifiable: if staging plus 384 waves does not cut the 112.6us median +// toward the ~20-40us band (and below the 73.3us Triton baseline), the +// binding cost is not per-step global latency and the next round must attack +// the staging prologue / combine / launch overheads instead. +// Measured result (accepted): median 15.65us / 206.96 GB/s / 6.43 TOPS at +// S=16. The exact-source ISA shows the remaining binding cost in the LDS-only +// K loop: per step per wavefront the row-major B fragments degrade to 16 +// ds_read_u8 (8 per fragment) at stride-64 addresses (offsets 0,64,128,...) +// with per-read lgkmcnt waits, and PMC reports 328,704 LDS bank conflicts on +// 60,672 LDS instructions (~5.4 conflicts/instr) plus 88,787 LDS waits - the +// same B-fragment-aliasing / byte-reassembly signature that the gate_up +// M=4096 lineage on this worker fixed with an [N,K] pack + col_major +// fragments + an 8-consecutive-byte fragment loader (one ds_read2_b64 per +// fragment; VALU 16.08M -> 6.25M, accepted iter 13). +// +// Iteration 5 = packed-weight round. One focused layout change plus its +// minimal load-path rewiring: +// * launch_pack_w8a8_weight transposes [K,N] -> [N,K] for the exact gate_up +// pair (k,n)==(4096,768) (one-time, out of the timed region and out of +// Graph capture; every other (K,N) keeps the identity copy); +// * the B fragment is declared col_major: du_mma.h's matrix_b col_major +// loader maps lane l to 8 CONSECUTIVE bytes (x[i] = p[(l&15)*ldm + +// ((l>>4)<<3) + i]), so n-major storage makes each lane's fragment bytes +// contiguous - the staged K loop replaces the 16 ds_read_u8 per step per +// wavefront with one 8-byte load per fragment (ds_read_b64); the direct +// S<16 path reads the packed global layout through the library col_major +// loader; +// * staged slices use a bank-skewed row stride 272 = 256+16 (16 B aligned +// so the ds_write_b128 staging stays aligned; 272 % 8 == 0 so every +// fragment b64 read is 8 B aligned), cutting the stride-256 8-way bank +// aliasing to ~2-way on the remaining reads; +// * the scalar generic fallback decodes the packed [N,K] layout whenever +// (k,n)==(4096,768) (this covers the paired M=2 validation shape with the +// same (N,K)); every other shape keeps the raw [K,N] decode. +// Per-element int32 accumulation order is unchanged (k0 ascending inside each +// split, splits summed in split order by the combine kernel), and the +// col_major loader fills the same fragment slots with the same logical (k,n) +// bytes, so results stay bit-identical to iteration 4. +// Expected: LDS instructions in the K loop drop 17 -> 3 per step per +// wavefront (A: 1 ds_read2_b32, B: 2 ds_read_b64), LDS waits and conflicts +// drop roughly 4-6x, so median should move from 15.65us toward ~10-13us. +// Falsifiable: if the median does not land below ~14us, LDS issue/latency is +// not the binding cost and the next round must attack the staging prologue / +// combine / launch overhead or the grid (finer one-wave zero-barrier grid or +// split-K probes [2,3,5,7,8,10,15], which now measure the packed direct +// path plus combine cost via METAINFER_W8A8_SPLIT_K). +// +// launch_pack_w8a8_weight is an out-of-timed-region pack op: an [N,K] +// transpose for the exact (k,n)==(4096,768) gate_up weight, identity +// device-to-device copy for every other (K,N). +// +// Iteration 8 = combine-round: specialize the S=16 combine for the default +// split count. The accepted S=16 operator (median 13.34us / 242.85 GB/s / +// 7.55 TOPS) spends ~22-25% of its time in the combine kernel (3.36us of the +// 15.04us profiled operator aggregate; GEMM kernel 11.68us). The exact-source +// ISA of the combine shows why: each thread's 16 plane loads (one scalar +// 4-B global_load_dword per plane, 3,456 vmem_read instructions per replay) +// are issued in small groups of ~4 with the per-plane address chain +// (v_mad_u64_u32 reusing the previous address register) and vmcnt waits +// serializing them, so a thread has only ~4-6 HBM round trips in flight. +// With 96 blocks x 128 threads = 192 waves and ~1.6 waves/CU the combine is +// latency-bound on those serialized round trips. This round makes ONE focused +// change: when split_k == 16 (the compile-time-known default), the combine +// loads all 16 planes into a register array with a fully unrolled loop BEFORE +// summing them, so the compiler schedules 16 independent loads in flight and +// the ascending int32 sum then consumes them. Per-element accumulation order +// (planes s=0..15 ascending, exactly the generic loop's order) is unchanged, +// so results stay bit-identical to iteration 5; the generic loop and every +// other split count are untouched. Falsifiable: if the median does not move +// by the combine's ~2-2.5us share (13.34us -> ~11-12us), the combine was not +// latency-bound on load MLP and iteration 9 must attack the GEMM kernel's +// remaining LDS conflicts / A-fragment repack VALU in the staged K loop. +// +// Include order is load-bearing on this DTK: HIP runtime first, then +// hip_bfloat16.h, then du_mma.h (not self-contained before the HIP headers). + +#include +#include +#include + +#include +#include + +namespace { + +// gfx928 wavefront is 64 lanes; blockDim must stay a multiple of 64. +constexpr int kScalarBlockThreads = 128; +constexpr int kCopyBlockThreads = 256; + +// Exact (k,n) pair that launch_pack_w8a8_weight transposes to [N,K] and that +// the scalar fallback decodes as n-major. Declared before the scalar kernel +// below (which references them for the exact gate_up pair). +constexpr int kPackedK = 4096; +constexpr int kPackedN = 768; + +// Scalar INT8 dot-product GEMM, one thread per output element. +// out[m, n] = bf16( (sum_k a[m,k] * b[k,n]) * x_scale[m] * weight_scale[n] ) +// Accumulation order is exactly k0..K-1 in int32 (identical for every output +// element), so results are bit-exact against the float/int32 reference. +__global__ void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (idx >= total) { + return; + } + // Adjacent lanes map to adjacent columns in the fastest-changing N + // dimension, so output stores (and A rows) stay coalesced. + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + + const int8_t* __restrict__ a_row = a + static_cast(row) * k; + // Iteration 5: the exact gate_up pair (k,n)==(4096,768) is packed [N,K] + // (n-major) by launch_pack_w8a8_weight; the fallback decodes that layout so + // it stays correct against the packed buffer (paired M=2 validation shape + // included). Every other (K,N) keeps the raw [K,N] row-major decode. + const bool packed = (k == kPackedK && n == kPackedN); + const int8_t* __restrict__ b_col = + b + (packed ? static_cast(col) * k : col); + const int64_t b_stride = packed ? 1 : static_cast(n); + + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * b_stride]); + } + + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Minimal gfx928 DUMMA INT8 path (assigned shape M=16, N=768, K=4096). +// Tile: 16x16x32, int32 accumulation, one 64-lane wavefront per block. +// --------------------------------------------------------------------------- + +// gfx928 DUMMA INT8 m16n16k32 tile constants. +constexpr int kDummaM = 16; +constexpr int kDummaN = 16; +constexpr int kDummaK = 32; +// Iteration 2 launch geometry: 2 wavefronts per block, each owning 2 +// adjacent 16-wide N tiles (a 16x32 column span) that share one A fragment. +// Block tile = 16 x 64 columns, grid = 768/64 = 12 blocks, 128 threads. +constexpr int kDummaWavesPerBlock = 2; +constexpr int kDummaTilesPerWave = 2; +constexpr int kDummaBlockThreads = kDummaWavesPerBlock * 64; // 128 +constexpr int kDummaWaveN = kDummaTilesPerWave * kDummaN; // 32 cols/wave +constexpr int kDummaBlockN = kDummaWavesPerBlock * kDummaWaveN; // 64 cols/block + +// Iteration 3 split-K: runtime-selectable split count (1..16, matching the +// caller workspace contract WORKSPACE_MAX_SPLIT_K=16), default 16 as of +// iteration 4 (CU-filling grid + LDS-staged slice, see below). The +// combine+scale kernel uses 128 threads (multiple of wavefront 64). +constexpr int kDummaSplitDefault = 16; +constexpr int kDummaSplitMax = 16; +constexpr int kDummaCombineThreads = 128; + +// Iteration 4 bounded LDS prefetch. At split_k == kDummaSplitMax every split's +// K slice is exactly 128/16 = 8 DUMMA steps = 256 rows, so the kernel stages +// the whole slice into LDS once per block (A 16x256 = 4 KiB, B 256x64 = +// 16 KiB) and the K loop then reads its fragments from LDS: no global loads +// and no barrier inside the loop (zero-barrier K loop; one __syncthreads +// after the staging copy, one in the drain epilogue). LDS per block = 4 KiB +// accumulator drain + 20 KiB staged = 24 KiB, keeping two blocks per CU +// feasible (49 KiB <= 64 KiB per CU) so the 192-block S=16 grid still covers +// all 120 CUs. Slices larger than this (S < 16) keep the iteration-3 +// direct-global path. +constexpr int kStagedSliceRows = 256; + +// Iteration 5 packed-weight layout. The gate_up weight (k,n)==(4096,768) is +// packed once (out of the timed region) as [N,K] n-major; B fragments are +// declared col_major so each lane's 8 fragment bytes are consecutive in +// memory and load with one ds_read_b64 instead of 8 ds_read_u8. Staged rows +// use a bank-skewed stride 272 = 256+16: 16 B aligned (ds_write_b128 staging +// stays aligned, 272 = 17*16), 8 B aligned (every fragment b64 read stays +// aligned), and not a multiple of 128 B, so rows no longer alias onto one LDS +// bank phase (stride 256 put all 16 rows on the same phase -> 8-way fragment +// conflicts). +constexpr int kStagedRowStride = 272; + +// Iteration 5 packed B-fragment loader. With n-major [N,K] storage and a +// col_major fragment, du_mma.h's matrix_b col_major address rule is +// x[i] = p[(lane&15)*ldm + ((lane>>4)<<3) + i] - lane l's 8 fragment bytes +// are consecutive in memory, so one 8-byte LDS read per fragment replaces the +// 8 ds_read_u8 at stride-64 addresses that the row-major loader emits (and +// du_mma_sync consumes the fragment as one 64-bit value per lane, so writing +// the packed bytes straight into x[0..7] is exactly the library's own +// storage convention). ldm is the n-major row stride in bytes; the pointer +// must be 8 B aligned (guaranteed: staged base is __align__(256), ldm % 8 == +// 0, tile offsets are multiples of 16, local_k is a multiple of 32). +__device__ __forceinline__ void load_b_frag_packed( + du::dumma::DUFragment& b_frag, + const signed char* p, int ldm, int lane) { + const int row = lane & 15; + const int col = (lane >> 4) << 3; + const uint64_t v = + *reinterpret_cast(p + static_cast(row) * ldm + col); + *reinterpret_cast(b_frag.x) = v; +} + +// M=16 W8A8 GEMM with the minimal 16x16x32 DUMMA tile and runtime split-K. +// Block layout is the iteration-2 geometry (12 N-groups of 64 columns, 2 +// wavefronts per block, 2 adjacent tiles per wave sharing one A fragment); +// blockIdx.x encodes (split, n_group) = (blockIdx.x / 12, blockIdx.x % 12). +// Each block computes its split's slice of the K loop (kDummaK-aligned) and +// drains the raw int32 accumulator tiles into workspace plane `split`; the +// combine kernel later sums the planes and applies the two float scales. +__global__ __launch_bounds__(kDummaBlockThreads) +void w8a8_gemm_m16_dumma_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int32_t* __restrict__ partials, + int n, + int k, + int split_k) { + const int wave = static_cast(threadIdx.x) >> 6; + const int lane = static_cast(threadIdx.x) & 63; + const int n_groups = n / kDummaBlockN; + const int split = static_cast(blockIdx.x) / n_groups; + const int block_n_base = + (static_cast(blockIdx.x) % n_groups) * kDummaBlockN; + const int n_base = block_n_base + wave * kDummaWaveN; + + // Split s owns DUMMA steps [(s*total_steps)/S, ((s+1)*total_steps)/S), so + // every split's K slice is an exact multiple of kDummaK for any legal S. + const int total_steps = k / kDummaK; + const int k_lo = ((split * total_steps) / split_k) * kDummaK; + const int k_hi = (((split + 1) * total_steps) / split_k) * kDummaK; + const int slice_rows = k_hi - k_lo; // bytes of K owned by this block + + __shared__ __align__(16) + int32_t acc_tile[kDummaWavesPerBlock][kDummaTilesPerWave] + [kDummaM * kDummaN]; + // Iteration 4/5 staged slices: A is [16, slice] row-major (reused by both + // N-tiles of each wavefront across every K step) and B is [64, slice] + // n-major (the block's 64-column span of the packed [N,K] weight). Both use + // the iteration-5 bank-skewed row stride 272 so fragment reads from LDS + // stay 8 B aligned and rows do not alias onto one bank phase. + __shared__ __align__(256) signed char lds_a[kDummaM * kStagedRowStride]; + __shared__ __align__(256) signed char lds_b[kDummaBlockN * kStagedRowStride]; + + du::dumma::DUFragment + a_frag; + // Iteration 5: col_major B fragments match the packed [N,K] storage so each + // lane's 8 fragment bytes are consecutive (one ds_read_b64 per fragment in + // the staged loop instead of 8 ds_read_u8). + du::dumma::DUFragment + b_frag[kDummaTilesPerWave]; + du::dumma::DUFragment + acc_frag[kDummaTilesPerWave]; + +#pragma unroll + for (int t = 0; t < kDummaTilesPerWave; ++t) { + du::dumma::du_fill_fragment(acc_frag[t], 0); + } + + if (slice_rows == kStagedSliceRows) { + // ---- Iteration 4 staged path (split_k == 16). Cooperative global->LDS + // prefetch of the whole split slice: 128 threads issue wide 16 B vector + // loads (A: 2 int4/thread, B: 8 int4/thread), one block-wide barrier, + // then the K loop is LDS-only: du_load_matrix_sync emits ds_read instead + // of global_load, so no HBM latency sits on the MMAC critical path and + // the loop is barrier-free. + const int4* __restrict__ a4 = reinterpret_cast(a); + const int4* __restrict__ b4 = reinterpret_cast(b); + int4* __restrict__ lds_a4 = reinterpret_cast(lds_a); + int4* __restrict__ lds_b4 = reinterpret_cast(lds_b); + const int tid = static_cast(threadIdx.x); + // A slice: 16 rows x kStagedSliceRows bytes = 256 int4s; each thread + // stages two int4s. All loads issue before any LDS store so the HBM + // round trips pipeline instead of serializing behind per-load waitcnts. + // int4 = 16 bytes, so flat int4 indices divide byte offsets by 16 + // (>>4): row stride k>>4 = 256 int4s (K=4096 bytes/row), slice column + // offset k_lo>>4, and one int4 per a_c (16 int4s per 256 B row). The LDS + // row stride is kStagedRowStride>>4 = 17 int4s (bank-skewed padding). + const int a_r = tid >> 4; // A row of the first int4 (0..7) + const int a_c = tid & 15; // 16 int4s per 256 B row + int4 a0 = a4[a_r * (k >> 4) + (k_lo >> 4) + a_c]; + int4 a1 = a4[(a_r + 8) * (k >> 4) + (k_lo >> 4) + a_c]; + // B slice (iteration 5): the packed [N,K] weight is n-major, so the + // block's 64 columns are 64 n-rows of 256 slice bytes = 1024 int4s; each + // thread stages eight. b4 row stride = k>>4 int4s (K=4096 bytes per + // n-row), slice column offset k_lo>>4, LDS row stride 17 int4s. + const int b_nr = tid >> 4; // first n-row of this thread (0..7) + const int b_k16 = tid & 15; // 16 int4s per 256 B slice row + int4 b0 = b4[(static_cast(block_n_base) + b_nr + 0) * (k >> 4) + + (k_lo >> 4) + b_k16]; + int4 b1 = b4[(static_cast(block_n_base) + b_nr + 8) * (k >> 4) + + (k_lo >> 4) + b_k16]; + int4 b2 = b4[(static_cast(block_n_base) + b_nr + 16) * (k >> 4) + + (k_lo >> 4) + b_k16]; + int4 b3 = b4[(static_cast(block_n_base) + b_nr + 24) * (k >> 4) + + (k_lo >> 4) + b_k16]; + int4 b4v = b4[(static_cast(block_n_base) + b_nr + 32) * (k >> 4) + + (k_lo >> 4) + b_k16]; + int4 b5 = b4[(static_cast(block_n_base) + b_nr + 40) * (k >> 4) + + (k_lo >> 4) + b_k16]; + int4 b6 = b4[(static_cast(block_n_base) + b_nr + 48) * (k >> 4) + + (k_lo >> 4) + b_k16]; + int4 b7 = b4[(static_cast(block_n_base) + b_nr + 56) * (k >> 4) + + (k_lo >> 4) + b_k16]; + lds_a4[a_r * (kStagedRowStride >> 4) + a_c] = a0; + lds_a4[(a_r + 8) * (kStagedRowStride >> 4) + a_c] = a1; + lds_b4[(b_nr + 0) * (kStagedRowStride >> 4) + b_k16] = b0; + lds_b4[(b_nr + 8) * (kStagedRowStride >> 4) + b_k16] = b1; + lds_b4[(b_nr + 16) * (kStagedRowStride >> 4) + b_k16] = b2; + lds_b4[(b_nr + 24) * (kStagedRowStride >> 4) + b_k16] = b3; + lds_b4[(b_nr + 32) * (kStagedRowStride >> 4) + b_k16] = b4v; + lds_b4[(b_nr + 40) * (kStagedRowStride >> 4) + b_k16] = b5; + lds_b4[(b_nr + 48) * (kStagedRowStride >> 4) + b_k16] = b6; + lds_b4[(b_nr + 56) * (kStagedRowStride >> 4) + b_k16] = b7; + __syncthreads(); + + // Fragment loads now come from LDS; ldm is the staged n-major row stride. + // A bytes are reused 2x per step (one fragment feeds both N-tiles); B + // bytes are consumed once per step from the staged copy, each B fragment + // as one 8-byte ds_read_b64 via load_b_frag_packed (col_major address + // rule). The trip count is compile-time (slice_rows == kStagedSliceRows), + // so the scheduler can overlap the ds_read streams of consecutive steps. +#pragma unroll 2 + for (int step = 0; step < kStagedSliceRows / kDummaK; ++step) { + const int local_k = step * kDummaK; + du::dumma::du_load_matrix_sync(a_frag, lds_a + local_k, + kStagedRowStride); +#pragma unroll + for (int t = 0; t < kDummaTilesPerWave; ++t) { + load_b_frag_packed( + b_frag[t], + lds_b + (wave * kDummaWaveN + t * kDummaN) * kStagedRowStride + + local_k, + kStagedRowStride, lane); + } +#pragma unroll + for (int t = 0; t < kDummaTilesPerWave; ++t) { + du::dumma::du_mma_sync(acc_frag[t], a_frag, b_frag[t], acc_frag[t]); + } + } + } else { + // ---- Iteration 3 direct-global path (any other legal split count). One + // A fragment serves both adjacent tiles of this wavefront. +#pragma unroll 2 + for (int k0 = k_lo; k0 < k_hi; k0 += kDummaK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); +#pragma unroll + for (int t = 0; t < kDummaTilesPerWave; ++t) { + // Iteration 5: b is the packed [N,K] n-major weight for this shape, + // so the tile's 16 columns live at (n_base + t*16)..+15 n-rows with + // row stride k; the col_major loader addresses p[(l&15)*ldm + + // ((l>>4)<<3) + i], i.e. n-major rows of k bytes. + du::dumma::du_load_matrix_sync( + b_frag[t], + b + (static_cast(n_base) + t * kDummaN) * k + k0, k); + } +#pragma unroll + for (int t = 0; t < kDummaTilesPerWave; ++t) { + du::dumma::du_mma_sync(acc_frag[t], a_frag, b_frag[t], acc_frag[t]); + } + } + } + +#pragma unroll + for (int t = 0; t < kDummaTilesPerWave; ++t) { + du::dumma::du_store_matrix_sync(acc_tile[wave][t], acc_frag[t], kDummaN, + du::dumma::mem_row_major); + } + __syncthreads(); + + // Drain raw int32 partials into workspace plane `split` (row-major + // [16, n]); scaling and bf16 conversion happen in the combine kernel. + int32_t* __restrict__ plane = + partials + static_cast(split) * (kDummaM * n); +#pragma unroll + for (int t = 0; t < kDummaTilesPerWave; ++t) { +#pragma unroll + for (int linear = lane; linear < kDummaM * kDummaN; linear += 64) { + const int row = linear >> 4; + const int col = linear & 15; + plane[static_cast(row) * n + n_base + t * kDummaN + col] = + acc_tile[wave][t][linear]; + } + } +} + +// Combine + scale: out[m,n] = bf16( (sum_s partials[s][m,n]) * x_scale[m] * +// weight_scale[n] ). The S planes are summed in ascending split order; int32 +// addition is exact and overflow-free (|full-K dot| <= 6.6e7 << 2^31), so the +// result is bit-identical to the single-chain k0-ascending accumulation of +// iterations 1-2. Runs on the caller stream right after the GEMM kernel, so +// its cost is inside the timed Graph. +__global__ void w8a8_gemm_m16_combine_kernel( + const int32_t* __restrict__ partials, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n, + int split_k) { + const int total = kDummaM * n; + const int idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int row = idx / n; + const int col = idx - row * n; + int32_t acc; + if (split_k == kDummaSplitMax) { + // Iteration 8 S=16 specialization: all 16 plane loads are independent, so + // load them into a register array with a fully unrolled loop FIRST and + // sum afterwards. The generic loop's exact-source ISA grouped ~4 scalar + // 4-B loads per HBM round trip behind a serialized address chain (v_mad + // reusing the previous address + vmcnt waits), leaving ~4-6 loads in + // flight per thread; the unrolled array version puts 16 loads in flight + // and collapses the combine to ~one memory round trip per wave. The sum + // is taken in the same ascending s order as the generic loop, so every + // element's int32 accumulation is bit-identical. + int32_t v[kDummaSplitMax]; +#pragma unroll + for (int s = 0; s < kDummaSplitMax; ++s) { + v[s] = partials[static_cast(s) * total + idx]; + } + acc = 0; +#pragma unroll + for (int s = 0; s < kDummaSplitMax; ++s) { + acc += v[s]; + } + } else { + // Generic path: any other legal split count (runtime loop bound). + acc = 0; + for (int s = 0; s < split_k; ++s) { + acc += partials[static_cast(s) * total + idx]; + } + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); +} + +// Identity device-to-device copies used by the bootstrap pack_weight op. +__global__ void w8a8_identity_copy_i8_kernel( + const int8_t* __restrict__ src, int8_t* __restrict__ dst, int64_t numel) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < numel) { + dst[idx] = src[idx]; + } +} + +__global__ void w8a8_identity_copy_f32_kernel( + const float* __restrict__ src, float* __restrict__ dst, int64_t numel) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < numel) { + dst[idx] = src[idx]; + } +} + +// Iteration 5 one-time pack for the exact gate_up pair (k,n)==(4096,768): +// transpose the raw [K,N] int8 weight to [N,K] n-major (element (kk,nn) -> +// dst[nn*k + kk]). Runs out of the timed region and out of Graph capture; +// performance is irrelevant (single weight prep per layer load). +__global__ void w8a8_transpose_i8_kernel( + const int8_t* __restrict__ src, int8_t* __restrict__ dst, int k, int n) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(k) * n; + if (idx >= total) { + return; + } + const int kk = static_cast(idx / n); + const int nn = static_cast(idx - static_cast(kk) * n); + dst[static_cast(nn) * k + kk] = src[idx]; +} + +} // namespace + +// Timed-region GEMM entry point (called from gemm_out on PyTorch's current +// HIP stream). Must stay allocation/sync/compile-free and launch only on the +// caller-provided stream. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // Shape-exact M=16 DUMMA specialization. Guarded by the exact (m,n,k) + // triple so the paired M=2 API shape with the same (N,K) and every other + // shape still reach the scalar fallback below. + if (m == 16 && n == 768 && k == 4096) { + // Iteration 3/4/5: runtime split-K selected by METAINFER_W8A8_SPLIT_K + // (default 16, legal 1..16; clamped to workspace planes and K/32 stage + // count). grid = 12*S blocks x 128 threads, block tile 16x64. b is the + // [N,K] n-major packed weight (launch_pack_w8a8_weight for (4096,768)); + // at S=16 the GEMM kernel stages its 256-row K slice into LDS and runs an + // LDS-only K loop with col_major B fragments (one ds_read_b64 per + // fragment); at S<16 it keeps the direct-global loop on the packed + // layout. Each block writes raw int32 partials to workspace plane + // `split`; the combine kernel then sums the S planes, scales and stores + // bf16. Both kernels launch on the caller stream inside the timed Graph. + int split_k = kDummaSplitDefault; + if (const char* env = getenv("METAINFER_W8A8_SPLIT_K")) { + char* end = nullptr; + const long parsed = strtol(env, &end, 10); + if (end != env && parsed >= 1 && parsed <= kDummaSplitMax) { + split_k = static_cast(parsed); + } + } + const int64_t plane_bytes = static_cast(16) * n * 4; + const int max_by_ws = static_cast(workspace_bytes / plane_bytes); + const int max_by_k = k / kDummaK; + if (split_k > max_by_ws) { + split_k = max_by_ws; + } + if (split_k > max_by_k) { + split_k = max_by_k; + } + if (split_k >= 1) { + int32_t* partials = static_cast(workspace); + const int n_groups = n / kDummaBlockN; // 12 + const dim3 grid(static_cast(n_groups * split_k)); + const dim3 block(kDummaBlockThreads); + hipLaunchKernelGGL(w8a8_gemm_m16_dumma_kernel, grid, block, 0, stream, + a, b, partials, n, k, split_k); + const int total_out = m * n; // 12288 + const dim3 combine_grid( + (total_out + kDummaCombineThreads - 1) / kDummaCombineThreads); + const dim3 combine_block(kDummaCombineThreads); + hipLaunchKernelGGL(w8a8_gemm_m16_combine_kernel, combine_grid, + combine_block, 0, stream, partials, x_scale, + weight_scale, static_cast(out), n, + split_k); + return; + } + // No workspace plane available: fall through to the scalar fallback. + } + + // Scalar generic fallback: valid for every (m, n, k). + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads)); + const dim3 block(kScalarBlockThreads); + hipLaunchKernelGGL(w8a8_gemm_scalar_kernel, grid, block, 0, stream, a, b, + x_scale, weight_scale, static_cast(out), m, + n, k); +} + +// Out-of-timed-region weight packing. Iteration 5: for the exact gate_up +// pair (k,n)==(4096,768) the raw [K,N] int8 weight is transposed to an +// [N,K] n-major packed layout so DUMMA col_major B fragments read 8 +// consecutive bytes per lane (one ds_read_b64 instead of 8 ds_read_u8); every +// other (K,N) keeps the identity device-to-device copy and the raw [K,N] +// decode. The [N] fp32 scales are always copied as-is. Runs outside the timed +// region and outside Graph capture. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_numel = static_cast(k) * n; + const int64_t scale_numel = static_cast(n); + if (weight_numel > 0) { + const dim3 grid(static_cast( + (weight_numel + kCopyBlockThreads - 1) / kCopyBlockThreads)); + const dim3 block(kCopyBlockThreads); + if (k == kPackedK && n == kPackedN) { + hipLaunchKernelGGL(w8a8_transpose_i8_kernel, grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else { + hipLaunchKernelGGL(w8a8_identity_copy_i8_kernel, grid, block, 0, stream, + raw_weight, packed_weight, weight_numel); + } + } + if (scale_numel > 0) { + const dim3 grid(static_cast( + (scale_numel + kCopyBlockThreads - 1) / kCopyBlockThreads)); + const dim3 block(kCopyBlockThreads); + hipLaunchKernelGGL(w8a8_identity_copy_f32_kernel, grid, block, 0, stream, + weight_scale, packed_weight_scale, scale_numel); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip new file mode 100644 index 00000000..94a1fee0 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip @@ -0,0 +1,480 @@ +// @@variant shape=hy3_tp4_o_proj_m4096 commit=7c36244e756baf1e8268807cb6481eb118a30133 added=2026-08-23 +// median_us=803.9 p90_us=805.1 tops=85.48 bandwidth_gb_s=62.65 speedup=24.73 baseline_us=1.988e+04 +// source=hy3-tp4-m4096-dsh-test1-bfdf3002 +// INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. +// +// Worker: worker_1 (physical GPU 1), assigned shape: +// hy3_tp4_o_proj_m4096 : M=4096, N=4096, K=2048 +// +// Iteration 1 - DUMMA throughput baseline with a 2-D macro-tile (128x64). +// * Large-prefill path (exact (m, n, k) == (4096, 4096, 2048)): +// native INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 +// output tile per block; four wavefronts (256 threads); each wave owns +// a 64x32 quadrant built from eight m16n16k32 fragments; block +// cooperatively vector-loads A[128,64] and packed B panels into 26 KiB +// of LDS with a padded 80-byte B row stride; 64-deep K double buffering +// with a register prefetch of the next stage (one loop-carried barrier +// per 64-K stage); fused dot * x_scale[m] * weight_scale[n] epilogue +// stored directly as bf16 from the accumulator fragments. +// Rationale vs the bootstrap 64x64x128 single-buffer tile (verified on +// the exact shape in the trusted ledger): 128x64 doubles B reuse across +// rows, halves the grid to 2048 blocks, stages each B stage as one +// coalesced 4 KiB stream from pre-packed [N/64, K, 64] panels, breaks +// the 64-byte LDS bank periodicity with an 80-byte stride, and hides +// global->LDS latency with double buffering. +// * launch_pack_w8a8_weight packs exactly (k, n) == (2048, 4096) into +// contiguous [N/64, K, 64] int8 panels; every other (k, n) keeps the +// identity copy. Packing never happens inside the timed GEMM. +// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), +// including the paired M=2 decode shape with the same (K, N); the +// fallback reads the packed panel layout whenever (k, n) == (2048, 4096) +// and the identity layout otherwise. +// +// Iteration 3 - A-tile LDS bank-conflict fix: +// PMC on the accepted iteration-1 source: 77,594,624 LDS bank conflicts +// (~6.88 per LDS instruction). Root cause in the fragment loads: the A +// tile still uses an unpadded 64-byte row stride (16 banks/row), so the +// bank pattern repeats every two rows and an m16n16k32 A-fragment load +// (64 lanes x 8 B over 16 rows) puts eight lanes on the same bank pair +// (8-way conflicts), while the already padded 80-byte B rows (20 +// banks/row, period 8) only reach 4-way over their 32-row span - matching +// the PMC mix (8 A-loads x 8-way + 4 B-loads x 4-way per 64-K stage). +// Focused change: pad the A LDS rows 64 -> 80 bytes (five bank phases, +// same as B) so each A-fragment load drops from 8-way to 2-way bank +// conflicts. LDS grows 26,624 -> 30,720 B and stays at 2 blocks/CU +// (61,440 B of the 64 KiB LDS budget); 80 is a multiple of 16 so every +// cooperative int4 stage store stays aligned; the packed [N/64, K, 64] +// global layout, exact int32 accumulation, grid, and fallback are +// untouched. +// +// Iteration 10 - stage-boundary restructure (early publish, one full stage of +// write->read distance): +// Rounds 7 (in-loop A-fragment prefetch, 845.0 us) and 8 (occupancy probe, +// 866.6 us) falsified the in-kk LDS latency and wave-count theories and +// indicted the structural stage-start chain: [last MMACs] -> [s_waitcnt +// vmcnt for the prefetch loads] -> [3x ds_write_b128] -> [s_waitcnt +// lgkmcnt(0)] -> [s_barrier] -> [stage-start fragment loads] (ISA 0x3BC0- +// 0x3BF4 on the accepted source). The next-stage LDS publish currently +// happens AFTER the MMAC burst, so the store issue+retire latency and the +// fresh-write visibility sit directly between the barrier and the next +// stage's first reads. Focused change: the loop-carried int4 payload for +// stage k0+kStageK is published into the idle ping-pong buffer at the TOP +// of stage k0 (before the burst) and the global prefetch for stage +// k0+2*kStageK now has a full extra stage of lead, so the store latency +// retires under the burst and the stage-end barrier publishes data that has +// been resident in LDS for a full stage before the next stage reads it. +// One barrier per stage, 30,720 B LDS (2 blocks/CU), grid, packed layout, +// exact int32 accumulation order, shape guards, and the scalar fallback are +// unchanged. + +#include +#include +#include + +#include + +namespace { + +using namespace du::dumma; + +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 +constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kTargetM = 4096; +constexpr int kTargetN = 4096; +constexpr int kTargetK = 2048; +constexpr int kBlockM = 128; +constexpr int kBlockN = 64; +constexpr int kStageK = 64; +constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) +constexpr int kBStride = kBlockN + kBPad; +constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride +constexpr int kBlockThreads = 4 * kWaveSize; + +using bf16_t = hip_bfloat16; + +// --------------------------------------------------------------------------- +// Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. +// Each wave owns a 64x32 quadrant (eight m16n16k32 accumulators); the block +// cooperatively stages A[128,64] and packed-B [64,64] panels in two ping-pong +// buffers of 64-K stages (26 KiB total LDS). One barrier per stage. +// --------------------------------------------------------------------------- +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int base_row, + int base_col, + int m, + int n, + int lane) { + // Verified gfx928 int8 m16n16k32 accumulator ownership (matches + // du_store_matrix_sync): lane & 15 selects the row, lane >> 4 selects + // col % 4, and x[i] maps to columns col%4 + 4*i. + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M masking: padded rows never write + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); + } +} + +__global__ __launch_bounds__(kBlockThreads) void +w8a8_dumma_128x64x64_packed_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + // Packed layout: panel p = B[:, p*64:(p+1)*64] stored as [K][64] rows. + const int b_panel_offset = + static_cast(blockIdx.x) * kTargetK * kBlockN; + const int local_row = wave_row * 64; + const int local_col = wave_col * 32; + + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[2][kStageK * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + + // Prologue: cooperative vectorized staging of stage 0. Each thread owns one + // int4 in A rows [0,64) and one in A rows [64,128), plus one int4 of packed + // B. A[128,64] is 512 int4s; B[64,64] is 256 int4s. + const int vector_byte_offset = tid * sizeof(int4); + const int stage_row = vector_byte_offset / kStageK; + const int stage_col = vector_byte_offset - stage_row * kStageK; + *reinterpret_cast(a_tile[0] + stage_row * kAStride + stage_col) = + *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + stage_col); + *reinterpret_cast(a_tile[0] + + (stage_row + kBlockM / 2) * kAStride + stage_col) = + *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + + stage_col); + *reinterpret_cast(b_tile[0] + stage_row * kBStride + stage_col) = + *reinterpret_cast( + weight + b_panel_offset + stage_row * kBlockN + stage_col); + __syncthreads(); + + // Loop-carried register payload for the stage after the one being computed + // (stage k0 + kStageK): loaded one full stage ahead and published into the + // idle ping-pong buffer at the TOP of the following stage. The publish's + // LDS store latency then retires under the MMAC burst, and the stage-end + // barrier publishes it a full stage before the next stage's first fragment + // loads - so the structural LDS write -> barrier -> read chain leaves the + // stage boundary (rounds 7/8 falsified in-loop A-fragment prefetch and + // occupancy; this is the prescribed stage-boundary restructure). + int4 cur_a0{}; + int4 cur_a1{}; + int4 cur_b{}; + if (k > kStageK) { + cur_a0 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + kStageK + stage_col); + cur_a1 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + + kStageK + stage_col); + cur_b = *reinterpret_cast( + weight + b_panel_offset + (kStageK + stage_row) * kBlockN + stage_col); + } + + int current = 0; + for (int k0 = 0; k0 < k; k0 += kStageK) { + const int next_k = k0 + kStageK; + const int next = current ^ 1; + // Publish the loop-carried payload (stage next_k) into the idle buffer + // BEFORE the MMAC burst; the stage-end barrier below publishes it a full + // stage before it is read, hiding the store latency under the burst. + if (next_k < k) { + *reinterpret_cast(a_tile[next] + stage_row * kAStride + + stage_col) = cur_a0; + *reinterpret_cast(a_tile[next] + + (stage_row + kBlockM / 2) * kAStride + + stage_col) = cur_a1; + *reinterpret_cast( + b_tile[next] + stage_row * kBStride + stage_col) = cur_b; + } + // Prefetch stage k0 + 2*kStageK into the payload registers (a full stage + // of extra global->LDS lead) so the publish never waits on the loads. + const int load_k = k0 + 2 * kStageK; + if (load_k < k) { + cur_a0 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + load_k + stage_col); + cur_a1 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + + load_k + stage_col); + cur_b = *reinterpret_cast( + weight + b_panel_offset + (load_k + stage_row) * kBlockN + + stage_col); + } + + // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + b_frag0, b_tile[current] + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile[current] + kk * kBStride + local_col + kTileN, + kBStride); + du_load_matrix_sync( + a_frag0, a_tile[current] + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile[current] + (local_row + kTileM) * kAStride + kk, + kAStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_load_matrix_sync( + a_frag0, a_tile[current] + (local_row + 2 * kTileM) * kAStride + kk, + kAStride); + du_load_matrix_sync( + a_frag1, a_tile[current] + (local_row + 3 * kTileM) * kAStride + kk, + kAStride); + du_mma_sync(acc20, a_frag0, b_frag0, acc20); + du_mma_sync(acc21, a_frag0, b_frag1, acc21); + du_mma_sync(acc30, a_frag1, b_frag0, acc30); + du_mma_sync(acc31, a_frag1, b_frag1, acc31); + } + + if (next_k < k) { + // The stage-end barrier publishes the payload written at the top of + // this stage (a full stage before the next stage reads it) and protects + // buffer `current` from the publish at the top of the next stage. + __syncthreads(); + current = next; + } + } + + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + store_prefill_fragment(acc00, x_scale, weight_scale, out, + base_row, base_col, m, n, lane); + store_prefill_fragment(acc01, x_scale, weight_scale, out, + base_row, base_col + kTileN, m, n, lane); + store_prefill_fragment(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_prefill_fragment(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, lane); + store_prefill_fragment(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, lane); + store_prefill_fragment(acc31, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + kTileN, m, n, lane); +} + +// --------------------------------------------------------------------------- +// Generic scalar fallback: one thread per output element. Exact int32 dot +// over K, then fused x_scale * weight_scale, then bf16 store. Handles every +// unmatched (m, n, k), including all small-M API cases. When (k, n) matches +// the packed pair (2048, 4096) the weight is read from the [N/64, K, 64] +// panel layout; otherwise from the identity layout. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + int32_t acc = 0; + if (n == kTargetN && k == kTargetK) { + const int panel = col / kBlockN; + const int panel_col = col - panel * kBlockN; + const int64_t panel_offset = + static_cast(panel) * kTargetK * kBlockN; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(x_q[static_cast(row) * k + kk]) * + static_cast(weight[panel_offset + + static_cast(kk) * kBlockN + + panel_col]); + } + } else { + const int8_t* a_row = x_q + static_cast(row) * k; + const int8_t* b_col = weight + col; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + } + const float scaled = static_cast(acc) * + x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Weight packing: exact (k, n) == (2048, 4096) becomes contiguous +// [N/64, K, 64] int8 panels (each int4 stays within one 64-column row, so +// loads are coalesced 4 KiB streams per 64-K stage); identity copy for any +// other (k, n). Scales are always copied identity. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_64n_panel_kernel( + const int8_t* __restrict__ raw_weight, + int8_t* __restrict__ packed_weight) { + const int vec = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + constexpr int kVectors = + kTargetK * kTargetN / static_cast(sizeof(int4)); + if (vec >= kVectors) { + return; + } + constexpr int kPanelBytes = kTargetK * kBlockN; + const int byte_offset = vec * sizeof(int4); + const int panel = byte_offset / kPanelBytes; + const int panel_offset = byte_offset - panel * kPanelBytes; + const int kk = panel_offset / kBlockN; + const int col = panel_offset - kk * kBlockN; + reinterpret_cast(packed_weight)[vec] = + *reinterpret_cast( + raw_weight + static_cast(kk) * kTargetN + + panel * kBlockN + col); +} + +__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int64_t count) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +__global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( + const float* __restrict__ raw, + float* __restrict__ packed, + int count) { + const int linear = blockIdx.x * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Stable host launch symbols consumed by csrc/bindings.cpp. +// Both launchers are pure dispatch: no allocation, no packing, no +// synchronization, no default-stream launch; they run on the caller-provided +// PyTorch stream and are CUDA-Graph safe. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; // no split-K: the GEMM does not use the workspace + (void)workspace_bytes; + auto* out_bf16 = reinterpret_cast(out); + + // Explicit dispatch. The assigned shape (M=4096, N=4096, K=2048) takes the + // packed-panel DUMMA 128x64 path; every other (m, n, k) - including + // small-M API cases and the paired M=2 shape with the same (K, N) - takes + // the scalar fallback, which understands the packed panel layout. + if (m == kTargetM && n == kTargetN && k == kTargetK) { + const dim3 grid(kTargetN / kBlockN, kTargetM / kBlockM); + const dim3 block(kBlockThreads); + hipLaunchKernelGGL(w8a8_dumma_128x64x64_packed_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + constexpr int kBlock = 256; + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + (void)hipGetLastError(); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kBlock = 256; + if (k == kTargetK && n == kTargetN) { + constexpr int kPackVectors = + kTargetK * kTargetN / static_cast(sizeof(int4)); + const dim3 weight_grid( + static_cast((kPackVectors + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_pack_64n_panel_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight); + } else { + const int64_t weight_count = static_cast(k) * n; + const dim3 weight_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_pack_identity_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, weight_count); + } + + const int64_t scale_count = n; + const dim3 scale_grid( + static_cast((scale_count + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, + scale_grid, block, 0, stream, + weight_scale, packed_weight_scale, n); + (void)hipGetLastError(); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/qkv_proj.hip new file mode 100644 index 00000000..919b94e1 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/qkv_proj.hip @@ -0,0 +1,1169 @@ +// @@variant shape=hy3_tp4_qkv_proj_m4096 commit=61a0f079ab9622cc6c83a9c44af6ecd626ef0640 added=2026-08-23 +// median_us=756.7 p90_us=760.2 tops=113.5 bandwidth_gb_s=63.78 speedup=31.78 baseline_us=2.405e+04 +// source=hy3-tp4-m4096-dsh-test1-bfdf3002 +// MetaInfer W8A8 INT8 GEMM bootstrap for gfx928 (K500SM_AI). +// +// Worker: worker_0 (physical GPU 0) +// Assigned shape: hy3_tp4_qkv_proj_m4096 (M=4096, N=2560, K=4096) +// +// Operator contract (fixed by the control plane): +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// with the packed weight layout equal to the logical [K, N] row-major int8 +// weight for the bootstrap (launch_pack_w8a8_weight is an identity D2D copy). +// +// Iteration 3 (harness) strategy: single-buffered 64x128 macro-tile with +// cooperative A+B LDS staging -- the matched control arm of the mandated +// pipeline comparison (single vs double buffering across K tiles). +// * Large-M path (M >= 128, N % 128 == 0, K % 128 == 0): native INT8 DUMMA +// m16n16k32 tiled kernel with a 64x128 output tile per block and 256 +// threads (4 wavefronts). Each wave owns a 32x64 quadrant (eight +// m16n16k32 int32 accumulator fragments). K is staged cooperatively in a +// SINGLE LDS buffer at 64-K granularity: the next stage's 3 x int4 +// (1 A + 2 B vectors) are loaded from global only after the current +// stage's 16-MMAC DUMMA burst and its all-consumption barrier, waited +// on, and committed as int32 LDS stores into the same buffer, followed +// by a store-visibility barrier -- TWO __syncthreads per 64-K stage +// (128 over K=4096). The global-load latency is fully exposed (no +// compute overlap), which isolates the value of the double-buffer +// overlap at the same 64-K stage size, odd-word LDS strides, and +// 2-blocks/CU residency as the accepted iteration-2 double-buffered +// kernel (1108.58 us median; iteration-1 single-buffered 128-K stage: +// 1296.50 us). LDS/block = 64x68 + 64x132 = 12,800 B -> 2 resident +// blocks/CU (VGPR-bound). The grid (N/128) x (ceil(M/64)) supplies 1280 +// blocks for the assigned shape. +// * Comparison evidence (fact ledger): double-64 vs single-128 improved +// the benchmark median 1296.50 -> 1108.58 us, but PMC moved against the +// 'reduced VMEM stalls' reading: vmem_read_instructions stayed flat at +// 1,177,600, lds_wait_instructions rose 19.48M -> 20.69M, bank conflicts +// rose 5.24M -> 7.86M, L2 hit fell 89.46% -> 88.25%. This round re-runs +// the single-buffered pipeline at the SAME 64-K stage to decide +// retention on matched evidence (benchmark median/P90 is authoritative +// per the interpretation guard; PMC counters are diagnostic). +// * Operand-reuse accounting per 64x128 macro-tile (A[64,64] + B[64,128] +// per 64-K stage): cooperative staging loads each A element from global +// once per stage and reuses it for 8 MMAs (512 MMA-uses per element over +// K=4096), and each B element for 2 MMAs per stage (128 MMA-uses over +// K=4096). Without staging, the two waves sharing each A row / B column +// segment would refetch everything (2x global traffic) with per-lane +// 8-byte fragment loads (4x more global-load instructions), so the +// cooperative A+B staging path is kept. +// * LDS row strides stay padded to odd word counts (A: 68 B = 17 words, +// B: 132 B = 33 words). The 64x64 bootstrap used unpadded 128/64-byte +// rows, so every INT8 m16n16k32 fragment load collapsed onto a handful +// of LDS banks (16-way conflicts; 199.2M conflict events over 24.9M LDS +// instructions in PMC). The odd word stride makes the lane row/k-group +// index contribute to the bank, spreading fragment loads across banks +// (A and B: 16-way -> 2-4-way per LDS instruction) at 2 blocks/CU. +// * The epilogue applies float(dot) * x_scale[m] * weight_scale[n] and +// stores bf16 directly from the accumulator fragments using the verified +// gfx928 int8 m16n16k32 lane mapping (row = lane & 15, col_mod4 = +// lane >> 4, frag.x[i] -> columns col_mod4 + 4*i). +// * Fallbacks are preserved unchanged: the 64x64 kernel for large-M shapes +// with N % 64 == 0 but N % 128 != 0, and the scalar int8/int32 +// grid-stride fallback for every other (m, n, k) including the paired +// M=2 API shape. launch_pack_w8a8_weight stays an identity +// device-to-device copy valid for every (K, N). +// +// Iteration 5 (packing round) strategy: keep the accepted iteration-3 +// pipeline (single-buffered 64-K stage, 64x128 macro-tile, 3 blocks/CU) and +// add ONE weight packing/swizzle for the exact assigned shape (K=4096, +// N=2560): launch_pack_w8a8_weight transposes the logical [K, N] weight to +// an n-major packed[n][k] layout once, outside the timed region and outside +// Graph capture (same byte count k*n, same allocation, graph-stable packed +// layout). The exact shape routes (before the generic 64x128 path) to a new +// w8a8_dumma_prefill_64x128_packedb_kernel that stages B into an n-major LDS +// tile with a 16-byte-aligned 80-byte row stride (20 words, five bank +// phases) and consumes it with col_major m16n16k32 B fragments via an +// explicit 8-byte loader (load_b_frag8): each lane's 8 fragment elements at +// p[row*ldm + col + i] are one contiguous 8-byte-aligned chunk (one +// ds_read2_b64 per B fragment instead of 8 byte-granular ds_read_u8 at +// stride 132 in the accepted kernel's ISA), and each 16-byte staging vector +// commits with one ds_write_b128 instead of four ds_write_b32. A stays +// row-major (already vectorized ds_read2_b32 in the accepted ISA); the int32 +// accumulation order (k0-outer, kk-inner) and the fused scale/bf16 epilogue +// are unchanged, so results stay bit-identical. All other shapes keep the +// identity pack and the accepted iteration-3 kernel byte-for-byte; the 64x64 +// kernel, the scalar fallback (which decodes packed[n][k] when (k, n) +// matches), M tails, and the exact shape guards are preserved. +// +// Iteration 6 (epilogue round) strategy: keep the accepted iteration-5 +// packedb kernel's pipeline byte-identical (single-buffered 64-K stage, +// n-major B with 80-B rows, load_b_frag8, 3 blocks/CU) and replace ONLY its +// final store. The scales and bf16 conversion are already fused in the +// kernel and launch_w8a8_gemm has no workspace/combine pass (workspace is +// unused), so the remaining epilogue inefficiency is the store pattern: the +// m16n16k32 accumulator lane mapping (row = lane & 15, column group +// c4 = lane >> 4, frag.x[i] -> column c4 + 4*i) makes the direct store four +// 2-byte scalar stores per lane whose wavefront addresses touch every 32-B +// sector at 25% utilization. store_prefill_fragment_coalesced transposes the +// 4-element groups within each 4-lane column group (lanes r, r+16, r+32, +// r+48) with two 2x2 shuffle steps (8 shfl_xor + 8 v_cndmask; this DTK +// lowers each shfl_xor to one ds_bpermute -- +64 LDS permutes per wavefront +// at the block tail, where the LDS pipe is otherwise idle), so lane (r, c4) +// owns the four CONTIGUOUS columns 4*c4 .. 4*c4+3, then converts to bf16, +// packs 4 bf16 (8 B), and issues ONE 8-byte store per lane: 64 lanes x 8 B = +// 512 B per fragment per wavefront in 16 fully-used 32-B sectors (100% +// store sector efficiency; the assigned shape's vmem_write_instructions drop +// 163,840 -> 40,960). Only the int32 values are re-routed between lanes; the +// per-element float scale/bf16 math is unchanged, so output bits are +// identical. The generic 64x128/64x64 kernels and the scalar fallback are +// untouched (their store_prefill_fragment epilogue is preserved +// byte-for-byte). +// +// The timed operator (launch_w8a8_gemm) performs no allocation, compilation, +// autotuning, packing, host/device synchronization, or default-stream launch: +// it only dispatches kernels on the caller-provided HIP stream. +// +// Include order is fixed by the control plane: hip_runtime, hip_bfloat16, +// then du_mma (this DTK's du_mma.h is not self-contained otherwise). + +#include +#include +#include + +#include + +namespace { + +constexpr int kWaveSize = 64; +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; +constexpr int kBlockM = 64; +constexpr int kBlockN = 64; +constexpr int kStageK = 128; +// 4 wavefronts; must remain a multiple of the gfx928 wavefront size (64). +constexpr int kThreadsPerBlock = 256; + +// Iteration 3 (pipeline round): 64x128 macro-tile, SINGLE-buffered 64-K K +// stage (matched control for the single-vs-double buffering comparison). +// LDS row strides stay padded to odd word counts (A: 68 B = 17 words, +// B: 132 B = 33 words) so the INT8 m16n16k32 fragment loads spread across +// banks instead of collapsing 16-way onto one bank (128-byte rows are 32 +// words, i.e. 0 mod 32 for every row -> single-bank collisions). +constexpr int kBlockM128 = 64; +constexpr int kBlockN128 = 128; +// Single-buffered K stage: one 64-K buffer = 12,800 B/block (the accepted +// iteration-2 double buffer was 2 x 64 = 25,600 B; the iteration-1 +// single-buffered 128-K stage was 25,344 B). Occupancy stays 2 blocks/CU +// (VGPR-bound at arch_vgpr ~112), identical to both prior kernels. +constexpr int kStageK128 = 64; +constexpr int kAStride128 = kStageK128 + 4; // 68 bytes per A row (17 words) +constexpr int kBStride128 = kBlockN128 + 4; // 132 bytes per B row (33 words) + +// Iteration 5 (packing round): packed n-major B layout for the exact +// assigned shape (K=4096, N=2560). B is transposed once, outside timing, to +// packed[n][k]; the packed kernel stages B into an n-major LDS tile with a +// 16-byte-aligned non-power-of-two row stride (80 B = 20 words, five bank +// phases) so every DUMMA m16n16k32 B fragment load (col_major) reads its +// lane's 8 bytes contiguously (ds_read2_b64) and every 16-byte staging +// vector commits with one ds_write_b128 (b_n*80 + b_k16 is 0 mod 16). +constexpr int kPackedBStride = 80; // 64 data + 16 pad bytes (20 words) +constexpr int kPackedBK = 4096; // exact K of the packed assigned shape +constexpr int kPackedBN = 2560; // exact N of the packed assigned shape + +// Iteration 18 (fragment-pipeline round): the packedb compute loop splits the +// 64-K stage into two explicit 32-K steps (kk = 0, kk = 32) and hoists ALL +// fragment loads for both steps before the first v_mmac. This requires +// exactly two kTileK steps per stage. +static_assert(kStageK128 == 2 * kTileK, + "pipelined kk loop requires exactly two 32-K steps per stage"); + +using namespace du::dumma; + +// Direct accumulator epilogue for gfx928 INT8 m16n16k32 (verified against +// du_store_matrix_sync): row = lane & 15, col_mod4 = lane >> 4, +// frag.x[i] maps to columns col_mod4 + 4*i. Stores +// bf16(float(dot) * x_scale[row] * weight_scale[col]); out-of-range rows are +// masked (tail-M handling). +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Iteration 6 (epilogue round): coalesced fragment store. The m16n16k32 +// accumulator lane mapping (row = lane & 15, column group c4 = lane >> 4, +// frag.x[i] -> column c4 + 4*i) gives each lane four elements strided by 4 +// columns, so the direct store (store_prefill_fragment) is four 2-byte +// scalar stores per lane whose wavefront addresses touch each 32-B sector at +// 25% utilization (8 B used of 32 B; 4 store instructions per fragment per +// wave). This epilogue transposes the 4-element groups within each 4-lane +// column group (lanes r, r+16, r+32, r+48 -- a 4x4 transpose, two 2x2 steps +// with shfl_xor 16 then 32, one v_cndmask per element per step; no staging +// tile round trip), so lane (r, c4) ends up holding the four CONTIGUOUS +// columns 4*c4 .. 4*c4+3, converts them to bf16, packs 4 bf16 (8 B), and +// writes ONE 8-byte store per lane: 64 lanes x 8 B = 512 B per fragment per +// wavefront in 16 fully-used 32-B sectors (100% store sector efficiency; +// vmem_write_instructions for the assigned shape drops 163,840 -> 40,960). +// Note: this DTK's __shfl_xor lowers to ds_bpermute (one LDS permute per +// shuffle, i.e. +64 LDS instructions per wavefront per block -- accepted +// because it runs at the block tail where the LDS pipe is otherwise idle; +// there is no global round trip and no staging tile). Only the int32 values +// are re-routed between lanes -- the per-element scale multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are +// unchanged, so the stored bits are identical to store_prefill_fragment. +// The row>=m guard is wavefront-uniform (lane & 15 cycles the same 16 rows +// in every 16-lane group), so the shuffles never mix active and inactive +// lanes; base_col is a multiple of 64 and n*2 a multiple of 8, so the float4 +// weight_scale load (col0 % 4 == 0) and the 8-byte store are aligned. +template +__device__ __forceinline__ void store_prefill_fragment_coalesced( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 64, n*2 is a multiple of 8). + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + row * n + col0) = packed; +} + +// Large-M prefill kernel: 64x64 output tile per block, K staged in LDS +// (single buffer), four waves each owning a 32x32 quadrant = four +// m16n16k32 int8->int32 DUMMA accumulators. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM * kStageK]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBlockN]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[64,128] (row-major, zero-filled past M). + constexpr int kAVectors = kBlockM * kStageK / static_cast(sizeof(int4)); + for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int local_row = byte_offset / kStageK; + const int kk = byte_offset - local_row * kStageK; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile)[vec] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + kk) + : int4{0, 0, 0, 0}; + } + + // Stage B[128,64] (row-major [K, N] logical weight layout). + constexpr int kBVectors = kStageK * kBlockN / static_cast(sizeof(int4)); + for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + reinterpret_cast(b_tile)[vec] = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + col); + } + __syncthreads(); + + // Consume the stage: each wave does four m16n16k32 MMAs per 32-K step. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kStageK + kk, kStageK); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kStageK + kk, kStageK); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBlockN + local_col, kBlockN); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBlockN + local_col + kTileN, kBlockN); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + // Single buffer: make sure every wave finished reading LDS before the + // next stage overwrites it. + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Large-M prefill kernel #2 (iteration 3, pipeline round): 64x128 output +// tile per block, 256 threads (4 wavefronts). Each wave owns a 32x64 +// quadrant = eight m16n16k32 int32 accumulators (2 row halves x 4 column +// quarters). K is staged cooperatively in a SINGLE LDS buffer at 64-K +// granularity as the matched control for the mandated single-vs-double +// buffering comparison: it shares the iteration-2 double-buffered kernel's +// 64-K stage size, odd-word LDS strides (A: 68 B = 17 words, B: 132 B = 33 +// words), 16-MMAC DUMMA burst, and 2-blocks/CU residency, but uses ONE +// 12,800-B buffer and TWO __syncthreads per stage (128 over K): the next +// stage's 3 x int4 (1 A + 2 B) are loaded from global only after the +// current burst and its all-consumption barrier, vmcnt-waited, and +// committed into the same buffer, then a store-visibility barrier. The +// global-load latency is fully exposed (no compute overlap), isolating the +// value of the double-buffer overlap at matched stage size and occupancy. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x128_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM128; + const int n0 = static_cast(blockIdx.x) * kBlockN128; + + __shared__ __align__(16) int8_t a_tile[kBlockM128 * kAStride128]; + __shared__ __align__(16) int8_t b_tile[kStageK128 * kBStride128]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; + DUFragment + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); + + const int num_stages = k / kStageK128; + + // Cooperative staging mapping (fixed per thread, reused every stage): + // A[64,64] -> 256 int4 vectors; thread `tid` owns vector `tid` + // (row = tid/4, 16-B column group = (tid%4)*16), so each + // wavefront covers 16 rows x 64 B contiguous per row. + // B[64,128] -> 512 int4 vectors; thread `tid` owns vectors `tid` and + // `tid+256` (kk = tid/8, column group = (tid%8)*16), so + // each wavefront covers 8 kk rows x 128 B contiguous. + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + const int b_kk0 = tid >> 3; + const int b_col16 = (tid & 7) << 4; + + // VGPR payload for one full 64-K stage (3 x int4 = 12 VGPR), loaded just + // before it is committed (no one-stage-ahead overlap in this control). + int4 vA, vB0, vB1; + + // Prologue: load stage 0 into registers, commit it to the single buffer, + // and make it visible before the first burst. + { + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + vB0 = *reinterpret_cast( + weight + b_kk0 * n + n0 + b_col16); + vB1 = *reinterpret_cast( + weight + (b_kk0 + kStageK128 / 2) * n + n0 + b_col16); + // Commit as int32 stores (skips the 4-byte pad column; 4 x ds_write_b32 + // per 16-byte vector because the odd strides are 4 B mod 16). + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + int32_t* bdst0 = reinterpret_cast( + b_tile + b_kk0 * kBStride128 + b_col16); + bdst0[0] = vB0.x; + bdst0[1] = vB0.y; + bdst0[2] = vB0.z; + bdst0[3] = vB0.w; + int32_t* bdst1 = bdst0 + (kStageK128 / 2) * (kBStride128 / 4); + bdst1[0] = vB1.x; + bdst1[1] = vB1.y; + bdst1[2] = vB1.z; + bdst1[3] = vB1.w; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + // Consume the single buffer: each wave does eight m16n16k32 MMAs per + // 32-K step over its 32x64 quadrant. + const int local_row = wave_row * 32; + const int local_col = wave_col * 64; +#pragma unroll + for (int kk = 0; kk < kStageK128; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride128 + kk, kAStride128); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride128 + kk, + kAStride128); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride128 + local_col, kBStride128); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride128 + local_col + kTileN, + kBStride128); + du_load_matrix_sync( + b_frag2, b_tile + kk * kBStride128 + local_col + 2 * kTileN, + kBStride128); + du_load_matrix_sync( + b_frag3, b_tile + kk * kBStride128 + local_col + 3 * kTileN, + kBStride128); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + } + // All consumption of the single buffer is complete only after every + // wavefront passes this barrier; the stage s+1 stores below overwrite + // that same buffer, so they must wait for it. + __syncthreads(); + + // Load stage s+1 and commit it into the single buffer. The global-load + // latency is fully exposed here: the loads are issued after the burst + // (nothing to overlap) and the compiler's vmcnt wait before the DS + // stores is on the critical path. `s + 1 < num_stages` is block-uniform, + // so the branch and its barrier are free of divergence. + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vB0 = *reinterpret_cast( + weight + (s1 + b_kk0) * n + n0 + b_col16); + vB1 = *reinterpret_cast( + weight + (s1 + b_kk0 + kStageK128 / 2) * n + n0 + b_col16); + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + int32_t* bdst0 = reinterpret_cast( + b_tile + b_kk0 * kBStride128 + b_col16); + bdst0[0] = vB0.x; + bdst0[1] = vB0.y; + bdst0[2] = vB0.z; + bdst0[3] = vB0.w; + int32_t* bdst1 = bdst0 + (kStageK128 / 2) * (kBStride128 / 4); + bdst1[0] = vB1.x; + bdst1[1] = vB1.y; + bdst1[2] = vB1.z; + bdst1[3] = vB1.w; + // Make the stage s+1 stores visible to every wavefront before the + // next burst. + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 64; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc02, x_scale, weight_scale, out, m, n, + base_row, base_col + 2 * kTileN, lane); + store_prefill_fragment( + acc03, x_scale, weight_scale, out, m, n, + base_row, base_col + 3 * kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); + store_prefill_fragment( + acc12, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + 2 * kTileN, lane); + store_prefill_fragment( + acc13, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + 3 * kTileN, lane); +} + +// Iteration 5 (packing round): explicit 8-byte loader for the col_major +// m16n16k32 B fragment (the gate_up-lineage-validated load_fragment8 +// pattern). The col_major fragment mapping is n = lane&15, +// k = 8*(lane>>4) + i; with the n-major LDS tile at row stride 80 the +// lane's 8 elements are contiguous (p[row*ldm + col .. +7], 8-byte aligned +// because ldm=80 and col are multiples of 8), so this compiles to ONE +// ds_read2_b64 instead of 8 ds_read_u8 + mask/OR reassembly. The byte +// placement is identical to du_load_matrix_sync, so the +// v_mmac_i32_16x16x32_i8 fragment registers receive the same values. +__device__ __forceinline__ void load_b_frag8( + DUFragment& f, + const int8_t* __restrict__ p, + int ldm) { + const unsigned row = static_cast(__lane_id()) & 0xfu; + const unsigned col = (static_cast(__lane_id()) >> 4) << 3; + const int64_t v = *reinterpret_cast(p + row * ldm + col); + *reinterpret_cast(&f.x[0]) = v; +} + +// Large-M prefill kernel #3 (iteration 5, packing round): 64x128 macro-tile +// for the exact assigned shape (K=4096, N=2560) with the weight packed +// n-major. The pipeline is byte-identical to the accepted iteration-3 kernel +// (single-buffered 64-K stage, TWO __syncthreads per stage with fully +// exposed global-load latency, 3 x int4 VGPR payload, A staged row-major +// with the 68-byte odd-word stride); ONLY the B operand path changes: +// * pack: the logical [K, N] weight is transposed once (outside the timed +// region and Graph capture) to packed[n][k] by launch_pack_w8a8_weight. +// * staging: thread tid owns rows n_local = tid>>2 and n_local+64, each +// with the 16-byte k-run (tid&3)*16; two global_load_dwordx4 are +// committed with ONE ds_write_b128 each (16-byte-aligned LDS address: +// row stride 80 and the k16 offset are both 0 mod 16), vs four +// ds_write_b32 per vector in the accepted kernel. +// * fragments: load_b_frag8 (col_major mapping n = local_col + lane&15, +// k = kk + 8*(lane>>4) + i) reads each lane's 8 elements as one +// 8-byte-aligned contiguous chunk -> one ds_read2_b64 per B fragment, +// vs 8 byte-granular ds_read_u8 at stride 132 in the accepted kernel's +// ISA. The mapping is the transpose of the accepted row_major mapping, +// so the v_mmac_i32_16x16x32_i8 accumulation order is unchanged and the +// int32 results are bit-identical. Bank safety: the 20-word row stride +// spreads each k-group's 16 n-lanes over 8 distinct banks with five +// bank phases. +// * A path, epilogue, tail masking, and occupancy are unchanged. +// LDS/block = 64x68 + 128x80 = 14,592 B -> 3 resident blocks/CU +// (VGPR-bound at arch_vgpr ~80), same as the accepted kernel. +// +// Iteration 18 (fragment-pipeline round) strategy: keep the accepted +// iteration-6 packedb kernel byte-identical in EVERYTHING except the inner +// fragment schedule of the 64-K-stage compute loop. The accepted kernel's +// exact gfx928 ISA (iteration 7 evidence) shows the compiler already +// software-pipelines the kk=32 A fragment loads under the kk=0 MMAC burst, +// but its four kk=32 B-fragment ds_read_b64 loads are issued only AFTER +// MMACs 0-7, with s_waitcnt lgkmcnt(4..1) interleaved into MMACs 8-11 (16+ +// lgkmcnt waits per stage; PMC lds_wait 6,560,572 = 1.18x lds_instructions, +// the dominant stall class of a SIMD ~36%-issue-busy kernel). This round +// hoists ALL twelve ds_read fragment loads of both 32-K steps before the +// first v_mmac of the stage (explicit two-step unroll of the kk loop with +// separate kk=32 fragment variables), so the LDS latency concentrates in one +// wait and the 16-MMAC stream has no lgkmcnt waits. The MMAC sequence per +// (m,n) output tile and the int32 k0-outer/kk-inner accumulation order are +// unchanged, so output bits stay identical. Occupancy target is retained: +// arch_vgpr 72 -> <= 85 at 3 blocks/CU (65,280 <= 65,536 VGPRs), LDS 14,592 +// B, grid 1280 x 256, TWO __syncthreads per stage, staging mapping, packed +// B layout, coalesced epilogue, and the exact shape guard (m >= 128 && +// n == 2560 && k == 4096) all unchanged; generic 64x128/64x64 kernels, the +// pack kernel, and the scalar fallback are untouched. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x128_packedb_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM128; + const int n0 = static_cast(blockIdx.x) * kBlockN128; + + __shared__ __align__(16) int8_t a_tile[kBlockM128 * kAStride128]; + // n-major B tile: 128 n-rows of 80 bytes (64 data + 16 pad). + __shared__ __align__(16) int8_t b_tile[kBlockN128 * kPackedBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; + DUFragment + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); + + const int num_stages = k / kStageK128; + + // A staging mapping unchanged (row-major x_q, odd-word stride 68). + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + // B staging mapping for packed[n][k]: thread tid owns the 16-byte k-run + // (b_k16..b_k16+15) of n-rows b_n and b_n + 64. + const int b_n = tid >> 2; + const int b_k16 = (tid & 3) << 4; + + // VGPR payload for one full 64-K stage (3 x int4 = 12 VGPR). + int4 vA, vB0, vB1; + + // Prologue: load stage 0 into registers, commit it to the single buffers, + // and make it visible before the first burst. + { + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + vB0 = *reinterpret_cast( + packed_b + (n0 + b_n) * k + b_k16); + vB1 = *reinterpret_cast( + packed_b + (n0 + b_n + kBlockN128 / 2) * k + b_k16); + // Commit A as int32 stores (68-byte rows are 4 mod 16 -> no b128). + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + // Commit B as 16-byte vector stores: b_n*80 + b_k16 is 0 mod 16. + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB0; + *reinterpret_cast( + b_tile + (b_n + kBlockN128 / 2) * kPackedBStride + b_k16) = vB1; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + // Consume the single buffer: each wave does eight m16n16k32 MMAs per + // 32-K step over its 32x64 quadrant. + const int local_row = wave_row * 32; + const int local_col = wave_col * 64; + const int8_t* abase = a_tile + local_row * kAStride128; + const int8_t* bbase = b_tile + local_col * kPackedBStride; + // Iteration 18 (fragment-pipeline round): software-pipeline the kk=32 + // fragment loads UNDER the kk=0 MMAC burst. The accepted kernel's exact + // ISA already prefetches the kk=32 A fragments this way, but its four + // kk=32 B ds_read_b64 loads are issued only AFTER MMACs 0-7, with + // s_waitcnt lgkmcnt(4..1) interleaved into MMACs 8-11 (16+ lgkmcnt waits + // per stage; PMC lds_wait 6,560,572 = 1.18x lds_instructions). Loading + // both 32-K steps' operands before the first v_mmac concentrates the LDS + // latency into one wait and leaves the 16-MMAC stream free of lgkmcnt + // waits. Same 12 ds_read instructions per stage, same fragment values, + // same int32 k0-outer/kk-inner accumulation order (the MMAC sequence per + // output tile is unchanged) -> bit-identical results. + du_load_matrix_sync(a_frag0, abase, kAStride128); + du_load_matrix_sync(a_frag1, abase + kTileM * kAStride128, kAStride128); + load_b_frag8(b_frag0, bbase, kPackedBStride); + load_b_frag8(b_frag1, bbase + kTileN * kPackedBStride, kPackedBStride); + load_b_frag8(b_frag2, bbase + 2 * kTileN * kPackedBStride, kPackedBStride); + load_b_frag8(b_frag3, bbase + 3 * kTileN * kPackedBStride, kPackedBStride); + DUFragment + a2_frag0, a2_frag1; + DUFragment + b2_frag0, b2_frag1, b2_frag2, b2_frag3; + du_load_matrix_sync(a2_frag0, abase + kTileK, kAStride128); + du_load_matrix_sync( + a2_frag1, abase + kTileM * kAStride128 + kTileK, kAStride128); + load_b_frag8(b2_frag0, bbase + kTileK, kPackedBStride); + load_b_frag8( + b2_frag1, bbase + kTileN * kPackedBStride + kTileK, kPackedBStride); + load_b_frag8( + b2_frag2, bbase + 2 * kTileN * kPackedBStride + kTileK, + kPackedBStride); + load_b_frag8( + b2_frag3, bbase + 3 * kTileN * kPackedBStride + kTileK, + kPackedBStride); + // kk = 0 burst (eight m16n16k32 MMAs). + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + // kk = 32 burst (same accumulators, same MMAC order as the accepted + // two-iteration kk loop). + du_mma_sync(acc00, a2_frag0, b2_frag0, acc00); + du_mma_sync(acc01, a2_frag0, b2_frag1, acc01); + du_mma_sync(acc02, a2_frag0, b2_frag2, acc02); + du_mma_sync(acc03, a2_frag0, b2_frag3, acc03); + du_mma_sync(acc10, a2_frag1, b2_frag0, acc10); + du_mma_sync(acc11, a2_frag1, b2_frag1, acc11); + du_mma_sync(acc12, a2_frag1, b2_frag2, acc12); + du_mma_sync(acc13, a2_frag1, b2_frag3, acc13); + // All consumption of the single buffer is complete only after every + // wavefront passes this barrier; the stage s+1 stores below overwrite + // that same buffer, so they must wait for it. + __syncthreads(); + + // Load stage s+1 and commit it into the single buffer (fully exposed + // global-load latency, exactly as in the accepted iteration-3 kernel). + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vB0 = *reinterpret_cast( + packed_b + (n0 + b_n) * k + s1 + b_k16); + vB1 = *reinterpret_cast( + packed_b + (n0 + b_n + kBlockN128 / 2) * k + s1 + b_k16); + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB0; + *reinterpret_cast( + b_tile + (b_n + kBlockN128 / 2) * kPackedBStride + b_k16) = vB1; + // Make the stage s+1 stores visible to every wavefront before the + // next burst. + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 64; + // Iteration 6 (epilogue round): coalesced stores -- register 4x4 + // transpose (8 shfl_xor -> ds_bpermute + 8 v_cndmask, no staging tile) + // then ONE 8-byte store per lane per fragment (4x fewer VMEM store + // instructions, 100% store sector utilization vs 25% for the scalar + // 2-byte stores). + store_prefill_fragment_coalesced( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment_coalesced( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment_coalesced( + acc02, x_scale, weight_scale, out, m, n, + base_row, base_col + 2 * kTileN, lane); + store_prefill_fragment_coalesced( + acc03, x_scale, weight_scale, out, m, n, + base_row, base_col + 3 * kTileN, lane); + store_prefill_fragment_coalesced( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment_coalesced( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); + store_prefill_fragment_coalesced( + acc12, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + 2 * kTileN, lane); + store_prefill_fragment_coalesced( + acc13, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + 3 * kTileN, lane); +} + +// Generic scalar int8/int32 fallback for unmatched shapes and small-M API +// cases. One output element per grid-stride step; exact int32 accumulation +// (K <= 4096 keeps the int8 dot well within int32 range), then the fused +// float scale and bf16 store. +__global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; + idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = a + static_cast(row) * k; + // Iteration 5 (packing round): the exact assigned shape (K=4096, + // N=2560) uses the packed n-major layout packed[n][k] (set up once by + // launch_pack_w8a8_weight); every other (k, n) keeps the logical + // [K, N] row-major layout. The branch is grid-uniform per launch. + const bool packed_b = (k == kPackedBK && n == kPackedBN); + const int8_t* b_col = + packed_b ? b + static_cast(col) * k : b + col; + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + const int32_t bw = packed_b + ? static_cast(b_col[kk]) + : static_cast( + b_col[static_cast(kk) * n]); + acc += static_cast(a_row[kk]) * bw; + } + const float scaled = static_cast(acc) * + x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// Identity device-to-device byte copy (pack_weight bootstrap, valid for any +// (K, N)). +__global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +// Identity device-to-device float copy for the weight scales. +__global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +// Iteration 5 (packing round): one-time n-major transpose of the logical +// [K, N] row-major weight into packed[n][k] for the exact assigned shape +// (K=4096, N=2560). Runs only from launch_pack_w8a8_weight, outside the +// timed region and outside Graph capture. The packed buffer keeps the same +// byte count (k*n) as the identity pack, so allocations and graph-stable +// addresses are unchanged. Element (k, n) of the logical weight lands at +// packed[n * K + k]; each thread copies one 16-byte k-run (coalesced read +// side; the strided write side is off the critical path). The guard +// guarantees n % 16 == 0, so every 16-byte chunk lies inside one logical +// row and the transpose is exact. +__global__ __launch_bounds__(256) void w8a8_pack_nmajor_b_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int k, + int n) { + const int64_t chunks = (static_cast(k) * n) >> 4; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t c = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + c < chunks; + c += stride) { + const int64_t off = c << 4; + const int kk = static_cast(off / n); + const int col = static_cast(off - static_cast(kk) * n); + const int4 v = *reinterpret_cast(src + off); + const int8_t* bytes = reinterpret_cast(&v); +#pragma unroll + for (int i = 0; i < 16; ++i) { + // Logical element (k = kk, n = col + i) lands at packed[(col+i)*K+kk]. + dst[static_cast(col + i) * k + kk] = bytes[i]; + } + } +} + +} // namespace + +// Stable host launch symbol consumed by csrc/bindings.cpp. Dispatches on the +// caller-provided stream only; never allocates, synchronizes, or touches the +// default stream. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + auto* out_ptr = static_cast(out); + + // Iteration 5 (packing round): exact assigned shape (K=4096, N=2560) + // routes to the packed n-major B variant of the 64x128 kernel. The guard + // is exact (m >= 128, n == 2560, k == 4096) and sits BEFORE the generic + // 64x128 path; every other shape keeps its existing path (generic 64x128 + // for n % 128 == 0, 64x64 for n % 64 == 0, scalar fallback otherwise) and + // the identity-packed layout. + if (m >= 128 && n == kPackedBN && k == kPackedBK) { + const dim3 grid( + static_cast(n / kBlockN128), + static_cast((m + kBlockM128 - 1) / kBlockM128)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x128_packedb_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Native INT8 DUMMA m16n16k32 prefill path (2-D macro-tile 64x128, + // single-buffered 64-K K stage; pipeline-round control arm) for every + // large-M shape with compatible geometry (covers the assigned M=4096, + // N=2560, K=4096 shape: 2560 % 128 == 0). The 128-K guard is written + // explicitly so the dispatch condition is bit-identical to the accepted + // iteration-1 kernel even though the internal K stage is now 64 + // (single-buffered). + if (m >= 128 && (n % kBlockN128) == 0 && (k % 128) == 0) { + const dim3 grid( + static_cast(n / kBlockN128), + static_cast((m + kBlockM128 - 1) / kBlockM128)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x128_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Native INT8 DUMMA m16n16k32 prefill path (64x64 tile) for large-M shapes + // whose N is not a multiple of 128 but is a multiple of 64. + if (m >= 128 && (n % kBlockN) == 0 && (k % kStageK) == 0) { + const dim3 grid( + static_cast(n / kBlockN), + static_cast((m + kBlockM - 1) / kBlockM)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x64_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Scalar int8/int32 fallback for unmatched (m, n, k) and small-M API cases. + const int64_t total = static_cast(m) * n; + constexpr int kFallbackThreads = 256; + int64_t blocks = (total + kFallbackThreads - 1) / kFallbackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + const dim3 grid(static_cast(blocks)); + const dim3 block(kFallbackThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); +} + +// Stable host launch symbol consumed by csrc/bindings.cpp. For the exact +// assigned shape (K=4096, N=2560) it performs the one-time n-major B pack +// (logical [K, N] -> packed[n][k]) outside the timed region and outside +// Graph capture; every other (K, N) keeps the identity device-to-device +// copy, valid for every (K, N). The packed buffer size is k*n in both +// cases, so the allocation and graph-stable packed layout are unchanged. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kPackThreads = 256; + const int64_t weight_count = static_cast(k) * n; + if (weight_count > 0) { + if (k == kPackedBK && n == kPackedBN) { + const int64_t chunks = weight_count >> 4; + int64_t blocks = (chunks + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_nmajor_b_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + int64_t blocks = (weight_count + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_identity_i8_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count); + } + } + if (n > 0) { + int64_t blocks = (static_cast(n) + kPackThreads - 1) / + kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_identity_f32_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/shared_down_proj.hip new file mode 100644 index 00000000..f65afbf3 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/shared_down_proj.hip @@ -0,0 +1,631 @@ +// @@variant shape=hy3_tp4_shared_down_proj_m4096 commit=62971c1acacf5866bcbbadc375ee87334100b9c7 added=2026-08-23 +// median_us=214.8 p90_us=215.1 tops=59.97 bandwidth_gb_s=171 speedup=20.39 baseline_us=4381 +// source=hy3-tp4-m4096-dsh-test1-bfdf3002 +// MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). +// +// This file is owned by worker_3 and provides the two stable host launch +// symbols consumed by csrc/bindings.cpp: +// * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's +// HIP stream; no allocation, packing, +// autotuning, or synchronization. +// * launch_pack_w8a8_weight(...) - optional out-of-timed-region packing; +// round 14: n-major transpose pack +// (packed[n][k], K-contiguous per output +// column); bootstrap was an identity copy. +// +// Large-M prefill strategy (2-D macro-tile family, native INT8 DUMMA +// m16n16k32): +// * A single templated kernel computes a (kBlockM x kBlockN) output tile +// from (kBlockM/32) x (kBlockN/32) wavefronts, each wave owning one +// 32x32 quadrant (four 16x16 int32 accumulators) and keeping its +// accumulation resident across the whole K loop. +// * Three macro-tiles are instantiated - 64x64, 64x128 and 128x64 - and +// the active one is selected at compile time (kActiveTileIndex). All +// three share the cooperative A+B LDS staging pipeline with 16-byte +// padded LDS row strides (A stride kStageK+16, B stride kBlockN+16). +// The padding keeps every LDS row start 16-byte aligned (vectorized +// int4 cooperative staging) while breaking the 128-byte bank aliasing +// that produced ~8-way LDS conflicts in the bootstrap (PMC: +// 29,884,416 conflicts over 3,735,552 LDS instructions); the padded +// strides keep the accepted kernel's ~2-way fragment-load conflict +// profile. +// * Operand feed (iteration 2): direct du_load_matrix_sync global loads +// are rejected for this shape (they would amplify L2 requests 32-64x +// per unique element and degrade to byte-granular strided requests), +// so the staged path is kept and made latency-tolerant: K is staged in +// kPrefillStageK=32 chunks into a double-buffered LDS pair while each +// stage's global int4 loads are issued a full stage early into VGPRs +// (software prefetch). One barrier per stage, so the vmcnt(0) wait + +// ds_write of stage s+1 overlap the MMACs of stage s instead of +// serializing before them (iteration-1 single buffering exposed the +// global latency and produced 6,885,116 LDS wait instructions over +// 3,686,400 LDS instructions). Footprint 17,408 B/block -> 3 blocks/CU +// (24 resident wavefronts), an occupancy probe from the trusted set +// [2, 3, 4, 5, 6] that fits the 64 KiB LDS and the stage-alignment +// constraints (kPrefillGuardK % kPrefillStageK == 0, stage % 32 == 0, +// stage % 16 == 0 for int4 staging). +// * B fragment transport (iteration 5): the exact-source ISA shows the +// matrix_b fragment path as byte-granular - 16 ds_read_u8 at k-row +// stride kBStride=80, ~40 VALU of byte assembly (v_or3_b32, +// v_lshlrev_b32, v_or_b32_sdwa) and ~5 s_waitcnt lgkmcnt(N) groups per +// wave per stage feeding only 4 v_mmac_i32_16x16x32_i8 - while the +// matrix_a path is already 2 vectorized ds_read2_b32 (8 contiguous +// bytes per lane). du_load_matrix_sync(matrix_b, row_major) fetches +// p[(col+i)*ldm + row] (8 bytes spread over 8 k-rows), so B is now +// staged into LDS in a transposed n-major layout (b_tile[n][k], row +// stride kBTStride = kStageK + 8 = 40, 8-byte-aligned K runs) and +// consumed through col_major matrix_b fragments, whose loader fetches +// p[row*ldm + col + i] = 8 contiguous bytes per lane and lowers to the +// same vectorized ds_read2_b32 path as A (verified against +// /opt/dtk/include/du_mma.hpp: x[i] maps to the same B[k0+col+i][n0+row] +// elements in both layouts and du_mma_sync passes the fragment +// registers through LayoutB-agnostically, so the v_mmac inputs - and +// therefore the exact int32 accumulation - are bit-identical). Staging +// cost moves from 1 ds_write_b128 to 16 bank-spread ds_write_b8 per +// staging thread: each int4 covers 16 consecutive n columns and all +// sixteen bytes are scattered into the transposed tile at +// (16*seg+j)*kBTStride + kk, so every n column is written every stage +// (~4-way write conflicts vs ~2-way before, but B fragment reads drop +// 16 -> 2 instructions per wave per stage with no byte assembly, ~3x +// fewer LDS instructions and ~5x fewer LDS wait groups per stage). LDS +// footprint unchanged at 17,408 B/block. +// * Pipeline depth (iteration 7): kPrefillStageK raised 32 -> 64. The +// probe keeps the double-buffered structure and the transposed-B +// staging from iteration 5 but halves the barrier count (12 -> 6 stage +// barriers + prologue) and doubles the in-flight global bytes. LDS +// footprint moves 17,408 -> 29,696 B/block. The PMC launch record +// showed the code object at 74 VGPRs (arch_vgpr 80), so residency is +// actually 1 block/CU / 8 waves - the __launch_bounds__(512, 2) pin +// was ignored (2 blocks/CU would need <= 64 VGPRs), and stage 128 +// (iteration 9) and 3 blocks/CU at 64x64 (iteration 10) both regressed. +// * Intra-wave ILP (iteration 11): the 2-trip kk loop is fully unrolled +// into two independent fragment register sets (all eight ds_read2 of a +// stage issue before the MMACs), so the second step's LDS latency +// overlaps the first step's swizzle + MMACs instead of serializing +// behind them. ~8 extra VGPRs (74 -> ~82) keep residency at the +// measured 1 block/CU / 8 waves; LDS, depth, barriers, and the launch +// config are unchanged. +// * Vectorized B staging (iteration 14): the B transpose moves out of the +// timed kernel into launch_pack_w8a8_weight, which now emits an n-major +// packed weight (packed[n][k]; the API contract allows any contiguous +// opaque packed layout and the pack runs once per weight tensor outside +// the timed region and outside Graph capture). The timed kernel therefore +// stages each n row's k-contiguous run with two 8-byte ds_write_b64 per +// thread (b_tile[n][kk] addresses and the col_major fragment loads are +// untouched, so every LDS byte and every MMAC input is identical), and +// the 16 ds_write_b8 + ~12 VALU byte extracts of stage_b_transposed per +// staging thread disappear from the per-stage critical path. The removed +// byte-extract register pressure (v25..v38 in the accepted ISA) is also +// what kept the code object at 69 VGPRs (> 64, so the +// __launch_bounds__(512, 2) minBlocks pin never held); with it gone the +// first clean 2 blocks/CU probe at the constant 128x64 tile is expected +// to materialize (LDS 2 * 29,696 = 59,392 B <= 64 KiB). +// * Fused direct fragment -> scale -> bf16 epilogue; no separate epilogue +// kernel, no split-K. +// * All other (m, n, k), including M < 128: scalar int8/int32 fallback. +// +// Header order is fixed by the control plane: hip_runtime first (du_mma.h is +// not self-contained before it), hip_bfloat16 second, du_mma.h last. + +#include +#include +#include + +#include + +namespace { + +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; +constexpr int kWaveSize = 64; + +// Staging stage length for the prefill pipeline. 64 = 2x the DUMMA K unit +// (iteration 7): with the 128x64 macro-tile the double-buffered footprint +// is 2*(128*80 + 64*72) = 29,696 B (A stride 80, transposed-B stride 72), +// which fits the 64 KiB LDS either way. Actual residency is 1 block/CU / +// 8 waves, capped by the 74-VGPR code object (2 blocks/CU at 512 threads +// would need <= 64 VGPRs; the __launch_bounds__(512, 2) minBlocks pin was +// ignored). Stage 64 halves the per-block barrier count vs 32 (12 -> 6 +// stage barriers) and doubles the software-prefetch distance (one full +// 64-k stage of global loads in flight instead of 32 k). Padded row strides +// stay aligned: kAStride 80 % 16 == 0 for int4 staging, kBTStride 72 % 8 +// == 0 for the col_major 8-byte fragment reads; the launch guard +// (k % 128 == 0) still implies k % kPrefillStageK == 0. +constexpr int kPrefillStageK = 64; + +// Launch-guard divisibility is kept at 128 (unchanged from iteration 1), +// which implies k % kPrefillStageK == 0 for every dispatched shape. +constexpr int kPrefillGuardK = 128; + +// Compile-time selection of the active 2-D macro-tile: +// 0 -> 64x64, 1 -> 64x128, 2 -> 128x64. +// All three instantiations are compiled so the family stays measurable; only +// the active one is launched inside the timed region. +constexpr int kActiveTileIndex = 2; +constexpr int kActiveBlockM = kActiveTileIndex == 2 ? 128 : 64; +constexpr int kActiveBlockN = kActiveTileIndex == 1 ? 128 : 64; + +// LDS padding in bytes added to each staged row. 16 keeps every row start +// 16-byte aligned for int4 staging while shifting fragment rows off the same +// 32-bank phase (128-byte strides alias every row onto identical banks). +constexpr int kLdsPad = 16; + +// LDS row stride of the transposed (n-major) B tile: each row holds one +// output column n with kStageK contiguous K bytes. kLdsPadB = 8 keeps every +// row start 8-byte aligned so the col_major matrix_b fragment loader can +// lower its 8 contiguous bytes per lane to a single vectorized ds_read2_b32 +// (like the A path) and so the round-14 staging can land each n row's +// k-contiguous 16-byte run as two 8-byte ds_write_b64 (72 % 8 == 0; note +// 72 % 16 == 8, so a single ds_write_b128 would be 16-byte misaligned for +// odd n rows - the explicit two-half store keeps the lowering legal). +constexpr int kLdsPadB = 8; +constexpr int kBTStride = kPrefillStageK + kLdsPadB; + +using namespace du::dumma; + +// gfx928 int8 m16n16k32 accumulator ownership, established against +// du_store_matrix_sync: lane % 16 selects the row, lane / 16 selects the +// column mod 4, and x[i] selects columns separated by four. +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = + static_cast(frag.x[i]) * xs * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// B staging (iteration 14): launch_pack_w8a8_weight emits the weight packed +// n-major (packed[n][k], one contiguous K run per output column), so the +// timed kernel stages each 16-byte int4 run straight into the transposed +// LDS tile as two 8-byte ds_write_b64 (the B_tile row address +// n_local * kBTStride + kchunk is 8-byte aligned but not 16-byte aligned, +// so a single ds_write_b128 would be illegal for odd n_local; the explicit +// two-half store keeps the lowering legal). This replaces the old 16 +// byte-granular ds_write_b8 + ~12 VALU byte extracts per staging thread of +// stage_b_transposed, whose (16*seg+j)*kBTStride + kk scatter wrote every n +// column every stage. The LDS tile layout, every staged byte, the col_major +// fragment reads, and the kk/acc ordering are unchanged, so the int32 +// accumulation is bit-identical. + +// --------------------------------------------------------------------------- +// Large-M prefill path. One block computes a kBlockM x kBlockN output tile +// with (kBlockM/32) x (kBlockN/32) wavefronts (one 32x32 quadrant per wave). +// A[kBlockM, K] and B[K, kBlockN] are cooperatively staged into a +// double-buffered, bank-padded LDS in kStageK chunks with VGPR prefetch: +// the global int4 loads for stage s+1 are issued before the MMAC sequence +// of stage s, and the alternate LDS buffer is written right after it, so +// the vmcnt(0) wait and ds_write overlap compute. One barrier per stage. +// A stays row-major in LDS (stride 80); B is staged transposed (n-major, +// stride 72, K-contiguous) and loaded through col_major matrix_b fragments +// so both operand fragment paths use vectorized 8-byte LDS reads (iteration +// 5; the iteration-2 B path was 16 byte-granular ds_read_u8 + ~40 VALU of +// assembly per wave per stage). Since iteration 14 the B transpose happens +// at pack time (launch_pack_w8a8_weight emits packed[n][k]), so the staged +// B runs land with two 8-byte ds_write_b64 per thread instead of the 16 +// ds_write_b8 + ~12 VALU byte extracts of stage_b_transposed. Each wave +// keeps four 16x16 int32 +// accumulators resident over the whole K loop, then a fused scale/bf16 +// store. The kk ordering (k0 outer, kk inner step kTileK) and the +// acc00/acc01/acc10/acc11 update order are unchanged from iteration 1, and +// the col_major loader places the same B[k][n] values in the same fragment +// registers as row_major, so int32 accumulation is bit-identical. +// kPrefillStageK = 64 (iteration 7): six 64-k stages (7 barriers per block +// including the prologue), one full stage of global loads in flight, and +// measured residency 1 block/CU / 8 waves (code object 69 VGPRs at +// iteration 12; the __launch_bounds__(512, 2) minBlocks pin is a floor, not +// a guarantee - round 14 expects the code object to drop to <= 64 VGPRs +// once the B byte-extract temps are gone, which would finally make 2 +// blocks/CU / 16 waves / 4 waves per SIMD satisfiable at the constant +// 128x64 tile). +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kBlockM * kBlockN / 1024 * kWaveSize, 2) void +w8a8_dumma_prefill_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWavesN = kBlockN / 32; + constexpr int kAStride = kStageK + kLdsPad; + static_assert(kAStride % 16 == 0, + "A LDS row stride must stay 16-byte aligned"); + static_assert(kBTStride % 8 == 0, + "transposed-B LDS row stride must stay 8-byte aligned"); + static_assert(kStageK % kTileK == 0, + "K stage must be a multiple of the DUMMA K unit"); + static_assert(kStageK % sizeof(int4) == 0, + "K stage must keep int4 staging aligned"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWavesN; + const int wave_col = wave - wave_row * kWavesN; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + // Transposed (n-major) B tile: element B[k0+kk][n] lives at + // b_tile[n][kk], so col_major fragment loads read 8 contiguous K bytes + // per lane (vectorized ds_read2_b32). The staged runs are written n-major + // (packed[n][k]) as two 8-byte ds_write_b64 per thread. + __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBTStride]; + + DUFragment + a_frag0, a_frag1, a_frag0_1, a_frag1_1; + DUFragment + b_frag0, b_frag1, b_frag0_1, b_frag1_1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + constexpr int kAInt4PerRow = kStageK / sizeof(int4); + constexpr int kAInt4 = kBlockM * kAInt4PerRow; + // B is packed n-major by launch_pack_w8a8_weight (packed[n][k], one + // contiguous K run per output column), so a stage row holds kStageK + // contiguous bytes per column and stages as kChunksPerBRow 16-byte int4 + // runs (two 8-byte ds_write_b64 each). kBlockN * kChunksPerBRow is 256 for + // 64x64 / 128x64 and 512 for 64x128, so one chunk per thread fits every + // instantiation of the 512-thread block (no loop, matching the flat + // staging issue order the accepted kernel depends on). + constexpr int kChunksPerBRow = kStageK / sizeof(int4); + constexpr int kBInt4 = kBlockN * kChunksPerBRow; + const int kStages = k / kStageK; + + // Prologue: prefetch stage 0 into VGPRs and land it in buffer 0. Every + // source address is 16-byte aligned (k % 16 == 0 and n % 16 == 0 by the + // launch guard) and every destination row start is 16-byte aligned. + int4 a_pref{0, 0, 0, 0}; + int4 b_pref{0, 0, 0, 0}; + { + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + const int global_row = m0 + row; + if (global_row < m) { + a_pref = *reinterpret_cast( + x_q + global_row * k + seg * sizeof(int4)); + } + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + b_pref = *reinterpret_cast( + weight + (n0 + n_local) * k + kchunk); + } + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + reinterpret_cast(a_tile[0])[row * (kAStride / sizeof(int4)) + + seg] = a_pref; + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + const int64_t* src64 = reinterpret_cast(&b_pref); + int64_t* dst64 = reinterpret_cast( + b_tile[0] + n_local * kBTStride + kchunk); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + __syncthreads(); + } + + for (int s = 0; s < kStages; ++s) { + const int k0 = s * kStageK; + const int buf = s & 1; + + // Issue the global loads for stage s+1 now so their latency hides + // behind the fragment loads and MMACs below. Out-of-range A rows are + // zero-filled (tail-M guard, matching iteration 1). + if (s + 1 < kStages) { + const int k1 = k0 + kStageK; + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + const int global_row = m0 + row; + if (global_row < m) { + a_pref = *reinterpret_cast( + x_q + global_row * k + k1 + seg * sizeof(int4)); + } else { + a_pref = int4{0, 0, 0, 0}; + } + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + b_pref = *reinterpret_cast( + weight + (n0 + n_local) * k + k1 + kchunk); + } + } + + // Consume stage s from buffer buf: four fragment loads and four MMACs + // per wave per kTileK step. kStageK = 64 = 2 kTileK steps, now fully + // unrolled (round 11) into two independent fragment register sets. + // The accepted ISA (exact-source kk loop) shows each rolled trip as + // 4 ds_read2_b32 -> 4 s_waitcnt lgkmcnt(3..0) -> ~24 VALU byte-swizzle + // (the du_load_matrix_sync lowering) -> 4 v_mmac_i32_16x16x32_i8, with + // VGPR reuse forcing the kk=32 loads to wait for the kk=0 MMACs, so + // each stage serializes two load-wait-swizzle-MMAC chains per wave. + // With both steps' eight ds_read2 issued up front into distinct + // registers, the second step's LDS latency overlaps the first step's + // swizzle + MMACs and the per-stage lgkmcnt wait window collapses. + // The kk ordering (k0 outer, kk = 0 then kTileK) and the + // acc00/acc01/acc10/acc11 update order are unchanged, so int32 + // accumulation is bit-identical. The second fragment set landed at 69 + // code-object VGPRs (measured iteration 12, fewer than the round-11 + // prediction); at 512 threads that keeps residency at 1 block/CU / 8 + // waves (2 blocks/CU would need <= 64 VGPRs), so occupancy, LDS + // (29,696 B), stage depth (64 k), barrier count, and launch config are + // pinned exactly where the accepted iteration-7 kernel runs. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + static_assert(kStageK == 2 * kTileK, + "round-11 unrolled body assumes kStageK == 2 * kTileK"); + du_load_matrix_sync( + a_frag0, a_tile[buf] + local_row * kAStride, kAStride); + du_load_matrix_sync( + a_frag1, a_tile[buf] + (local_row + 16) * kAStride, kAStride); + du_load_matrix_sync( + b_frag0, b_tile[buf] + local_col * kBTStride, kBTStride); + du_load_matrix_sync( + b_frag1, b_tile[buf] + (local_col + 16) * kBTStride, kBTStride); + // Second 32-k step (kk = kTileK): independent registers, issued while + // the first step's data is still in flight. + du_load_matrix_sync( + a_frag0_1, a_tile[buf] + local_row * kAStride + kTileK, kAStride); + du_load_matrix_sync( + a_frag1_1, a_tile[buf] + (local_row + 16) * kAStride + kTileK, + kAStride); + // col_major: element B[k0+kk+k][n] at b_tile[n][kk+k], so the loader's + // p[row*ldm + col + i] reads 8 contiguous bytes per lane. + du_load_matrix_sync( + b_frag0_1, b_tile[buf] + local_col * kBTStride + kTileK, kBTStride); + du_load_matrix_sync( + b_frag1_1, b_tile[buf] + (local_col + 16) * kBTStride + kTileK, + kBTStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc00, a_frag0_1, b_frag0_1, acc00); + du_mma_sync(acc01, a_frag0_1, b_frag1_1, acc01); + du_mma_sync(acc10, a_frag1_1, b_frag0_1, acc10); + du_mma_sync(acc11, a_frag1_1, b_frag1_1, acc11); + + // Land the prefetched stage s+1 in the alternate buffer. The compiler + // inserts the vmcnt(0) wait here via the data dependency, so the wait + // and the ds_write overlap the MMACs above instead of stalling before + // them. + if (s + 1 < kStages) { + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + reinterpret_cast(a_tile[buf ^ 1])[row * + (kAStride / sizeof(int4)) + + seg] = a_pref; + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + const int64_t* src64 = reinterpret_cast(&b_pref); + int64_t* dst64 = reinterpret_cast( + b_tile[buf ^ 1] + n_local * kBTStride + kchunk); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + } + // One barrier per stage: makes the alternate-buffer writes visible to + // the next stage's reads and retires this stage's reads before the + // parity-flipped buffer is overwritten two stages later. + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + 16, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + 16, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + 16, base_col + 16, lane); +} + +// --------------------------------------------------------------------------- +// Generic scalar int8 x int8 -> int32 fallback for unmatched (m, n, k) and +// small-M API cases. One thread per output element, coalesced along N. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_scalar_fallback_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int total = m * n; + if (linear >= total) { + return; + } + const int row = linear / n; + const int col = linear - row * n; + const int8_t* a_row = x_q + row * k; + // The packed weight is n-major since round 14 (packed[n][k]), so the + // column stride is k; the products and their int32 summation order are + // unchanged, keeping the fallback bit-identical to the pre-round-14 code. + const int8_t* b_col = weight + col * k; + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[kk]); + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// n-major transpose pack (round 14). pack_weight runs once per weight tensor +// outside the timed region and outside CUDA/HIP Graph capture, so a +// byte-granular kernel is fine: packed[col * k + row] = raw[row * n + col] +// gives every output column a contiguous K run, which the timed GEMM stages +// with two vectorized ds_write_b64 per thread instead of the 16 +// byte-granular ds_write_b8 + ~12 VALU byte extracts per staging thread of +// the identity-layout path. The layout is opaque to the API contract. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_transpose_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int k, + int n) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(k) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + dst[static_cast(col) * k + row] = src[linear]; +} + +__global__ __launch_bounds__(256) void w8a8_identity_copy_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t numel) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < numel) { + dst[linear] = src[linear]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Host launch symbols (extern "C", consumed by csrc/bindings.cpp). +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + auto* out_bf16 = reinterpret_cast<__hip_bfloat16*>(out); + + // Native INT8 DUMMA m16n16k32 tiled path for large-M prefill. All three + // macro-tile instantiations share the exact (m, n, k) guard; the active + // one is fixed at compile time, so the timed region never tunes. The + // guard divisibility (k % 128 == 0) is unchanged from iteration 1 and + // implies k % kPrefillStageK == 0 for the pipelined kernel. + if (m >= 128 && (k % kPrefillGuardK) == 0 && (n % kActiveBlockN) == 0) { + const dim3 grid( + static_cast(n / kActiveBlockN), + static_cast((m + kActiveBlockM - 1) / kActiveBlockM)); + const dim3 block(static_cast( + kActiveBlockM * kActiveBlockN / 1024 * kWaveSize)); + if (kActiveTileIndex == 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 64, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else if (kActiveTileIndex == 1) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 128, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<128, 64, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + return; + } + + // Generic scalar fallback for every unmatched (m, n, k) and M < 128. + const int total = m * n; + constexpr int kFallbackThreads = 256; + const dim3 grid(static_cast( + (total + kFallbackThreads - 1) / kFallbackThreads)); + const dim3 block(static_cast(kFallbackThreads)); + hipLaunchKernelGGL( + w8a8_scalar_fallback_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kCopyThreads = 256; + // Round 14: n-major transpose pack (packed[n][k], K-contiguous per output + // column). The API contract (int8_w8a8_gemm_api.py) allows any contiguous + // opaque packed layout; the timed GEMM consumes exactly this layout, so + // its B staging drops from 16 byte-granular ds_write_b8 + ~12 VALU byte + // extracts per staging thread to two 8-byte ds_write_b64. The scale copy + // below is unchanged (N-length, order-independent). + const int64_t weight_numel = static_cast(k) * n; + const dim3 weight_grid(static_cast( + (weight_numel + kCopyThreads - 1) / kCopyThreads)); + hipLaunchKernelGGL( + w8a8_transpose_i8_kernel, + weight_grid, dim3(kCopyThreads), 0, stream, + raw_weight, packed_weight, k, n); + + const dim3 scale_grid(static_cast( + (n + kCopyThreads - 1) / kCopyThreads)); + hipLaunchKernelGGL( + w8a8_identity_copy_f32_kernel, + scale_grid, dim3(kCopyThreads), 0, stream, + weight_scale, packed_weight_scale, static_cast(n)); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/shared_gate_up_proj.hip new file mode 100644 index 00000000..af3b0397 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/shared_gate_up_proj.hip @@ -0,0 +1,552 @@ +// @@variant shape=hy3_tp4_shared_gate_up_proj_m4096 commit=5a1e581a09dc66853a94d300b508d6db6bd6731e added=2026-08-23 +// median_us=220.7 p90_us=221.5 tops=116.8 bandwidth_gb_s=118.9 speedup=30.26 baseline_us=6678 +// source=hy3-tp4-m4096-dsh-test1-bfdf3002 +// INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. +// +// Worker 2 assigned shape: hy3_tp4_shared_gate_up_proj_m4096 +// (M=4096, N=768, K=4096). +// +// DUMMA 2-D macro-tile family (iteration 1, refined iteration 3): +// * Large-prefill path (M >= 128, N % 64 == 0, K % 128 == 0): native INT8 +// DUMMA m16n16k32 tiled kernel with int32 accumulation. One templated +// kernel covers the 64x64 (4 waves), 64x128 (8 waves) and 128x64 +// (8 waves) output tiles; every wavefront owns a 32x32 quadrant of four +// m16n16k32 DUMMA accumulators, so per-wave DUMMA/LDS work is identical +// across the family and only grid, LDS footprint and occupancy differ. +// Operands are staged cooperatively with 16-byte vector loads into +// bank-skewed LDS (A row stride kStageK+16, B row stride kBlockN+16): +// the +16-byte skew keeps every int4 store aligned while breaking the +// 16-way fragment bank conflicts of the power-of-two strides. No +// split-K, no raw asm. +// * The exact assigned shape (4096, 768, 4096) is dispatched to the +// 64x128 tile: 6x64 = 384 blocks, 8 waves/block, grid 384. Iteration 3 +// switches this shape to a double-buffered (software-pipelined) K loop +// with K stage 64: staging of stage s+1 is issued into registers before +// the MMA compute of stage s and committed to the second LDS buffer after +// it, so the global-load round trips overlap the LDS-wait-bound compute +// phase instead of serializing at the top of every stage (the iteration-1 +// ISA shows each global_load_dwordx4 followed by an immediate +// s_waitcnt vmcnt(0) before its ds_write). Buffer footprint is 2 x +// (64x80 A + 64x144 B) = 28,672 B/block -> still 2 blocks/CU -> 16 waves/CU, +// i.e. the pipeline is added without sacrificing the 16-wave occupancy +// that the 64x128 tile established. Barrier count stays ~64/block (one +// prologue + one per stage vs. two per stage in the single-buffer loop). +// All other routed shapes keep the proven single-buffer K stage 128 loop. +// Iteration 10 (exact shape only) fixes the dominant LDS cost of this +// path: the DUMMA int8 B-fragment loader reads 8 k-rows per lane, and +// with any 16-byte-aligned [K, N] row stride the four k-octets alias onto +// one LDS bank group (16-way conflict, eight ds_read_u8 per fragment; +// PMC: 16.5M bank conflicts, 23.15M LDS waits on the stage-64 winner). +// pack_weight now stores the exact shape's weight transposed [N, K], the +// double-buffered path stages B in [N, K] LDS order and loads B fragments +// with the col_major loader, so each B fragment becomes one ds_read2_b32 +// of 8 consecutive k-values (~4-way floor). LDS grows 28,672 -> 30,720 +// B/block, still 2 blocks/CU (16 waves); same MMAC sequence and operand +// values -> exact int32 accumulation unchanged; grid and barriers +// unchanged. Generic fallback and all other shapes are untouched. +// Iteration 13 (exact shape only) removes the residual per-byte fragment +// reassembly: du_mma_sync passes each int8 fragment to the v_mmac builtin +// as one packed 8-byte operand, and du_load_matrix_sync assigns x[0..7] = +// 8 separate byte loads, so the compiler emits a 5-op mask/OR chain +// (v_and 0xff00/0xff0000/0xff000000 + v_or_b32_sdwa + v_or3_b32) for every +// loaded dword (~50-66 VALU per wave per stage, ~13M of the 16.08M VALU +// per replay; PMC of the accepted source). The exact-shape k-chunk now +// writes the same 8 consecutive bytes directly into the fragment storage +// (same element-to-slot mapping, same byte order), so the v_mmac operand +// bit patterns and the exact int32 accumulation are unchanged while the +// reassembly VALU disappears and the LDS->MMA dependency chain shortens. +// LDS reads stay one ds_read2_b32 per fragment at the 4-way floor; grid, +// occupancy, barriers and the generic paths are untouched. +// * Everything else (small M, unmatched shapes): simple scalar int8/int32 +// fallback with exact int32 accumulation. +// * pack_weight: identity copy for every shape except the exact assigned +// shape (K=4096, N=768), which is stored transposed [N, K]. The packed +// layout is opaque per the API contract (prepare_weight may return any +// contiguous layout); the correctness reference computes from raw +// weights, and pack runs outside the timed/CUDA-Graph region. +// +// Mathematical contract (exact int32 dot before float scaling): +// out[m, n] = bf16(int32_dot(a[m, :], b[:, n]) * x_scale[m] * weight_scale[n]) +// +// Header order is fixed by the control plane for this DTK: +// hip_runtime.h -> hip_bfloat16.h -> du_mma.h + +#include +#include +#include + +#include + +namespace { + +// gfx928 INT8 DUMMA primitive: m16n16k32, int8 x int8 -> int32. +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; + +// gfx928 wavefront is 64 lanes; every blockDim must be a multiple of 64. +constexpr int kWaveSize = 64; + +// Large-prefill dispatch geometry. Each macro-tile config uses a 32x32 +// quadrant per wave, so waves per block = (BlockM/32) * (BlockN/32). +constexpr int kPrefillMinM = 128; +constexpr int kPrefillStageK = 128; +constexpr int kDefaultBlockM = 64; +constexpr int kDefaultBlockN = 64; + +constexpr int kScalarThreads = 256; + +using namespace du::dumma; + +// Verified gfx928 INT8 DUMMA accumulator ownership: +// row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> columns +// col_mod4 + 4*i. Scale by x_scale[row] and weight_scale[col] and store +// bf16 directly from the fragment (no accumulator LDS round trip). +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Direct 8-byte LDS fragment load for the double-buffered exact-shape path. +// du_load_matrix_sync's int8 loaders assign x[0..7] = 8 consecutive bytes at +// (lane & 15) * ldm + ((lane >> 4) << 3) for both matrix_a row_major and +// matrix_b col_major, and du_mma_sync passes the fragment to the v_mmac +// builtin as one packed 8-byte operand, so the compiler emits a per-byte +// mask/OR reassembly chain (v_and 0xff00/0xff0000/0xff000000 + v_or_b32_sdwa +// + v_or3_b32) for every loaded dword. Writing the same 8 bytes directly into +// the fragment storage keeps the operand bit pattern identical (exact int32 +// accumulation unchanged) and lets the compiler feed the ds_read2_b32 pair +// straight to the v_mmac (still one ds_read2_b32 per fragment, 4-way floor). +__device__ __forceinline__ void load_fragment8( + DUFragmentBase& frag, + const signed char* __restrict__ base, + int stride, + int lane) { + const int off = (lane & 15) * stride + ((lane >> 4) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// Large-prefill path: one block computes a kBlockM x kBlockN output tile. +// kWaveRows x kWaveCols wavefronts each own a 32x32 quadrant (four 16x16 +// DUMMA accumulators), while the block cooperatively stages A[kBlockM, kStageK] +// and B[kStageK, kBlockN] in bank-skewed LDS. Tail M rows are zero-filled on +// load and masked on store, so any M >= 128 is supported. +// +// kBuffers == 1: single-buffered K loop (K stage 128, two barriers per stage): +// stage s+1's global loads cannot start until the stage-s barrier, so each +// stage serializes one global-load round trip (global_load -> vmcnt(0) -> +// ds_write -> lgkmcnt(0) -> barrier) in front of the compute. +// kBuffers == 2: software-pipelined K loop (K stage 64, one barrier per stage): +// stage s+1's loads are issued into registers before the stage-s MMA compute +// and committed to the other LDS buffer after it, overlapping the global +// round trip with the LDS-wait-bound compute. Requires at most one int4 per +// thread per operand (static_asserted), which holds for the exact shape. +template +__global__ __launch_bounds__((kBlockM / 32) * (kBlockN / 32) * kWaveSize) void +w8a8_gemm_prefill_tiled_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWaveRows = kBlockM / 32; + constexpr int kWaveCols = kBlockN / 32; + constexpr int kThreads = kWaveRows * kWaveCols * kWaveSize; + // 16-byte bank skew: keeps int4 staging stores aligned while spreading the + // fragment rows across distinct LDS banks (4-way floor for 8-byte/lane + // reads instead of the 16-way aliasing of power-of-two strides). + constexpr int kAStride = kStageK + sizeof(int4); + // Double-buffered path (exact shape) stages B transposed [N, K] so the + // DUMMA B fragments read 8 consecutive k-values per lane (one ds_read2_b32, + // ~4-way) instead of 8 strided k-rows (eight ds_read_u8, 16-way: with any + // 16-byte-aligned row stride S, B k-rows 8 apart alias onto one bank group + // because (S/4)*8*g == 0 mod 32 for all g). The generic single-buffer path + // keeps [K, N] B and the row-major B loader. + constexpr int kBStride = (kBuffers == 2) ? (kStageK + sizeof(int4)) + : (kBlockN + sizeof(int4)); + constexpr int kBRows = (kBuffers == 2) ? kBlockN : kStageK; + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = + (kBuffers == 2) ? (kStageK / sizeof(int4)) : (kBlockN / sizeof(int4)); + constexpr int kBVectors = kBRows * kBVectorsPerRow; + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWaveCols; + const int wave_col = wave - wave_row * kWaveCols; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBuffers][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBuffers][kBRows * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + if constexpr (kBuffers == 2) { + static_assert(kAVectors <= kThreads && kBVectors <= kThreads, + "double-buffered path stages at most one int4 per thread"); + // The exact shape's packed weight is [N, K] (transposed by + // launch_pack_w8a8_weight), so B fragments load 8 consecutive k-values + // per lane from the [N, K] LDS tile via the col_major loader; the + // element-to-slot mapping (slot i = B[k = kk + 8*g + i][n = col + row]) + // is identical to the row-major loader on a [K, N] tile, so the v_mmac + // operand values and the exact int32 accumulation are unchanged. + DUFragment + b_frag_t0, b_frag_t1; + // Prologue: stage K tile 0 into buffer 0, then one barrier. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + // B is packed [N, K] for the exact shape: each thread stages 16 + // consecutive k-values of one n row (n-stride in global is k). + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + n_local * kBStride)[v] = + *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + const int n_stages = k / kStageK; + for (int s = 0; s < n_stages; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + const bool has_next = (s + 1) < n_stages; + // Issue the stage s+1 global loads now; the data is consumed by the + // ds_write after the compute, so the vmcnt wait lands after the MMA + // loop instead of stalling the front of the stage. + int4 a_reg{0, 0, 0, 0}; + int4 b_reg{0, 0, 0, 0}; + if (has_next) { + const int k1 = (s + 1) * kStageK; + if (tid < kAVectors) { + const int local_row = tid / kAVectorsPerRow; + const int v = tid - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + a_reg = global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k1 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + const int n_local = tid / kBVectorsPerRow; + const int v = tid - n_local * kBVectorsPerRow; + b_reg = *reinterpret_cast( + weight + (n0 + n_local) * k + k1 + + v * static_cast(sizeof(int4))); + } + + // Compute stage s from the buffer staged last iteration. + // Iteration 13: load the four int8 fragments as raw 8-byte LDS reads + // into the fragment storage instead of du_load_matrix_sync's per-byte + // assignments (same bytes, same element-to-slot mapping), so the + // compiler feeds the ds_read2_b32 pair straight to the v_mmac and the + // per-dword mask/OR reassembly VALU disappears from the steady state. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + load_fragment8( + a_frag0, a_tile[cur] + local_row * kAStride + kk, kAStride, + lane); + load_fragment8( + a_frag1, a_tile[cur] + (local_row + kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8( + b_frag_t0, b_tile[cur] + local_col * kBStride + kk, kBStride, + lane); + load_fragment8( + b_frag_t1, b_tile[cur] + (local_col + kTileN) * kBStride + kk, + kBStride, lane); + du_mma_sync(acc00, a_frag0, b_frag_t0, acc00); + du_mma_sync(acc01, a_frag0, b_frag_t1, acc01); + du_mma_sync(acc10, a_frag1, b_frag_t0, acc10); + du_mma_sync(acc11, a_frag1, b_frag_t1, acc11); + } + + // Commit the prefetched stage s+1 into the other LDS buffer, then one + // barrier: it orders both this iteration's compute reads of buffer cur + // (against the next-next prefetch, which reuses cur) and the prefetch + // writes of buffer nxt (against the next iteration's compute reads). + if (has_next) { + if (tid < kAVectors) { + const int local_row = tid / kAVectorsPerRow; + const int v = tid - local_row * kAVectorsPerRow; + reinterpret_cast(a_tile[nxt] + local_row * kAStride)[v] = + a_reg; + } + const int n_local = tid / kBVectorsPerRow; + const int v = tid - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[nxt] + n_local * kBStride)[v] = b_reg; + } + __syncthreads(); + } + } else { + for (int k0 = 0; k0 < k; k0 += kStageK) { + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int kk = vec / kBVectorsPerRow; + const int v = vec - kk * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + kk * kBStride)[v] = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + du_load_matrix_sync( + a_frag0, a_tile[0] + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile[0] + (local_row + kTileM) * kAStride + kk, + kAStride); + du_load_matrix_sync( + b_frag0, b_tile[0] + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile[0] + kk * kBStride + local_col + kTileN, + kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Generic scalar fallback: one output element per grid-stride iteration with +// exact int32 accumulation. Correct for any (m, n, k), including the small-M +// API shapes (M=2/16) and any unmatched geometry. +__global__ __launch_bounds__(kScalarThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + int32_t acc = 0; + const int8_t* a_row = x_q + static_cast(row) * k; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast( + weight[static_cast(kk) * n + col]); + } + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// pack_weight. The packed layout is opaque per the API contract; for the +// exact assigned shape (K=4096, N=768) the weight is stored transposed +// [N, K] (n-major) so the double-buffered GEMM stages B in [N, K] LDS order +// and the DUMMA B fragments read 8 consecutive k-values per lane. Every +// other shape keeps the identity [K, N] copy (generic fallback unchanged). +__global__ __launch_bounds__(kScalarThreads) void +w8a8_pack_kernel( + const int8_t* __restrict__ raw_weight, + const float* __restrict__ weight_scale, + int8_t* __restrict__ packed_weight, + float* __restrict__ packed_weight_scale, + int64_t weight_elems, + int n, + int k) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (k == 4096 && n == 768) { + // Transpose [K, N] -> [N, K]: packed[n * k + kk] = raw[kk * n + n]. + if (idx < weight_elems) { + const int64_t n_idx = idx / static_cast(k); + const int64_t k_idx = idx - n_idx * static_cast(k); + packed_weight[idx] = raw_weight[k_idx * n + n_idx]; + } + } else if (idx < weight_elems) { + packed_weight[idx] = raw_weight[idx]; + } + if (idx < n) { + packed_weight_scale[idx] = weight_scale[idx]; + } +} + +} // namespace + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // No split-K: the provided workspace is not touched. The timed operator + // performs no allocation, synchronization, packing, or default-stream + // launch; only the caller-provided out is written. + (void)workspace; + (void)workspace_bytes; + auto* out_bf16 = static_cast(out); + + if (m >= kPrefillMinM && (n % kDefaultBlockN == 0) && + (k % kPrefillStageK == 0)) { + if (m == 4096 && n == 768 && k == 4096) { + // Exact assigned shape: 64x128 macro-tile, 8 waves, grid 6x64 = 384. + // Double-buffered K loop, K stage 64, B staged transposed [N, K] from + // the transposed packed weight (col_major B fragments): 30,720 B LDS + // -> 2 blocks/CU -> 16 waves/CU (occupancy preserved), global staging + // overlapped with the MMA compute, B fragment loads cut from 8 strided + // ds_read_u8 (16-way conflicts) to one ds_read2_b32 (~4-way) each. + const dim3 grid(768 / 128, 4096 / 64); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel<64, 128, 64, 2>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (n % 128 == 0) { + // 64x128 tile, 8 waves; A rows are re-read by only 6 N-blocks. + const dim3 grid(n / 128, (m + 63) / 64); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel<64, 128, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (m % 128 == 0) { + // 128x64 tile, 8 waves. + const dim3 grid(n / 64, m / 128); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel<128, 64, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + { + // 64x64 tile, 4 waves (generic default). + const dim3 grid(n / 64, (m + 63) / 64); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel<64, 64, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + } + + const int64_t total = static_cast(m) * n; + const unsigned blocks = static_cast( + (total + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_elems = static_cast(k) * n; + const int64_t total = weight_elems > n ? weight_elems : n; + const unsigned blocks = static_cast( + (total + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_pack_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + raw_weight, weight_scale, packed_weight, packed_weight_scale, + weight_elems, n, k); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/o_proj.hip new file mode 100644 index 00000000..40e1cdef --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/o_proj.hip @@ -0,0 +1,830 @@ +// @@variant shape=hy3_tp8_o_proj_m16 commit=4ddcb92e2032eb19fef17c7d4e19ac6739780ba9 added=2026-08-27 +// median_us=12.43 p90_us=12.59 +// source=hy3-dsh-tp8-m16-2-dc53295a +// MetaInfer W8A8 INT8 GEMM HIP implementation for gfx928 (worker_1). +// +// Iteration 4: bounded LDS prefetch pipeline for the assigned shape +// hy3_tp8_o_proj_m16 (M=16, N=4096, K=1024). +// * Iteration 1 (accepted): 256 blocks x 1 wavefront x 1 16x16 N tile, +// direct global fragment loads, 32 K-steps; official median +// 97.3785 us / p90 99.7009 us. +// * Iteration 2 (rejected): 128 blocks x 2 wavefronts x 2 N tiles, +// median 168.109 us / p90 177.434 us -- intra-block wave pairing with +// the direct byte-load B path regressed, so launch-geometry co- +// residency alone is not the lever. +// * Iteration 3 (accepted official best): split-K=2 pipeline with int32 +// partials in the caller workspace -- (1) +// w8a8_dumma_m16_n16_splitk_partial_kernel<2>: grid = 2*256 = 512 +// one-wave blocks (one 16x16 N tile per block; split s owns the K-half +// [s*512, s*512+512), 16 ascending m16n16k32 steps with the +// iteration-1 direct-load double-buffered K loop), each block storing +// its int32 tile to workspace plane s via the API LDS plane +// (mem_row_major); (2) w8a8_dumma_m16_combine_kernel<2>: grid = 256 +// one-wave blocks summing the planes in ascending split order and +// applying the exact reference fp32 scale/bf16 chain. Official median +// 31.446 us / p90 31.509 us (beats the fixed Triton Graph baseline +// 32.137/35.969 us). +// * This round (mandated architecture/pipeline axis: bounded +// register/LDS prefetch): the exact code object of the accepted +// iteration-3 partial kernel shows per K-step ONE global_load_dwordx2 +// per lane for A (8 B, immediate K-offsets -- already vectorized and +// L2-hot) but EIGHT per-lane global_load_ubyte for B (one byte per +// k-row, k-rows strided by N=4096) with s_waitcnt vmcnt(N) interleaved +// with the byte reassembly -- 16 dwordx2 + 128 ubyte + 150 waits + 16 +// mmac per block, i.e. ~9 serialized global-latency rounds per K-step. +// A-only staging would remove only ~1 of those 9 rounds, and +// multi-N-tile reuse would add 8 B-ubyte rounds per extra tile per +// step (per-step chain grows while blocks/CU halves), so the load path +// itself is replaced with a BOUNDED depth-1 LDS prefetch pipeline: the +// partial kernel keeps the exact split-K=2 geometry (512 one-wave +// blocks = 4.27/CU, zero barriers, same combine kernel, same workspace +// planes, same ascending-K int32 order) and only the B load path +// changes. Each block stages its 8-KiB B tile slice into a +// double-buffered LDS window (2 x 64 rows x 48-B padded stride; one +// window = 2 m16n16k32 steps; 8 windows for kLen=512) with ONE 16-B +// global_load_dwordx4 per lane per window, issued one full window +// ahead of consumption; the K loop then reads B fragments from LDS +// (ds_read_ubyte x8 + reassembly, ~30-cycle LDS latency, no vmcnt) +// while A stays on its proven single-dwordx2 direct path. Byte reuse: +// A (8 KiB/block slice; 16 KiB total, L2-hot) is re-read per K-step +// exactly as before; B (8 KiB per block) is read from global exactly +// once per block through the vectorized prefetch and every staged byte +// is consumed by exactly one m16n16k32 step (B is inherently +// single-use per block -- the change converts 128 scattered +// ubyte-load+wait rounds into 8 vectorized prefetch rounds plus LDS +// reads). +// * Falsifiable gate: median_us < 31.446 AND p90_us < 31.509 (strict +// improvement over the iteration-3 official best with the p90 noise +// guard); expected effect ~2x (partial kernel ~3-4x faster, combine +// ~2-3 us unchanged). +// * The CU-aligned non-power-of-two candidate (SPLIT_K=3, grid = 768 +// blocks = 6.4 blocks/CU, non-uniform 32-aligned K slices 384/384/256) +// remains implemented in the same templated partial+combine pair and +// explicitly instantiated; this round launches SPLIT_K=2 as the +// primary candidate. +// * Workspace contract unchanged: partials occupy SPLIT_K * 65536 int32 +// planes (512 KiB for S=2) and are overwritten on every launch; stream +// order makes combine read the planes written by the same replay. +// Graph capture/replay safe: two kernel launches, no allocation, no +// sync, no host reads, no default stream. +// * Iteration 6 (this round): the exact code object of the accepted +// iteration-4 partial kernel shows the B fragment is still reassembled +// per m16n16k32 step from EIGHT per-lane ds_read_u8 (k-rows strided by +// the 48-B padded window stride, one byte per k-row) with ~8 +// s_waitcnt lgkmcnt + ~15 VALU byte-reassembly ops per step, while A is +// already one dwordx2 per lane per step and software-pipelined 4-6 +// steps ahead by the compiler (offsets 128..288 issued early). The +// remaining per-step serial chain is therefore the B-from-LDS byte +// path, so this round replaces the B data path with the validated +// gate_up/down_proj recipe (int8-w8a8-gemm-decode skill): pack the +// exact-shape weight OUT of the timed region into n-major fragment +// layout packed[n*K + k] = raw[k*N + n] (same 4 MiB byte count, same +// buffer, graph-stable addresses; identity pack retained for every +// other (K, N)), keep the depth-1 double-buffered LDS window pipeline, +// but store each window n-major (16 n-cols x 64 k-rows, 80-B padded +// k-stride = 64+16, 16-B aligned non-power-of-two bank skew) so the +// col_major B fragment load (du_mma.hpp: lane l reads +// p[(l&15)*ldm + (l>>4)*8 + i], i=0..7 -- 8 CONTIGUOUS bytes) becomes +// one 8-B vectorized read per lane per step (emitted as two +// ds_read_b32) with no reassembly. Stage-in stays +// one 16-B global_load_dwordx4 per lane per window (now from the packed +// n-major buffer), issued one window ahead of consumption. A path, +// split-K=2 geometry (512 + 256 one-wave blocks), combine kernel, +// workspace planes, ascending-K exact int32 order, and zero-barrier +// structure are held fixed; the scalar fallback decodes the n-major +// pack for (k,n)==(1024,4096) (keeps the paired M=2 API shape exact). +// * Iteration 7 (rejected, reverted): single-buffering leg of the +// mandated double-vs-single buffer comparison -- all 8 B windows staged +// in a prologue burst (LDS 11264 B, zero B-related vmcnt in the K loop, +// 525 -> 371 instructions). Median 14.7632 / p90 14.7728 us vs the +// iteration-6 gate 14.7233/14.7345 -> NOT accepted. Conclusion: the B +// stage-in latency was not the binding constraint; the residual +// per-wave serial latency is the direct A-path global loads waited +// inside the K loop (vmcnt(0..6) pipelined waits) with only 1.07 +// waves/SIMD (4.27 one-wave blocks/CU) to overlap them. +// * Iteration 8 (this round, mandated HIP-only occupancy probe): PMC for +// the accepted iteration-6 build (arch_vgpr 40, sgpr 16, LDS 3584 B, +// scratch 0, 512 blocks, partial kernel 12.8 us) shows LDS allows 18 +// blocks/CU and VGPR 24 blocks/CU with zero spills -- neither resource +// binds at 4.27 blocks/CU, so the only occupancy limiter is the grid +// (total one-wave blocks). This round flips ONLY the split factor to +// SPLIT_K=4: 1024 one-wave partial blocks = 8.53 blocks/CU = 2.13 +// wavefronts/SIMD (first candidate in the trusted [2..8] probe set that +// crosses 2 waves/SIMD), uniform 64-window-aligned K slices 256/256/ +// 256/256 (numWindows = 4), all blocks resident (LDS 18.3 >= 8.53, +// VGPR 24 >= 8.53). The kernel body is bit-identical to the accepted +// iteration-6 code: depth-1 double-buffered LDS window pipeline +// (b_lds[2][1280], one 16-B global_load_dwordx4 per lane per window +// issued one window ahead), col_major B fragment (one 8-B vectorized +// LDS read per lane per step), direct A dwordx2 path, zero barriers, +// same combine/scale chain, same n-major pack, same fallbacks. A and B +// per-replay HBM totals stay constant (A: 1024 x 4 KiB, B: 1024 x +// 4 KiB); only the combine kernel now sums 4 freshly-written +// L2-resident planes (+512 KiB reads, no repeated A/B HBM reads). +// Workspace: 4 planes x 256 KiB = 1 MiB <= the 16-plane (4 MiB) caller +// budget; the launcher workspace guard keeps the scalar fallback for an +// undersized workspace. Expected effect: the ~2x resident waves overlap +// the A-load global stalls (and the 4 instead of 8 per-block stage-in +// rounds halve the per-wave exposed-latency chain), partial kernel +// ~12.8 us -> ~7-10 us, combine ~2 us -> ~2.5-3.5 us, total ~10-13 us. +// Falsifiable gate: median_us < 14.72327470779419 AND p90_us < +// 14.734469652175903 (strict improvement over the iteration-6 official +// best with the p90 noise guard). If the median stays ~14.7 us, the +// kernel is not occupancy-limited on this axis and the next round must +// attack the A-path itself (LDS-stage A, or A layout) or probe SPLIT_K=8 +// (4.27 waves/SIMD, combine reads 2 MiB). +// * Iteration 11 (rejected, reverted): SPLIT_K=8 occupancy probe (2048 +// one-wave partial blocks = 17.07 blocks/CU = 4.27 waves/SIMD, kernel +// body bit-identical to iteration 8) regressed to 17.3119/17.3270 vs the +// 14.4528/14.4672 gate -- the occupancy axis is falsified (1.07 through +// 4.27 waves/SIMD all flat or worse; 2.13 is the peak). Per-kernel +// profile of the accepted build: partial kernel 12.16 us + combine +// kernel 2.88 us (profiled; operator 15.04 us), i.e. the combine is a +// ~20% SERIAL second phase plus its launch gap. +// * Iteration 12 (this round, HIP-only consolidation): fuse the combine +// into the partial kernel with the qkv_proj-lineage per-tile +// last-arrival protocol -- monotonic counters in the workspace tail +// (256 x int32 after plane 3; one async zero per workspace pointer on +// its first eager use, before capture), each block does one +// __threadfence() + atomicAdd(&counters[tile], 1), and the arrival with +// (arrived % SPLIT_K) == SPLIT_K - 1 sums the tile's 4 planes in +// ascending split order and writes the scaled bf16 output with the exact +// combine-kernel math (bit-identical bytes). The operator becomes ONE +// launch per replay: the ~2.9-us serial combine phase and its launch gap +// disappear, each tile's combine runs as soon as its 4 splits land +// (overlapped with sibling blocks), and the plane reads happen while the +// planes are L2-fresh. SPLIT_K=4 geometry, kernel body, n-major pack, +// workspace planes, exact int32 order, and fallbacks are unchanged; the +// two-kernel path remains for workspaces that fit the planes but not the +// counters. Expected: partial kernel ~12 us (combine tail on the +// last-arrivers only) + no second kernel -> total ~12-13 us. Falsifiable +// gate: median_us < 14.452790021896362 AND p90_us < 14.467190504074097 +// (strict improvement over the iteration-8 official best with the p90 +// noise guard). If the median stays ~14.4 us, the serial combine phase +// was not the binding cost and the next round must attack the partial +// kernel's own issue chain (dual accumulator fragments) or geometry. +// * Iteration 21 (this round, HIP-only): A-LDS double-buffered window +// staging -- A is staged through the same depth-1 double-buffered LDS +// window pipeline as B (a_lds[2][1280], one 16-B dwordx4 per lane per +// window issued one window ahead) and the per-step a_frag loads read LDS +// (ldm = 80) instead of global (ldm = 1024). Iteration 7 showed removing +// every B-related vmcnt wait from the K loop does not move the time, and +// identified the DIRECT A-PATH global loads waited inside the loop as the +// residual per-wave serial latency; iteration 10 hoisted all A dwordx2 +// loads into a prologue burst and regressed (14.8639) because the 8 +// steps of A sat in VGPRs. Staging A in LDS (not VGPRs) keeps the +// prefetch in registers only (the compiler hoists all 8 A/B window +// dwordx4 prefetches into the prologue with progressive vmcnt waits; +// arch_vgpr 40 -> 50, 2-3 waves/SIMD -- lockstep evidence says extra +// co-residency does not change the serial-chain makespan), LDS 3588 -> +// 6148 B/block (64 KiB / 6.15 KiB = 10.4 blocks/CU >= 8.53 needed, grid +// stays fully resident), and moves the A vmcnt wait out of the per-step +// mmac chain to the top-of-window ds_write where it overlaps a full +// compute window -- the exact recipe that took the B path from 31.45 to +// 19.01 us in iteration 4. Fragment values are byte-identical (same +// global bytes staged, same col_major-style 8-B fragment mapping), so +// the k-ascending int32 accumulation, the ascending-split sum, the fp32 +// scale/bf16 chain and every output byte are unchanged. Geometry +// (SPLIT_K=4, 1024 one-wave blocks, one launch per replay), B path, +// fused last-arrival counter protocol, workspace/planes/counters budget, +// Graph safety, two-kernel fallback (counters = nullptr) and generic +// scalar fallback are UNCHANGED. Expected: per-block serial latency +// drops by the exposed A-load round trips (~8 waited dwordx2 -> 4 +// prologue-hidden dwordx4), operator median toward ~8-11 us. +// Falsifiable gate: median_us < 13.084909915924072 AND p90_us < +// 13.115299940109253 (strict improvement over the iteration-12 official +// best with the p90 noise guard). If the median stays ~13 us, the A path +// is falsified as the binding cost and the residual is the fused +// atomic/combine tail, irreducible in this protocol. +// +// The guarded exact-shape branch keeps every other (m, n, k) -- including the +// paired M=2 API shape with the same (N, K) -- on the generic scalar fallback +// (which decodes the n-major pack only for (k,n)==(1024,4096) and otherwise +// reads the row-major weight, matching the identity pack below). +// Later optimization rounds may replace only the exact-shape branch +// implementation (and the pack op) while this dispatch structure stays. + +#include +#include +#include + +#include + +namespace { + +// gfx928 wavefront is 64; block dimensions must be multiples of 64. +constexpr int kScalarBlockThreads = 256; +constexpr int kWaveSize = 64; + +// Exact-shape DUMMA geometry for hy3_tp8_o_proj_m16 (M=16, N=4096, K=1024). +constexpr int kTargetM = 16; +constexpr int kTargetN = 4096; +constexpr int kTargetK = 1024; +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; + +// Split-K planning for the exact shape: one 16x16 N tile per block means +// kNTiles = N/16 = 256 tiles; a split-K=S partial grid has S*kNTiles +// one-wave blocks. Every K boundary is aligned to the 32-element DUMMA step +// and to the 64-k-row LDS window (uniform 512/512 for S=2; uniform 256/256/ +// 256/256 for S=4; non-uniform 384/384/256 for S=3). +constexpr int kNTiles = kTargetN / kDummaTileN; // 256 +constexpr int kPlaneInts = kTargetM * kTargetN; // 65536 + +__device__ __forceinline__ constexpr int kSplitStart(int split_k, int s) { + if (split_k == 2) return s * 512; + if (split_k == 4) return s * 256; + return s * 384; // S=3 +} +__device__ __forceinline__ constexpr int kSplitLen(int split_k, int s) { + if (split_k == 2) return 512; + if (split_k == 4) return 256; + return (s < 2) ? 384 : 256; // S=3 +} + +// --------------------------------------------------------------------------- +// Generic scalar int8 dot-product kernel: one thread per output element. +// Fallback for every (m, n, k) not handled by an exact-shape specialization. +// For the exact (k, n) == (1024, 4096) the packed weight buffer holds the +// n-major transpose packed[n*K + k] = raw[k*N + n] (iteration 6), so the +// fallback decodes that layout there (this keeps the paired M=2 API shape +// with the same (N, K) byte-exact); every other (k, n) keeps the identity +// pack (raw row-major [K, N]). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + + const int8_t* a_row = a + static_cast(row) * k; + const bool packed_nmajor = (k == kTargetK && n == kTargetN); + const int8_t* b_col = b + col; + int32_t acc = 0; + if (packed_nmajor) { + // packed[col * 1024 + kk] == raw[kk * 4096 + col] (iteration-6 pack). + const int8_t* b_pack_col = b + static_cast(col) * kTargetK; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_pack_col[kk]); + } + } else { + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Split-K partial kernel: one 64-thread wavefront per block, one 16x16 N +// tile per block, one K-slice per split. Iteration 6 replaces the B data +// path with the validated n-major packed + col_major fragment recipe: the +// exact-shape weight is packed (out of the timed region) as +// packed[n*K + k] = raw[k*N + n], and each block stages its B slice +// (kLen k-rows x 16 n-cols) into a double-buffered LDS window stored +// n-major (16 n-cols x 64 k-rows, 80-B padded k-stride) with ONE 16-B +// global_load_dwordx4 per lane per window, issued one full window (2 +// m16n16k32 steps) ahead of consumption; the per-step K loop then reads B +// fragments from LDS with the col_major load (du_mma.hpp lane mapping: +// p[(l&15)*ldm + (l>>4)*8 + i], i=0..7), i.e. one 8-B vectorized read per +// lane per step -- emitted by the compiler as two ds_read_b32 (8 contiguous +// bytes) -- so the 8 ds_read_u8 + ~8 lgkmcnt + ~15 VALU byte-reassembly +// chain of iteration 4 is gone (validated on gate_up/down_proj in this +// lineage). +// Iteration 21: A takes the same depth-1 double-buffered LDS window path as +// B (a_lds[2][1280], one 16-B dwordx4 per lane per window issued one window +// ahead; per-step a_frag reads from LDS at ldm = kWindowKStride), so the +// direct A-path vmcnt wait leaves the per-step mmac chain (iteration 7 +// identified those waited A loads as the residual per-wave serial latency; +// iteration 10's register hoist regressed on VGPR pressure -- LDS staging +// avoids it). Explicit int32 accumulation (max |dot| per slice = +// 16,516,096 << 2^31, k-ascending). +// Each block publishes its int32 tile through the API LDS plane +// (mem_row_major) and stores it to workspace plane `split` with one 16-B +// vector store per lane. Single wavefront -> no barrier anywhere (LDS +// store->read ordering is per-wave program order; buffer w&1 is stored at +// the top of iteration w and read only after that store). Guarded in the +// launcher: only this exact (m, n, k) reaches it. +// Iteration 12 (fused last-arrival combine, qkv_proj lineage recipe): with +// the counters pointer the tile's LAST arrival (monotonic per-tile counters +// in the workspace tail, (arrived % SPLIT_K) == SPLIT_K - 1) sums the +// tile's SPLIT_K planes in ascending split order and emits the scaled bf16 +// output in-kernel, byte-identical to the separate combine kernel, so the +// operator becomes ONE launch per replay; counters = nullptr keeps the +// iteration-8 two-kernel behavior for undersized workspaces. +// --------------------------------------------------------------------------- +constexpr int kWindowKRows = 64; // one window = 2 m16n16k32 steps +constexpr int kWindowNCols = 16; // block tile width (n-cols per window) +constexpr int kWindowKStride = 80; // 64 k-rows + 16 pad; 16-B aligned, + // non-power-of-two bank skew (<=2-way) +constexpr int kNumBStages = 2; // double-buffered (depth-1 prefetch) + +template +__global__ __launch_bounds__(kWaveSize) void +w8a8_dumma_m16_n16_splitk_partial_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int32_t* __restrict__ partials, + int32_t* __restrict__ counters, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out) { + static_assert(SPLIT_K == 2 || SPLIT_K == 3 || SPLIT_K == 4, + "exact-shape split-K plan supports SPLIT_K=2, 3 or 4"); + const int lane = static_cast(threadIdx.x); + const int split = static_cast(blockIdx.x) / kNTiles; + const int tile = static_cast(blockIdx.x) % kNTiles; + const int n0 = tile * kDummaTileN; + const int kStart = kSplitStart(SPLIT_K, split); + const int kLen = kSplitLen(SPLIT_K, split); + const int numWindows = kLen / kWindowKRows; // 8 (S=2), 6/4 (S=3) + + __shared__ __align__(16) int8_t b_lds[kNumBStages][kWindowNCols * + kWindowKStride]; + // Iteration 21: A is staged through the same double-buffered LDS window + // pipeline as B (16 m-rows x 64 k-rows per window, 80-B padded stride, + // one 16-B dwordx4 per lane per window issued one window ahead), so the + // per-step a_frag loads read LDS (ldm = kWindowKStride) instead of global + // (ldm = kTargetK). The A-path vmcnt wait leaves the per-step mmac chain + // and lands at the top-of-window ds_write alongside B's (iteration 7 + // identified the direct A-path global loads as the residual per-wave + // serial latency; iteration 10's register hoist was rejected -- staging in + // LDS, not VGPRs, keeps occupancy and the windowed depth-1 prefetch). + __shared__ __align__(16) int8_t a_lds[kNumBStages][kWindowNCols * + kWindowKStride]; + + du::dumma::DUFragment + a_frag[2]; + du::dumma::DUFragment + b_frag[2]; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Depth-1 double-buffered LDS prefetch. Window w covers B k-rows + // [kStart + w*64, kStart + w*64 + 64) at the block's 16 n-cols. The + // packed buffer is n-major, so lane l fetches 16 CONTIGUOUS k-rows + // (16 B, one dwordx4) at n-col (n0 + (l&15)) and k-offset + // 16*(l>>4) within the window; the data for window w is issued one + // iteration ahead (prologue for w=0) so the vmcnt wait lands at the + // top-of-iteration ds_write_b128, a full compute window (~2 mmac steps) + // after issue. + const int8_t* b_packed_col = + b + static_cast(n0 + (lane & 15)) * kTargetK; + uint4 prefetch = *reinterpret_cast( + b_packed_col + kStart + 16 * (lane >> 4)); + // A prefetch mirrors B: lane l fetches 16 contiguous k-rows (16 B, one + // dwordx4) of its A row (lane & 15) at k-offset 16*(lane >> 4) within the + // window; every staged byte is consumed by exactly one m16n16k32 step via + // the row_major col_major-style fragment (the 16-B chunk covers two + // 8-B k-groups, exactly like the B window). + const int8_t* a_row = a + static_cast(lane & 15) * kTargetK; + uint4 a_prefetch = *reinterpret_cast( + a_row + kStart + 16 * (lane >> 4)); +#pragma unroll + for (int w = 0; w < numWindows; ++w) { + // Store the previously-prefetched window w into buffer w&1 (its vmcnt + // wait has had the prior iteration's compute to complete), then issue + // the next window's prefetches before computing this window from LDS. + *reinterpret_cast( + &b_lds[w & 1][(lane & 15) * kWindowKStride + 16 * (lane >> 4)]) = + prefetch; + *reinterpret_cast( + &a_lds[w & 1][(lane & 15) * kWindowKStride + 16 * (lane >> 4)]) = + a_prefetch; + if (w + 1 < numWindows) { + prefetch = *reinterpret_cast( + b_packed_col + kStart + (w + 1) * kWindowKRows + 16 * (lane >> 4)); + a_prefetch = *reinterpret_cast( + a_row + kStart + (w + 1) * kWindowKRows + 16 * (lane >> 4)); + } + + const int8_t* bwin = &b_lds[w & 1][0]; + const int8_t* awin = &a_lds[w & 1][0]; + du::dumma::du_load_matrix_sync(a_frag[0], awin, kWindowKStride); + du::dumma::du_load_matrix_sync(b_frag[0], bwin, kWindowKStride); + du::dumma::du_load_matrix_sync(a_frag[1], awin + kDummaTileK, + kWindowKStride); + du::dumma::du_load_matrix_sync(b_frag[1], bwin + kDummaTileK, + kWindowKStride); + du::dumma::du_mma_sync(acc_frag, a_frag[0], b_frag[0], acc_frag); + du::dumma::du_mma_sync(acc_frag, a_frag[1], b_frag[1], acc_frag); + } + + // Publish the int32 tile through the API (mem_row_major), then one 16-B + // vector store per lane. Plane layout is per-tile contiguous + // workspace[split][tile][256] (e = row*16 + col), so each block's 1-KiB + // int32 tile is one coalesced 64-lane x 16-B store. Single wavefront per + // block: the LDS write->read dependency is ordered by the compiler's + // lgkmcnt wait, no barrier needed. + __shared__ __align__(16) int32_t tile_lds[kDummaTileM * kDummaTileN]; + du::dumma::du_store_matrix_sync(tile_lds, acc_frag, kDummaTileN, + du::dumma::mem_row_major); + + const int e = lane * 4; + const uint4 packed = + make_uint4(tile_lds[e], tile_lds[e + 1], tile_lds[e + 2], + tile_lds[e + 3]); + *reinterpret_cast(partials + split * kPlaneInts + + tile * (kDummaTileM * kDummaTileN) + e) = packed; + + // Fused last-arrival combine tail (iteration 12): when counters != nullptr + // the separate combine kernel is removed and the operator is ONE launch per + // replay. Every block signals its arrival for its tile with one lane-0 + // __threadfence() + atomicAdd(&counters[tile], 1); the arrival whose + // pre-increment value satisfies (arrived % SPLIT_K) == SPLIT_K - 1 is the + // tile's LAST arrival of this replay (counters are monotonic: each replay + // adds exactly SPLIT_K to every tile counter) and it sums the tile's + // SPLIT_K planes in ascending split order and writes the scaled bf16 + // output -- byte-identical to the separate combine kernel (same plane + // chunk reads, same ascending int32 sum, same + // float(acc) * x_scale[row] * weight_scale[col] -> __float2bfloat16 chain, + // same store addresses). Each tile's combine now runs as soon as its 4 + // splits land, while sibling blocks are still computing, so the ~2.9-us + // serial combine phase and its launch gap disappear and the plane reads + // happen while the planes are L2-fresh (per-tile, immediately after the + // tile's last write). Exact int32 accumulation unchanged: per-slice max + // |dot| << 2^31, ascending-split sum order (same as the reference's + // ascending-K order). The counters live in the workspace tail (last 256 + // int32 = 1 KiB, past plane SPLIT_K-1) and are zeroed once per workspace + // pointer by the launcher before the first eager use (never inside Graph + // capture); the two-kernel fallback passes counters = nullptr and the tail + // branch is dead there. + if (counters != nullptr) { + __syncthreads(); + __shared__ int s_is_last; + if (lane == 0) { + __threadfence(); // release: this block's plane store is visible to the + // observer of this block's arrival atomic + const int arrived = atomicAdd(&counters[tile], 1); + __threadfence(); // acquire: the reads below (after the barrier) see + // every sibling plane store released before its + // arrival atomic + s_is_last = ((arrived % SPLIT_K) == SPLIT_K - 1); + } + __syncthreads(); + if (s_is_last) { + const int row = lane >> 2; // 0..15 + const int cbase = (lane & 3) * 4; // 0, 4, 8, 12 within the tile + // 16-B aligned reads of the per-tile-contiguous plane layout + // workspace[split][tile][256]; plane stride is kPlaneInts/4 uint4s. + const uint4* plane = + reinterpret_cast(partials) + + tile * (kDummaTileM * kDummaTileN / 4) + row * (kDummaTileN / 4) + + (lane & 3); + int32_t acc[4] = {0, 0, 0, 0}; +#pragma unroll + for (int s = 0; s < SPLIT_K; ++s) { + const uint4 v = plane[s * (kPlaneInts / 4)]; + acc[0] += v.x; + acc[1] += v.y; + acc[2] += v.z; + acc[3] += v.w; + } + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = n0 + cbase + i; + const float scaled = + static_cast(acc[i]) * xs * weight_scale[col]; + out[static_cast(row) * kTargetN + col] = + __float2bfloat16(scaled); + } + } + } +} + +// --------------------------------------------------------------------------- +// Combine+scale kernel: one 64-thread wavefront per block, one 16x16 N tile +// per block. Each lane loads its 4-element chunk from every split plane with +// one 16-B vector load per plane, sums the SPLIT_K planes in ascending split +// order (exact int32; ascending-split sum == the reference's ascending-K +// order), then applies the exact reference fp32 chain +// float(acc) * x_scale[row] * weight_scale[n0 + col] -> __float2bfloat16 +// (RN). Same expression order as the reference -> bit-identical output for +// the same int32 dot. Zero barriers. Reads only planes written by the +// partial kernel of the same stream-ordered launch. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kWaveSize) void w8a8_dumma_m16_combine_kernel( + const int32_t* __restrict__ partials, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out) { + const int lane = static_cast(threadIdx.x); + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + const int row = lane >> 2; // 0..15 + const int cbase = (lane & 3) * 4; // 0, 4, 8, 12 within the tile + + // 16-B aligned reads of the per-tile-contiguous plane layout + // workspace[split][tile][256]; plane stride is kPlaneInts/4 uint4s. + const uint4* plane = + reinterpret_cast(partials) + + blockIdx.x * (kDummaTileM * kDummaTileN / 4) + row * (kDummaTileN / 4) + + (lane & 3); + int32_t acc[4] = {0, 0, 0, 0}; +#pragma unroll + for (int s = 0; s < SPLIT_K; ++s) { + const uint4 v = plane[s * (kPlaneInts / 4)]; + acc[0] += v.x; + acc[1] += v.y; + acc[2] += v.z; + acc[3] += v.w; + } + + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = n0 + cbase + i; + const float scaled = static_cast(acc[i]) * xs * weight_scale[col]; + out[static_cast(row) * kTargetN + col] = + __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// Identity device-to-device packing: byte-for-byte copy. This is the generic +// fallback for any (K, N) and is valid for every shape. +// pack_weight runs outside the timed region and outside Graph capture. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_pack_identity_bytes_kernel( + const uint8_t* __restrict__ src, + uint8_t* __restrict__ dst, + int64_t num_bytes) { + const int64_t stride = + static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + i < num_bytes; i += stride) { + dst[i] = src[i]; + } +} + +// --------------------------------------------------------------------------- +// Exact-shape n-major pack (iteration 6): for (K, N) == (1024, 4096) the +// weight is transposed one-time (outside the timed region / Graph) into +// packed[n * K + k] = raw[k * N + n], so each lane's 8-byte col_major B +// fragment (du_mma.hpp lane mapping) is contiguous and every m16n16k32 +// fragment is a single vectorized load. Same byte count and buffer as the +// identity pack, so captured addresses are unchanged. One thread per output +// byte; runs once during weight prep. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_pack_nmajor_bytes_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t total = static_cast(k) * n; + const int64_t stride = + static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; idx += stride) { + const int col = static_cast(idx / k); + const int kk = static_cast(idx - static_cast(col) * k); + packed[idx] = raw[static_cast(kk) * n + col]; + } +} + +// Explicit instantiations: the CU-aligned non-power-of-two candidate +// (SPLIT_K=3, grid = 3*kNTiles = 768 blocks = 6.4 blocks/CU, non-uniform +// 32-aligned K slices 384/384/256) and the SPLIT_K=4 occupancy candidate +// (grid = 4*kNTiles = 1024 blocks = 8.53 blocks/CU, uniform 64-window-aligned +// K slices 256/256/256/256) are compiled so a future round flips one +// constant; this round launches SPLIT_K=4 as the primary mandated candidate. +template __global__ void w8a8_dumma_m16_n16_splitk_partial_kernel<3>( + const int8_t*, const int8_t*, int32_t*, int32_t*, const float*, + const float*, hip_bfloat16*); +template __global__ void w8a8_dumma_m16_combine_kernel<3>( + const int32_t*, const float*, const float*, hip_bfloat16*); +template __global__ void w8a8_dumma_m16_n16_splitk_partial_kernel<4>( + const int8_t*, const int8_t*, int32_t*, int32_t*, const float*, + const float*, hip_bfloat16*); +template __global__ void w8a8_dumma_m16_combine_kernel<4>( + const int32_t*, const float*, const float*, hip_bfloat16*); + +} // namespace + +// --------------------------------------------------------------------------- +// Host launchers (stable symbols consumed by csrc/bindings.cpp). +// --------------------------------------------------------------------------- + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + hip_bfloat16* out_bf16 = static_cast(out); + + // Guarded exact-shape dispatch: only hy3_tp8_o_proj_m16 (16, 4096, 1024) + // reaches the split-K DUMMA pipeline. Every other (m, n, k) -- including + // the paired M=2 API shape with the same (N, K) -- takes the generic + // scalar fallback (which reads the row-major weight, matching the + // identity pack). + if (m == kTargetM && n == kTargetN && k == kTargetK) { + // Split-K pipeline (iteration 8: occupancy probe -- SPLIT_K=4, grid = + // 4*kNTiles = 1024 one-wave partial blocks = 8.53 blocks/CU = 2.13 + // wavefronts/SIMD. PMC evidence for the accepted iteration-6 SPLIT_K=2 + // build (arch_vgpr 40, sgpr 16, LDS 3584 B, scratch 0, 512 blocks) shows + // LDS allows 18 blocks/CU and VGPR 24 blocks/CU -- neither binds at the + // 4.27 blocks/CU needed, so the only occupancy limiter is the grid + // itself (1.07 waves/SIMD). Iteration 7 proved removing every B-related + // vmcnt wait from the K loop does not move the time (14.763 vs 14.723 + // us), i.e. the per-wave serial latency is the direct A-path global + // loads waited inside the loop with no sibling wave to overlap them. + // SPLIT_K=4 doubles resident waves to 2.13/SIMD (all 1024 blocks stay + // resident: LDS 18.3 >= 8.53, VGPR 24 >= 8.53) while keeping the exact + // iteration-6 depth-1 double-buffered window pipeline, A/B totals and + // per-replay HBM traffic constant (only the combine reads 4 instead of + // 2 freshly-written L2-resident planes, +512 KiB). The partial kernel + // reads B from the n-major packed buffer through the col_major fragment + // path (one 8-B vectorized LDS read per lane per step, staged via the + // double-buffered LDS window prefetched one window ahead): int32 + // partials in the caller workspace, combine+scale kernel in the timed + // Graph. Both launches are stream ordered inside the same captured + // region; the partial kernel overwrites every partial element on each + // launch (no workspace clear needed). + constexpr int kSplitK = 4; + constexpr int kPartialBlocks = kSplitK * kNTiles; // 1024 + constexpr int kCombineBlocks = kNTiles; // 256 + const int64_t plane_bytes = + static_cast(kPlaneInts) * sizeof(int32_t); // 256 KiB + const int64_t counters_bytes = + static_cast(kNTiles) * sizeof(int32_t); // 1 KiB + int32_t* partials = static_cast(workspace); + if (workspace_bytes >= kSplitK * plane_bytes + counters_bytes) { + // Fused single-launch path (iteration 12): the combine runs inside the + // partial kernel (per-tile last arrival), so the operator is ONE kernel + // launch per replay and the separate combine kernel + its launch gap + // disappear. The per-tile arrival counters (256 x int32 = 1 KiB) live + // immediately after plane kSplitK-1 (the partial kernel writes planes + // 0..kSplitK-1 only; the two-kernel fallback below never touches the + // counters). Counters are monotonic -- each replay adds exactly SPLIT_K + // to every tile counter -- so the last-arrival test is + // (arrived % SPLIT_K) == SPLIT_K - 1 and NO per-replay reset is needed. + // One async zero per workspace pointer is issued on the workspace's + // first eager use (the Graph capture flow always warms up eagerly + // first): the guarded hipMemsetAsync is a stream op on the caller's + // stream, no host sync, and is NOT part of the captured graph -- the + // qkv_proj lineage validated this exact protocol. + int32_t* counters = partials + kSplitK * kPlaneInts; + static const void* s_fused_counters_ws = nullptr; + if (s_fused_counters_ws != workspace) { + hipMemsetAsync(counters, 0, static_cast(counters_bytes), + stream); + s_fused_counters_ws = workspace; + } + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_n16_splitk_partial_kernel), + dim3(kPartialBlocks), + dim3(kWaveSize), + 0, + stream, + a, + b, + partials, + counters, + x_scale, + weight_scale, + out_bf16); + return; + } + if (workspace_bytes >= kSplitK * plane_bytes) { + // Two-kernel fallback for workspaces that fit the partial planes but + // not the counters: partial kernel without the fused tail + // (counters = nullptr -> the tail branch is dead) then the combine + // kernel, byte-identical to the accepted iteration-8 behavior. + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_n16_splitk_partial_kernel), + dim3(kPartialBlocks), + dim3(kWaveSize), + 0, + stream, + a, + b, + partials, + static_cast(nullptr), + x_scale, + weight_scale, + out_bf16); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_combine_kernel), + dim3(kCombineBlocks), + dim3(kWaveSize), + 0, + stream, + partials, + x_scale, + weight_scale, + out_bf16); + return; + } + // Workspace smaller than the split-K=4 partial planes: fall through to + // the generic scalar fallback (correct for every (m, n, k)); never run + // the tuned pipeline with an undersized workspace. + } + + const int64_t total = static_cast(m) * n; + const unsigned blocks = static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_scalar_gemm_kernel), + dim3(blocks), + dim3(kScalarBlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_bf16, + m, + n, + k); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_bytes = static_cast(k) * n; + const int64_t scale_bytes = static_cast(n) * sizeof(float); + constexpr int kPackThreads = kScalarBlockThreads; + + // Exact-shape n-major pack (iteration 6): only (k, n) == (1024, 4096) + // gets the transposed fragment layout consumed by the tuned partial + // kernel; every other (K, N) keeps the byte-identical identity pack so + // the generic scalar fallback stays correct for all other shapes. + if (k == kTargetK && n == kTargetN) { + const unsigned pack_blocks = static_cast( + (weight_bytes + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_pack_nmajor_bytes_kernel), + dim3(pack_blocks), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + const unsigned weight_blocks = static_cast( + (weight_bytes + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_pack_identity_bytes_kernel), + dim3(weight_blocks), + dim3(kPackThreads), + 0, + stream, + reinterpret_cast(raw_weight), + reinterpret_cast(packed_weight), + weight_bytes); + } + + const unsigned scale_blocks = static_cast( + (scale_bytes + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_pack_identity_bytes_kernel), + dim3(scale_blocks), + dim3(kPackThreads), + 0, + stream, + reinterpret_cast(weight_scale), + reinterpret_cast(packed_weight_scale), + scale_bytes); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/qkv_proj.hip new file mode 100644 index 00000000..7ff2b06e --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/qkv_proj.hip @@ -0,0 +1,941 @@ +// @@variant shape=hy3_tp8_qkv_proj_m16 commit=467827c4a33a7e06fe4b3d9871cad428d75648fe added=2026-08-27 +// median_us=23.17 p90_us=23.21 +// source=hy3-dsh-tp8-m16-2-dc53295a +// INT8 W8A8 GEMM for Hygon K500SM_AI / gfx928. +// +// Iteration 1 (DUMMA bootstrap): replace the correctness-first scalar kernel +// for the assigned shape hy3_tp8_qkv_proj_m16 (M=16, N=1280, K=4096) with the +// minimal native gfx928 DUMMA INT8 m16n16k32 tile: one 64-lane wavefront per +// block, one 16x16 output tile per block (grid = N/16 = 80 blocks), direct +// global-to-fragment loads (no LDS staging), explicit int32 accumulation, and +// a single-wave LDS epilogue with no cross-wave barrier. Measured 176.98 us. +// +// Iteration 2 (architecture round: launch geometry): keep the direct-load +// data path unchanged and change the launch geometry to 2 wavefronts per +// block (128 threads) with in-block split-K=2. Each wavefront accumulates +// exactly one K-half (2048 = 64 m16n16k32 steps) in int32 k-ascending order, +// publishes its int32 partial to its own LDS plane, and after one END-of-K +// barrier the two wavefronts add partial 0 + partial 1 (ascending split +// order, bit-identical to the unsplit k-ascending int32 accumulation) and +// emit the scaled bf16 tile. grid stays N/16 = 80 blocks. Measured +// 102.12 us median / 105.59 us p90 (vs the 74.782 us Triton Graph baseline); +// the ~1.73x geometry gain shows each active CU still runs one serial +// load->wait->MMA chain per wave (64 steps after split-K=2), with the two +// chains of one block co-scheduled on the same CU. +// +// Iteration 3 (architecture round: workspace split-K pair). Mandate: the +// unsplit grid (80 blocks) has fewer than two blocks per device CU (120 CUs) +// and K=4096 >= 1024, so implement split-K=2 plus at least one CU-aligned +// candidate (non-power-of-two allowed), writing int32 partials into the +// caller workspace and including the combine+scale kernel in the timed Graph. +// Replace the iteration-2 single-launch in-block-combine kernel with a +// workspace split-K pair (the pattern validated on the hy3 TP4 qkv lineage, +// which took the same architecture round at its iteration 3): +// +// 1. w8a8_dumma_m16_splitk_partial_kernel: one 64-thread +// wavefront per block, grid = (N/16 = 80 tiles, SPLIT_K K-slices). +// Each block computes the int32 partial dot for one 16x16 output tile +// over one 32-aligned K slice and du_store_matrix_sync's the partial +// straight into the caller workspace. ZERO barriers, ZERO LDS, ZERO +// atomics in the main kernel. +// 2. w8a8_dumma_m16_combine_scale_kernel: 80 blocks x 256 threads, one +// thread per output element; exact int32 sum of the SPLIT_K partials +// (ascending split order), then x_scale/weight_scale multiply and the +// bf16 (RN-even) epilogue. Barrier-free (each thread reads only its +// own 16x16 element's S partial values, written by the previous kernel +// on the same stream). +// +// Both kernels launch on the caller's stream inside the timed Graph, so the +// combine cost is included in the measured operator latency. +// +// Iteration 4 (register prefetch, rejected): double-buffered a_frag/b_frag in +// registers one K step ahead. Regressed to 120.5 us median: the exact code +// object kept the full vmcnt drain before each v_mmac, so prefetching could +// not remove the per-step global round trip from the MMA critical path. +// Iterations 5-6 were killed by the agent infrastructure (no evidence). +// +// Iteration 7 (architecture round: LDS slice staging). The exact gfx928 code +// object of the iteration-3 kernel shows every one of the 64 K steps compiles +// to 16 narrow global_load_ubyte + byte-reassembly VALU + a full progressive +// vmcnt drain before the single v_mmac: the whole global round trip is paid +// serially per step with only ~1.3 waves/CU to hide it. This round applies +// the TP4 qkv lineage remedy (validated 26.05 us on N=2560) scaled to this +// shape: stage the block's entire 32-aligned K slice into LDS once with +// coalesced dword global loads, one __syncthreads, then a zero-barrier +// LDS-only K loop (du_load_matrix_sync reads both fragments from LDS, so the +// MMA path sees only ~30-cycle LDS latency instead of a global round trip +// per step): +// +// w8a8_dumma_m16_slicestage_partial_kernel: one 64-thread +// wavefront per block, grid = (N/16 = 80 tiles, SPLIT_K K-slices), same +// per-tile 16x16 output and the same workspace int32 plane layout, so the +// combine+scale kernel below is unchanged. A (16 x slice) and B +// (slice x 16) are staged into LDS with the same 32-aligned non-uniform +// slice bounds and the same k-ascending int32 MMA order as iteration 3, so +// outputs stay bit-identical (0 mismatches) and Graph capture/replay with +// changed contents must pass. +// +// Round-8 geometry (default): split-K=16 -> grid = (80 tiles, 16 slices) = +// 1,280 one-wave blocks = 10.67 blocks/CU on the 120-CU device. Each block +// stages exactly 256 K (A 16x256 -> 4,160 B + B 256x16 -> 5,120 B = 9,280 B +// LDS = 7 resident blocks/CU by the 64 KiB LDS) and runs exactly 8 uniform +// zero-barrier LDS MMAs (128/16 = 8, so no 11-step stragglers). This is an +// occupancy (LDS-footprint) tune against the round-7 PMC evidence: at +// split-K=12 the profile shows grid 960 = 8 waves/CU available but +// 12,736 B LDS/block caps residency at 5 blocks/CU (62.5%), with VGPR 24 / +// scratch 0 / SGPR 32 proving registers and spills are NOT limiters -- the +// per-block LDS footprint is the binding occupancy limiter, and S=16 is the +// largest split that fits the 16-plane contract workspace (1,310,720 B +// exactly), so residency rises to 7 of 10.67 waves/CU. A bytes are reused +// across the 80 tile blocks of each slice (1,280 x 16 x 256 = 5.24 MB total +// A reads vs 64 KB unique, L2-hot); B 16-column strips are read exactly once +// (5.24 MB unique), with the 64-B L2 sectors shared with the adjacent +// 16-column tiles, so HBM A/B traffic is byte-identical to split-K=12 +// (~5.30 MB unique; no repeated HBM reads traded for occupancy -- only the +// split-K combine workspace round trip grows from 12 to 16 planes, +0.33 MB +// write + read). vmem_read_instructions stay ~40,960/replay (slightly fewer +// than S=12's 42,240: each 256-K slice needs 32 wavefront loads vs 44). +// +// Iteration 11 (packed col-major B; the validated TP4 gate_up/down_proj +// recipe). The accepted S=16 staged kernel's fresh PMC shows +// lds_bank_conflicts 81,920 == the 8 scalar ds_read_u8 per matrix_b fragment +// (8 steps x 1,280 blocks), i.e. every B fragment read in the zero-barrier K +// loop conflicts at the 20-B k-major pitch, while round 9 falsified +// staging-load MLP batching (58.56 us), so the K-loop B fragment reads are +// the remaining hot-path lever. launch_pack_w8a8_weight now transposes the +// exact (k,n)==(4096,1280) weight to [N,K] n-major once, outside the timed +// region and out of Graph capture (every other (K,N) keeps the identity +// copy); the staged kernel stages its 16-column strip as 16-B int4 straight +// copies into an n-major LDS tile b_s[16][kMaxKlen+16] (row stride +// kMaxKlen+16: 16-B aligned for the int4 stores, 8-B aligned for the +// fragment b64 reads, not a multiple of 128 B -> no bank-phase aliasing), so +// each lane's 8 fragment bytes are consecutive and load with ONE ds_read_b64 +// (load_b_frag_packed) instead of 8 conflicting ds_read_u8. The direct +// S in {2,3,4,5} path reads the packed global layout through the library +// col_major loader, and the scalar fallback decodes the packed layout for +// (k,n)==(4096,1280) (covers the paired M=2 validation shape). Fragment +// contents, the k-ascending int32 du_mma order, the slice-major int32 plane +// layout, the combine+scale kernel and the generic fallback are all +// unchanged, so outputs stay bit-identical (0 mismatches) and Graph +// capture/replay with changed contents must pass. Per-block LDS drops +// 9,280 -> 8,512 B (65,536/8,512 = 7.70, still 7 resident/CU); A/B HBM bytes +// are unchanged (B stays read exactly once: 16-B contiguous per n-row, each +// 64-B sector read once at the uniform S=16 default). +// +// Iteration 12 (staging-load MLP batching, alignment-safe). The fresh +// exact-source ISA of the accepted iteration-11 object shows both staging +// loops still compile to a fully serialized load -> s_waitcnt vmcnt(0) -> +// store chain (A: 16 sequential global round trips per block; B: 4), i.e. +// every block pays ~20 serial global round trips before the single barrier, +// while the K loop is already lean (2 LDS reads + 2 lgkmcnt waits + 1 v_mmac +// per step). Iteration 11 already validated the B-side fix (16 serialized +// dword loads -> 4 aligned int4 loads as part of the accepted 40.38 -> +// 30.66 us round). This round batches both staging loops 4-deep in source: +// 4 independent loads are issued before their 4 stores, so one vmcnt drain +// covers 4 in-flight round trips (A 16 -> 4, B 4 -> 1; ~20 -> ~5 per block). +// A stays dword-wide because kAStride = kMaxKlen+4 is not 16-B aligned -- +// round 9's 58.56 us regression was misaligned int4 staging on the old +// layouts, not an MLP falsification; B stays int4 (aligned). No layout, +// stride, fragment, slice, plane, combine, guard or fallback changes: +// outputs stay bit-identical (0 mismatches) and Graph capture/replay with +// changed contents must pass. vmem_read/lds instruction counts are +// unchanged (25,600 / 46,080); the wait structure is what changes. +// +// The ZTH_W8A8_QKV_SPLIT_K environment override keeps its semantics but now +// selects the data path too: S in {6,8,9,12,16} run the staged kernel (their +// slices fit the 64 KiB LDS), S in {2,3,4,5} keep the iteration-3 direct-load +// kernel (S=2's 2,048-K slice alone would need 64 KiB for B, and S=3's +// 1,408-K slice does not fit either). +// +// Round-3 geometry (direct-load path, still selectable for S in +// {2,3,4,5}): split-K=2 -> 160 +// one-wave zero-barrier blocks +// (160 independent wavefronts = the same total parallelism as iteration 2's +// 80 blocks x 2 waves, but with no in-block barrier, no LDS partial planes, +// and 160 independent block-level load streams that the block scheduler can +// co-resident 2-deep on 40 of the 120 CUs). The trusted occupancy-probe +// sweep SPLIT_K in {2,3,4,5,6,8,9,12,16} is implemented as template +// instantiations and selectable at launch time with the +// ZTH_W8A8_QKV_SPLIT_K environment variable, so the control plane can +// measure split-K=2 and the CU-aligned candidates without source edits. +// For the direct path only, the CU-aligned (integer blocks/CU) candidates +// for N/16=80 tiles were S=3 -> 240 blocks = exactly 2 blocks/CU +// (non-power-of-two), S=6 -> 480 = 4/CU, S=9 -> 720 = 6/CU, S=12 -> 960 = +// 8/CU; with the iteration-7/8 routing those S values now run the staged +// kernel instead (block counts 480/720/960/1,280 for S=6/9/12/16), and +// S in {2,3,4,5} keep the direct path. Every candidate fits the contract +// workspace (1,310,720 bytes = capacity 16 planes of 81,920 B; S=16 uses +// exactly all 16 planes) and keeps +// every K slice a multiple of the DUMMA K tile 32 (non-uniform slices differ +// by at most one 32-K step; S=16 is uniform: 128/16 = 8 steps). +// +// Exact int32 accumulation is preserved: each slice accumulates int32 over an +// ascending 32-aligned K range; slices tile [0, K) exactly once; the combine +// sums the partials in int32 ascending split order (per-tile max |dot| = +// 4096*127*127 = 66,064,384 << 2^31), so every per-tile dot total is +// bit-identical to the scalar fallback and to iterations 1-2 regardless of +// SPLIT_K. The scale + bf16 (RN-even) epilogue runs once per output element +// in the combine kernel. +// +// Launch is guarded by the exact (m,n,k) = (16,1280,4096) shape; every other +// shape (including the paired M=2 shape with the same (N,K)) falls back to +// the generic scalar kernel below. launch_pack_w8a8_weight is an identity +// device-to-device copy for every (K,N) except the exact pair +// (k,n)==(4096,1280), which is transposed once to an [N,K] n-major packed +// layout (iteration 11) so DUMMA col_major B fragments read 8 consecutive +// bytes per lane; the scalar fallback decodes that packed layout for the +// same pair, so every path stays correct against the packed buffer. +// +// Graph-safety: this file only launches kernels on the caller-provided stream +// (PyTorch's current HIP stream). It performs no allocation, compilation, +// autotuning, packing, host synchronization, or device synchronization, and +// touches no buffer other than the caller-provided out/workspace pointers. +// The split-K choice is a host-side static dispatch decision (env override or +// default 16) made once per launch call, never inside the Graph replay, so +// capture/replay determinism holds. +// +// Header order is deliberate: hip_runtime.h, then hip_bfloat16.h, then +// du_mma.h (du_mma.h on this DTK is not self-contained before the HIP runtime +// headers). + +#include +#include +#include + +#include +#include + +namespace { + +// gfx928 native wavefront is 64 lanes; every block size below is a multiple +// of 64. +constexpr int kScalarBlockThreads = 128; +constexpr int kCopyBlockThreads = 256; + +// Minimal DUMMA tile constants (gfx928 INT8 support is m16n16k32). +constexpr int kDummaM = 16; +constexpr int kDummaN = 16; +constexpr int kDummaK = 32; +constexpr int kDummaThreads = 64; // one wavefront (64 lanes) per block +constexpr int kCombineThreads = 256; // 16x16 = 256 output elements/block + +// Exact (k,n) pair whose weight launch_pack_w8a8_weight transposes to the +// [N,K] n-major packed layout and that the scalar fallback decodes as +// n-major (iteration 11; covers the paired M=2 validation shape). +constexpr int kPackedK = 4096; +constexpr int kPackedN = 1280; + +// Trusted occupancy-probe split-K candidates (control plane): all fit the +// contract workspace for this shape (1,310,720 B = 16 partial planes of +// 81,920 B each) and produce 32-aligned K slices. Block counts = 80*S: +// S=3 -> 240 blocks = exactly 2 blocks/CU (CU-aligned, non-power-of-two), +// S=6 -> 480 = 4/CU, S=9 -> 720 = 6/CU, S=12 -> 960 = 8/CU, +// S=16 -> 1,280 = 10.67/CU (16 is outside the old probe list but is the +// largest split that fits the 16-plane workspace, divides 128 K-steps +// exactly into uniform 256-K slices, and is the iteration-8 default: it +// lowers the binding occupancy limiter -- per-block LDS footprint -- from +// 12,736 B (5 resident/CU) to 9,280 B (7 resident/CU)). +constexpr int kTrustedSplitK[] = {2, 3, 4, 5, 6, 8, 9, 12, 16}; + +// --------------------------------------------------------------------------- +// Scalar GEMM kernel: one thread computes one output element (generic +// fallback for every unmatched shape, including the paired M=2 shape). +// --------------------------------------------------------------------------- +// Adjacent lanes own adjacent N columns (the fastest-changing dimension), so +// output stores are coalesced across each wavefront. +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * kScalarBlockThreads + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + + // Complete K loop, k ascending, accumulated exactly in int32. The maximum + // API K (<= 6144) keeps the int8 dot well inside int32 range. + const int8_t* a_row = a + static_cast(row) * k; + // Iteration 11: the exact (k,n)==(4096,1280) pair is packed [N,K] n-major + // by launch_pack_w8a8_weight; the fallback decodes that layout so it stays + // correct against the packed buffer (paired M=2 validation shape included). + // Every other (K,N) keeps the raw [K,N] row-major decode. + const bool packed = (k == kPackedK && n == kPackedN); + const int8_t* b_col = + b + (packed ? static_cast(col) * k : col); + const int64_t b_stride = packed ? 1 : static_cast(n); + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * b_stride]); + } + + // Same float32 evaluation order as the exact reference: + // (float(dot) * x_scale[m]) * weight_scale[n], then bf16 (RN). + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +void launch_scalar_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + hip_bfloat16* out, + int m, + int n, + int k, + hipStream_t stream) { + const int64_t total = static_cast(m) * n; + const unsigned grid = static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + dim3(grid), + dim3(kScalarBlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out, + m, + n, + k); +} + +// --------------------------------------------------------------------------- +// Workspace split-K partial kernel (template over SPLIT_K), M == 16: +// partial[tile, slice][m, n] = int32_dot(x_q[m, k0:k1], weight[k0:k1, n]) +// --------------------------------------------------------------------------- +// Layouts (all contiguous): +// x_q [16, K] int8 row-major (ldm = K) +// weight [K, N] int8 row-major (ldm = N, identity pack layout) +// partials [(SPLIT_K * num_tiles) * 256] int32, plane = slice*num_tiles + +// tile, row-major 16x16 tile per plane (caller workspace) +// One block = one wavefront = one (tile, slice). The K slice is 32-aligned +// (non-uniform: slices differ by at most one 32-K DUMMA step) and covers +// [k0, k0 + steps*32) in ascending order. No LDS, no __syncthreads, no +// atomics: each block writes exactly one private int32 partial plane. +template +__global__ __launch_bounds__(kDummaThreads) void +w8a8_dumma_m16_splitk_partial_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int32_t* __restrict__ partials, + int n, + int k) { + const int tile = static_cast(blockIdx.x); + const int slice = static_cast(blockIdx.y); + const int n0 = tile * kDummaN; + + // 32-aligned non-uniform slice bounds over the K dimension. + const int total_steps = k / kDummaK; // 4096 / 32 = 128 + const int base = total_steps / SPLIT_K; + const int rem = total_steps - base * SPLIT_K; + const int steps = base + (slice < rem ? 1 : 0); + const int k0 = (slice * base + (slice < rem ? slice : rem)) * kDummaK; + + du::dumma::DUFragment + a_frag; + // Iteration 11: b is the packed [N,K] n-major weight for this shape, so + // the B fragment is col_major (each lane's 8 fragment bytes are consecutive + // in memory -> one 8-B read instead of 8 byte reads). + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + for (int i = 0; i < steps; ++i) { + const int kk = k0 + i * kDummaK; + du::dumma::du_load_matrix_sync(a_frag, a + kk, k); + // Packed n-major row n starts at b[n*k + k0]; the library col_major + // loader addresses p[(lane&15)*ldm + ((lane>>4)<<3) + i], i.e. n-major + // rows of k bytes, so the tile's 16 columns live at n0..n0+15 n-rows + // with row stride k. + du::dumma::du_load_matrix_sync( + b_frag, b + static_cast(n0) * k + kk, k); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Private int32 partial plane: slice-major so a tile's S planes are + // stride-uniform in the combine kernel. + const int num_tiles = static_cast(gridDim.x); + du::dumma::du_store_matrix_sync( + partials + (slice * num_tiles + tile) * (kDummaM * kDummaN), + acc_frag, + kDummaN, + du::dumma::mem_row_major); +} + +// Template launch helper for the partial kernel (one wavefront per block, +// grid = (num_tiles, SPLIT_K)). +template +void launch_splitk_partial( + const int8_t* a, + const int8_t* b, + int32_t* partials, + int n, + int k, + hipStream_t stream, + int num_tiles) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk_partial_kernel), + dim3(static_cast(num_tiles), static_cast(SPLIT_K)), + dim3(static_cast(kDummaThreads)), + 0, + stream, + a, + b, + partials, + n, + k); +} + +// Iteration 11 packed B-fragment loader (validated on the TP4 gate_up +// (4096,768) lineage): with n-major [N,K] storage and a col_major fragment, +// du_mma.h's matrix_b col_major address rule is x[i] = p[(lane&15)*ldm + +// ((lane>>4)<<3) + i] -- lane l's 8 fragment bytes are consecutive in +// memory, so one 8-byte LDS read per fragment replaces the 8 ds_read_u8 at +// stride-64 addresses that the row-major loader emits. du_mma_sync consumes +// the fragment as one 64-bit value per lane, so writing the packed bytes +// straight into x[0..7] is exactly the library's own storage convention. +// p must be 8 B aligned (guaranteed: b_s is __align__(256), kBStride % 8 == +// 0, step offsets are multiples of 32, and (lane&15)*kBStride % 8 == 0). +__device__ __forceinline__ void load_b_frag_packed( + du::dumma::DUFragment& b_frag, + const int8_t* p, int ldm, int lane) { + const int row = lane & 15; + const int col = (lane >> 4) << 3; + const uint64_t v = *reinterpret_cast( + p + static_cast(row) * ldm + col); + *reinterpret_cast(b_frag.x) = v; +} + +// --------------------------------------------------------------------------- +// Whole-K-slice LDS staging partial kernel (M == 16, iteration 7): +// partial[tile, slice][m, n] = int32_dot(x_q[m, k0:k1], weight[k0:k1, n]) +// --------------------------------------------------------------------------- +// Same 32-aligned non-uniform K slices and the same k-ascending int32 +// du_mma accumulation as w8a8_dumma_m16_splitk_partial_kernel, so the +// combine+scale kernel below sums bit-identical partials and the outputs are +// bit-identical to the accepted iteration-3 kernel. Differences: +// - A (16 x slice) is staged ONCE into LDS with coalesced dword global +// loads (4 B/lane, 256 B per wavefront instruction); B (16 x slice) is +// staged as 16-B int4 loads straight from the packed [N,K] n-major +// weight (iteration 11) into an n-major LDS tile. +// - One __syncthreads() after staging, then a zero-barrier LDS-only K +// loop: du_load_matrix_sync reads the A fragment from LDS and +// load_b_frag_packed reads each B fragment with one ds_read_b64 (~30-cycle +// LDS latency) instead of a global round trip per K step. +// Layouts (all contiguous): +// x_q [16, K] int8 row-major (ldm = K) +// weight [K, N] int8 row-major for every (K,N) except the exact pair +// (k,n)==(4096,1280), which is packed [N,K] n-major (element +// (k,n) at P[n*K+k]; iteration 11) +// partials [(SPLIT_K * num_tiles) * 256] int32, plane = slice*num_tiles + +// tile, row-major 16x16 tile per plane (caller workspace) +// LDS: a_s[16][kMaxKlen + 4] (compile-time row stride kMaxKlen+4, an odd +// dword count -> the library's matrix_a row-major lane pattern (row = +// lane&15, 8 consecutive bytes per lane) is bank-conflict-free); +// b_s[16][kMaxKlen + 16] (iteration 11: n-major, row stride kMaxKlen+16 = +// 16-B aligned for the int4 staging stores, 8-B aligned for the col_major +// fragment b64 reads, not a multiple of 128 B -> the 16 staged n-rows do not +// alias onto one LDS bank phase; each lane's 8 fragment bytes are +// consecutive, so every B fragment loads with ONE ds_read_b64 instead of 8 +// ds_read_u8). Per-template LDS fits: S=6 22,848 B, S=8 16,704 B, S=9 +// 15,680 B, S=12 11,584 B, S=16 8,512 B (all <= 64 KiB; resident blocks/CU +// by LDS: 2/3/4/5/7). Only launched +// under the exact (m,n,k) == (16,1280,4096) guard, so total_steps = k/32 = +// 128 is exact for the compile-time LDS sizing below. +template +__global__ __launch_bounds__(kDummaThreads) void +w8a8_dumma_m16_slicestage_partial_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int32_t* __restrict__ partials, + int n, + int k) { + const int tile = static_cast(blockIdx.x); + const int slice = static_cast(blockIdx.y); + const int n0 = tile * kDummaN; + + // 32-aligned non-uniform slice bounds (identical formula to the direct + // kernel: slices tile [0, K) exactly once in ascending order). + const int total_steps = k / kDummaK; // 4096 / 32 = 128 under the guard + const int base = total_steps / SPLIT_K; + const int rem = total_steps - base * SPLIT_K; + const int steps = base + (slice < rem ? 1 : 0); + const int k0 = (slice * base + (slice < rem ? slice : rem)) * kDummaK; + const int klen = steps * kDummaK; + + // Compile-time LDS sizing for the longest slice this instantiation can + // receive (base+1 steps, or base when SPLIT_K divides 128 evenly). + constexpr int kMaxSteps = 128 / SPLIT_K + (128 % SPLIT_K == 0 ? 0 : 1); + constexpr int kMaxKlen = kMaxSteps * kDummaK; + __shared__ int8_t a_s[16][kMaxKlen + 4]; + // Iteration 11: B is staged n-major (16 columns x kMaxKlen+16 bytes, row + // stride kMaxKlen+16) so col_major B fragments read 8 consecutive bytes + // per lane (one ds_read_b64). kMaxKlen is a multiple of 32, so the row + // stride is 16-B aligned (int4 staging stores stay aligned), 8-B aligned + // (fragment b64 reads stay aligned) and not a multiple of 128 B (no LDS + // bank-phase aliasing across the 16 staged n-rows). + __shared__ __align__(256) int8_t b_s[16][kMaxKlen + 16]; + + // Row strides are compile-time constants equal to the array row strides + // (a_s: kMaxKlen+4 bytes = odd dword count; b_s: kMaxKlen+16 bytes), so + // the staging stores and the K-loop fragment reads always agree even for + // short slices (klen < kMaxKlen). The A stride stays an odd dword count + // (bank-conflict-free row-major lane pattern); the B n-major stride keeps + // every fragment b64 read 8 B aligned and off the single bank phase. + constexpr int kAStride = kMaxKlen + 4; + constexpr int kBStride = kMaxKlen + 16; + + // Stage A: 16 rows x klen bytes = 4*klen dwords; consecutive lanes copy + // consecutive dwords of one row (coalesced 256 B per wavefront load). + // Iteration 12 (staging MLP, alignment-safe): the fresh exact-source ISA + // (iteration-12 profile, digest ab7fc108...) shows both staging loops + // compile to a fully serialized load -> s_waitcnt vmcnt(0) -> store chain + // (A: 16 sequential global round trips per block, B: 4), and the accepted + // iteration-11 round already validated this fix on the B side (16 + // serialized dword loads -> 4 aligned int4 loads). Both loops below batch + // 4-deep in source: 4 independent loads are issued before their 4 stores, + // so one vmcnt drain covers 4 in-flight round trips (A 16 -> 4, B 4 -> 1). + // A stays dword-wide: 4-B loads AND stores are aligned at every a_s row, + // while kAStride = kMaxKlen+4 is not 16-B aligned, so int4 A staging would + // repeat round 9's misaligned-store regression (58.56 us) -- that round + // falsified misaligned int4 staging, not the MLP idea. B stays int4 (its + // global rows are 16-B aligned and kBStride % 16 == 0). Fragment reads, + // LDS strides, the K loop and the int32 accumulation order are untouched, + // so outputs stay bit-identical (0 mismatches) and Graph capture/replay + // with changed contents must pass. klen % 32 == 0 => a_dwords % 128 == 0, + // so the batch loop is exact whenever klen % 128 == 0 (S=16: klen = 256 -> + // 1024 dwords = 4 batches exactly; S=6/9/12's odd slices fall through to + // the scalar tail, which preserves correctness for every S in the sweep). + const int a_dwords = 4 * klen; + const int a_dw_per_row = klen / 4; + constexpr int kABatch = 4 * kDummaThreads; + { + const auto ld_a = [&](int e) { + const int row = e / a_dw_per_row; + return *reinterpret_cast( + a + static_cast(row) * k + k0 + (e - row * a_dw_per_row) * 4); + }; + const auto st_a = [&](int e, int32_t v) { + const int row = e / a_dw_per_row; + *reinterpret_cast(&a_s[row][(e - row * a_dw_per_row) * 4]) = v; + }; + int e = static_cast(threadIdx.x); + for (; e + 3 * kDummaThreads < a_dwords; e += kABatch) { + const int32_t v0 = ld_a(e); + const int32_t v1 = ld_a(e + kDummaThreads); + const int32_t v2 = ld_a(e + 2 * kDummaThreads); + const int32_t v3 = ld_a(e + 3 * kDummaThreads); + st_a(e, v0); + st_a(e + kDummaThreads, v1); + st_a(e + 2 * kDummaThreads, v2); + st_a(e + 3 * kDummaThreads, v3); + } + for (; e < a_dwords; e += kDummaThreads) { + st_a(e, ld_a(e)); + } + } + + // Stage B (iteration 11): the exact (k,n)==(4096,1280) weight is packed + // [N,K] n-major (element (k,n) at P[n*K+k]) by launch_pack_w8a8_weight, so + // the block's 16-column strip is 16 n-rows of klen contiguous bytes. Each + // lane stages 16-B int4s straight into the n-major LDS tile b_s[16][ + // kMaxKlen+16] (4 loads per lane at the uniform S=16 default: rows 0..15 x + // 16 int4 columns). Iteration 12: the same 4-deep batching as A, so the 4 + // aligned int4 round trips collapse to one vmcnt drain per block. + const int k16_per_row = klen / 16; + const int b_int4s = kDummaN * k16_per_row; + const int4* __restrict__ b4 = reinterpret_cast(b); + int4* __restrict__ lds_b4 = reinterpret_cast(b_s); + constexpr int kBBatch = 4 * kDummaThreads; + { + const auto ld_b = [&](int e) { + const int row = e / k16_per_row; + return b4[(static_cast(n0) + row) * (k >> 4) + (k0 >> 4) + + (e - row * k16_per_row)]; + }; + const auto st_b = [&](int e, const int4& v) { + const int row = e / k16_per_row; + lds_b4[row * (kBStride >> 4) + (e - row * k16_per_row)] = v; + }; + int e = static_cast(threadIdx.x); + for (; e + 3 * kDummaThreads < b_int4s; e += kBBatch) { + const int4 v0 = ld_b(e); + const int4 v1 = ld_b(e + kDummaThreads); + const int4 v2 = ld_b(e + 2 * kDummaThreads); + const int4 v3 = ld_b(e + 3 * kDummaThreads); + st_b(e, v0); + st_b(e + kDummaThreads, v1); + st_b(e + 2 * kDummaThreads, v2); + st_b(e + 3 * kDummaThreads, v3); + } + for (; e < b_int4s; e += kDummaThreads) { + st_b(e, ld_b(e)); + } + } + + __syncthreads(); + + du::dumma::DUFragment + a_frag; + // Iteration 11: col_major B fragments match the n-major staged tile, so + // each lane's 8 fragment bytes are consecutive (one ds_read_b64 per + // fragment in the K loop instead of 8 ds_read_u8). + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Zero-barrier LDS-only K loop: every fragment read comes from LDS, so no + // global latency sits on the MMA critical path after the single staging + // barrier above. + const int lane = static_cast(threadIdx.x) & 63; + for (int i = 0; i < steps; ++i) { + du::dumma::du_load_matrix_sync(a_frag, &a_s[0][i * kDummaK], kAStride); + // One 8-B ds_read_b64 per B fragment (col_major: each lane's 8 fragment + // bytes are consecutive in the n-major tile) instead of 8 ds_read_u8 at + // the 20-B k-major pitch. + load_b_frag_packed(b_frag, &b_s[0][i * kDummaK], kBStride, lane); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Private int32 partial plane: same slice-major layout as the direct + // kernel (plane = slice*num_tiles + tile). + const int num_tiles = static_cast(gridDim.x); + du::dumma::du_store_matrix_sync( + partials + (slice * num_tiles + tile) * (kDummaM * kDummaN), + acc_frag, + kDummaN, + du::dumma::mem_row_major); +} + +// Template launch helper for the staged partial kernel (one wavefront per +// block, grid = (num_tiles, SPLIT_K)); only instantiated for the SPLIT_K +// values whose longest slice fits the 64 KiB LDS (6, 8, 9, 12, 16). +template +void launch_slicestage_partial( + const int8_t* a, + const int8_t* b, + int32_t* partials, + int n, + int k, + hipStream_t stream, + int num_tiles) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_slicestage_partial_kernel), + dim3(static_cast(num_tiles), static_cast(SPLIT_K)), + dim3(static_cast(kDummaThreads)), + 0, + stream, + a, + b, + partials, + n, + k); +} + +// --------------------------------------------------------------------------- +// Combine + scale kernel (M == 16): one 16x16 tile per block, one thread per +// output element. Exact int32 sum of the SPLIT_K workspace partials +// (ascending split order), then x_scale/weight_scale scaling and the bf16 +// (RN-even) store. No barrier is needed: every thread reads only its own +// 16x16 element's S partial values (written by the previous kernel on the +// same stream). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kCombineThreads) void +w8a8_dumma_m16_combine_scale_kernel( + const int32_t* __restrict__ partials, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int split_k, + int num_tiles, + int n) { + const int tile = static_cast(blockIdx.x); + const int linear = static_cast(threadIdx.x); // 0..255 + const int row = linear >> 4; + const int col = linear & 15; + const int out_col = tile * kDummaN + col; + + int32_t sum = 0; + for (int s = 0; s < split_k; ++s) { + sum += partials[(s * num_tiles + tile) * (kDummaM * kDummaN) + linear]; + } + + const float scaled = + static_cast(sum) * x_scale[row] * weight_scale[out_col]; + out[row * n + out_col] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Weight pack op (outside the timed region). +// --------------------------------------------------------------------------- +// For every (K,N) except the exact pair (k,n)==(4096,1280) the packed weight +// stays a contiguous [K, N] int8 buffer with an [N, 1] fp32 scale, +// byte-identical to the raw inputs. For (k,n)==(4096,1280) (iteration 11) +// launch_pack_w8a8_weight transposes the raw [K,N] weight to an [N,K] +// n-major packed layout (element (kk,nn) -> dst[nn*k + kk]) so DUMMA +// col_major B fragments read 8 consecutive bytes per lane (one ds_read_b64 +// instead of 8 ds_read_u8); the staged/direct DUMMA kernels and the scalar +// fallback (for the paired M=2 validation shape) all decode that layout. +__global__ __launch_bounds__(kCopyBlockThreads) void +w8a8_identity_copy_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t linear = + static_cast(blockIdx.x) * kCopyBlockThreads + threadIdx.x; + if (linear < count) { + dst[linear] = src[linear]; + } +} + +__global__ __launch_bounds__(kCopyBlockThreads) void +w8a8_identity_copy_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t linear = + static_cast(blockIdx.x) * kCopyBlockThreads + threadIdx.x; + if (linear < count) { + dst[linear] = src[linear]; + } +} + +// Iteration 11 one-time pack for the exact pair (k,n)==(4096,1280): +// transpose the raw [K,N] int8 weight to [N,K] n-major (element (kk,nn) -> +// dst[nn*k + kk]). Runs out of the timed region and out of Graph capture; +// performance is irrelevant (single weight prep per layer load). +__global__ __launch_bounds__(kCopyBlockThreads) void +w8a8_transpose_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int k, + int n) { + const int64_t linear = + static_cast(blockIdx.x) * kCopyBlockThreads + threadIdx.x; + const int64_t total = static_cast(k) * n; + if (linear >= total) { + return; + } + const int kk = static_cast(linear / n); + const int nn = static_cast(linear - static_cast(kk) * n); + dst[static_cast(nn) * k + kk] = src[linear]; +} + +} // namespace + +// --------------------------------------------------------------------------- +// Stable host launch symbols (declared in csrc/bindings.cpp). +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + auto* out_bf16 = static_cast(out); + + // Exact-shape guard for the assigned shape hy3_tp8_qkv_proj_m16 + // (M=16, N=1280, K=4096). The DUMMA specialization can never capture + // other shapes (in particular the paired M=2 shape with the same (N, K)). + if (m == 16 && n == 1280 && k == 4096) { + const int num_tiles = n / kDummaN; // 1280 / 16 = 80 + const int64_t plane_bytes = + static_cast(num_tiles) * kDummaM * kDummaN * + static_cast(sizeof(int32_t)); // 80 * 256 * 4 = 81,920 + + // Split-K selection: default 16 (the iteration-8 staged geometry: 1,280 + // one-wave blocks = 10.67 blocks/CU; per-block LDS drops from 12,736 B + // (5 resident/CU, the binding occupancy limiter at S=12 per PMC: grid + // 960 = 8 waves/CU available, VGPR 24 / 0 scratch not limiting) to + // 9,280 B = 7 resident/CU, with uniform 256-K slices = exactly 8 MMAs + // per block; A/B HBM bytes unchanged). The ZTH_W8A8_QKV_SPLIT_K + // environment variable selects any trusted probe candidate + // {2,3,4,5,6,8,9,12,16} so the sweep can be measured without source + // edits (S=6 -> 480 blocks = 4/CU, S=9 -> 720 = 6/CU, S=12 -> 960 = 8/CU, + // S=16 -> 1,280 = 10.67/CU, S=3 -> 240 direct-load blocks, ...). The + // read happens once per launch call on the host, never during Graph + // replay, so capture/replay determinism is preserved. + int split_k = 16; + if (const char* env = std::getenv("ZTH_W8A8_QKV_SPLIT_K")) { + char* end = nullptr; + const long parsed = std::strtol(env, &end, 10); + if (end != env && *end == '\0') { + for (int t : kTrustedSplitK) { + if (static_cast(parsed) == t) { + split_k = t; + break; + } + } + } + } + + // Workspace fit: clamp to the largest trusted candidate that fits the + // caller workspace (contract guarantees 16 planes = 1,310,720 B for this + // shape, so the clamp never fires; kept defensive). + int best_fit = 0; + for (int t : kTrustedSplitK) { + if (static_cast(t) * plane_bytes <= workspace_bytes) { + best_fit = t; + } + } + if (best_fit == 0) { + launch_scalar_gemm(a, b, x_scale, weight_scale, out_bf16, m, n, k, + stream); + return; + } + if (split_k > best_fit) { + split_k = best_fit; + } + + // Int32 partial planes in the caller workspace, then the combine+scale + // kernel, both on the caller's stream inside the timed Graph. + // Iteration-7/8 staged kernels: S in {6,8,9,12,16} (their longest slice + // fits the 64 KiB LDS; see w8a8_dumma_m16_slicestage_partial_kernel). + // Iteration-3 direct-load kernels: S in {2,3,4,5} (their slices exceed + // the staging capacity). Both families produce the same slice-major + // int32 planes, so the combine kernel is shared and unchanged. + int32_t* partials = static_cast(workspace); + switch (split_k) { + case 6: + launch_slicestage_partial<6>(a, b, partials, n, k, stream, + num_tiles); + break; + case 8: + launch_slicestage_partial<8>(a, b, partials, n, k, stream, + num_tiles); + break; + case 9: + launch_slicestage_partial<9>(a, b, partials, n, k, stream, + num_tiles); + break; + case 12: + launch_slicestage_partial<12>(a, b, partials, n, k, stream, + num_tiles); + break; + case 16: + launch_slicestage_partial<16>(a, b, partials, n, k, stream, + num_tiles); + break; + case 2: + launch_splitk_partial<2>(a, b, partials, n, k, stream, num_tiles); + break; + case 3: + launch_splitk_partial<3>(a, b, partials, n, k, stream, num_tiles); + break; + case 4: + launch_splitk_partial<4>(a, b, partials, n, k, stream, num_tiles); + break; + case 5: + launch_splitk_partial<5>(a, b, partials, n, k, stream, num_tiles); + break; + default: + launch_scalar_gemm(a, b, x_scale, weight_scale, out_bf16, m, n, k, + stream); + return; + } + + hipLaunchKernelGGL( + w8a8_dumma_m16_combine_scale_kernel, + dim3(static_cast(num_tiles)), + dim3(static_cast(kCombineThreads)), + 0, + stream, + partials, + x_scale, + weight_scale, + out_bf16, + split_k, + num_tiles, + n); + return; + } + + // Generic scalar fallback for every unmatched (m, n, k). + launch_scalar_gemm(a, b, x_scale, weight_scale, out_bf16, m, n, k, stream); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Out-of-timed-region weight prep: [N,K] n-major transpose pack for the + // exact pair (k,n)==(4096,1280) (iteration 11), identity copy otherwise. + const int64_t weight_count = static_cast(k) * n; + const unsigned weight_grid = static_cast( + (weight_count + kCopyBlockThreads - 1) / kCopyBlockThreads); + if (k == kPackedK && n == kPackedN) { + // Iteration 11: [N,K] n-major transpose pack for the exact pair so DUMMA + // col_major B fragments read 8 consecutive bytes per lane (one + // ds_read_b64 instead of 8 ds_read_u8). + hipLaunchKernelGGL( + w8a8_transpose_i8_kernel, + dim3(weight_grid), + dim3(kCopyBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + hipLaunchKernelGGL( + w8a8_identity_copy_i8_kernel, + dim3(weight_grid), + dim3(kCopyBlockThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count); + } + + const unsigned scale_grid = static_cast( + (static_cast(n) + kCopyBlockThreads - 1) / kCopyBlockThreads); + hipLaunchKernelGGL( + w8a8_identity_copy_f32_kernel, + dim3(scale_grid), + dim3(kCopyBlockThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + static_cast(n)); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_down_proj.hip new file mode 100644 index 00000000..01194376 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_down_proj.hip @@ -0,0 +1,537 @@ +// @@variant shape=hy3_tp8_shared_down_proj_m16 commit=6f8e396e4eb7b316acecbb928ea5db9f2f3ffe4d added=2026-08-27 +// median_us=7.696 p90_us=8.174 +// source=hy3-dsh-tp8-m16-2-dc53295a +// INT8 W8A8 GEMM implementation for gfx928 (K500SM_AI). +// +// Worker: worker_3, physical GPU 3. +// Assigned shape: hy3_tp8_shared_down_proj_m16 -> (M, N, K) = (16, 4096, 192). +// +// Iteration 1: minimal DUMMA bootstrap. The exact assigned shape +// (M=16, N=4096, K=192) is served by a native gfx928 DUMMA INT8 +// m16n16k32 kernel with one 64-lane wavefront per block, one 16x16 output +// tile per block, and explicit int32 accumulation: +// - A fragment: du::dumma::DUFragment, loaded directly from global memory each K step with +// leading dimension K. +// - B fragment: the analogous matrix_b fragment, loaded directly from +// global memory at (k0, n0) with leading dimension N. +// - C: du::dumma::DUFragment, filled to 0 and +// accumulated by du_mma_sync over the 6 K-tiles (K=192 = 6*32). +// - Epilogue: du_store_matrix_sync (mem_row_major) into a 1 KiB shared +// tile, one intra-block __syncthreads (single wavefront, no cross-wave +// barrier), then fp32 scale multiply in the reference order +// (dot * x_scale[m] * weight_scale[n]) and a coalesced bf16 store. +// +// Iteration 2 (rejected): 2-wave/128-block geometry probe regressed to +// 18.165 us median, falsifying the launch-geometry hypothesis and pointing +// the search at the load path. +// +// Iteration 3 (this revision): load path. The iteration-1 code object showed +// the poison pattern - du_load_matrix_sync expands to ~16 global_load_ubyte +// per K=32 step with a full vmcnt drain before each single +// v_mmac_i32_16x16x32_i8 (6 serialized chains per block). This revision +// removes global traffic from the K loop entirely: +// - A [16,192] and the block's B tile [192,16] are staged ONCE into LDS +// with 16-B vector loads (global_load_dwordx4 -> ds_write_b128, 3 per +// lane per tile), rows padded 192 -> 208 (kPaddedK = 192 + 16 bank-skew +// stride, the TP4-validated down_proj stride), +// - B is consumed from the n-major packed layout P[n*K+kk] = W[kk*N+n] +// produced out of the timed region / Graph by launch_pack_w8a8_weight, +// so each lane's B fragment is 8 contiguous bytes of one packed row, +// - every m16n16k32 step then reads both fragments as one 8-B LDS read +// (ds_read_b64) per lane and issues one v_mmac_i32_16x16x32_i8; no +// vmcnt dependency exists inside the K loop (single wavefront, LDS only). +// The library fragment lane layout (row = lane&15, col = (lane>>4)*8) and +// byte order are reproduced exactly, so the int32 accumulation order is +// unchanged and outputs stay bit-identical to the accepted iteration-1 +// kernel. +// +// Iteration 4 (this revision): architecture/pipeline round - multi-N-tile +// reuse. Each block now covers TWO adjacent 16-column N tiles (16x32 output) +// with the same single 64-lane wavefront: grid = 4096/32 = 128 blocks, still +// >= the 120 physical CUs so every CU keeps at least one block. The A tile +// [16,192] is staged ONCE per block (3 global_load_dwordx4 per lane, same as +// before) and every lane's 8-B A fragment for each of the 6 K steps is read +// from LDS ONCE and consumed by BOTH m16n16k32 MMACs (tile 0 and tile 1): +// A bytes are reused 2x per step in registers, so A global traffic and A LDS +// staging/read traffic halve vs two separate 16-col blocks. B bytes are NOT +// reused - each of the block's 32 packed-B rows is read exactly once per K +// step (8 B per lane per tile per step); B traffic per output column is +// unchanged. The two accumulator fragments are independent, so each K step +// issues two v_mmac_i32_16x16x32_i8 with no accumulator dependency between +// them (2x MMAC ILP vs the iteration-3 single serialized chain), while the +// dispatch block count halves (256 -> 128). The trusted occupancy-probe +// split candidates [2, 3] were evaluated for this axis: 2 fits (128 blocks +// >= 120 CUs, 32 divides N=4096, 16-B stage alignment preserved); 3 would +// give ceil(4096/48) = 86 blocks < 120 CUs and 48 does not divide N=4096, +// so 3 is excluded by the cover-all-CUs constraint. (The alternative reading +// of [2,3] as waves-per-block was already falsified in iteration 2: +// 2-wave/128-block regressed to 18.165 us median; 3 waves/block would give +// 86 blocks < 120 CUs.) Everything else is unchanged: same staged +// [16][208]-padded LDS layout, same 6-step k-ascending int32 accumulation +// per tile (bit-identical fragments, so outputs stay bit-identical to the +// accepted iteration-1/3 kernels), same library lane layout, same LDS +// accumulator round-trip epilogue. +// +// Iteration 11 (accepted, current best 8.4951 us): HIP-only consolidation - +// isolate the ONE untested component of the iteration-9 epilogue bundle +// (measured 10.7085 us, -1.13%, inside the +/-2% band): the 17-word (68-B) +// padded accumulator rows. Iteration 10 isolated the bundle's scale hoisting +// alone (10.7575 us, -1.58%, inside the band); this revision isolates the +// bank-conflict removal ALONE: the two 16x16 int32 tiles are stored manually +// with the library's verified fragment layout (word = (lane&15)*17 + +// (lane>>4) + 4*i, the exact mapping du_store_matrix_sync writes with stride +// 16) into 17-word rows so row r starts at bank r (17 is coprime with 32) +// instead of du_store_matrix_sync into the natural 16-word rows, where every +// row starts at bank 0 and every epilogue LDS access hits a 16-way conflict +// per phase (PMC of the accepted code object measures 32768 +// lds_bank_conflicts = ~256/block, overwhelmingly from this round trip). The +// epilogue read mapping becomes row*17 + col (same values, same (row, col) +// positions); the accepted masked-branch scale loop, its in-loop x_scale +// loads (falsified as a separate axis in iteration 10), the scale order, +// output addresses, staging, K loop, int32 accumulation and the exact-shape +// guard are untouched, so outputs stay bit-identical. LDS 12032 -> 12160 B +// (still 5 blocks/CU of the 64 KiB), 0 s_barrier preserved (single +// wavefront, program-order stores then same-wave reads). +// +// Iteration 14 (this revision): HIP-only consolidation - merge the accepted +// two per-tile scale loops into ONE runtime masked loop that shares each +// pass's x_scale load across BOTH tiles. Exact-source ISA of the accepted +// object shows the epilogue as two separate runtime loops per block, each +// pass issuing ONE x_scale global_load_dword and then s_waitcnt vmcnt(0) +// before its store (0x2780/0x2798, 0x28A0/0x28B8) - 8 fully exposed L2-hot +// (64-B x_scale line) load latencies per block, and the per-pass x_scale row +// set (rows (lane>>4)+4*i) is IDENTICAL for both tiles (it does not depend +// on the tile index). The compiler cannot CSE across the two runtime loops, +// so the second tile re-loads the same four rows. The iteration-13 +// straight-line rewrite (batched up-front loads + unrolled bodies) regressed +// to 11.0035 us (-22.8%) and iteration-10 register hoisting was flat on the +// old base - both falsified forms changed the per-pass wait structure and +// added live registers; this revision instead keeps the EXACT accepted +// per-pass structure (in-loop x_scale load, per-pass vmcnt(0) wait, masked +// 4-pass loop, no new registers) and only removes the cross-tile redundancy: +// one 4-pass loop, each pass loading x_scale[row_m] once and feeding both +// tiles' (acc * x_scale * weight_scale) bodies. Exposed epilogue vmcnt(0) +// waits halve 8 -> 4 per block, loop/exec-mask machinery halves, vmem read +// instructions 2432 -> 1920 (19 -> 15 loads/block: 9 staging + 4 x_scale + 2 +// weight_scale), weight_scale stays LICM-hoisted (2/block), LDS 12160 B and +// 0 s_barrier unchanged. Same linear -> (row_m, col_n) mapping, same 17-word +// acc rows, same left-to-right fp32 order, same output addresses: outputs +// stay bit-identical. +// +// The scalar fallback kernel below remains the generic path for every +// unmatched (m, n, k), including the paired M=2 API shape with the same +// (N, K); it never enters the guarded DUMMA branch. Because the exact-shape +// pack is n-major, the scalar fallback decodes the packed B with a +// b_transposed flag whenever (n, k) == (4096, 192) (TP4 down_proj +// precedent); all other shapes keep the identity pack. +// +// Layout contract (logical, contiguous): +// x_q [M, K] int8, row-major, stride (K, 1) +// packed_weight[K, N] int8 (identity pack for all shapes EXCEPT the exact +// (k, n) == (192, 4096), where it is the n-major transpose +// P[n*192 + kk] = W[kk*4096 + n]; same byte count, same +// graph-stable buffer, produced once out of the timed region) +// x_scale [M, 1] fp32 +// packed_weight_scale [N, 1] fp32 +// out [M, N] bf16, row-major +// +// Mathematical reference (bit-exact for every supported shape because the +// full int8 dot fits exactly in int32 and every intermediate fp32 value is +// an exact integer below 2^24): +// out[m, n] = bf16((int32_dot(x_q[m, :], packed_weight[:, n])) +// * x_scale[m] * packed_weight_scale[n]) +// +// The timed operator (launch_w8a8_gemm) performs no allocation, compilation, +// autotuning, weight packing, host/device synchronization, or default-stream +// launch. It only launches kernels on the caller-provided stream and uses the +// caller-provided out tensor. Weight packing lives in the out-of-timed-region +// launch_pack_w8a8_weight (identity device-to-device copy). + +#include +#include +#include + +#include +#include + +namespace { + +// gfx928 native wavefront is 64 lanes; the block size must be a multiple of 64. +constexpr int kBlockThreads = 256; + +// DUMMA INT8 m16n16k32 tile constants (gfx928 native tensor core tile). +constexpr int kDummaWaveSize = 64; // one wavefront per block +constexpr int kDummaM = 16; +constexpr int kDummaN = 16; +constexpr int kDummaK = 32; + +// LDS row stride for the staged exact-shape tiles: K=192 padded by 16 to +// 208 (a 16-B-aligned non-power-of-two bank-skew stride; the same value the +// TP4 shared_down_proj lineage validated for its K=192 staging halves). +constexpr int kPaddedK = kDummaK * 6 + 16; // 208 = 192 + 16 + +// --------------------------------------------------------------------------- +// Generic scalar GEMM kernel: one thread per output element. +// +// Thread mapping keeps adjacent lanes on adjacent addresses in the +// fastest-changing N dimension: threads with consecutive threadIdx.x handle +// consecutive columns of the same row, so the B column loads (stride 1) are +// fully coalesced and A row loads are broadcast within the warp. +// +// b_transposed != 0 selects the n-major packed B layout +// P[n*K + kk] = W[kk*N + n] used by the exact-shape (n, k) == (4096, 192) +// pack; the per-column load then becomes P[col*K + kk] with the same +// k-ascending int32 accumulation order (bit-identical results). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kBlockThreads) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, // [M, K] row-major + const int8_t* __restrict__ b, // [K, N] row-major, or [N, K] + // n-major pack when transposed + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + __hip_bfloat16* __restrict__ out, // [M, N] row-major + const int n, + const int k, + const int b_transposed) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; + if (col >= n) { + return; + } + + const int8_t* a_row = a + static_cast(row) * k; + + // Complete K loop accumulated exactly in int32. The maximum supported K + // keeps the exact int8 dot well inside the int32 range, so this matches the + // fp32 reference bit-for-bit. + int32_t acc = 0; + if (b_transposed != 0) { + // Packed n-major B: b[kk * n + col] == packed[col * k + kk]. + const int8_t* b_col = b + static_cast(col) * k; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[kk]); + } + } else { + const int8_t* b_col = b + col; // column n index, stride n + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + } + + // Convert to float only for the two scale multiplies (same left-to-right + // order as the reference `(A@B) * x_scale * weight_scale.T`), then store + // bf16 with the header's round-to-nearest-even float conversion. + const float scaled = static_cast(acc) * x_scale[row] * weight_scale[col]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// DUMMA INT8 m16n16k32 kernel for the exact assigned shape (M=16, N=4096, +// K=192). One 64-lane wavefront per block; each block computes TWO adjacent +// 16x16 output tiles (16x32 output, n0 = blockIdx.x * 32) for all 16 rows, +// accumulating the six 32-wide K tiles explicitly in int32 per tile. +// +// Iteration 4 multi-N-tile reuse: A and the block's two packed-B tiles are +// staged once into LDS with 16-B vector loads (3 per lane per tile); every K +// step then reads the A fragment once and each tile's B fragment as one 8-B +// LDS read (ds_read_b64) per lane and issues two independent +// v_mmac_i32_16x16x32_i8 (A fragment register-reused across both MMACs). No +// vmcnt dependency exists inside the K loop. See the file header for the +// full mechanism and the bit-identical fragment byte order. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDummaWaveSize) void +w8a8_gemm_dumma_m16n32k32_kernel( + const int8_t* __restrict__ a, // [16, K] row-major + const int8_t* __restrict__ b, // [N, K] n-major pack + // P[n*K+kk] = W[kk*N+n] + const float* __restrict__ x_scale, // [16] + const float* __restrict__ weight_scale, // [N] + __hip_bfloat16* __restrict__ out, // [16, N] row-major + const int n, + const int k) { + const int lane = static_cast(threadIdx.x) & (kDummaWaveSize - 1); + // Two adjacent 16-column tiles per block: grid = N / 32 = 128 blocks. + const int n0 = static_cast(blockIdx.x) * (kDummaN * 2); + + // Exact-shape kernel (guard in launch_w8a8_gemm): K = 192 = 6 * 32. The + // compile-time step count is required for full unrolling of the K loop. + constexpr int kExactKSteps = kDummaK * 6 / kDummaK; // 6 + if (k != kDummaK * kExactKSteps) { + return; + } + + __shared__ __align__(16) int8_t a_smem[kDummaM][kPaddedK]; + __shared__ __align__(16) int8_t b_smem[2][kDummaN][kPaddedK]; + __shared__ __align__(16) int32_t acc_tile[2][kDummaM * (kDummaN + 1)]; + + // Cooperative global -> LDS staging with 16-B vector loads. All staged + // tiles use the same [16][208] layout (rows padded 192 -> 208 bank skew): + // a_smem[m][kk] = a[m][kk] (A is shared by all blocks) + // b_smem[0][cn][kk] = P[n0 + cn][kk] (packed B tile, rows n0..n0+15) + // b_smem[1][cn][kk] = P[n0 + 16 + cn][kk] (packed B tile, rows n0+16..n0+31) + // Each row splits into 12 chunks of 16 B; 64 lanes = 16 rows x 4 lanes per + // row, each lane copies 3 chunks ({0,1,2}, {3,4,5}, {6,7,8} or {9,10,11}). + const int srow = lane >> 2; // 0..15 + const int schunk = (lane & 3) * 3; // 0, 3, 6, 9 + const int8_t* a_src = a + srow * k + schunk * 16; + const int8_t* b0_src = b + (static_cast(n0) + srow) * k + + schunk * 16; + const int8_t* b1_src = b + (static_cast(n0) + 16 + srow) * k + + schunk * 16; + int8_t* a_dst = &a_smem[srow][schunk * 16]; + int8_t* b0_dst = &b_smem[0][srow][schunk * 16]; + int8_t* b1_dst = &b_smem[1][srow][schunk * 16]; + + uint4 a_chunk[3]; + uint4 b0_chunk[3]; + uint4 b1_chunk[3]; +#pragma unroll + for (int p = 0; p < 3; ++p) { + a_chunk[p] = *reinterpret_cast(a_src + p * 16); + b0_chunk[p] = *reinterpret_cast(b0_src + p * 16); + b1_chunk[p] = *reinterpret_cast(b1_src + p * 16); + } +#pragma unroll + for (int p = 0; p < 3; ++p) { + *reinterpret_cast(a_dst + p * 16) = a_chunk[p]; + *reinterpret_cast(b0_dst + p * 16) = b0_chunk[p]; + *reinterpret_cast(b1_dst + p * 16) = b1_chunk[p]; + } + __syncthreads(); + + du::dumma::DUFragment a_frag; + du::dumma::DUFragment b0_frag; + du::dumma::DUFragment b1_frag; + du::dumma::DUFragment acc0_frag; + du::dumma::DUFragment acc1_frag; + du::dumma::du_fill_fragment(acc0_frag, 0); + du::dumma::du_fill_fragment(acc1_frag, 0); + + // K loop: 192 / 32 = 6 DUMMA steps per tile, accumulated in the int32 + // fragments. The library fragment lane layout is row = lane&15, + // col = (lane>>4)*8 with a.x[i] = a[row][k0+col+i] and + // b.x[i] = b[k0+col+i][row]; all sequences are 8 contiguous bytes in the + // staged/packed layouts, so one 8-B LDS read per lane reproduces each + // fragment byte-identically. The A fragment is read once per step and + // register-reused by BOTH independent MMAC chains (A bytes reused 2x per + // step); the two accumulators have no dependency, so the two + // v_mmac_i32_16x16x32_i8 per step pipeline against each other. All loads + // are LDS: no vmcnt dependency and no global traffic in the loop. + const int row = lane & 15; + const int col = (lane >> 4) * 8; + const int frag_off = row * kPaddedK + col; +#pragma unroll + for (int s = 0; s < kExactKSteps; ++s) { + const int off = frag_off + s * kDummaK; + uint64_t a_val; + uint64_t b0_val; + uint64_t b1_val; + memcpy(&a_val, a_smem[0] + off, sizeof(a_val)); + memcpy(&b0_val, b_smem[0][0] + off, sizeof(b0_val)); + memcpy(&b1_val, b_smem[1][0] + off, sizeof(b1_val)); + memcpy(a_frag.x, &a_val, sizeof(a_val)); + memcpy(b0_frag.x, &b0_val, sizeof(b0_val)); + memcpy(b1_frag.x, &b1_val, sizeof(b1_val)); + du::dumma::du_mma_sync(acc0_frag, a_frag, b0_frag, acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag, b1_frag, acc1_frag); + } + + // Materialize both 16x16 int32 tiles in shared memory with the library's + // verified fragment layout - lane l owns acc[row = lane&15][col = (lane>>4) + // + 4*i] for i = 0..3, the exact word mapping du_store_matrix_sync writes + // (word = row*16 + col in the accepted 16-word rows) - but stored into + // 17-word (68-B) padded rows so row r starts at bank r: the natural 64-B + // row stride starts every row at bank 0, forcing a 16-way bank conflict per + // phase (the measured 32768 lds_bank_conflicts = ~256/block on the accepted + // code object come overwhelmingly from this round trip); a 17-word stride + // (17 coprime with 32) spreads the 16 rows over 16 distinct banks and drops + // every epilogue LDS access (4 ds_write2_b32 + 8 ds_read_b32 per block) to + // the ~2-4-way bandwidth floor. Same stored int32 values, same (row, col) + // positions, same single-wavefront program order (stores, then + // __syncthreads that elides to 0 s_barrier, then same-wave reads), so the + // scale loop below reads back bit-identical values. + constexpr int kAccRowWords = kDummaN + 1; // 17-word padded rows + const int acc_row = lane & 15; // fragment row + const int acc_col4 = lane >> 4; // fragment col group +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc_tile[0][acc_row * kAccRowWords + acc_col4 + 4 * i] = acc0_frag.x[i]; + acc_tile[1][acc_row * kAccRowWords + acc_col4 + 4 * i] = acc1_frag.x[i]; + } + __syncthreads(); + + // Epilogue (iteration 14): ONE runtime masked loop for both tiles. Each + // pass loads x_scale[row_m] once (the per-pass row set (lane>>4)+4*i is + // tile-independent) and feeds both tiles' scale/store bodies, halving the + // exposed global-load-latency waits (8 -> 4 per block) and the loop/ + // exec-mask machinery vs the accepted two separate tile loops, while + // keeping the exact accepted per-pass structure (in-loop x_scale load, + // per-pass s_waitcnt vmcnt(0), no new live registers - the iteration-10 + // register hoisting and iteration-13 straight-line forms were both + // falsified). Same linear -> (row_m, col_n) mapping, same 17-word acc + // rows, same left-to-right (acc * x_scale * weight_scale) fp32 order and + // same output addresses: outputs stay bit-identical. +#pragma unroll + for (int linear = lane; linear < kDummaM * kDummaN; + linear += kDummaWaveSize) { + const int row_m = linear >> 4; // / kDummaN + const int col_n = linear & 15; // % kDummaN + const int acc_off = row_m * kAccRowWords + col_n; + const float xs = x_scale[row_m]; + const float scaled0 = + static_cast(acc_tile[0][acc_off]) * xs * + weight_scale[n0 + col_n]; + out[static_cast(row_m) * n + n0 + col_n] = + __float2bfloat16(scaled0); + const float scaled1 = + static_cast(acc_tile[1][acc_off]) * xs * + weight_scale[n0 + kDummaN + col_n]; + out[static_cast(row_m) * n + n0 + kDummaN + col_n] = + __float2bfloat16(scaled1); + } +} + +// --------------------------------------------------------------------------- +// N-major transpose pack for the exact assigned shape (k == 192, n == 4096): +// packed[n_idx * k + kk] = weight[kk * n + n_idx] +// Runs once out of the timed region / Graph (weight preprocessing; same byte +// count, graph-stable buffer). Each 256-thread block covers one 16-column +// tile: threads 0-15 own the 16 columns and each 16-thread group copies 12 +// consecutive K rows (12 * 16 = 192 = K exactly). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kBlockThreads) void w8a8_pack_nmajor_kernel( + const int8_t* __restrict__ w, // [K, N] row-major raw weight + int8_t* __restrict__ packed, // [N, K] n-major packed weight + const int n, + const int k) { + const int n0 = static_cast(blockIdx.x) * 16; + const int n_idx = n0 + (static_cast(threadIdx.x) & 15); + const int k_base = (static_cast(threadIdx.x) >> 4) * 12; // 0..180 +#pragma unroll + for (int i = 0; i < 12; ++i) { + const int kk = k_base + i; + packed[static_cast(n_idx) * k + kk] = + w[static_cast(kk) * n + n_idx]; + } +} + +// --------------------------------------------------------------------------- +// Identity device-to-device copy used by the bootstrap pack_weight op. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kBlockThreads) void w8a8_identity_copy_kernel( + const T* __restrict__ src, + T* __restrict__ dst, + const int64_t numel) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < numel) { + dst[i] = src[i]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Timed GEMM launcher (stable host symbol required by csrc/bindings.cpp). +// +// Exact-shape guard first: only (m, n, k) == (16, 4096, 192) enters the +// assigned-shape DUMMA branch. The paired M=2 API shape with the same (N, K) +// (tp8_shared_down_proj_m2) never matches the guard and reaches the generic +// scalar path below, as required. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + + auto* out_bf16 = reinterpret_cast<__hip_bfloat16*>(out); + + if (m == 16 && n == 4096 && k == 192) { + // Exact assigned shape: hy3_tp8_shared_down_proj_m16. + // Grid: 4096 / 32 = 128 one-wave blocks, each computing a 16x32 output + // (two adjacent 16x16 N tiles sharing one staged A tile). + // `b` is the n-major packed weight produced by launch_pack_w8a8_weight + // for this exact (k, n). + const dim3 block(kDummaWaveSize); + const dim3 grid(n / (kDummaN * 2)); + hipLaunchKernelGGL(HIP_KERNEL_NAME(w8a8_gemm_dumma_m16n32k32_kernel), + grid, block, 0, stream, a, b, x_scale, weight_scale, + out_bf16, n, k); + return; + } + + // Generic scalar fallback for every unmatched (m, n, k), including the + // paired M=2 API shape with the same (N, K). Whenever (n, k) == (4096, 192) + // the packed B is the n-major transpose, so the scalar path decodes it via + // the b_transposed flag; every other shape keeps the identity pack. + const dim3 block(kBlockThreads); + const dim3 grid((n + kBlockThreads - 1) / kBlockThreads, m); + const int b_transposed = (n == 4096 && k == 192) ? 1 : 0; + hipLaunchKernelGGL(HIP_KERNEL_NAME(w8a8_gemm_scalar_kernel), grid, block, 0, + stream, a, b, x_scale, weight_scale, out_bf16, n, k, + b_transposed); +} + +// --------------------------------------------------------------------------- +// Optional out-of-timed-region weight packing launcher (stable host symbol +// required by csrc/bindings.cpp). +// +// For the exact assigned (k, n) == (192, 4096) it packs raw_weight [K, N] +// int8 -> packed_weight [N, K] n-major transpose P[n*K+kk] = W[kk*N+n] (the +// layout consumed by the exact-shape DUMMA kernel and the b_transposed +// scalar fallback). Every other (K, N) keeps the identity device-to-device +// copy raw_weight [K, N] -> packed_weight [K, N]. weight_scale [N] fp32 -> +// packed_weight_scale [N] fp32 is always the identity copy. Both run once +// out of the timed region / Graph with graph-stable buffer addresses. +// --------------------------------------------------------------------------- +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_numel = static_cast(k) * n; + if (k == 192 && n == 4096) { + // Exact assigned shape: n-major transpose pack (grid = N/16 tiles). + const dim3 block(kBlockThreads); + const dim3 grid(n / 16); + hipLaunchKernelGGL(HIP_KERNEL_NAME(w8a8_pack_nmajor_kernel), grid, block, + 0, stream, raw_weight, packed_weight, n, k); + } else if (weight_numel > 0) { + const dim3 block(kBlockThreads); + const dim3 grid(static_cast( + (weight_numel + kBlockThreads - 1) / kBlockThreads)); + hipLaunchKernelGGL(HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + grid, block, 0, stream, raw_weight, packed_weight, + weight_numel); + } + if (n > 0) { + const dim3 block(kBlockThreads); + const dim3 grid( + static_cast((n + kBlockThreads - 1) / kBlockThreads)); + hipLaunchKernelGGL(HIP_KERNEL_NAME(w8a8_identity_copy_kernel), grid, + block, 0, stream, weight_scale, packed_weight_scale, + n); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_gate_up_proj.hip new file mode 100644 index 00000000..3b0c1c88 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_gate_up_proj.hip @@ -0,0 +1,1182 @@ +// @@variant shape=hy3_tp8_shared_gate_up_proj_m16 commit=2321c42b8618aae23700bc2cf3532b59c8f6c8d8 added=2026-08-27 +// median_us=10.38 p90_us=10.42 +// source=hy3-dsh-tp8-m16-2-dc53295a +// MetaInfer W8A8 INT8 GEMM - HIP implementation (gfx928 / K500SM_AI). +// +// worker_2 iteration 8: split-K=15 one-wave LDS-staged partial GEMM + +// split-15 combine for the assigned shape hy3_tp8_shared_gate_up_proj_m16: +// M = 16, N = 384, K = 4096 (M*N = 6144 output elements) +// +// Operator contract: +// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) +// * x_scale[m, 0] * weight_scale[n, 0] ) +// +// Iteration 1 = minimal gfx928 DUMMA INT8 bootstrap (exact mandatory tile): +// * one 64-lane wavefront per block, one 16x16 output tile per wave, +// explicit int32 accumulation, grid = N/16 = 24 blocks, direct global +// fragment loads, LDS-drain epilogue; measured 117.91 us median with +// only 24 of 120 CUs active (24 waves, 1 per CU, no co-residency). +// +// Iteration 4 = accepted split-K=10 grid-parallelism probe (19.84 us): +// * K=4096 split into 10 non-uniform 32-aligned slices, grid = 24 N-tiles +// x 10 splits = 240 one-wave blocks (2 per CU); direct global fragment +// loads in the K loop; separate 24-block combine kernel sums the 10 +// int32 planes in ascending split order (bit-identical int32 order); +// both launches inside the timed Graph region. +// * PMC: 49,152 vmem_read instructions for the partial kernel (~205 per +// block = byte-level fragment loads + reassembly, one vmcnt wait per +// load before each mmac), lds_instructions = 0 -> the K loop is +// serialized on global load latency (issue/latency-bound). +// +// Iteration 6 = accepted split-K=16 LDS-staged kernel (18.73 us median / +// 19.04 us p90): block tile 16x32 (2 waves x 1 N-tile of 16 columns), 128 +// thr, grid = 12 N-groups x split-K=16 = 192 blocks = 1.6 blocks/CU = 3.2 +// waves/CU; whole uniform 256-row slice staged once (A 16x256 stride 272, +// B 32x256 stride 48), one __syncthreads, zero-barrier LDS-only K loop +// (8 steps x 1 v_mmac per wave); workspace = 16 int32 planes = 393,216 B = +// the caller's workspace_split_k_capacity(16,384,4096) budget; combine = +// 24 blocks x 128 thr, ascending split sum; PMC: 192 blocks < 2/CU target. +// +// Iteration 7 = measured split-K=10 one-wave grid (this file's round-7 +// candidate, REJECTED by the p90 guard): block tile 16x16 (1 wavefront, 64 +// thr), grid = 24 N-groups x split-K=10 = 240 one-wave blocks = exactly +// 2 blocks/CU; K=4096 -> 10 non-uniform 32-aligned slices (8 x 416 + 2 x +// 384); whole slice staged once (A stride 432, B stride 48; 26,880 B/block +// = 2 x 26,880 = 53,760 B/CU at 2 blocks, fits 64 KiB), 0 s_barrier (the +// single-wave __syncthreads is eliminated; lgkmcnt waits cover the ds +// hazard, verified in the exact gfx928 code object); 10 planes = 245,760 B; +// combine = 24 blocks x 128 thr. Measured median 15.432 us (-21% vs +// iteration 6) but p90 19.434 > best 19.040 -> p90_guard_passed = false; +// the 30 samples were bimodal (4 slow windows 19.4-23.2 us then 26 fast +// windows 15.4-15.9 us): a machine-state artifact, not kernel behavior. +// +// Iteration 8 (this file's active M=16 path) = mandatory HIP-only +// occupancy round: tune one occupancy limiter (LDS footprint -> blocks/CU) +// with PMC evidence. The split-K=10 kernel's 26,880 B LDS/block caps it at +// exactly 2 blocks/CU (3 x 26,880 = 80,640 > 64 KiB). split-K=15 (trusted +// probe set [2,5,10,15]) shortens every slice (288/256 vs 416/384 rows), +// shrinking the per-block LDS footprint to 18,688 B -> 3 x 18,688 = 56,064 +// B/CU <= 64 KiB, so 3 one-wave blocks/CU become legal, and the grid = +// 24 N-groups x 15 splits = 360 one-wave blocks = EXACTLY 3 blocks/CU on +// 120 CUs (above the 2-blocks/CU latency-hiding target, integer-balanced, +// no tail wave). No repeated HBM reads: B stays once-read (1,572,864 B), +// A stays L2-hot, every block still stages its own slice exactly once. +// * block tile 16x16 (1 wavefront x 1 N-tile), 64 thr; grid = 24 x 15 = +// 360 one-wave blocks = exactly 3 blocks/CU; +// * K=4096 -> 15 non-uniform 32-aligned slices: 8 x 288 (9 DUMMA steps) +// + 7 x 256 (8 steps); k0-ascending within each slice and ascending +// split sum in the combine preserve the bit-identical int32 order; +// * whole slice staged once into bank-skewed LDS (A 16x288 stride 304, +// B 288x16 stride 48; 18,688 B/block), batched 16-B vector loads, one +// __syncthreads (0 s_barrier for the 1-wave block), then a zero-barrier +// LDS-only K loop (9/8 steps x 1 A-frag + 1 B-frag ds_read + 1 v_mmac); +// * workspace = 15 int32 planes [15][16][384] = 368,640 B (<= the +// caller's 393,216 B capacity; guard updated, unsplit DUMMA fallback +// when workspace < 15 planes); combine = 24 blocks x 128 thr, all 15 +// plane loads issued into a register array before the ascending sum; +// * total replay traffic ~2.32 MB (B 1.57 MB + planes RW 737,280 B + +// out 12,288 B) vs 2.47 MB for the accepted S=16 kernel (-6%); vs S=10 +// the plane RW grows +245,760 B (+50%) - the combine cost of more +// splits being benchmarked against the occupancy gain. +// +// Iteration 9 (this file's active M=16 path) = HIP-only two-launch attack. +// Round 8 measured 15.529 us median / 16.024 us p90 and its falsifiable +// branch ("if median stays >= 15.4, the next round attacks the two-launch +// structure - fused/single-kernel combine - instead of geometry") fired. +// Change: fuse the 24-block combine INTO the partial kernel as a +// last-arrival tail using the validated TP4 qkv lineage protocol - every +// block writes its 16x16 int32 plane tile, __threadfence() (release), then +// atomicAdd(&counters[tile], 1); the block whose atomicAdd returns +// kSplitK-1 (the 15th arriver of its N-group) __threadfence()s (acquire), +// sums the 15 planes in ascending split order (exact int32 order preserved +// -> outputs bit-identical), applies both scales, stores bf16, and resets +// its counter to 0 (atomicExch) so the next replay restarts from 0. The 24 +// counters live in the workspace tail (96 B right after the 15 planes; +// guard updated to 368,736 B) and are zeroed once by the host before the +// first launch (a static-flag hipMemsetAsync on the caller stream - issued +// in the eager precheck, outside the timed Graph in the normal flow). +// Result: ONE kernel launch per replay (grid 360 one-wave blocks, 3/CU, +// LDS unchanged 18,688 B); the combine no longer pays a second dispatch +// and each N-group's combine tail overlaps the still-running partials of +// other groups instead of waiting for full grid completion + relaunch. +// +// Repair 1/4 (correctness): the round-9 draft let ALL 64 lanes of every +// block execute atomicAdd(&counters[tile], 1) (64 increments per block = +// 960 per tile per replay), so the kSplitK-1 condition fired on the 15th +// increment inside the FIRST-arriving block - before its 14 siblings had +// stored their planes - and the unguarded atomicExch reset made several +// blocks combine per replay. Exact check failed: 6048/6144 mismatched, +// first_mismatch (0,0) actual=0.0 (premature combiner summed unwritten, +// zero planes). Fix: exactly one arrival per block (lane 0 fence + +// atomicAdd, LDS s_is_last flag, barriers) + lane-0-only atomicExch reset, +// mirroring the validated TP4 qkv fused tail. GEMM/plane-store body and the +// ascending 15-plane sum are unchanged (bit-identical int32 order). +// +// Iteration 10 (this file's active M=16 path) = final HIP-only +// consolidation round (raw inline asm not allowed by the ISA policy: +// plateau=false, valid HIP rounds < 8). Round 9 measured 14.516 us median / +// 14.544 us p90 and its falsifiable branch ("if median stays >= 14.5 ... the +// single-wave combine tail (60 loads/lane) is the new critical path - the +// next round vectorizes the combiner's plane loads (int4) or shrinks the +// plane count instead") fired (14.516 >= 14.5). Change: vectorize ONLY the +// fused combine tail's plane loads - each lane now owns 4 CONSECUTIVE +// columns (row = lane>>2, col0 = (lane&3)*4) and reads each of the 15 +// planes with one aligned int4 (16-B) load instead of 4 scalar 4-B loads +// (60 -> 15 vmem loads per lane, 4x fewer tail load instructions, same +// bytes, same L2-request footprint); weight_scale is read once as float4 +// and the 4 adjacent bf16 stores merge into one 8-B store. Per-element +// int32 accumulation order is unchanged (ascending split sum per output +// element) -> outputs remain bit-identical (0 mismatches). Everything else +// is untouched: grid = 24 N-groups x 15 splits = 360 one-wave blocks = 3 +// blocks/CU, per-block LDS 18,688 B, zero-barrier LDS-only K loop, +// workspace = 15 planes + 24 counters = 368,736 B, one launch per replay, +// exact (m,n,k) guard, unsplit DUMMA fallback, generic scalar fallback. +// Falsifiable: (1) median_us < 14.516 (target ~13.6-14.2) with p90 <= 14.544 +// (guard vs the accepted best); (2) 0 mismatches + graph_capture_passed; +// (3) PMC: grid 360 / workgroup 64 / LDS 18,688 / s_barrier 0 unchanged, +// arch_vgpr stays ~40-60 with 0 spills, the tail portion of vmem_read drops +// ~4x (staging ~3,264 unchanged; tail 60 -> 15 loads/lane); (4) replay +// traffic unchanged ~2.07 MB. If median stays >= 14.516 with p90 <= 14.544, +// the tail was not the critical path and the kernel sits at its +// streaming/launch floor -> the round documents the plateau evidence. +// +// Iteration 11 (this file's active M=16 path) = final conditional inline-asm +// round -> mandatory HIP-only consolidation (raw asm is NOT allowed: isa_policy +// says phase=hip_only, plateau=false, raw_inline_asm_allowed=false, +// valid_hip_rounds=7 < 8 - the control plane has not confirmed a HIP plateau). +// Round 10 measured 13.925 us median / 13.972 us p90 (< 14.516 target), so the +// int4-vectorized tail branch did NOT fire and the tail is not the critical +// path anymore. The accepted kernel's fresh exact gfx928 code object exposes +// the next concrete compiler limitation: the slice staging is SERIALIZED - +// every global_load_dwordx4 is followed by s_waitcnt vmcnt(0) before its +// ds_write_b128 (load_wait0_ds_write_windows = 18 across the two slice +// instantiations, 10 long + 8 short), so each block pays ~10/8 full +// global-memory latencies one after another before the K loop's first +// ds_read (the last write is followed by lgkmcnt(0) waits). Change: in +// split15_slice_body, issue ALL staging global loads into one register array +// (int4 s_buf[10/8]) BEFORE any ds_write - the compiler is then forced to put +// the loads in flight back-to-back and move the vmcnt waits to the write +// side, collapsing the staging critical path to one memory latency. Same +// addresses, same per-lane order, same LDS layout, same __syncthreads, same +// zero-barrier LDS-only K loop, same fused last-arrival tail: the int32 +// accumulation order is untouched (outputs bit-identical, 0 mismatches +// expected). Everything else is untouched: grid = 24 N-groups x 15 splits = +// 360 one-wave blocks = exactly 3 blocks/CU (trusted probe set member 15), +// per-block LDS 18,688 B (still the occupancy binder: 3 x 18,688 = 56,064 B +// <= 64 KiB), workspace = 15 int32 planes + 24 arrival counters = 368,736 B +// (<= 393,216 B budget; unsplit DUMMA fallback when smaller), one launch per +// replay, fence/atomic/barrier protocol unchanged, exact (m,n,k) guard, +// generic scalar fallback for every unmatched shape including the paired M=2. +// Falsifiable predictions: (1) median_us < 13.925 (target ~11.5-12.8) with +// p90_us <= 13.972 (guard vs the accepted best); (2) correctness 0 mismatches +// + graph_capture_passed (single-launch structure, write-in-place planes, +// self-resetting counters, no clear); (3) ISA: the first staging vmcnt(0) +// now covers ~7 outstanding global_load_dwordx4 (old kernel: exactly 1) - +// memory-level parallelism per block rises 1 -> ~7-10 and the staging +// critical path collapses from ~10 serialized latencies to ~1-2; vmcnt(0) +// 20 -> ~19, arch_vgpr ~69-75 (locally verified 69 VGPR / 40 SGPR) with 0 +// spills, grid 360 / LDS 18,692 / s_barrier 0 / v_mmac 17 / ds_write_b128 18 +// / ds_read2_b32 17 / ds_read_u8 136 unchanged; (4) PMC vmem_read ~4,392 and +// replay traffic ~2.06 MB unchanged (same loads, same bytes). If +// median stays >= 13.925 with p90 <= 13.972, the 3-blocks/CU overlap already +// hid the staging latency and the kernel sits at its streaming/launch floor - +// the round then documents the plateau evidence for this shape. +// +// Iteration 15 (this file's active M=16 path) = final conditional inline-asm +// round -> mandatory HIP-only consolidation (raw asm is NOT allowed: isa_policy +// says phase=hip_only, plateau=false, raw_inline_asm_allowed=false - the +// control plane has NOT confirmed a HIP plateau: the three recent valid +// improvements were -28.62% / -0.033% / +3.86% (iterations 12/13/14), the +// -28.62% being the rejected round-12 regression and +3.86% the round-14 +// packed-B median gain, and no prior ISA-guided round recorded a compiler +// limitation plus target instructions). Round 14 (n-major-packed B staging +// loads only) measured median 12.512 us / p90 13.064 us - best median ever, +// rejected only by the p90 guard (13.064 > 13.052) - and the source was +// restored to the accepted iteration-11 digest. The accepted kernel's exact +// gfx928 code object exposes the remaining half of the B path: the K loop +// still reads each lane's 8 B-fragment elements as 8 ds_read_u8 from the +// k-major stride-48 LDS (4-way bank conflicted; the persistent PMC signature +// since iteration 9), while the A fragment already reads its 8 consecutive +// bytes as ONE ds_read2_b32. The TP4 gate_up/down_proj lineage validated the +// fix for exactly this ("one ds_read_b64 per step, was 8 ds_read_u8"; rule 3: +// 'Pack the cold once-read weight (B) ... into the fragment's layout'). +// Change (one consolidation): complete the round-14 B-pack into a per-tile +// n-major transpose packed[tile][n][kk] = raw[kk*384 + tile*16 + n] +// ((k,n)==(4096,384) ONLY; every other (k,n) keeps the identity pack), stage +// B into an n-major LDS layout b_lds[n][k] (element (k,n) at n*304 + k; +// column stride 304 = 288 + 16 bank skew, 16-B aligned so the staging writes +// stay ds_write_b128), and load the B fragment with the library col_major +// loader: a.x[i] = p[row*304 + col + i] reads the lane's 8 elements as 8 +// CONSECUTIVE bytes -> one ds_read2_b32 per step (was 8 ds_read_u8 at stride +// 48). Same elements, same x[i] register assignment (x[i] = element +// (k0+col+i, n_base+row) in both layouts), same v_mmac, so the int32 +// accumulation order is untouched (k-ascending within each slice, ascending +// split sum - outputs bit-identical, 0 mismatches expected). The B staging +// loads also get the round-14 locality win: each lane holds one 16-B +// row-chunk of 16 consecutive k values of one n column (chunk c -> n = c&15, +// k-chunk = c>>4) - 64 lanes cover 16 x 64-B contiguous sectors per wave-load +// (was 64 scattered 16-B requests), sequential HBM pages, no straddles. LDS +// per block drops 18,688 -> 9,728 B (A 16x304 = 4,864 + B 16x304 = 4,864), +// so 3 x 9,728 = 29,184 B/CU << 64 KiB (LDS is no longer the occupancy +// binder, but the grid stays 360 = exactly 3 blocks/CU, so residency is +// unchanged); per-block LDS instructions drop ~4x (136 ds_read_u8 + 17 +// ds_read2_b32 -> 34 ds_read2_b32), attacking the per-CU LDS-pipe pressure +// that round 13's evidence identified as the shared per-CU resource. +// Everything else is untouched: grid = 24 N-groups x 15 splits = 360 +// one-wave blocks = exactly 3 blocks/CU (trusted probe-set member 15), +// workspace = 15 int32 planes + 24 arrival counters = 368,736 B (<= the +// caller's 393,216 B budget; unsplit DUMMA fallback when smaller), one launch +// per replay, fence/atomic/barrier protocol unchanged (lane-0 atomicAdd, LDS +// s_is_last flag, lane-0 atomicExch reset), exact (m,n,k) guard, generic +// scalar fallback for every unmatched shape including the paired M=2 (both +// fallbacks decode the packed layout for (k,n)==(4096,384): the scalar path +// reads element (kk, col) at (col>>4)*65536 + (col&15)*4096 + kk; the unsplit +// DUMMA path loads B with the row_major loader at ldm = K from +// packed + tile*65536 + k0, giving p[(col+i)*4096 + row] = element +// (k0+col+i, n_base+row) - identical values to the raw layout). +// Falsifiable predictions: (1) median_us < 12.995 (target ~12.0-12.7; the +// round-14 pack alone measured 12.512 and this also cuts the per-block LDS +// read instructions 8x and the per-CU LDS-pipe work ~4x) with p90_us <= +// 13.052 (guard vs the accepted best p90); (2) correctness 0 mismatches +// (bit-identical int32 order) and graph_capture_passed (single-launch +// structure unchanged, write-in-place planes, self-resetting counters, pack +// out-of-timed/out-of-Graph, no clear); (3) PMC/ISA: grid_blocks 360 / +// workgroup_size 64 / s_barrier 0 / v_mmac 17 / ds_write_b128 18 unchanged, +// ds_read_u8 136 -> 0, ds_read2_b32 17 -> 34 (A + B), group_segment_fixed_size +// 18,692 -> ~9,732, arch_vgpr ~69 +/- a few with 0 spills, occupancy 3 +// blocks/CU (now grid-bound), lds_instructions ~31,632 -> ~20k, +// lds_bank_conflicts ~63,744 -> ~15-25k (the B u8 4-way reads disappear), +// vmem_read_instructions ~4,392 unchanged (same loads, same bytes); +// (4) memory_traffic.total_bytes_per_operator_replay unchanged ~2.06 MB (no +// HBM delta). If median stays >= 12.995 with p90 <= 13.052, the LDS pipe was +// not on the critical path either (the kernel is at its launch/memory floor) +// and this final round documents the plateau evidence for this shape instead +// of claiming a gain. +// +// Iteration 17 (this file's active M=16 path) = HIP-only B-pack ORDER +// consolidation (isa_policy: phase=hip_only, plateau=false, +// raw_inline_asm_allowed=false, skill_allowed=false - the three recent valid +// HIP rounds are NOT all within [-2%, +2%): -0.033% / +3.86% / +8.35%, so the +// plateau is not proven and one more HIP-only change is mandated; the +// round-16 attempt died in the agent infrastructure before a valid proposal, +// so this round returns to the accepted iteration-15 source (digest +// b39a985c...) and starts ONE new bounded experiment). Round 15 (n-major B +// pack + n-major LDS) measured median 11.994 us / p90 12.293 us - accepted - +// and left the remaining half of the B memory path on the table: the n-major +// pack places chunk (n, k-chunk) at n*4096 + k-chunk*16, so each staged B +// wave-load touches 16 x 64-B sectors at 4,096-B stride (16 scattered L2 +// requests / up to 16 DRAM pages) - the round-14 contiguous-load locality +// was traded away to buy the one-ds_read2_b32 B fragment read. Change (one +// consolidation, pack ORDER only): the pack kernel now emits the SAME 16-B +// chunks (16 consecutive k of one n column) in ascending (slice, n, k-chunk) +// order - chunk index = slice_chunk_base(s) + n*cps(s) + k-chunk, byte +// address = tile*65536 + chunk*16 - so each staged wave-load (64 consecutive +// chunks) covers 1,024 CONTIGUOUS bytes (sequential L2 sectors and +// sequential DRAM pages; was 16 scattered 64-B sectors at 4,096-B stride). +// Chunk content, the n-major LDS destination b_lds[n][k] (n*304 + k), the +// one-ds_read2_b32 B fragment read, the v_mmac count, the zero-barrier K +// loop and the fused last-arrival tail are all UNTOUCHED: every staged byte +// and every fragment register is identical to iteration 15, so the int32 +// accumulation order is bit-identical (0 mismatches expected). The +// long-slice staging write index changes from (c&15, c>>4) to (c/18, c%18) +// (the short 256-row instantiation keeps (c>>4, c&15); divisions resolved +// by the compiler, staging path only); the generic scalar fallback and the +// unsplit DUMMA fallback decode the new pack order through a shared +// packed384_byte_offset helper (per-element gather in the unsplit path, +// fallback-only: x[i] = element (k0 + (lane>>4)*8 + i, n_base + (lane&15)), +// the same fragment registers as the library row_major loader). +// Everything else is untouched: grid = 24 N-groups x 15 splits = 360 +// one-wave blocks = exactly 3 blocks/CU (trusted probe-set member 15), +// per-block LDS 9,732 B, workspace = 15 int32 planes + 24 arrival counters +// = 368,736 B (<= 393,216 B budget; unsplit DUMMA fallback when smaller), +// one launch per replay, fence/atomic/barrier protocol unchanged, exact +// (m,n,k) guard, generic scalar fallback for every unmatched shape +// including the paired M=2. +// Falsifiable predictions: (1) median_us < 11.994 (target ~11.2-11.8; the B +// staging wave-loads become fully sequential, removing the 16-scattered- +// sector DRAM page pattern) with p90_us <= 12.293 (guard vs the accepted +// best p90); (2) correctness 0 mismatches (bit-identical int32 order) and +// graph_capture_passed (single-launch structure unchanged, write-in-place +// planes, self-resetting counters, pack out-of-timed/out-of-Graph, no +// clear); (3) PMC/ISA: grid_blocks 360 / workgroup_size 64 / lds_bytes +// 9,732 / s_barrier 0 / v_mmac 17 / ds_read2_b32 34 / ds_write_b128 18 +// unchanged, arch_vgpr ~69-72 with 0 spills, vmem_read ~4,392 unchanged +// (same loads, same bytes), the B staging address stream becomes sequential +// ((k_base + c)*16 per chunk) so l2_misses ~32,321 may drop (fewer sector +// straddles / page effects) while lds_bank_conflicts may shift (the +// long-slice staging write bank pattern changes; instruction count and +// bytes identical); (4) memory_traffic.total_bytes_per_operator_replay +// unchanged ~2.06 MB (same bytes, better DRAM locality). If median stays +// >= 11.994 with p90 <= 12.293, the B staging DRAM pattern was not on the +// critical path either (the kernel is at its launch/latency floor) and this +// round documents the next plateau datum for this shape. +// +// Iteration 21 (this file's active M=16 path) = HIP-only, ONE consolidation: +// A register-only transport (the validated TP4 o_proj lineage recipe +// "register-only K-loop transport with depth-1 prefetch"; raw asm is NOT +// allowed: isa_policy says phase=hip_only, plateau=false, +// raw_inline_asm_allowed=false, skill_allowed=false - the three recent valid +// HIP improvements +3.86% / +8.35% / +5.35% are not within [-2%, +2%), so the +// HIP plateau is not proven and one more HIP-only change is mandated; rounds +// 18/19/20 died in the agent infrastructure before valid proposals, so this +// round returns to the accepted iteration-17 source (digest +// 88ea95a832fbc79e4d625b2aa22df3a0d533d6eb2c6aff08529c7e60a92942c4, verified +// sha256-identical before editing) and starts ONE new bounded experiment - no +// round-18/19/20 partial candidate was repaired or replayed). The accepted +// kernel (11.385 us median / 11.954 us p90) still stages the small L2-hot A +// slice (16 x 288/256 int8) through LDS: per block that is 5/4 16-B global +// loads + 5/4 ds_write_b128 + 9/8 ds_read2_b32 (2-way bank conflicted, +// structural for any 16-B-aligned stride) that the K loop then reads back. +// Every block reads its own A rows exactly once, so the LDS round trip only +// adds instructions and a serialization point; the A values are 64 KiB total +// and L2-resident (re-read by all 360 blocks), so direct global fragment +// loads are cheap. Change (one consolidation): (1) split15_slice_body no +// longer stages A; each lane preloads its own 8-B fragment slice for EVERY +// DUMMA step into a per-lane register array a_slice[kStepsT] (9/8 +// global_load_dwordx2 per lane, issued back-to-back BEFORE any ds_write - +// the iteration-11 register-array batching pattern - so the A slice costs +// ONE memory latency and zero LDS instructions); (2) the K loop consumes the +// A fragments from registers (du_mma_sync(acc, a_slice[s], b_frag, acc)) - +// same elements per lane per step as the LDS row_major loader (verified +// against du_mma.hpp: x[i] = p[row*ldm + col + i], row = lane&15, col = +// (lane>>4)*8, ldm = k), so the int32 accumulation order is untouched +// (k-ascending per split, ascending split sum -> outputs bit-identical, 0 +// mismatches expected); (3) the shared a_lds array is deleted -> per-block +// LDS drops 9,732 -> 4,868 B (B 4,864 + s_is_last 4). B staging, the B +// col_major LDS reads, the zero-barrier structure (one __syncthreads before +// the B reads, unchanged), the slice-major pack, the fused last-arrival tail +// and every fallback are UNTOUCHED. Everything else is unchanged: grid = 24 +// N-groups x 15 splits = 360 one-wave blocks = exactly 3 blocks/CU (trusted +// probe-set member 15; LDS is no longer the binder at 3 x 4,868 = 14,604 +// B/CU), workspace = 15 int32 planes + 24 arrival counters = 368,736 B (<= +// the caller's 393,216 B budget; unsplit DUMMA fallback when smaller), one +// launch per replay, fence/atomic/barrier protocol unchanged (lane-0 +// atomicAdd, LDS s_is_last flag, lane-0 atomicExch reset), exact (m,n,k) +// guard, generic scalar fallback for every unmatched shape including the +// paired M=2 (the packed384_byte_offset decode is untouched). +// Falsifiable predictions: (1) median_us < 11.385 (target ~10.8-11.3; the +// per-block staging loses the A LDS round trip and the K loop loses 9/8 +// LDS reads) with p90_us <= 11.954 (guard vs the accepted best p90); +// (2) correctness 0 mismatches (bit-identical int32 order; M=2 paired +// fallback exact via the packed decode) and graph_capture_passed (single- +// launch structure unchanged, write-in-place planes, self-resetting +// counters, no clear); (3) PMC/ISA: grid_blocks 360 / workgroup_size 64 / +// s_barrier 0 / v_mmac 17 unchanged; group_segment_fixed_size 9,732 -> +// ~4,868 B; ds_write_b128 18 -> ~9/8 (B staging only); ds_read2_b32 34 -> +// ~17 (B only); ds_read_u8 stays 0; the A slice appears as 9/8 batched +// global_load_dwordx2 before the first B ds_write (explicit int2 loads; if +// instead the ISA shows the A loads sunk into the K loop or scalarized to +// global_load_ubyte, the compiler refused the batching/vectorization and +// this round documents that finding); arch_vgpr ~69 -> ~85-95 (a_slice +// holds 18/16 extra VGPR) with 0 spills (per-wave ~90 x 64 = 5,760 VGPR << +// 65,536/CU, occupancy stays 3 blocks/CU = grid-bound); lds_instructions +// ~10,128 -> ~5-6k and lds_bank_conflicts ~29,184 -> ~15k (the A 2-way +// fragment-read conflicts disappear); vmem_read_instructions ~4,392 -> +// ~10,464 (A loads become 8-B per step instead of 16-B staged chunks - SAME +// bytes, more instructions, all issued in one batch); (4) +// memory_traffic.total_bytes_per_operator_replay unchanged ~2.06 MB (same +// A+B bytes; A stays L2-hot). If median stays >= 11.385 with p90 <= 11.954, +// the A LDS round trip was not on the critical path either (the kernel is +// at its launch/latency floor with 3 waves/CU) and this round documents the +// next plateau datum for this shape. +// +// The generic scalar path remains the fallback for every unmatched +// (m,n,k), including the paired M=2 shape with the same (N,K); the M=16 +// DUMMA path is guarded by the exact (m,n,k) triple. +// +// Host symbols (matched against csrc/bindings.cpp): +// launch_w8a8_gemm(...) +// launch_pack_w8a8_weight(...) +// +// Known-good include order for this DTK: HIP runtime first, then the HIP +// bfloat16 header, then du_mma.h. + +#include +#include +#include + +#include + +namespace { + +// Byte offset of element (k, col) inside the exact (k,n)==(4096,384) +// slice-major-chunk weight pack (iteration 17; see +// w8a8_pack_weight_gateup_slicemajor_kernel): each 16-B chunk holds 16 +// CONSECUTIVE k values of one n column, and chunks are stored in ascending +// (slice, n, k-chunk) order. With tile = col>>4, n_local = col&15 and +// kchunk = k>>4: +// kchunk < 144 (the 8 long 288-k slices): s = kchunk/18, kk' = kchunk%18, +// chunk base = s*288, chunks per n = 18; +// kchunk >= 144 (the 7 short 256-k slices): c2 = kchunk-144, +// s = 8 + c2/16, kk' = c2%16, chunk base = 2304 + (c2/16)*256, +// chunks per n = 16; +// offset = tile*65536 + (chunk_base + n_local*cps + kk')*16 + (k & 15). +// Fallback paths only (the timed M=16 staging path uses the same mapping +// implicitly through the pack's chunk order). 144 = 8 long slices * 18 +// k-chunks per 288-k slice. +__device__ __forceinline__ int64_t packed384_byte_offset(int k, int col) { + const int64_t tile = static_cast(col) >> 4; // col / 16 + const int64_t n_local = static_cast(col) & 15; + const int kchunk = k >> 4; + int64_t base; + int kk_p, cps; + if (kchunk < 144) { + const int s = kchunk / 18; + base = static_cast(s) * 288; + kk_p = kchunk % 18; + cps = 18; + } else { + const int c2 = kchunk - 144; + const int s = 8 + c2 / 16; + base = 2304 + static_cast(c2 / 16) * 256; + kk_p = c2 % 16; + cps = 16; + } + return tile * 65536 + (base + n_local * cps + kk_p) * 16 + (k & 15); +} + +// Scalar W8A8 GEMM: one thread computes one output element (m, n). +// +// Adjacent threads map to adjacent N columns (the fastest-changing dimension +// of both `out` and the [K, N] weight), so global reads coalesce in N. +// The K loop is executed exactly once per output element in ascending k and +// accumulated in int32; every assigned K keeps the exact int8 dot within +// int32 range (max |dot| = 128*128*4096 = 67,108,864 < 2^31). The result is +// converted to float only for `dot * x_scale[m] * weight_scale[n]`, then +// stored as bfloat16, matching the contract reference exactly. +// +// gfx928: wavefront = 64, blockDim must be a multiple of 64 (128 here). +__global__ __launch_bounds__(128) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, // x_q [M, K] row-major + const int8_t* __restrict__ b, // weight [K, N] row-major + const float* __restrict__ x_scale, // [M, 1] contiguous + const float* __restrict__ weight_scale, // [N, 1] contiguous + hip_bfloat16* __restrict__ out, // [M, N] row-major + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + + const int8_t* __restrict__ a_row = a + static_cast(row) * k; + // For (k,n)==(4096,384) the weight buffer holds the per-tile + // slice-major-chunk pack (iteration 17; see packed384_byte_offset); every + // other (k,n) keeps the raw [K, N] row-major layout (identity pack). + // Decode so this fallback stays exact for the paired M=2 shape and all + // unmatched m. + const bool packed_384 = (k == 4096 && n == 384); + const int64_t b_row_stride = packed_384 ? 0 : static_cast(n); + + int32_t acc = 0; +#pragma unroll 8 + for (int kk = 0; kk < k; ++kk) { + const int8_t b_val = + packed_384 ? b[packed384_byte_offset(kk, col)] + : b[static_cast(col) + + static_cast(kk) * b_row_stride]; + acc += static_cast(a_row[kk]) * static_cast(b_val); + } + + // Reference order: (float(dot) * x_scale[m]) * weight_scale[n], then bf16. + // hip_bfloat16's float constructor performs the supported RNE conversion. + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[idx] = hip_bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Minimal gfx928 DUMMA INT8 path (assigned shape M=16, N=384, K=4096). +// Tile: 16x16x32, int32 accumulation, one 64-lane wavefront per block, +// one 16x16 output tile per wave (grid = N/16 = 24 blocks). +// --------------------------------------------------------------------------- + +// gfx928 DUMMA INT8 m16n16k32 tile constants. +constexpr int kDummaM = 16; +constexpr int kDummaN = 16; +constexpr int kDummaK = 32; +constexpr int kDummaBlockThreads = 64; // one wavefront per block + +// M=16 W8A8 GEMM with the minimal 16x16x32 DUMMA tile: one wavefront owns one +// independent 16x16 output tile. A (x_q [16, K] row-major) and B (weight +// [K, N] row-major) fragments are loaded directly from global memory every +// DUMMA step with the library loaders; the int32 accumulator is explicit and +// drained through LDS in a coalesced epilogue that applies x_scale and +// weight_scale before the bf16 store. No LDS staging, no double buffering, +// and no cross-wave barrier in the K loop (a single wavefront per block; +// the one __syncthreads in the epilogue is reached by every block thread). +__global__ __launch_bounds__(kDummaBlockThreads) void +w8a8_gemm_m16_dumma_kernel( + const int8_t* __restrict__ a, // x_q [16, K] row-major + const int8_t* __restrict__ b, // weight [K, N] row-major + const float* __restrict__ x_scale, // [16, 1] contiguous + const float* __restrict__ weight_scale, // [N, 1] contiguous + hip_bfloat16* __restrict__ out, // [16, N] row-major + int n, + int k) { + const int lane = static_cast(threadIdx.x); // 0..63 + const int n_base = static_cast(blockIdx.x) * kDummaN; + + __shared__ __align__(16) int32_t acc_tile[kDummaM * kDummaN]; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // K is divisible by kDummaK for every assigned shape. A fragment rows are + // x_q rows (ldm = k); B fragment rows are weight k-rows (ldm = n), with the + // tile's 16 columns at n_base. The explicit int32 accumulation is + // k0-ascending over exact 32-element dot products, so the final int32 total + // is bit-identical to the scalar chain (integer arithmetic is exact). + // Iteration 17: this fallback is only reachable for (m,n,k)==(16,384,4096), + // whose weight buffer holds the per-tile slice-major-chunk pack. The + // library row_major loader (a.x[i] = p[(col+i)*ldm + row], row = lane&15, + // col = (lane>>4)*8) cannot decode that non-linear order, so the B + // fragment is filled with per-element gathers from the pack: x[i] = + // element (k0 + (lane>>4)*8 + i, n_base + (lane&15)) - the SAME fragment + // registers the raw-layout loader would produce (the round-15 n-major + // decode asserted the same mapping). Fallback path only - never on the + // timed M=16 split-K kernel. + const bool packed_384 = (k == 4096 && n == 384); +#pragma unroll 4 + for (int k0 = 0; k0 < k; k0 += kDummaK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); + if (packed_384) { + const int n_local = lane & 15; + const int k_off = (lane >> 4) * 8; +#pragma unroll + for (int i = 0; i < 8; ++i) { + b_frag.x[i] = + b[packed384_byte_offset(k0 + k_off + i, n_base + n_local)]; + } + } else { + du::dumma::du_load_matrix_sync( + b_frag, b + static_cast(n_base >> 4) * (kDummaN * k) + k0, + k); + } + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + du::dumma::du_store_matrix_sync(acc_tile, acc_frag, kDummaN, + du::dumma::mem_row_major); + __syncthreads(); + + // Coalesced epilogue: each lane drains 256/64 = 4 accumulator elements, + // applies the per-row/per-column float scales and stores bf16. +#pragma unroll + for (int linear = lane; linear < kDummaM * kDummaN; linear += 64) { + const int row = linear >> 4; + const int col = linear & 15; + const float scaled = static_cast(acc_tile[linear]) * + x_scale[row] * weight_scale[n_base + col]; + out[static_cast(row) * n + n_base + col] = hip_bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// Iteration 8 (HIP round): split-K=15 one-wave LDS-staged partial GEMM + +// split-15 combine (mandatory occupancy round - one limiter: LDS footprint +// -> blocks/CU; trusted probe set [2,5,10,15]). +// +// Round 7 measured the split-K=10 one-wave grid (240 blocks = exactly +// 2 blocks/CU, 26,880 B LDS/block = 53,760 B/CU) at 15.43 us median but was +// rejected by the p90 guard (19.43 > 19.04; bimodal samples - machine +// state, not kernel behavior). Its 26,880 B/block LDS caps occupancy at +// exactly 2 blocks/CU (3 x 26,880 = 80,640 B > 64 KiB). split-K=15 shortens +// every slice (288/256 vs 416/384 rows), shrinking the per-block LDS +// footprint to 18,688 B -> 3 x 18,688 = 56,064 B/CU <= 64 KiB, so 3 one-wave +// blocks/CU become legal, and the grid = 24 N-groups x 15 splits = 360 +// one-wave blocks = EXACTLY 3 blocks/CU on 120 CUs (above the 2-blocks/CU +// latency-hiding target, integer-balanced, no tail wave). B stays once-read +// from HBM and A stays L2-hot (no repeated global reads - each block stages +// its own slice exactly once). +// * block tile 16x16 (1 wavefront x 1 N-tile), 64 thr; grid = 24 x 15 = +// 360 one-wave blocks = exactly 3 blocks/CU; +// * K=4096 -> 15 non-uniform 32-aligned slices: 8 x 288 (9 DUMMA steps) +// + 7 x 256 (8 steps); k0-ascending within each slice, ascending split +// sum -> bit-identical int32 order (0 mismatches expected); +// * whole slice staged once into bank-skewed LDS (A 16x288 stride 304, +// B 288x16 stride 48; 18,688 B/block), batched 16-B vector loads, one +// __syncthreads (emitted as 0 s_barrier for the 1-wave block), then a +// zero-barrier LDS-only K loop (9/8 steps x 1 A-frag + 1 B-frag ds_read +// + 1 v_mmac, 9/8 v_mmac per block); +// * workspace = 15 int32 planes [15][16][384] = 368,640 B (<= the +// caller's 393,216 B workspace_split_k_capacity budget; every launch +// overwrites every partial tile - no workspace clear); +// * combine kernel: 24 blocks x 128 thr, all 15 plane loads issued into a +// register array before the ascending split sum (exact int32 order), +// then (float(acc) * x_scale[row]) * weight_scale[col] -> bf16; +// * replay traffic ~2.32 MB (B 1.57 MB + planes RW 737,280 B + out 12,288 +// B) vs 2.47 MB for the accepted S=16 kernel (-6%); vs S=10 the plane +// RW grows +245,760 B (+50%) - the combine cost of more splits is the +// explicit trade being benchmarked against the occupancy gain. +// --------------------------------------------------------------------------- + +constexpr int kSplitK = 15; // 15 planes; workspace 368,640 B +constexpr int kNumNGroups = 24; // N=384 / 16 (partial block tiles) +constexpr int kNumLongSplits = 8; // 8 x 288-row slices + 7 x 256-row +constexpr int kSliceLong = 288; // 9 x kDummaK steps (non-uniform) +constexpr int kSliceShort = 256; // 8 x kDummaK steps +constexpr int kStepsLong = kSliceLong / kDummaK; // 9 +constexpr int kStepsShort = kSliceShort / kDummaK; // 8 +constexpr int kPartialBlockThreads = 64; // one wavefront per block +constexpr int kCounterInts = kNumNGroups; // one arrival counter per N-group +constexpr int kStagedAStride = 304; // 288 + 16 bank skew (16-B aligned) +// Iteration 15: B is staged n-major (element (k,n) at n*kStagedBStride + k, +// the transpose of the A staging convention), so the column stride must hold +// the longest k-slice (288) plus bank skew; 304 = 288 + 16 keeps every +// staging ds_write_b128 and every 8-byte fragment read 16-/8-B aligned. +constexpr int kStagedBStride = 304; // 288 + 16 bank skew (16-B aligned) + +// Partial GEMM: block (tile, split) computes one 16x16 output tile for the +// split's contiguous K slice and stores the int32 partial to +// partials[split][row][n]. The whole slice is staged into LDS once with +// batched 16-B vector loads, one barrier, then a zero-barrier LDS-only K +// loop. k0-ascending within the slice; the combine kernel sums slices in +// ascending split order, so the total int32 accumulation order is +// bit-identical to the scalar chain. The slice body is a compile-time +// template (kSliceKT, kStepsT) so the staging loops and the K loop are +// fully unrolled; the grid-uniform `split < kNumLongSplits` branch picks the +// 288- or 256-row instantiation for each block (one launch total). +template +__device__ __forceinline__ void split15_slice_body( + const int8_t* __restrict__ a, // x_q [16, K] row-major + const int8_t* __restrict__ b, // weight [K, N] row-major + int32_t* __restrict__ partials, // [kSplitK][16][N] int32 + int n, + int k, + int split, + int tile, + int n_base, + int k_base, + int8_t (*b_lds)[kStagedBStride]) { + const int tid = static_cast(threadIdx.x); + + // Iteration 21 (HIP-only, A register-only transport; the validated TP4 + // o_proj lineage recipe "register-only K-loop transport"). The A slice is + // small (16 x kSliceKT = 4,608 / 4,096 B, L2-hot 64-KiB activation) and + // every block reads its own rows exactly once, so staging it through LDS + // only buys the LDS round trip: per block it costs kALoads 16-B global + // loads + kALoads ds_write_b128 + kStepsT ds_read2_b32 (2-way bank + // conflicted). This round removes the A LDS stage entirely: each lane + // loads its own 8-B fragment slice for EVERY DUMMA step into a per-lane + // register array a_slice[kStepsT] (the library matrix_a row_major loader + // with ldm = k reads x[i] = p[row*k + k_base + s*kDummaK + col + i], row = + // lane&15, col = (lane>>4)*8 - the SAME 8 elements per lane per step the + // LDS version read, so every fragment register is identical and the int32 + // accumulation order is bit-identical). All kStepsT A loads are issued in + // one register-array block BEFORE any ds_write (the iteration-11 batching + // pattern that forced the staging loads back-to-back in the accepted code + // object), so the A slice costs one memory latency and zero LDS + // instructions; the K loop then consumes the fragments straight from + // registers. Guarded lanes are impossible (kStepsT * 64 lanes * 8 B = + // kSliceKT * 16 B exactly, no tail). + du::dumma::DUFragment + a_slice[kStepsT]; + // Explicit vector loads: the library loader assigns the 8 bytes + // individually, which LLVM scalarizes to 8 global_load_ubyte per step per + // lane from GLOBAL memory (verified in the round-21 draft code object: + // 136 ubyte loads = the old "direct fragment byte-load" poison). The + // fragment's 8 elements are 8 CONSECUTIVE bytes of one row (m = lane&15) + // at k = k_base + s*kDummaK + (lane>>4)*8, 8-B aligned (k_base, kDummaK + // and the col offset are all multiples of 8; x_q is tensor-aligned), so + // one explicit int2 load per lane per step forces ONE global_load_dwordx2 + // and the __builtin_memcpy (8 bytes, register-to-register) fills the same + // fragment registers the LDS row_major loader produced. + const int a_row = tid & 15; // m index (library loader row) + const int a_col = (tid >> 4) * 8; // k offset (library loader col) + const int8_t* a_base = + a + static_cast(a_row) * k + k_base + a_col; +#pragma unroll + for (int s = 0; s < kStepsT; ++s) { + const int2 raw = *reinterpret_cast(a_base + s * kDummaK); + __builtin_memcpy(a_slice[s].x, &raw, 8); + } + + // Stage the B slice (kSliceKT rows x 16 cols) with coalesced 16-B vector + // loads, then one barrier (B only - the A stage is gone). B chunks + // (iteration 17) come from the per-tile slice-major-chunk pack (see + // w8a8_pack_weight_gateup_slicemajor_kernel): chunk c holds 16 CONSECUTIVE + // k values of one n column (n = c/18 long, c/16 short; k-chunk = c%18 / + // c%16) and the chunks are stored in ascending c order, so the slice's + // chunks occupy the CONTIGUOUS byte range [(k_base + c)*16] of the tile - + // 64 lanes cover 1,024 CONTIGUOUS bytes per wave-load (the round-14 + // locality win). Same chunk content, same n-major LDS destination, same + // per-lane order -> the fragment reads and the int32 accumulation stay + // bit-identical. + // + // Iteration 11: issue ALL staging global loads into one register array + // BEFORE any ds_write. The accepted kernel's exact gfx928 code object + // serialized the slice staging: every global_load_dwordx4 was followed by + // s_waitcnt vmcnt(0) before its ds_write_b128 (load_wait0_ds_write_windows + // = 18 across the two slice instantiations, 10 long + 8 short), so each + // block's staging paid ~10 (long) / 8 (short) full global-memory latencies + // one after another before the K loop could issue its first ds_read + // (lgkmcnt(0)). Holding the 16-B chunks in a register array first forces + // the compiler to issue the loads back-to-back (all in flight + // concurrently) and move the vmcnt waits to the write side, collapsing the + // staging critical path to one memory latency. Guarded lanes (the unrolled + // tail beyond the chunk count) write a zero placeholder so no + // uninitialized register is ever consumed. + constexpr int kChunksPerBRow = kSliceKT / 16; // 18 / 16 (per n column) + constexpr int kBLoads = (kSliceKT + 63) / 64; // 5 / 4 + int4 s_buf[kBLoads]; +#pragma unroll + for (int i = 0; i < kBLoads; ++i) { + const int c = tid + 64 * i; + // Iteration 17: slice-major-chunk pack - chunk c (n = c/kChunksPerBRow, + // k-chunk = c%kChunksPerBRow) lives at byte (k_base + c)*16 of the + // tile, so a wave-load reads 1,024 contiguous bytes. + s_buf[i] = (c < kSliceKT) + ? *reinterpret_cast( + b + static_cast(tile) * (kDummaN * k) + + (static_cast(k_base) + c) * 16) + : int4{0, 0, 0, 0}; + } +#pragma unroll + for (int i = 0; i < kBLoads; ++i) { + const int c = tid + 64 * i; + if (c < kSliceKT) { + // n-major LDS destination (iteration 15, UNCHANGED): element (k,n) at + // n*kStagedBStride + k. Chunk c maps to (n = c/kChunksPerBRow, + // k-chunk = c%kChunksPerBRow): for the 256-row instantiation this is + // the same (c&15, c>>4) as iteration 15; for the 288-row + // instantiation the chunk order is n-major (18 chunks per column) and + // the division is resolved by the compiler (staging path only). Still + // one ds_write_b128 per 16-B chunk. + const int b_row = c / kChunksPerBRow; + *reinterpret_cast( + &b_lds[b_row][(c - b_row * kChunksPerBRow) * 16]) = s_buf[i]; + } + } + + __syncthreads(); + + // The block's one wavefront owns the whole 16x16 tile. Zero-barrier K + // loop: kStepsT steps of kDummaK=32; the A fragment comes from the + // preloaded register array a_slice[s] (iteration 21, no LDS), the B + // fragment from b_lds (iteration 15: col_major on the n-major layout, ldm + // 304), one v_mmac per step (9 or 8 v_mmac per block). The col_major + // loader computes b.x[i] = p[row*304 + col + i] (row = n_local = lane&15, + // col = k offset (lane>>4)*8), i.e. the lane's 8 elements B[k0+col+i] + // [n_base+row] are 8 CONSECUTIVE bytes of one n column -> ONE ds_read2_b32 + // per step (was 8 ds_read_u8 at stride 48). x[i] gets the same element as + // the old row_major loader, so the fragment registers (and the int32 + // accumulation order) are bit-identical. + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + +#pragma unroll + for (int s = 0; s < kStepsT; ++s) { + const int k0 = s * kDummaK; + du::dumma::du_load_matrix_sync(b_frag, &b_lds[0][k0], kStagedBStride); + du::dumma::du_mma_sync(acc_frag, a_slice[s], b_frag, acc_frag); + } + + // Direct int32 partial store: partials[split][row][n_base .. n_base+15]. + // Every launch overwrites every partial tile, so no workspace clear is + // needed before capture or between replays. + int32_t* partial_base = + partials + static_cast(split) * kDummaM * n + n_base; + du::dumma::du_store_matrix_sync(partial_base, acc_frag, n, + du::dumma::mem_row_major); +} + +// Fused last-arrival combine tail (iteration 9): removes the separate +// 24-block combine kernel -> exactly ONE kernel launch per replay. +// +// Protocol (the validated TP4 qkv lineage pattern): +// * every block writes its 16x16 int32 plane tile, then __threadfence() +// (release) and atomicAdd(&counters[tile], 1); +// * the block whose atomicAdd returns kSplitK-1 is the 15th (last) +// arriver of its N-group: it __threadfence()s (acquire), sums the 15 +// planes in ascending split order (exact int32 order), applies both +// scales, stores bf16, and resets its counter to 0 (atomicExch) so the +// next replay restarts from 0. The host zeroes the 24 counters once +// before the first launch; the self-reset keeps every later replay +// deterministic (stream-ordered, no racing increments). +__device__ __forceinline__ void split15_combine_tail( + const int32_t* __restrict__ partials, // [kSplitK][16][N] int32 + const float* __restrict__ x_scale, // [16, 1] contiguous + const float* __restrict__ weight_scale, // [N, 1] contiguous + hip_bfloat16* __restrict__ out, // [16, N] row-major + int32_t* __restrict__ counters, // [kNumNGroups] monotonic + int tile, + int n) { + const int lane = static_cast(threadIdx.x); // 0..63 (one wavefront) + const int n_base = tile * kDummaN; + const int plane = kDummaM * n; // int32 stride per split + + // Repair 1/4 (iteration 9): exactly ONE arrival per block. Only lane 0 + // fences, atomically increments the tile counter and publishes the result + // through LDS; the other 63 lanes must NOT atomicAdd. The round-9 draft had + // all 64 lanes incrementing -> 64 x 15 = 960 increments per tile per + // replay, so the 15th increment (arrived == kSplitK-1) fired inside the + // FIRST-arriving block, before its 14 siblings had stored their planes: + // the premature combiner summed mostly-zero planes (first_mismatch + // actual=0.0 at (0,0)) and the unguarded atomicExch reset made several + // blocks combine per replay (6048/6144 mismatched). This mirrors the + // validated TP4 qkv fused tail: tid-0 fence + atomicAdd, LDS flag, barriers. + __syncthreads(); + __shared__ int s_is_last; + if (lane == 0) { + __threadfence(); // release plane stores + const int arrived = atomicAdd(&counters[tile], 1); + __threadfence(); // acquire other planes + s_is_last = (arrived == kSplitK - 1) ? 1 : 0; + } + __syncthreads(); + if (s_is_last == 0) { + return; // not the last arriver + } + + // 256 outputs / 64 lanes = 4 CONSECUTIVE columns per lane (row = lane>>2, + // col0 = (lane & 3) * 4): one aligned int4 (16-B) plane load per split + // replaces the round-9 loop's 4 scalar 4-B loads per element, cutting the + // tail's vmem load instruction count 4x (60 -> 15 per lane) with no byte + // or L2-request increase (each 16-B load is fully consumed; the 4 lanes + // of one row cover a contiguous 64-B segment). The per-element int32 + // accumulation order is UNCHANGED: for each output element the 15 plane + // values are still summed in ascending split order, so the outputs stay + // bit-identical to the scalar chain. weight_scale is read once as a + // float4 (n_base + col0 is 16-B aligned) and the 4 adjacent bf16 stores + // merge into one 8-B store. Fence/atomic/barrier protocol untouched. + const int row = lane >> 2; // 0..15 (4 lanes per row) + const int col0 = (lane & 3) * 4; // 0,4,8,12 (16-B aligned) + const int32_t* p = partials + row * n + n_base + col0; + const float4 ws = + *reinterpret_cast(weight_scale + n_base + col0); + const float xs = x_scale[row]; + int32_t acc0 = 0, acc1 = 0, acc2 = 0, acc3 = 0; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + const int4 v = *reinterpret_cast( + p + static_cast(s) * plane); + acc0 += v.x; + acc1 += v.y; + acc2 += v.z; + acc3 += v.w; + } + hip_bfloat16* o = out + static_cast(row) * n + n_base + col0; + o[0] = hip_bfloat16(static_cast(acc0) * xs * ws.x); + o[1] = hip_bfloat16(static_cast(acc1) * xs * ws.y); + o[2] = hip_bfloat16(static_cast(acc2) * xs * ws.z); + o[3] = hip_bfloat16(static_cast(acc3) * xs * ws.w); + + // All kSplitK arrivals of this replay have happened; restart for the next + // replay (the kernel's completion on the stream orders this reset before + // the next replay's first atomicAdd). One lane is enough; all 15 blocks + // have already passed their atomicAdd, so no racing increment survives. + if (lane == 0) { + atomicExch(&counters[tile], 0); + } +} + +__global__ __launch_bounds__(kPartialBlockThreads) void +w8a8_gemm_m16_dumma_splitk15_lds_staged_partial_kernel( + const int8_t* __restrict__ a, // x_q [16, K] row-major + const int8_t* __restrict__ b, // weight [K, N] row-major + int32_t* __restrict__ partials, // [kSplitK][16][N] int32 + const float* __restrict__ x_scale, // [16, 1] contiguous + const float* __restrict__ weight_scale, // [N, 1] contiguous + hip_bfloat16* __restrict__ out, // [16, N] row-major + int32_t* __restrict__ counters, // [kNumNGroups] monotonic + int n, + int k) { + const int tile = static_cast(blockIdx.x) / kSplitK; // 0..kNumNGroups-1 + const int split = static_cast(blockIdx.x) % kSplitK; // 0..kSplitK-1 + const int n_base = tile * kDummaN; + const int k_base = + (split < kNumLongSplits) + ? split * kSliceLong + : kNumLongSplits * kSliceLong + (split - kNumLongSplits) * kSliceShort; + + // Iteration 15: n-major B stage, element (k,n) at n*kStagedBStride + k + // (16 n columns x 304 B = 4,864 B; the k-slice length is carried by the + // template instantiation, not by the shared array). Iteration 21: the A + // slice is no longer staged in LDS (register-only transport), so the + // per-block LDS is B + the combine flag only. + __shared__ __align__(16) int8_t b_lds[kDummaN][kStagedBStride]; + + if (split < kNumLongSplits) { + split15_slice_body(a, b, partials, n, k, split, + tile, n_base, k_base, b_lds); + } else { + split15_slice_body(a, b, partials, n, k, split, + tile, n_base, k_base, b_lds); + } + + // Fused combine tail (iteration 9): one launch per replay; the last + // arriver of each N-group sums the 15 planes and stores bf16. + split15_combine_tail(partials, x_scale, weight_scale, out, counters, tile, + n); +} + +// Identity device-to-device weight pack (iteration 1: the DUMMA kernel reads +// B with the raw [K, N] row-major layout, so identity packing is exact). +// Kept valid for every (K, N) including unmatched shapes; the exact +// (k,n)==(4096,384) shape uses the per-tile slice-major-chunk pack below +// instead (iteration 17). +__global__ void w8a8_pack_weight_identity_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t total) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + i < total; i += stride) { + dst[i] = src[i]; + } +} + +// Per-tile slice-major-chunk weight pack for the exact (k,n)==(4096,384) +// shape (iteration 17): each 16-B chunk holds 16 CONSECUTIVE k values of one +// n column, and the 4,096 chunks of a tile are stored in ascending (slice, +// n, k-chunk) order: +// chunk index = slice_chunk_base(s) + n_local*chunks_per_slice(s) + kk' +// with the 8 long 288-k slices (s < 8: base s*288, 18 chunks per n) followed +// by the 7 short 256-k slices (s >= 8: base 8*288 + (s-8)*256, 16 chunks per +// n); byte address = tile*65536 + chunk_index*16, chunk content = +// raw[k*384 + tile*16 + n_local] for k = k_base(s) + kk'*16 .. +15 (k_base +// numerically equals the chunk base: 288 k and 288 chunks per long slice). +// This is the iteration-15 n-major pack with the same chunk CONTENT but +// reordered memory (n*4096 + kk*16 -> slice-major chunk order), so a staged +// wave-load (64 consecutive chunks) covers 1,024 CONTIGUOUS bytes +// (sequential L2 sectors and sequential DRAM pages; was 16 x 64-B sectors +// at 4,096-B stride) while the n-major LDS destination, the +// one-ds_read2_b32 B fragment read and the int32 accumulation order are all +// untouched (outputs bit-identical to iteration 15). One-time, +// out-of-timed-region, out-of-Graph, same byte count and same buffer +// (graph-stable addresses unchanged). Every other (K, N) keeps the identity +// pack. Fallbacks decode this order via packed384_byte_offset. +__global__ void w8a8_pack_weight_gateup_slicemajor_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + constexpr int64_t kTotal = 4096 * 384; // 1,572,864 + if (idx >= kTotal) { + return; + } + const int64_t tile = idx >> 16; // / 65536 + const int64_t rem = idx - (tile << 16); // % 65536 (byte in tile) + const int64_t chunk = rem >> 4; // 16-B chunk index in tile + const int64_t in_chunk = rem & 15; // byte within the chunk + int64_t base; // chunk index of the slice start + int64_t n_local, kk_p; + if (chunk < 8 * 288) { // long slice (18 chunks per n) + const int64_t r = chunk % 288; + base = (chunk / 288) * 288; + n_local = r / 18; + kk_p = r % 18; + } else { // short slice (16 chunks per n) + const int64_t c2 = chunk - 8 * 288; + const int64_t r = c2 % 256; + base = 8 * 288 + (c2 / 256) * 256; + n_local = r / 16; + kk_p = r % 16; + } + // base == k_base (k index of the slice start) numerically: 288 k / 288 + // chunks per long slice, 256 / 256 per short slice. + const int64_t k = base + kk_p * 16 + in_chunk; + dst[idx] = src[k * 384 + tile * 16 + n_local]; +} + +// Identity device-to-device scale pack (bootstrap). +__global__ void w8a8_pack_scale_identity_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int total) { + const int stride = gridDim.x * blockDim.x; + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < total; i += stride) { + dst[i] = src[i]; + } +} + +} // namespace + +// Timed GEMM entry point (inside the CUDA/HIP Graph region). +// +// Graph-safety contract: no allocation, no compilation, no autotuning, no +// weight packing, no host/device synchronization, no default-stream launch. +// Launches only on the caller-provided stream and uses only caller-provided +// tensors. The iteration-9 split-K=15 path uses the caller's workspace as the +// int32 partial planes (write-in-place every launch, no clear needed) plus 24 +// arrival counters in the workspace tail and issues exactly ONE kernel per +// replay (LDS-staged partial GEMM with a fused last-arrival combine tail); +// the iteration-1 unsplit kernel is the workspace-too-small fallback. +// +// Exact-shape specializations are guarded HERE, before the generic scalar +// fallback, so the paired M=2 API shape with the same (N, K) and every other +// shape still reach the scalar path below. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + const int64_t total = static_cast(m) * n; + if (total <= 0) { + return; + } + + // Exact-shape M=16 specialization (iteration 9): split-K=15 LDS-staged + // partial GEMM with a FUSED last-arrival combine tail -> exactly one + // launch per replay (the round-8 falsifiable branch fired: median + // 15.53 >= 15.4, so the two-launch structure is attacked instead of + // geometry). grid = 24 N-groups (16x16 tiles) x 15 K-slices = 360 one-wave + // blocks = exactly 3 blocks/CU on 120 CUs; K=4096 -> 15 non-uniform + // 32-aligned slices (8 x 288 = 9 DUMMA steps + 7 x 256 = 8 steps); each + // block stages its whole slice in LDS once (18,688 B/block -> 3 x 18,688 = + // 56,064 B/CU <= 64 KiB), runs a zero-barrier LDS-only K loop, stores its + // int32 plane, then the 15th arriver of its N-group sums the 15 planes in + // ascending k order, applies both scales and stores bf16. The 24 arrival + // counters live in the 96 B right after the 15 planes and are zeroed once + // before the first launch (outside the timed Graph in the normal + // eager-first flow). The iteration-1 unsplit DUMMA kernel remains the + // fallback when the caller's workspace cannot hold 15 partial planes plus + // the counters (368,640 + 96 B <= the workspace_split_k_capacity budget of + // 393,216 B); the paired M=2 API shape with the same (N,K) and every + // other shape still reach the scalar fallback below. + if (m == 16 && n == 384 && k == 4096) { + constexpr int64_t kPartialBytes = + static_cast(kSplitK) * kDummaM * 384 * 4; // 368,640 B + constexpr int64_t kCounterBytes = + static_cast(kCounterInts) * sizeof(int32_t); // 96 B + if (workspace != nullptr && + workspace_bytes >= kPartialBytes + kCounterBytes) { + auto* partials = reinterpret_cast(workspace); + int32_t* counters = partials + kSplitK * kDummaM * 384; + // Zero the 24 arrival counters once before the first launch. Issued on + // the caller stream in the eager precheck, so the captured Graph + // contains no memset node; if the first call is captured, the memset + // becomes a replayable Graph node (still exact). The kernel's + // self-reset keeps every later replay deterministic. + static bool counters_zeroed = false; + if (!counters_zeroed) { + counters_zeroed = true; + hipMemsetAsync(counters, 0, kCounterBytes, stream); + } + hipLaunchKernelGGL( + w8a8_gemm_m16_dumma_splitk15_lds_staged_partial_kernel, + dim3(kNumNGroups * kSplitK), + dim3(kPartialBlockThreads), + 0, + stream, + a, + b, + partials, + x_scale, + weight_scale, + reinterpret_cast(out), + counters, + n, + k); + } else { + hipLaunchKernelGGL(w8a8_gemm_m16_dumma_kernel, dim3(n / kDummaN), + dim3(kDummaBlockThreads), 0, stream, a, b, x_scale, + weight_scale, reinterpret_cast(out), + n, k); + } + return; + } + + constexpr int kBlockThreads = 128; // multiple of the gfx928 wavefront (64) + const int grid = static_cast((total + kBlockThreads - 1) / kBlockThreads); + + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + dim3(grid), + dim3(kBlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + reinterpret_cast(out), + m, + n, + k); +} + +// Optional out-of-timed-region weight packing. Iteration 1: identity +// device-to-device copy of raw_weight -> packed_weight and weight_scale -> +// packed_weight_scale, valid for any (K, N) and any unmatched shape (the +// DUMMA kernel reads the raw [K, N] row-major layout). Iteration 17: the +// exact (k,n)==(4096,384) shape is packed per-tile slice-major-chunk +// instead (see w8a8_pack_weight_gateup_slicemajor_kernel: same 16-B chunks +// as iteration 15, chunk order (slice, n, k-chunk) so staged wave-loads are +// 1,024-B contiguous); every other (K, N) keeps the identity copy. Same +// byte count, same buffer, one-time and outside the timed Graph region. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t total = static_cast(k) * n; + if (total > 0) { + constexpr int kPackBlock = 256; + const int grid = static_cast((total + kPackBlock - 1) / kPackBlock); + if (k == 4096 && n == 384) { + hipLaunchKernelGGL(w8a8_pack_weight_gateup_slicemajor_kernel, + dim3(grid), dim3(kPackBlock), 0, stream, raw_weight, + packed_weight); + } else { + hipLaunchKernelGGL( + w8a8_pack_weight_identity_kernel, + dim3(grid), + dim3(kPackBlock), + 0, + stream, + raw_weight, + packed_weight, + total); + } + } + if (n > 0) { + constexpr int kScaleBlock = 64; + const int grid = (n + kScaleBlock - 1) / kScaleBlock; + hipLaunchKernelGGL( + w8a8_pack_scale_identity_kernel, + dim3(grid), + dim3(kScaleBlock), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/o_proj.hip new file mode 100644 index 00000000..8cf04a9f --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/o_proj.hip @@ -0,0 +1,666 @@ +// @@variant shape=hy3_tp8_o_proj_m4096 commit=710edb90ef1f462cd0751356a542a14e8bbba8c6 added=2026-08-27 +// median_us=323 p90_us=325.4 speedup=43.82 baseline_us=1.415e+04 +// source=hy3-dsh-tp8-m4096-1-0ebb994d +// @@variant shape=hy3_tp8_o_proj_m4096 added=2026-08-27 (bootstrap iteration 1) +// INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. +// +// Worker: worker_1 (physical GPU 1), assigned shape: +// hy3_tp8_o_proj_m4096 : M=4096, N=4096, K=1024 +// +// Bootstrap strategy (iteration 1, correctness-first but a usable profiling +// baseline - a large-Prefill scalar K loop would not be): +// * Large-prefill path (exact (m, n, k) == (4096, 4096, 1024)): +// native INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 +// output tile per block; four wavefronts (256 threads); each wave owns +// a 64x32 quadrant built from eight m16n16k32 accumulator fragments; +// the block cooperatively vector-loads A[128,64] and B[64,64] into one +// single-buffered 64-K LDS stage (15,360 B total: 128*80 + 64*80 with +// 16 B padding per row for bank skew); two __syncthreads per stage; +// fused dot * x_scale[m] * weight_scale[n] epilogue stored directly as +// bf16 from the accumulator fragments. Grid dim3(64, 32) = 2048 blocks +// dwarfs the 120 CUs, so no split-K is needed. +// The weight is read in the packed n-major [N, K] int8 layout +// (packed[n*K + kk] == raw[kk*N + n], a full transpose; identity +// [K, N] layout is retained for every other (k, n)). Every B tile of +// one (n-tile, K-stage) is one contiguous 4 KiB vector stream in both +// global and LDS, so each thread's B staging load/store is a single +// 16-byte-aligned int4. launch_pack_w8a8_weight performs the +// permutation once, outside the timed region and outside Graph +// capture, into the same byte-count buffer (graph-stable). +// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), +// including all small-M API cases (M=2, M=16) and M=3072 with the same +// (K, N) = (1024, 4096); the fallback decodes the packed n-major [N, K] +// layout for (k, n) == (1024, 4096) and the identity [K, N] layout for +// every other (k, n). +// * launch_pack_w8a8_weight: for (k, n) == (1024, 4096) transposes the raw +// [K, N] weight into the [N, K] int8 layout; for any other (k, n) it +// is an identity device-to-device copy. Scales are copied identity. +// Packing never happens inside the timed GEMM. +// +// Iteration 6 (epilogue round): the scales and bf16 conversion were already +// fused in the compute kernel and the launcher already ignores the workspace +// (no split-K combine pass exists), so the remaining epilogue inefficiency +// is the store pattern: the m16n16k32 lane mapping (row = lane & 15, +// c4 = lane >> 4, frag.x[i] -> column c4 + 4*i) makes the direct store four +// 2-byte scalar stores per lane whose wavefront addresses touch each 32-B +// sector at 25% utilization (32 global_store_short_d16_hi per wavefront per +// block in the gfx928 ISA; vmem_write_instructions 262,144). The new +// store_prefill_fragment_coalesced transposes the 4-element groups within +// each 4-lane column group (two __shfl_xor steps, 16 then 32; this DTK +// lowers them to ds_bpermute at the block tail where the LDS pipe is idle), +// so lane (r, c4) holds the four CONTIGUOUS columns 4*c4..4*c4+3 and issues +// ONE 8-byte store per lane: 64 lanes x 8 B = 512 B per fragment per +// wavefront in 16 fully-used 32-B sectors (100% sector efficiency; +// vmem_write_instructions -> 65,536). The MMAC loop, LDS staging, barriers +// and int32 accumulation order are untouched, and the per-element scale +// multiply order plus bf16 rounding are unchanged => stored bf16 bits are +// identical (mismatch 0, max_abs_error 0.0 expected). This is the exact +// mechanism validated on the worker-29 TP4 qkv_proj lineage (accepted +// iteration 6, 770.14 -> 721.87 us). +// +// Iteration 10 (B-fragment load round): the gfx928 code object of the +// accepted kernel (419.05 us median, 81.995 TOPS, arch_vgpr 80, LDS 15,360 B, +// grid 2048, workgroup 256) shows the DUMMA B-fragment path compiled as ~32 +// per-wave ds_read_u8 byte loads (the row_major int8 loader assigns +// x[i] = p[col*ldm + row + i], i.e. 8 k-rows strided by ldm per lane) plus a +// mask/OR byte-reassembly VALU chain (~6 ops per operand dword) with long +// lgkmcnt wait chains; PMC: lds_bank_conflicts 15,728,640 = 3 x (4,194,304 B +// ds_read_u8 + 1,048,576 A ds_read2_b32) and lds_wait_instructions +// 11,718,821. This round stages B n-major in LDS (b_tile[n][k], 64 rows x 80 +// B - same 5,120 B) and reads each B fragment as ONE 8-byte LDS load written +// directly into the fragment storage (the same 8 bytes in the same x[0..7] +// order; validated load_fragment8 pattern from the worker-29 TP4 gate_up +// lineage iteration 13, VALU 16.08M -> 6.25M), so the per-byte loads, the +// pack VALU and their waits disappear; du_mma_sync casts b.x straight to the +// v_mmac operand, so the operand bit pattern and the int32 accumulation are +// bit-identical (mismatch 0, max_abs_error 0.0 expected). The global pack for +// (k, n) == (1024, 4096) becomes n-major [N, K] so staging stays one +// coalesced int4 per thread (same trade as the worker-29 TP4 down_proj +// iteration 14), and the scalar fallback decodes the same n-major layout. +// Tile, A staging/A fragments, barriers, staging bytes, per-accumulator MMAC +// order and the iteration-6 epilogue are untouched. +// +// Iteration 11 (repair round): the iteration-10 candidate measured the fast +// mapping (327.36 us median, 104.96 TOPS, ~43.23x vs the fixed Triton +// baseline) but failed correctness (mismatch 16,775,923, max_abs_error 86.5) +// because its pack kernel vectorized along the wrong axis: it copied +// raw[kk][n0..n0+15] (16 consecutive n at fixed k) into packed[n0][kk..kk+15] +// (which must hold 16 consecutive k at fixed n). That wrote only the packed +// rows with n0 % 16 == 0 (93.66% of the buffer unwritten) and did misaligned +// int4 stores whenever kk % 16 != 0. The DUMMA kernel itself is verified +// correct against the gfx928 code object: the B-fragment path compiles to +// four ds_read2_b64 at exactly b_tile[(local_col + lane&15)*80 + kk + +// (lane>>4)*8] for (kk, frag) = (0,0), (32,0), (1280,1), (1312,1) (8-byte +// offset units), the A-fragment ds_read2_b32 instructions are byte-identical +// to the accepted kernel, and the v_mmac operand order (D, A, B) is +// unchanged, so with a correct n-major pack the int32 accumulation and the +// stored bf16 bits are the reference bits (mismatch 0, max_abs_error 0.0 +// expected). This round changes ONLY w8a8_pack_o_proj_panels_kernel: each +// thread now gathers its 16-byte destination vector from 16 raw rows +// (aligned byte loads strided by N) and stores one aligned int4; the pack is +// one-time, out-of-timed-region and out-of-Graph, so the timed GEMM kernel +// (staging, load_fragment8 fragments, MMAC order, epilogue, grid 2048, +// workgroup 256, LDS 15,360 B, arch_vgpr 82, zero scratch) is bit-identical +// to iteration 10. The exact (m, n, k) guard, the scalar fallback decode +// (packed[col*K + kk] == raw[kk*N + col]) and the identity packs are +// untouched. +// +// Iteration 12 (A-fragment load round): the iteration-11 repair was accepted +// (348.27 us median, 98.66 TOPS, mismatch 0). The exact gfx928 code object of +// the accepted kernel still shows the library du_load_matrix_sync path for the +// four matrix_a fragments emitting, per wave per stage, 8 ds_read2_b32 (A is +// m-major in LDS, so each lane's eight k-bytes are contiguous and the compiler +// already vectorized the byte loads) followed by a ~52-instruction per-dword +// mask/OR byte-reassembly VALU chain (v_and 0xff00/0xff0000/0xff000000 + +// v_or_b32_sdwa src0_sel:BYTE_0 + v_or3_b32) plus 7 lgkmcnt waitcnts that is +// arithmetic IDENTITY on every dword (the loaded bytes reach the v_mmac A +// operand unchanged): dead issue and dead latency on the A-load -> v_mmac +// dependency path. This round applies the exact iteration-10 mechanism +// (load_fragment8, validated on B; the worker-29 TP4 gate_up lineage measured +// VALU 16.08M -> 6.25M, and the TP4 qkv lineage accepted iteration 18 replaced +// the library row_major byte-reassembly with an explicit A-fragment loader) to +// the four A du_load_matrix_sync calls: each A fragment is filled by one +// 8-byte LDS read at a_tile[(local_row + (lane&15))*80 + kk + ((lane>>4)<<3)] +// writing the same 8 bytes to the same x[0..7] slots the row_major loader +// produced => the v_mmac A operand bit pattern and the int32 accumulation are +// bit-identical (mismatch 0, max_abs_error 0.0 expected). Tile, staging, the +// B load_fragment8 fragments, barriers, per-accumulator MMAC order, the +// iteration-6 coalesced epilogue, the pack layout, the exact (m, n, k) guard +// and the scalar fallback are untouched. + +#include +#include +#include + +#include + +namespace { + +using namespace du::dumma; + +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 +constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kTargetM = 4096; +constexpr int kTargetN = 4096; +constexpr int kTargetK = 1024; +constexpr int kBlockM = 128; +constexpr int kBlockN = 64; +constexpr int kStageK = 64; +constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) +constexpr int kBStride = kBlockN + kBPad; +constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride +constexpr int kBlockThreads = 4 * kWaveSize; + +using bf16_t = hip_bfloat16; + +// --------------------------------------------------------------------------- +// Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. +// Each wave owns a 64x32 quadrant (eight m16n16k32 int32 accumulators); the +// block cooperatively stages A[128,64] from x_q (row-major, stride k) and +// B[64,64] from the packed n-major [N, K] weight into a single +// 64-K LDS buffer. Two barriers per stage: one after the cooperative load, +// one before the next stage overwrites LDS. Direct fragment epilogue. +// --------------------------------------------------------------------------- +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int base_row, + int base_col, + int m, + int n, + int lane) { + // Verified gfx928 int8 m16n16k32 accumulator ownership (matches + // du_store_matrix_sync): lane & 15 selects the row, lane >> 4 selects + // col % 4, and x[i] maps to columns col%4 + 4*i. + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M masking: padded rows never write + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); + } +} + +// Iteration 6 (epilogue round): coalesced fragment store. The m16n16k32 +// accumulator lane mapping (row = lane & 15, column group c4 = lane >> 4, +// frag.x[i] -> column c4 + 4*i) gives each lane four elements strided by 4 +// columns, so the direct store (store_prefill_fragment) is four 2-byte +// scalar stores per lane whose wavefront addresses touch each 32-B sector at +// 25% utilization (8 B used of 32 B; 4 store instructions per fragment per +// wave, i.e. 32 global_store_short_d16_hi per wavefront per block in the +// current gfx928 ISA). This epilogue transposes the 4-element groups within +// each 4-lane column group (lanes r, r+16, r+32, r+48 -- a 4x4 transpose, +// two 2x2 steps with shfl_xor 16 then 32, one v_cndmask per element per +// step; no staging tile round trip), so lane (r, c4) ends up holding the +// four CONTIGUOUS columns 4*c4 .. 4*c4+3, converts them to bf16, packs 4 +// bf16 (8 B), and writes ONE 8-byte store per lane: 64 lanes x 8 B = 512 B +// per fragment per wavefront in 16 fully-used 32-B sectors (100% store +// sector efficiency; vmem_write_instructions drops 262,144 -> 65,536 for +// this shape). This is the exact mechanism validated on the worker-29 TP4 +// qkv_proj lineage (accepted iteration 6, 770.14 -> 721.87 us): this DTK's +// __shfl_xor lowers to ds_bpermute (one LDS permute per shuffle, +64 LDS +// instructions per wavefront per block -- accepted because it runs at the +// block tail where the LDS pipe is otherwise idle; there is no global round +// trip and no staging tile). Only the int32 values are re-routed between +// lanes -- the per-element scale multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are +// unchanged, so the stored bits are identical to store_prefill_fragment. +// The row>=m guard is wavefront-uniform (all 64 lanes of a wave share the +// same 16-row window base_row + (lane & 15)), so the shuffles never mix +// active and inactive lanes; base_col is a multiple of 16 and n*2 a multiple +// of 8, so the float4 weight_scale load (col0 % 4 == 0) and the 8-byte store +// ((row*n + col0) % 4 == 0 elements -> byte offset % 8 == 0) are aligned. +template +__device__ __forceinline__ void store_prefill_fragment_coalesced( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int base_row, + int base_col, + int m, + int n, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M masking: padded rows never write (wavefront-uniform) + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 16, n*2 is a multiple of 8). + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + static_cast(row) * n + col0) = + packed; +} + +// Iteration 10: direct 8-byte LDS fragment fill for matrix_b on the n-major +// [N, K] B tile. du_load_matrix_sync's int8 matrix_b loaders assign +// x[i] = p[col*ldm + row + i] (row = lane & 15, col = (lane >> 4) << 3), i.e. +// eight consecutive k-values per lane on an [N, K] tile, and du_mma_sync +// passes b.x straight to the v_mmac builtin as one packed 8-byte operand; +// the library path makes the compiler emit eight per-byte ds_read_u8 plus a +// mask/OR reassembly chain per loaded dword. Writing the same 8 bytes +// directly into the fragment storage keeps the operand bit pattern identical +// (exact int32 accumulation unchanged) and lets the compiler feed a single +// 8-byte LDS read straight to the v_mmac. This is the exact pattern +// validated on the worker-29 TP4 gate_up lineage (accepted iteration 13). +template +__device__ __forceinline__ void load_fragment8( + Frag& frag, + const signed char* __restrict__ base, + int stride, + int lane) { + const int off = (lane & 15) * stride + ((lane >> 4) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row * 64; + const int local_col = wave_col * 32; + + // Single-buffered 64-K stage: A[128, 80] + B[64, 80] = 15,360 B/block + // (4 blocks/CU fit the 64 KiB LDS budget; the padded 80-byte strides are + // 16-byte-aligned and break the 64-byte LDS bank periodicity). + // Iteration 10: B is staged n-major b_tile[n][k] (64 n-rows x 80 B - the + // same 5,120 B as the old [k][n] tile), so each lane's eight B-fragment + // k-values are contiguous and load as ONE 8-byte LDS read (load_fragment8) + // instead of eight ds_read_u8 + mask/OR pack VALU. + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + + // Cooperative staging: each thread owns one int4 in A rows [0,64), one in + // A rows [64,128), and one int4 of B. B comes from the packed n-major + // [N, K] layout (iteration 10): B[k0+stage_kv, n0+b_n] is at + // packed[(blockIdx.x*64 + b_n)*K + k0 + stage_kv], so the whole + // (n-tile, K-stage) B tile is one contiguous 4 KiB vector stream and the + // thread's global load and LDS store are each a single 16-byte-aligned + // int4. A[128,64] is 512 int4s; B[64,64] is 256 int4s. + const int vector_byte_offset = tid * static_cast(sizeof(int4)); + const int stage_row = vector_byte_offset / kStageK; + const int stage_col = vector_byte_offset - stage_row * kStageK; + const int b_n = tid >> 2; // n column of this thread's B int4 + const int b_kv = (tid & 3) * 16; // k offset (stage-local) of the B int4 + + for (int k0 = 0; k0 < k; k0 += kStageK) { + *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = + *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); + *reinterpret_cast(a_tile + + (stage_row + kBlockM / 2) * kAStride + stage_col) = + *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + + k0 + stage_col); + *reinterpret_cast(b_tile + b_n * kBStride + b_kv) = + *reinterpret_cast( + weight + (static_cast(blockIdx.x) * kBlockN + b_n) * k + + k0 + b_kv); + __syncthreads(); + + // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + // Iteration 10: B is n-major in LDS, so each lane's eight fragment + // k-values are contiguous; one 8-byte load fills the fragment (the + // same 8 bytes the old row_major loader gathered with eight ds_read_u8 + // plus mask/OR pack VALU), keeping the v_mmac operand bit pattern and + // the int32 accumulation identical. + load_fragment8(b_frag0, b_tile + local_col * kBStride + kk, kBStride, + lane); + load_fragment8(b_frag1, b_tile + (local_col + kTileN) * kBStride + kk, + kBStride, lane); + // Iteration 12: A gets the same direct 8-byte fill (load_fragment8). + // A is m-major in LDS, so lane (r, c4) owns the eight CONSECUTIVE + // k-values a_tile[(local_row + r)*80 + kk + c4*8 .. +7]; the code + // object shows du_load_matrix_sync already vectorizes these into + // ds_read2_b32 but still executes the per-dword mask/OR reconstruction + // (~52 identity VALU + 7 lgkmcnt waits per wave per stage) before the + // v_mmac A operands. Writing the same 8 bytes straight into the + // fragment storage removes that dead chain; the operand bit pattern + // and the int32 accumulation are unchanged. + load_fragment8(a_frag0, a_tile + local_row * kAStride + kk, kAStride, + lane); + load_fragment8(a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, + kAStride, lane); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + load_fragment8( + a_frag0, a_tile + (local_row + 2 * kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8( + a_frag1, a_tile + (local_row + 3 * kTileM) * kAStride + kk, + kAStride, lane); + du_mma_sync(acc20, a_frag0, b_frag0, acc20); + du_mma_sync(acc21, a_frag0, b_frag1, acc21); + du_mma_sync(acc30, a_frag1, b_frag0, acc30); + du_mma_sync(acc31, a_frag1, b_frag1, acc31); + } + + // Protect the LDS buffer from the next stage's cooperative overwrite. + __syncthreads(); + } + + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + // Iteration 6: coalesced epilogue store (one 8-byte store per lane per + // fragment; scales + bf16 conversion fused in-kernel, no workspace pass). + store_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, + base_row, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, + base_row, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, + lane); + store_prefill_fragment_coalesced(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc31, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + kTileN, + m, n, lane); +} + +// --------------------------------------------------------------------------- +// Generic scalar fallback: one thread per output element. Exact int32 dot +// over K, then fused x_scale * weight_scale, then bf16 store. Handles every +// unmatched (m, n, k), including all small-M API cases (M=2, M=16) and any +// M in (0, 4096] with the same (K, N). For (k, n) == (1024, 4096) it decodes +// the packed n-major [N, K] layout (raw[kk, col] == packed[col*K + kk]); +// every other (k, n) reads the identity [K, N] layout. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + int32_t acc = 0; + const int8_t* a_row = x_q + static_cast(row) * k; + if (k == kTargetK && n == kTargetN) { + // Iteration 10: the packed layout for (k, n) == (1024, 4096) is n-major + // [N, K] (packed[col*K + kk] == raw[kk*N + col]). + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast( + weight[static_cast(col) * k + kk]); + } + } else { + const int8_t* b_col = weight + col; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Weight packing (outside the timed region and outside Graph capture; the +// packed buffer keeps the same byte count K*N and the same allocated +// address, so the layout is graph-stable). For the target (k, n) == +// (1024, 4096) the raw [K, N] row-major weight is transposed into the +// n-major [N, K] layout: +// packed[n*K + kk] == raw[kk*N + n] +// so every DUMMA B tile (one 64-column n-tile x one 64-K stage) is a single +// contiguous 4 KiB stream of 16-byte vector loads and each lane's eight +// B-fragment k-values are contiguous in LDS (one 8-byte fragment read). +// Iteration 11 repair: the n-major transpose's 16-byte destination vectors +// run along K, so each thread GATHERS its 16 bytes from 16 raw rows (byte +// loads strided by N) instead of copying a contiguous raw vector (the +// iteration-10 pack copied 16 consecutive n at fixed k, leaving 93.66% of +// the packed buffer unwritten and issuing misaligned int4 stores). The +// pack is one-time and out-of-timed-region, so the gather cost is free. +// For every other (k, n) the pack is an identity device-to-device copy. +// Scales are copied identity in both cases. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_o_proj_panels_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + // Iteration 11 repair: one thread per 16-byte vector of the packed [N, K] + // layout, but the vector is a GATHER, not a memcpy. The destination vector + // packed[n*K + kk0 .. kk0+15] holds 16 CONSECUTIVE k-values of raw column n + // (packed[n*K + kk0 + j] == raw[(kk0 + j)*N + n], j = 0..15), so its 16 + // source bytes sit in 16 DIFFERENT raw rows (strided by N). The + // iteration-10 pack instead vectorized along the source axis (it copied + // raw[kk][n0..n0+15] - 16 consecutive n at fixed k - into + // packed[n0][kk..kk+15], which must hold 16 consecutive k at fixed n); + // that wrote only the packed rows with n0 % 16 == 0 (93.66% of the buffer + // left unwritten, rows 1..15 of every 16 never touched) and issued + // misaligned int4 stores whenever kk % 16 != 0. Here each thread gathers + // its 16 bytes with aligned byte loads and stores one aligned int4. + const int kk_vectors = k / 16; // 16-k vectors per packed row (k % 16 == 0) + const int64_t total4 = static_cast(n) * kk_vectors; + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear >= total4) { + return; + } + const int kk0 = static_cast(linear % kk_vectors) * 16; // multiple of 16 + const int nn = static_cast(linear / kk_vectors); // packed row = raw col + uint32_t d[4] = {0, 0, 0, 0}; +#pragma unroll + for (int j = 0; j < 16; ++j) { + const uint8_t b = raw[static_cast(kk0 + j) * n + nn]; + d[j >> 2] |= static_cast(b) << (8 * (j & 3)); + } + int4 v; + v.x = static_cast(d[0]); + v.y = static_cast(d[1]); + v.z = static_cast(d[2]); + v.w = static_cast(d[3]); + // Destination offset nn*K + kk0 is 16-byte aligned (K % 16 == 0 and + // kk0 % 16 == 0), so the int4 store is aligned. + *reinterpret_cast(packed + static_cast(nn) * k + kk0) = v; +} + +__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int64_t count) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +__global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( + const float* __restrict__ raw, + float* __restrict__ packed, + int count) { + const int linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Stable host launch symbols consumed by csrc/bindings.cpp. +// Both launchers are pure dispatch: no allocation, no packing, no +// synchronization, no default-stream launch; they run on the caller-provided +// PyTorch stream and are CUDA-Graph safe. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; // no split-K: the GEMM does not use the workspace + (void)workspace_bytes; + auto* out_bf16 = reinterpret_cast(out); + + // Explicit dispatch. The assigned shape (M=4096, N=4096, K=1024) takes the + // DUMMA 128x64 single-buffered path; every other (m, n, k) - including + // small-M API cases (M=2, M=16) and M=3072 with the same (K, N) - takes the + // scalar fallback, which decodes the packed n-major [N, K] layout for + // (k, n) == (1024, 4096) and the identity layout otherwise. + if (m == kTargetM && n == kTargetN && k == kTargetK) { + const dim3 grid(kTargetN / kBlockN, kTargetM / kBlockM); + const dim3 block(kBlockThreads); + hipLaunchKernelGGL(w8a8_dumma_128x64x64_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + constexpr int kBlock = 256; + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + (void)hipGetLastError(); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kBlock = 256; + const dim3 block(kBlock); + if (k == kTargetK && n == kTargetN) { + const int64_t total4 = static_cast(k) * n / 16; + const dim3 pack_grid( + static_cast((total4 + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_o_proj_panels_kernel, + pack_grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else { + const int64_t weight_count = static_cast(k) * n; + const dim3 weight_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_identity_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, weight_count); + } + + const int64_t scale_count = n; + const dim3 scale_grid( + static_cast((scale_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, + scale_grid, block, 0, stream, + weight_scale, packed_weight_scale, n); + (void)hipGetLastError(); +} +// @@end +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/qkv_proj.hip new file mode 100644 index 00000000..8149eb73 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/qkv_proj.hip @@ -0,0 +1,1211 @@ +// @@variant shape=hy3_tp8_qkv_proj_m4096 commit=54bce94f966dd5cb4dc5962db5258abd9fb392dc added=2026-08-27 +// median_us=418.2 p90_us=431.6 speedup=46.34 baseline_us=1.938e+04 +// source=hy3-dsh-tp8-m4096-1-0ebb994d +// MetaInfer W8A8 INT8 GEMM for gfx928 (K500SM_AI). +// +// Worker: worker_0 (physical GPU 0) +// Assigned shape: hy3_tp8_qkv_proj_m4096 (M=4096, N=1280, K=4096) +// +// Operator contract (fixed by the control plane): +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// +// Iteration 1 (DUMMA throughput baseline round): port the TP4-validated +// 64x128 packed-n-major-B architecture (113.5 TOPS on the sibling M=4096, +// N=2560, K=4096 shape) to the assigned TP8 shape as a templated 2-D +// macro-tile family, and benchmark the three mandated block tiles (64x64, +// 64x128, 128x64). +// * The timed exact shape (m >= 128, n == 1280, k == 4096) dispatches the +// <64,128> instantiation: 256 threads = 4 wavefronts, each wave owns a +// 32x64 quadrant = eight m16n16k32 int32 accumulator fragments, +// single-buffered 64-K LDS stage with TWO __syncthreads per stage, +// 3 resident blocks/CU (LDS 14,592 B/block, ~80 VGPR). +// * B is packed once (outside the timed region and Graph capture) to an +// n-major packed[n][k] layout for the exact (K=4096, N=1280): the LDS B +// tile rows are 80 B (64 data + 16 pad, five bank phases) and every +// m16n16k32 col_major B fragment is loaded with the explicit 8-byte +// loader load_b_frag8 (one ds_read2_b64 per fragment instead of 8 +// byte-granular reads at stride 132), while every 16-byte staging vector +// commits with one ds_write_b128. +// * The 64-K stage's two 32-K steps have ALL fragment loads hoisted before +// the first v_mmac (the TP4 iter-18 schedule), concentrating LDS latency +// into one wait per stage and leaving the 16-MMAC stream free of +// lgkmcnt waits. +// * The epilogue is the TP4 iter-6 coalesced store: a register 4x4 +// transpose (8 shfl_xor + 8 v_cndmask) puts four contiguous bf16 columns +// per lane, converted and stored as ONE 8-byte store per lane per +// fragment (100% store sector efficiency vs 25% for the scalar 2-byte +// stores). Iteration 6 (epilogue round) keeps this structure but peels +// the final K stage so the per-row/per-column scales are prefetched into +// registers once per lane (2 scalar + 4 float4 loads instead of 8 + 8 +// per-fragment reloads) after the last MMAC burst and the provably dead +// final all-consumption barrier is dropped (127 barriers per block); the +// per-element multiply order and bf16 rounding are unchanged, so output +// bits are identical. The timed path was audited for this round: scales, +// bf16 conversion and the final coalesced store are already fused inside +// the compute kernel and launch_w8a8_gemm ignores the workspace argument +// (no workspace/combine pass exists to remove; the only D2D copy is the +// one-time pack-time weight_scale copy, outside the timed region and +// outside Graph capture, required by the immutable API). +// * The <64,64> (LDS 9,472 B, ~4 blocks/CU) and <128,64> (LDS 13,824 B, +// ~2-3 blocks/CU) instantiations are compiled alongside and are +// benchmark candidates for later rounds; the exact-shape dispatch stays +// on <64,128> this round. +// * Generic arms preserved byte-for-byte: the 64x128 identity-layout +// kernel for large-M shapes with N % 128 == 0 (k % 64 == 0), the 64x64 +// identity-layout kernel for N % 64 == 0, and the scalar int8/int32 +// grid-stride fallback for every other (m, n, k), including the paired +// M=2 API shape (which decodes the packed n-major layout when +// (k, n) == (4096, 1280)). +// * launch_pack_w8a8_weight transposes the logical [K, N] weight to +// packed[n][k] for the exact (k, n) == (4096, 1280) and keeps the +// identity D2D copy for every other (K, N). Same byte count k*n, same +// allocation, graph-stable addresses. +// +// The timed operator (launch_w8a8_gemm) performs no allocation, compilation, +// autotuning, packing, host/device synchronization, or default-stream launch: +// it only dispatches kernels on the caller-provided HIP stream. +// +// Include order is fixed by the control plane: hip_runtime, hip_bfloat16, +// then du_mma (this DTK's du_mma.h is not self-contained otherwise). + +#include +#include +#include + +#include + +namespace { + +constexpr int kWaveSize = 64; +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; +constexpr int kBlockM = 64; +constexpr int kBlockN = 64; +constexpr int kStageK = 128; // K stage of the 64x64 kernel +// 4 wavefronts; must remain a multiple of the gfx928 wavefront size (64). +constexpr int kThreadsPerBlock = 256; + +// 64x128 macro-tile, single-buffered 64-K K stage. LDS row strides are +// padded to odd word counts (A: 68 B = 17 words, B: 132 B = 33 words) so the +// INT8 m16n16k32 fragment loads spread across LDS banks instead of +// collapsing onto a single bank phase (128-byte rows would be 0 mod 32 +// words -> single-bank collisions). +constexpr int kBlockM128 = 64; +constexpr int kBlockN128 = 128; +constexpr int kStageK128 = 64; +constexpr int kAStride128 = kStageK128 + 4; // 68 bytes per A row (17 words) +constexpr int kBStride128 = kBlockN128 + 4; // 132 bytes per B row (33 words) + +// Iteration 1 (throughput baseline): packed n-major B layout for the exact +// assigned shape (K=4096, N=1280). B is transposed once, outside timing, to +// packed[n][k]; the packed kernel stages B into an n-major LDS tile with a +// 16-byte-aligned non-power-of-two row stride (80 B = 20 words, five bank +// phases) so every DUMMA m16n16k32 B fragment load (col_major) reads its +// lane's 8 bytes contiguously (ds_read2_b64) and every 16-byte staging +// vector commits with one ds_write_b128. +constexpr int kPackedBStride = 80; // 64 data + 16 pad bytes (20 words) +constexpr int kPackedBK = 4096; // exact K of the packed assigned shape +constexpr int kPackedBN = 1280; // exact N of the packed assigned shape + +using namespace du::dumma; + +// Direct accumulator epilogue for gfx928 INT8 m16n16k32 (verified against +// du_store_matrix_sync): row = lane & 15, col_mod4 = lane >> 4, +// frag.x[i] maps to columns col_mod4 + 4*i. Stores +// bf16(float(dot) * x_scale[row] * weight_scale[col]); out-of-range rows are +// masked (tail-M handling). The float multiply order (dot * x_scale, then +// * weight_scale) matches the reference exactly. +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Iteration 1 (throughput baseline): coalesced fragment store. The m16n16k32 +// accumulator lane mapping (row = lane & 15, column group c4 = lane >> 4, +// frag.x[i] -> column c4 + 4*i) gives each lane four elements strided by 4 +// columns, so the direct store (store_prefill_fragment) is four 2-byte +// scalar stores per lane whose wavefront addresses touch each 32-B sector at +// 25% utilization. This epilogue transposes the 4-element groups within each +// 4-lane column group (lanes r, r+16, r+32, r+48 -- a 4x4 transpose, two 2x2 +// steps with shfl_xor 16 then 32, one v_cndmask per element per step; no +// staging tile round trip), so lane (r, c4) ends up holding the four +// CONTIGUOUS columns 4*c4 .. 4*c4+3, converts them to bf16, packs 4 bf16 +// (8 B), and writes ONE 8-byte store per lane: 64 lanes x 8 B = 512 B per +// fragment per wavefront in 16 fully-used 32-B sectors (100% store sector +// efficiency; vmem_write_instructions for the assigned shape drop from +// 81,920 to 20,480). Only the int32 values are re-routed between lanes -- +// the per-element scale multiply order (float(dot) * x_scale[row] * +// weight_scale[col]) and the bf16 rounding are unchanged, so the stored bits +// are identical to store_prefill_fragment. The row>=m guard is +// wavefront-uniform (lane & 15 cycles the same 16 rows in every 16-lane +// group), so the shuffles never mix active and inactive lanes; base_col is a +// multiple of 64 and n*2 a multiple of 8, so the float4 weight_scale load +// (col0 % 4 == 0) and the 8-byte store are aligned. +template +__device__ __forceinline__ void store_prefill_fragment_coalesced( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 64, n*2 is a multiple of 8). + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + row * n + col0) = packed; +} + +// Iteration 6 (epilogue round): register-scale variant of the coalesced +// store. The per-row scale (xs) and the 4 per-column scales (ws) are loaded +// ONCE per lane by the peeled final stage of the packed kernel (2 scalar + 4 +// float4 loads instead of 8 + 8 per-fragment reloads) and passed in; the +// per-element multiply order (float(dot) * xs * ws.i) and the bf16 rounding +// are byte-identical to store_prefill_fragment_coalesced, so the stored +// output bits are unchanged. The row>=m guard stays inside (tail-M blocks: +// the guard is uniform across each 4-lane row group, so the shuffle +// transpose never mixes active and inactive lanes). +template +__device__ __forceinline__ void store_prefill_fragment_coalesced_rs( + const AccFragment& frag, + float xs, + const float4 ws, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 64, n*2 is a multiple of 8). + const int col0 = base_col + 4 * c4; + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + row * n + col0) = packed; +} + +// Iteration 1 (throughput baseline): explicit 8-byte loader for the +// col_major m16n16k32 B fragment. The col_major fragment mapping is +// n = lane&15, k = 8*(lane>>4) + i; with the n-major LDS tile at row stride +// 80 the lane's 8 elements are contiguous (p[row*ldm + col .. +7], 8-byte +// aligned because ldm=80 and col are multiples of 8), so this compiles to +// ONE ds_read2_b64 instead of 8 ds_read_u8 + mask/OR reassembly. The byte +// placement is identical to du_load_matrix_sync, so the +// v_mmac_i32_16x16x32_i8 fragment registers receive the same values. +__device__ __forceinline__ void load_b_frag8( + DUFragment& f, + const int8_t* __restrict__ p, + int ldm) { + const unsigned row = static_cast(__lane_id()) & 0xfu; + const unsigned col = (static_cast(__lane_id()) >> 4) << 3; + const int64_t v = *reinterpret_cast(p + row * ldm + col); + *reinterpret_cast(&f.x[0]) = v; +} + +// Iteration 1 (throughput baseline): templated 2-D macro-tile packed-B +// prefill kernel (the TP4-validated 64x128 architecture generalized to the +// mandated tile set {64x64, 64x128, 128x64}). 256 threads = 4 wavefronts; +// each wave owns a (BM/2) x (BN/2) quadrant of m16n16k32 int32 accumulator +// fragments. K is staged cooperatively in a SINGLE LDS buffer at 64-K +// granularity with TWO __syncthreads per stage (iteration 6 peels the final +// stage and drops its provably dead all-consumption barrier: 127 barriers +// per block instead of 128): +// * A is staged row-major into a_tile[BM][68] (odd-word stride: 17 words, +// five bank phases); global-load latency is fully exposed, exactly as in +// the accepted TP4 pipeline (co-resident blocks hide it). +// * B is staged n-major from packed[n][k] into b_tile[BN][80] (20 words, +// five bank phases): each 16-byte staging vector commits with one +// ds_write_b128 (b_n*80 + b_k16 is 0 mod 16). +// * The 64-K stage is split into two explicit 32-K steps and ALL fragment +// loads (kWaveM16 A + kWaveN16 B per step, two steps) are hoisted before +// the first v_mmac of the stage, so LDS latency concentrates in one wait +// and the MMAC stream has no lgkmcnt waits. +// * The int32 accumulation order (k0-outer over 64-K stages, kk-inner +// kk=0 then kk=32 within a stage) and the element-to-slot fragment +// mapping are unchanged, so the results are bit-identical to the scalar +// reference. +// * The epilogue is the coalesced store: iteration 1 +// (store_prefill_fragment_coalesced); iteration 6 hoists the per-row and +// per-column scales into registers loaded once per lane after the peeled +// final stage's MMAC burst and stores via +// store_prefill_fragment_coalesced_rs (same math order, same stored +// bits). +template +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_packedb_tiled_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + // m16n16k32 accumulator tiles per wave quadrant: rows BM/32, cols BN/32. + constexpr int kWaveM16 = BM / 32; + constexpr int kWaveN16 = BN / 32; + // 16-byte staging vectors per thread (1 for 64-row/64-col sides, 2 for + // 128-row/128-col sides). + constexpr int kAVecsPerThread = + (BM * kStageK128 / static_cast(sizeof(int4))) / kThreadsPerBlock; + constexpr int kBVecsPerThread = + (BN * kStageK128 / static_cast(sizeof(int4))) / kThreadsPerBlock; + static_assert(BM % 32 == 0 && BN % 32 == 0, + "block tile must be a multiple of the wave quadrant"); + static_assert(kAVecsPerThread >= 1 && kBVecsPerThread >= 1, + "staging must fit the block thread count"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * BM; + const int n0 = static_cast(blockIdx.x) * BN; + + __shared__ __align__(16) int8_t a_tile[BM * kAStride128]; + __shared__ __align__(16) int8_t b_tile[BN * kPackedBStride]; + + DUFragment + a_frag[kWaveM16][2]; + DUFragment + b_frag[kWaveN16][2]; + DUFragment + acc[kWaveM16][kWaveN16]; +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_fill_fragment(acc[j][i], 0); + } + } + + const int num_stages = k / kStageK128; + + // Cooperative staging mapping (fixed per thread, reused every stage): + // A[BM,64]: vector v = tid + t*256 owns row v>>2 and 16-B column group + // (v&3)*16; each wavefront covers 16 rows x 64 contiguous B. + // B[BN,64]: n-major; vector v = tid + t*256 owns packed row (n0 + v>>2) + // and the 16-byte k-run (v&3)*16. + int4 vA[kAVecsPerThread]; + int4 vB[kBVecsPerThread]; + + // Prologue: load stage 0 into registers, commit it to the single buffers, + // and make it visible before the first burst. + { +#pragma unroll + for (int t = 0; t < kAVecsPerThread; ++t) { + const int v = tid + t * kThreadsPerBlock; + const int a_row = v >> 2; + const int a_k16 = (v & 3) << 4; + const int g_row = m0 + a_row; + vA[t] = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_k16) + : int4{0, 0, 0, 0}; + // Commit A as int32 stores (68-byte rows are 4 mod 16 -> no b128). + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_k16); + adst[0] = vA[t].x; + adst[1] = vA[t].y; + adst[2] = vA[t].z; + adst[3] = vA[t].w; + } +#pragma unroll + for (int t = 0; t < kBVecsPerThread; ++t) { + const int v = tid + t * kThreadsPerBlock; + const int b_n = v >> 2; + const int b_k16 = (v & 3) << 4; + vB[t] = *reinterpret_cast( + packed_b + (n0 + b_n) * k + b_k16); + // Commit B as 16-byte vector stores: b_n*80 + b_k16 is 0 mod 16. + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = + vB[t]; + } + __syncthreads(); + } + + // Stage consumption shared by the loop body and the peeled final stage: + // hoist all fragment loads of both 32-K steps before the first v_mmac, + // then run the MMAC bursts. + auto consume_stage = [&]() { + const int local_row = wave_row * (BM / 2); + const int local_col = wave_col * (BN / 2); + const int8_t* abase = a_tile + local_row * kAStride128; + const int8_t* bbase = b_tile + local_col * kPackedBStride; +#pragma unroll + for (int st = 0; st < 2; ++st) { + const int kk = st * kTileK; +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { + du_load_matrix_sync( + a_frag[j][st], abase + j * kTileM * kAStride128 + kk, + kAStride128); + } +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + load_b_frag8( + b_frag[i][st], + bbase + i * kTileN * kPackedBStride + kk, kPackedBStride); + } + } + // kk = 0 then kk = 32 bursts (same accumulators, same MMAC order as the + // scalar k-ascending reference). +#pragma unroll + for (int st = 0; st < 2; ++st) { +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_mma_sync(acc[j][i], a_frag[j][st], b_frag[i][st], acc[j][i]); + } + } + } + }; + + for (int s = 0; s < num_stages - 1; ++s) { + consume_stage(); + // All consumption of the single buffer is complete only after every + // wavefront passes this barrier; the stage s+1 stores below overwrite + // that same buffer, so they must wait for it. + __syncthreads(); + + // Load stage s+1 and commit it into the single buffer. The global-load + // latency is fully exposed here (co-resident blocks hide it); the + // compiler's vmcnt wait before the DS stores is on this path. + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; +#pragma unroll + for (int t = 0; t < kAVecsPerThread; ++t) { + const int v = tid + t * kThreadsPerBlock; + const int a_row = v >> 2; + const int a_k16 = (v & 3) << 4; + const int g_row = m0 + a_row; + vA[t] = (g_row < m) + ? *reinterpret_cast( + x_q + g_row * k + s1 + a_k16) + : int4{0, 0, 0, 0}; + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_k16); + adst[0] = vA[t].x; + adst[1] = vA[t].y; + adst[2] = vA[t].z; + adst[3] = vA[t].w; + } +#pragma unroll + for (int t = 0; t < kBVecsPerThread; ++t) { + const int v = tid + t * kThreadsPerBlock; + const int b_n = v >> 2; + const int b_k16 = (v & 3) << 4; + vB[t] = *reinterpret_cast( + packed_b + (n0 + b_n) * k + s1 + b_k16); + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = + vB[t]; + } + // Make the stage s+1 stores visible to every wavefront before the + // next burst. + __syncthreads(); + } + } + + // Iteration 6 (epilogue round): peel the final stage. After its MMAC + // burst no wavefront reads LDS again -- the epilogue touches only + // registers and global memory -- so the final all-consumption barrier is + // provably dead and is dropped (127 barriers per block instead of 128). + // The per-row (xs) and per-column (ws) scale registers are prefetched + // right after the last MMACs: 2 scalar + 4 float4 loads per lane instead + // of the 8 + 8 per-fragment reloads, with the load latency hidden behind + // the transpose/shuffle chains. The per-element multiply order and bf16 + // rounding are unchanged, so the stored bits are identical. + consume_stage(); + + const int base_row = m0 + wave_row * (BM / 2); + const int base_col = n0 + wave_col * (BN / 2); + const int lane_row = lane & 15; + const int c4 = lane >> 4; + float xs[kWaveM16]; +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { + const int row = base_row + j * kTileM + lane_row; + // Tail rows are clamped to 0.0f (never dereferenced); the store guard + // below skips them, and the guard is uniform across each 4-lane row + // group so the shuffle transpose never mixes active and inactive lanes. + xs[j] = (row < m) ? x_scale[row] : 0.0f; + } + float4 ws[kWaveN16]; +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + // col0 = base_col + i*kTileN + 4*c4 is a multiple of 4 (base_col is a + // multiple of 64), so the float4 load is 16-byte aligned. + ws[i] = *reinterpret_cast( + weight_scale + base_col + i * kTileN + 4 * c4); + } +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + store_prefill_fragment_coalesced_rs( + acc[j][i], xs[j], ws[i], out, m, n, + base_row + j * kTileM, base_col + i * kTileN, lane); + } + } +} + +// Large-M prefill kernel #1: 64x128 output tile per block, 256 threads +// (4 wavefronts). Each wave owns a 32x64 quadrant = eight m16n16k32 int32 +// accumulators (2 row halves x 4 column quarters). K is staged cooperatively +// in a SINGLE LDS buffer at 64-K granularity with TWO __syncthreads per +// stage (128 over K=4096): the next stage's 3 x int4 (1 A + 2 B vectors) are +// loaded from global only after the current stage's 16-MMAC DUMMA burst and +// its all-consumption barrier, vmcnt-waited, and committed into the same +// buffer, then a store-visibility barrier. The global-load latency is fully +// exposed (no compute overlap) -- the accepted direct/single-buffered +// bootstrap design. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x128_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM128; + const int n0 = static_cast(blockIdx.x) * kBlockN128; + + __shared__ __align__(16) int8_t a_tile[kBlockM128 * kAStride128]; + __shared__ __align__(16) int8_t b_tile[kStageK128 * kBStride128]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; + DUFragment + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); + + const int num_stages = k / kStageK128; + + // Cooperative staging mapping (fixed per thread, reused every stage): + // A[64,64] -> 256 int4 vectors; thread `tid` owns vector `tid` + // (row = tid/4, 16-B column group = (tid%4)*16), so each + // wavefront covers 16 rows x 64 B contiguous per row. + // B[64,128] -> 512 int4 vectors; thread `tid` owns vectors `tid` and + // `tid+256` (kk = tid/8, column group = (tid%8)*16), so + // each wavefront covers 8 kk rows x 128 B contiguous. + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + const int b_kk0 = tid >> 3; + const int b_col16 = (tid & 7) << 4; + + // VGPR payload for one full 64-K stage (3 x int4 = 12 VGPR), loaded just + // before it is committed (no one-stage-ahead overlap in this bootstrap). + int4 vA, vB0, vB1; + + // Prologue: load stage 0 into registers, commit it to the single buffer, + // and make it visible before the first burst. + { + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + vB0 = *reinterpret_cast( + weight + b_kk0 * n + n0 + b_col16); + vB1 = *reinterpret_cast( + weight + (b_kk0 + kStageK128 / 2) * n + n0 + b_col16); + // Commit as int32 stores (skips the 4-byte pad column; 4 x ds_write_b32 + // per 16-byte vector because the odd strides are 4 B mod 16). + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + int32_t* bdst0 = reinterpret_cast( + b_tile + b_kk0 * kBStride128 + b_col16); + bdst0[0] = vB0.x; + bdst0[1] = vB0.y; + bdst0[2] = vB0.z; + bdst0[3] = vB0.w; + int32_t* bdst1 = bdst0 + (kStageK128 / 2) * (kBStride128 / 4); + bdst1[0] = vB1.x; + bdst1[1] = vB1.y; + bdst1[2] = vB1.z; + bdst1[3] = vB1.w; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + // Consume the single buffer: each wave does eight m16n16k32 MMAs per + // 32-K step over its 32x64 quadrant. + const int local_row = wave_row * 32; + const int local_col = wave_col * 64; +#pragma unroll + for (int kk = 0; kk < kStageK128; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride128 + kk, kAStride128); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride128 + kk, + kAStride128); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride128 + local_col, kBStride128); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride128 + local_col + kTileN, + kBStride128); + du_load_matrix_sync( + b_frag2, b_tile + kk * kBStride128 + local_col + 2 * kTileN, + kBStride128); + du_load_matrix_sync( + b_frag3, b_tile + kk * kBStride128 + local_col + 3 * kTileN, + kBStride128); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + } + // All consumption of the single buffer is complete only after every + // wavefront passes this barrier; the stage s+1 stores below overwrite + // that same buffer, so they must wait for it. + __syncthreads(); + + // Load stage s+1 and commit it into the single buffer. The global-load + // latency is fully exposed here: the loads are issued after the burst + // (nothing to overlap) and the compiler's vmcnt wait before the DS + // stores is on the critical path. `s + 1 < num_stages` is block-uniform, + // so the branch and its barrier are free of divergence. + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vB0 = *reinterpret_cast( + weight + (s1 + b_kk0) * n + n0 + b_col16); + vB1 = *reinterpret_cast( + weight + (s1 + b_kk0 + kStageK128 / 2) * n + n0 + b_col16); + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + int32_t* bdst0 = reinterpret_cast( + b_tile + b_kk0 * kBStride128 + b_col16); + bdst0[0] = vB0.x; + bdst0[1] = vB0.y; + bdst0[2] = vB0.z; + bdst0[3] = vB0.w; + int32_t* bdst1 = bdst0 + (kStageK128 / 2) * (kBStride128 / 4); + bdst1[0] = vB1.x; + bdst1[1] = vB1.y; + bdst1[2] = vB1.z; + bdst1[3] = vB1.w; + // Make the stage s+1 stores visible to every wavefront before the + // next burst. + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 64; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc02, x_scale, weight_scale, out, m, n, + base_row, base_col + 2 * kTileN, lane); + store_prefill_fragment( + acc03, x_scale, weight_scale, out, m, n, + base_row, base_col + 3 * kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); + store_prefill_fragment( + acc12, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + 2 * kTileN, lane); + store_prefill_fragment( + acc13, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + 3 * kTileN, lane); +} + +// Large-M prefill kernel #2: 64x64 output tile per block, K staged in LDS +// (single 128-K buffer), four waves each owning a 32x32 quadrant = four +// m16n16k32 int8->int32 DUMMA accumulators. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM * kStageK]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBlockN]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[64,128] (row-major, zero-filled past M). + constexpr int kAVectors = kBlockM * kStageK / static_cast(sizeof(int4)); + for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int local_row = byte_offset / kStageK; + const int kk = byte_offset - local_row * kStageK; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile)[vec] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + kk) + : int4{0, 0, 0, 0}; + } + + // Stage B[128,64] (row-major [K, N] logical weight layout). + constexpr int kBVectors = kStageK * kBlockN / static_cast(sizeof(int4)); + for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + reinterpret_cast(b_tile)[vec] = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + col); + } + __syncthreads(); + + // Consume the stage: each wave does four m16n16k32 MMAs per 32-K step. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kStageK + kk, kStageK); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kStageK + kk, kStageK); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBlockN + local_col, kBlockN); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBlockN + local_col + kTileN, kBlockN); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + // Single buffer: make sure every wave finished reading LDS before the + // next stage overwrites it. + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Generic scalar int8/int32 fallback for unmatched shapes and small-M API +// cases (including the paired M=2 shape). One output element per grid-stride +// step; exact int32 accumulation (K <= 4096 keeps the int8 dot well within +// int32 range), then the fused float scale and bf16 store. The exact packed +// shape (K=4096, N=1280) uses the packed n-major layout packed[n][k] (set up +// once by launch_pack_w8a8_weight); every other (k, n) keeps the logical +// [K, N] row-major layout. The branch is grid-uniform per launch. +__global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const bool packed_b = (k == kPackedBK && n == kPackedBN); + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; + idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = a + static_cast(row) * k; + const int8_t* b_col = + packed_b ? b + static_cast(col) * k : b + col; + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + const int32_t bw = packed_b + ? static_cast(b_col[kk]) + : static_cast( + b_col[static_cast(kk) * n]); + acc += static_cast(a_row[kk]) * bw; + } + const float scaled = static_cast(acc) * + x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// Identity device-to-device byte copy (pack_weight bootstrap, valid for any +// (K, N)). +__global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +// Identity device-to-device float copy for the weight scales. +__global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +// Iteration 1 (throughput baseline): one-time n-major transpose of the +// logical [K, N] row-major weight into packed[n][k] for the exact assigned +// shape (K=4096, N=1280). Runs only from launch_pack_w8a8_weight, outside +// the timed region and outside Graph capture. The packed buffer keeps the +// same byte count (k*n) as the identity pack, so allocations and +// graph-stable addresses are unchanged. Element (k, n) of the logical weight +// lands at packed[n * K + k]; each thread copies one 16-byte k-run (coalesced +// read side; the strided write side is off the critical path). The guard +// guarantees n % 16 == 0, so every 16-byte chunk lies inside one logical row +// and the transpose is exact. +__global__ __launch_bounds__(256) void w8a8_pack_nmajor_b_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int k, + int n) { + const int64_t chunks = (static_cast(k) * n) >> 4; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t c = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + c < chunks; + c += stride) { + const int64_t off = c << 4; + const int kk = static_cast(off / n); + const int col = static_cast(off - static_cast(kk) * n); + const int4 v = *reinterpret_cast(src + off); + const int8_t* bytes = reinterpret_cast(&v); +#pragma unroll + for (int i = 0; i < 16; ++i) { + // Logical element (k = kk, n = col + i) lands at packed[(col+i)*K+kk]. + dst[static_cast(col + i) * k + kk] = bytes[i]; + } + } +} + +} // namespace + +// Stable host launch symbol consumed by csrc/bindings.cpp. Dispatches on the +// caller-provided stream only; never allocates, synchronizes, or touches the +// default stream. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + auto* out_ptr = static_cast(out); + + // Iteration 1 (throughput baseline): exact assigned shape (K=4096, + // N=1280) routes to the packed n-major B variant of the 2-D macro-tile + // kernel, dispatched as the <64,128> instantiation. The guard is exact + // (m >= 128, n == 1280, k == 4096) and sits BEFORE the generic 64x128 + // path; every other shape keeps its existing path (generic 64x128 for + // n % 128 == 0, 64x64 for n % 64 == 0, scalar fallback otherwise) and the + // identity-packed layout. + if (m >= 128 && n == kPackedBN && k == kPackedBK) { + const dim3 grid( + static_cast(n / kBlockN128), + static_cast((m + kBlockM128 - 1) / kBlockM128)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_tiled_kernel<64, 128>), + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Mandated 2-D macro-tile benchmark family: keep the <64,64> and <128,64> + // instantiations of the packed kernel compiled (and correct) so later + // rounds can flip the exact-shape dispatch without source surgery. This + // branch can never run (m <= 0 already returned above); it only forces + // template instantiation of the two sibling tiles. + if (m < 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_tiled_kernel<64, 64>), + dim3(1), + dim3(kThreadsPerBlock), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_tiled_kernel<128, 64>), + dim3(1), + dim3(kThreadsPerBlock), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Native INT8 DUMMA m16n16k32 prefill path #1 (64x128 macro-tile, + // single-buffered 64-K K stage) for every large-M shape with compatible + // geometry. Covers large-M shapes with N % 128 == 0 and K % 64 == 0 that + // are NOT the packed exact shape (which returned above). Grid = + // (N/128) x (ceil(M/64)). + if (m >= 128 && (n % kBlockN128) == 0 && (k % kStageK128) == 0) { + const dim3 grid( + static_cast(n / kBlockN128), + static_cast((m + kBlockM128 - 1) / kBlockM128)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x128_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Native INT8 DUMMA m16n16k32 prefill path #2 (64x64 tile) for large-M + // shapes whose N is not a multiple of 128 but is a multiple of 64 (or + // whose K is a multiple of 128 while not a multiple of 64). + if (m >= 128 && (n % kBlockN) == 0 && (k % kStageK) == 0) { + const dim3 grid( + static_cast(n / kBlockN), + static_cast((m + kBlockM - 1) / kBlockM)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x64_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Scalar int8/int32 fallback for unmatched (m, n, k) and small-M API cases. + const int64_t total = static_cast(m) * n; + constexpr int kFallbackThreads = 256; + int64_t blocks = (total + kFallbackThreads - 1) / kFallbackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + const dim3 grid(static_cast(blocks)); + const dim3 block(kFallbackThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); +} + +// Stable host launch symbol consumed by csrc/bindings.cpp. For the exact +// assigned shape (K=4096, N=1280) it performs the one-time n-major B pack +// (logical [K, N] -> packed[n][k]) outside the timed region and outside +// Graph capture; every other (K, N) keeps the identity device-to-device +// copy, valid for every (K, N). The packed buffer size is k*n in both +// cases, so the allocation and graph-stable packed layout are unchanged. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kPackThreads = 256; + const int64_t weight_count = static_cast(k) * n; + if (weight_count > 0) { + if (k == kPackedBK && n == kPackedBN) { + const int64_t chunks = weight_count >> 4; + int64_t blocks = (chunks + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_nmajor_b_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + int64_t blocks = (weight_count + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_identity_i8_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count); + } + } + if (n > 0) { + int64_t blocks = (static_cast(n) + kPackThreads - 1) / + kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_identity_f32_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_down_proj.hip new file mode 100644 index 00000000..6c5d9375 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_down_proj.hip @@ -0,0 +1,654 @@ +// @@variant shape=hy3_tp8_shared_down_proj_m4096 commit=426ea470e7ee02bcdbe1428e583651b7e518219e added=2026-08-27 +// median_us=171.2 p90_us=171.5 speedup=62.87 baseline_us=1.076e+04 +// source=hy3-dsh-tp8-m4096-1-0ebb994d +// MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). +// worker_3 iteration 1 (architecture round), assigned shape: +// hy3_tp8_shared_down_proj_m4096 : (M, N, K) = (4096, 4096, 192) +// +// This file provides the two stable host launch symbols consumed by +// csrc/bindings.cpp: +// * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's +// HIP stream; no allocation, packing, +// autotuning, or synchronization. +// * launch_pack_w8a8_weight(...) - out-of-timed-region packing. For the +// exact assigned pair (k, n) == (192, +// 4096) this round emits an n-major +// transpose pack (packed[n][k], one +// contiguous K run per output column); +// every other (k, n) pair keeps the +// identity device-to-device copy so the +// generic fallback and paired small-M +// shapes stay on the raw [K, N] layout. +// +// Round 1 strategy (DUMMA throughput baseline, 2-D macro-tile family): +// * Native INT8 DUMMA m16n16k32 tiled kernel with int32 accumulation. +// * 2-D macro-tile family (64x64, 64x128, 128x64); the active tile is +// selected at compile time (kActiveTileIndex = 1 -> 64x128 (round 4 +// aspect-ratio flip; rounds 1-3 ran 128x64), 512 threads +// = 8 wavefronts, one 32x32 quadrant per wave with four 16x16 int32 +// accumulators kept resident over the whole K loop). All three tiles +// share the templated double-buffered staging pipeline below. +// * K = 192 is pipelined in kPrefillStageK = 64 stages (3 stages) into a +// double-buffered LDS pair with software prefetch: the global int4 +// loads of stage s+1 are issued before the MMAC sequence of stage s and +// the alternate LDS buffer is written right after it, so the vmcnt(0) +// wait and ds_write overlap compute. One barrier per stage. +// * A stays row-major in LDS (stride 88, 8-byte-aligned); B is packed +// n-major at pack time (packed[n][k]) and staged n-major into LDS +// (b_tile[n][k], row stride 72, K-contiguous) with two 8-byte +// ds_write_b64 per thread, consumed through col_major matrix_b +// fragments; both operand paths use the round-8 explicit load_frag8 +// 8-byte ds_read_b64 reads (the bootstrap's row_major matrix_b path was +// byte-granular: 16 ds_read_u8 + ~40 VALU of byte assembly per wave per +// k-step, the measured LDS-issue bottleneck of the 634 us bootstrap). +// * kk body fully unrolled into two independent fragment register sets +// (kStageK == 2 * kTileK) so the second step's LDS latency overlaps the +// first step's swizzle + MMACs. +// * Fused direct fragment -> x_scale -> weight_scale -> bf16 epilogue, +// using the verified gfx928 accumulator ownership (row = lane & 15, +// col_mod4 = lane >> 4, x[i] -> columns col_mod4 + 4*i). The int32 dot +// is exact (max |dot| = 192*127*127 << 2^31), the int32 -> float +// conversion is exact, and the single fp32 multiply followed by one +// bfloat16 rounding matches the CPU int64 reference bit-for-bit. The +// col_major loader places the same B[k][n] values in the same fragment +// registers as row_major, so int32 accumulation is bit-identical to the +// bootstrap path. +// * All other (m, n, k) - including every M < 128 API case and the paired +// M=2 shape with the same (N, K) - go to a scalar int8/int32 fallback +// that decodes the packed layout only for (k, n) == (192, 4096) and the +// raw [K, N] layout for every other pair. +// +// Header order is fixed by the control plane: hip_runtime first (du_mma.h +// is not self-contained before it), hip_bfloat16 second, du_mma.h last. + +#include +#include +#include + +#include + +namespace { + +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; +constexpr int kWaveSize = 64; + +// Staging stage length for the prefill pipeline. 64 = 2x the DUMMA K unit: +// with the 128x64 macro-tile the double-buffered footprint is +// 2*(128*88 + 64*72) = 31,744 B; with the round-4 active 64x128 tile it is +// 2*(64*88 + 128*72) = 29,696 B (A stride 88, transposed-B stride 72; both +// footprints leave 2 blocks/CU LDS-fit at 512 threads (8 waves per block, +// 16 resident wavefronts): 59,392 B <= 64 KiB for 64x128, the same 29,696 B +// per block that round 1 ran resident at). Stage 64 keeps +// the barrier +// count low for the short K = 192 (3 stages -> 4 barriers per block +// including the prologue) while keeping one full 64-k stage of global loads +// in flight. Padded row strides stay aligned: kAStride 88 % 8 == 0 for the +// two-half int64 A staging, kBTStride 72 % 8 == 0 for the col_major 8-byte +// fragment reads and the two-half ds_write_b64 B staging. The A pad is 24 +// (stride 22 dwords): the row_major m16k32 fragment maps lane l to row +// l&15, k-quarter l/16 (8 contiguous bytes per lane), so the 16 fragment +// rows read banks (22*r mod 32) - 16 distinct even bank-pairs - and every +// A ds_read2 hits the 4-phase 32-bank floor with zero conflicts. Stride 20 +// dwords (pad 16) aliases rows r and r+8 onto identical bank pairs +// (20*r mod 32 has period 8), doubling each A fragment read to 8 phases; +// stride 22 is the smallest even dword stride >= 16 whose 16 multiples are +// pairwise distinct mod 32. +constexpr int kPrefillStageK = 64; + +// The only assigned shape for this worker has K = 192, so the optimized +// launch guard is shape-exact on this value (k == kAssignedK). 192 = 3 +// stage-64 chunks. +constexpr int kAssignedK = 192; +constexpr int kAssignedN = 4096; + +// Compile-time selection of the active 2-D macro-tile: +// 0 -> 64x64, 1 -> 64x128, 2 -> 128x64. +// All three instantiations are compiled so the family stays measurable; only +// the active one is launched inside the timed region. Round 4 (tile-shape +// round) flips the active tile from 128x64 to 64x128: the first constant- +// occupancy aspect-ratio A/B in this session (same 512 threads / 8 waves / +// 2 blocks/CU / grid 2048; only the (M,N) quadrant partition and grid +// geometry (N/128, M/64) swap). The flip halves per-A-tile global re-reads +// (64 -> 32) and doubles B re-reads (32 -> 64), both 768 KiB L2-resident; +// the active footprint is 29,696 B/block (59,392 B/CU at 2 blocks <= 64 KiB, +// the same per-block footprint round 1 ran resident at). +constexpr int kActiveTileIndex = 1; +constexpr int kActiveBlockM = kActiveTileIndex == 2 ? 128 : 64; +constexpr int kActiveBlockN = kActiveTileIndex == 1 ? 128 : 64; + +// LDS padding in bytes added to each staged A row. 24 keeps every row start +// 8-byte aligned for the two-half int64 staging while making the 16 +// m16k32-fragment row bases pairwise distinct mod the 32 LDS banks +// (stride 22 dwords: 22*r mod 32 is distinct for r = 0..15), so the A +// fragment ds_read2 instructions run at the conflict-free 4-phase floor. +// Strides of 20 dwords (pad 16) alias rows r and r+8 onto identical bank +// pairs (8 two-way conflicts per 16-lane group -> 8 phases per read); +// strides >= 16 that are multiples of 128 B alias every row onto one bank +// phase (up to 8-16-way conflicts). +constexpr int kLdsPad = 24; + +// LDS row stride of the transposed (n-major) B tile: each row holds one +// output column n with kStageK contiguous K bytes. kLdsPadB = 8 keeps every +// row start 8-byte aligned so the col_major matrix_b fragment loader can +// lower its 8 contiguous bytes per lane to a single vectorized ds_read2_b32 +// and the B staging can land each n row's k-contiguous 16-byte run as two +// 8-byte ds_write_b64 (72 % 8 == 0; 72 % 16 == 8, so a single ds_write_b128 +// would be 16-byte misaligned for odd n rows - the explicit two-half store +// keeps the lowering legal). +constexpr int kLdsPadB = 8; +constexpr int kBTStride = kPrefillStageK + kLdsPadB; + +using namespace du::dumma; + +// gfx928 int8 m16n16k32 accumulator ownership, established against +// du_store_matrix_sync: lane % 16 selects the row, lane / 16 selects the +// column mod 4, and x[i] selects columns separated by four. +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = + static_cast(frag.x[i]) * xs * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// Round 8 (LDS read-width round): explicit 8-byte fragment loader replacing +// the library du_load_matrix_sync for the int8 m16k32 fragments. The library +// row_major matrix_a and col_major matrix_b loaders both fetch the same +// per-lane bytes - lane l -> row (l & 15), k-quarter (l >> 4), eight +// contiguous bytes at p[row*ldm + (l>>4)*8 .. +7] - and the compiler lowers +// them to ds_read2_b32 (two 4-byte dword phases per half-wave). A +// ds_read2_b32 phase groups 32 lanes (rows 0-15 of k-quarter q with rows +// 0-15 of k-quarter q+1) into one 32-bank cycle, so the kq and kq+1 row sets +// share every phase: {22*r mod 32} and {22*r+2 mod 32} (A stride 22 dwords; +// 18*r and 18*r+2 for B stride 18) are the same 16 even banks, giving a +// 2-way conflict in every phase - 8 phases per read at ANY even dword row +// stride. This is why the exact-source PMC lds_bank_conflicts (2,162,688) +// is byte-identical between the round-5 stride-80 and round-6 stride-88 +// sources. Loading the same 8 bytes as one int64 forces the 8-byte +// ds_read_b64 form, whose phases are 16-lane groups (16 lanes x 8 B = 128 B +// = one bank cycle): within each group the 16 rows land on 16 distinct +// bank-pairs (A: 22*r mod 32 distinct for r = 0..15, B: 18*r mod 32 +// distinct), so every fragment read runs at the 4-phase conflict-free floor. +// The element-to-slot mapping is unchanged (f.x[0..7] = the same 8 bytes in +// the same little-endian order), the MMACs and the exact int32 k-ascending +// accumulation sequence are untouched, and the staging stores, LDS strides, +// barriers, tile, occupancy and epilogue are byte-identical. +// --------------------------------------------------------------------------- +template +__device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, + unsigned ldm) { + const unsigned row = __lane_id() & 0xf; + const unsigned kq = __lane_id() >> 4; + const int64_t v = + *reinterpret_cast(p + row * ldm + (kq << 3)); + reinterpret_cast(f.x)[0] = v; +} + +// --------------------------------------------------------------------------- +// Large-M prefill path. One block computes a kBlockM x kBlockN output tile +// with (kBlockM/32) x (kBlockN/32) wavefronts (one 32x32 quadrant per wave). +// A[kBlockM, K] and B[K, kBlockN] are cooperatively staged into a +// double-buffered, bank-padded LDS in kStageK chunks with VGPR prefetch: +// the global int4 loads for stage s+1 are issued before the MMAC sequence +// of stage s, and the alternate LDS buffer is written right after it, so +// the vmcnt(0) wait and ds_write overlap compute. One barrier per stage. +// A stays row-major in LDS (stride 88); B is packed n-major by +// launch_pack_w8a8_weight (packed[n][k], one contiguous K run per output +// column) and staged n-major (stride 72, K-contiguous) with two 8-byte +// ds_write_b64 per thread, loaded through the round-8 explicit load_frag8 +// 8-byte ds_read_b64 reads (row_major A and col_major B fragments place the +// same 8 contiguous bytes per lane in the same x[0..7] slots, so both +// operand paths are vectorized 8-byte LDS reads). Each wave keeps four +// 16x16 int32 accumulators resident over the whole K loop, then a fused +// scale/bf16 store. The kk ordering (k0 outer, kk inner step kTileK) and +// the acc00/acc01/acc10/acc11 update order are unchanged from the +// bootstrap, and the loaders place the same B[k][n] values in the same +// fragment registers as row_major, so int32 accumulation is bit-identical. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kBlockM * kBlockN / 1024 * kWaveSize, 2) void +w8a8_dumma_prefill_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWavesN = kBlockN / 32; + constexpr int kAStride = kStageK + kLdsPad; + static_assert(kAStride % 8 == 0, + "A LDS row stride must stay 8-byte aligned"); + static_assert(kBTStride % 8 == 0, + "transposed-B LDS row stride must stay 8-byte aligned"); + static_assert(kStageK % kTileK == 0, + "K stage must be a multiple of the DUMMA K unit"); + static_assert(kStageK % sizeof(int4) == 0, + "K stage must keep int4 staging aligned"); + static_assert(kStageK == 2 * kTileK, + "unrolled kk body assumes kStageK == 2 * kTileK"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWavesN; + const int wave_col = wave - wave_row * kWavesN; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + constexpr int kAInt4PerRow = kStageK / sizeof(int4); + constexpr int kAInt4 = kBlockM * kAInt4PerRow; + constexpr int kChunksPerBRow = kStageK / sizeof(int4); + constexpr int kBInt4 = kBlockN * kChunksPerBRow; + static_assert(kAInt4 <= kThreads && kBInt4 <= kThreads, + "one int4 chunk per thread requires chunks <= blockDim"); + const int kStages = k / kStageK; + + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + // Transposed (n-major) B tile: element B[k0+kk][n] lives at + // b_tile[n][kk], so col_major fragment loads read 8 contiguous K bytes + // per lane (vectorized ds_read2_b32). The staged runs are written n-major + // (packed[n][k]) as two 8-byte ds_write_b64 per thread. + __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBTStride]; + + DUFragment + a_frag0, a_frag1, a_frag0_1, a_frag1_1; + DUFragment + b_frag0, b_frag1, b_frag0_1, b_frag1_1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + // Prologue: prefetch stage 0 into VGPRs and land it in buffer 0. Every + // source address is 16-byte aligned (k % 16 == 0 by the exact k guard and + // n % 16 == 0 by the launch guard); A destination rows are 8-byte aligned + // (stride 88) and B destination rows 8-byte aligned (stride 72), so all + // staging stores are the two-half int64 ds_write2_b64 form. + int4 a_pref{0, 0, 0, 0}; + int4 b_pref{0, 0, 0, 0}; + { + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + const int global_row = m0 + row; + if (global_row < m) { + a_pref = *reinterpret_cast( + x_q + global_row * k + seg * sizeof(int4)); + } + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + b_pref = *reinterpret_cast( + weight + (n0 + n_local) * k + kchunk); + } + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + // Two 8-byte halves: kAStride 88 % 16 == 8 makes odd rows 16-byte + // misaligned, so the single-int4 store is illegal; the int64 pair + // lowers to one ds_write2_b64 per thread, exactly like the B staging. + const int64_t* src64 = reinterpret_cast(&a_pref); + int64_t* dst64 = reinterpret_cast( + a_tile[0] + row * kAStride + seg * sizeof(int4)); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + const int64_t* src64 = reinterpret_cast(&b_pref); + int64_t* dst64 = reinterpret_cast( + b_tile[0] + n_local * kBTStride + kchunk); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + __syncthreads(); + } + + for (int s = 0; s < kStages; ++s) { + const int k0 = s * kStageK; + const int buf = s & 1; + + // Issue the global loads for stage s+1 now so their latency hides + // behind the fragment loads and MMACs below. Out-of-range A rows are + // zero-filled (tail-M guard). + if (s + 1 < kStages) { + const int k1 = k0 + kStageK; + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + const int global_row = m0 + row; + if (global_row < m) { + a_pref = *reinterpret_cast( + x_q + global_row * k + k1 + seg * sizeof(int4)); + } else { + a_pref = int4{0, 0, 0, 0}; + } + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + b_pref = *reinterpret_cast( + weight + (n0 + n_local) * k + k1 + kchunk); + } + } + + // Consume stage s from buffer buf: four fragment loads and four MMACs + // per wave per kTileK step. kStageK = 64 = 2 kTileK steps, fully + // unrolled into two independent fragment register sets: all eight + // 8-byte fragment reads of a stage issue before the MMACs, so the + // second step's LDS latency overlaps the first step's MMACs. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + // Round 8: custom load_frag8 replaces du_load_matrix_sync for all eight + // fragment reads. Both library loaders (row_major A and col_major B) + // place the same 8 contiguous bytes per lane into the same x[0..7] + // slots; only the LDS instruction changes, from ds_read2_b32 (2-way + // conflict in every 32-lane dword phase at any even row stride) to the + // 8-byte ds_read_b64 (16-lane phases, distinct bank-pairs per group -> + // 4-phase conflict-free floor). Exact int32 accumulation is unchanged. + load_frag8(a_frag0, a_tile[buf] + local_row * kAStride, kAStride); + load_frag8(a_frag1, a_tile[buf] + (local_row + 16) * kAStride, kAStride); + load_frag8(b_frag0, b_tile[buf] + local_col * kBTStride, kBTStride); + load_frag8(b_frag1, b_tile[buf] + (local_col + 16) * kBTStride, + kBTStride); + // Second 32-k step (kk = kTileK): independent registers, issued while + // the first step's data is still in flight. + load_frag8(a_frag0_1, a_tile[buf] + local_row * kAStride + kTileK, + kAStride); + load_frag8(a_frag1_1, + a_tile[buf] + (local_row + 16) * kAStride + kTileK, kAStride); + // col_major: element B[k0+kk+k][n] at b_tile[n][kk+k], so the loader's + // p[row*ldm + col + i] reads 8 contiguous bytes per lane. + load_frag8(b_frag0_1, b_tile[buf] + local_col * kBTStride + kTileK, + kBTStride); + load_frag8(b_frag1_1, + b_tile[buf] + (local_col + 16) * kBTStride + kTileK, + kBTStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc00, a_frag0_1, b_frag0_1, acc00); + du_mma_sync(acc01, a_frag0_1, b_frag1_1, acc01); + du_mma_sync(acc10, a_frag1_1, b_frag0_1, acc10); + du_mma_sync(acc11, a_frag1_1, b_frag1_1, acc11); + + // Land the prefetched stage s+1 in the alternate buffer. The compiler + // inserts the vmcnt(0) wait here via the data dependency, so the wait + // and the ds_write overlap the MMACs above instead of stalling before + // them. + if (s + 1 < kStages) { + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + const int64_t* src64 = reinterpret_cast(&a_pref); + int64_t* dst64 = reinterpret_cast( + a_tile[buf ^ 1] + row * kAStride + seg * sizeof(int4)); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + const int64_t* src64 = reinterpret_cast(&b_pref); + int64_t* dst64 = reinterpret_cast( + b_tile[buf ^ 1] + n_local * kBTStride + kchunk); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + } + // One barrier per stage: makes the alternate-buffer writes visible to + // the next stage's reads and retires this stage's reads before the + // parity-flipped buffer is overwritten two stages later. + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + 16, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + 16, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + 16, base_col + 16, lane); +} + +// --------------------------------------------------------------------------- +// Generic scalar int8 x int8 -> int32 fallback for unmatched (m, n, k) and +// small-M API cases. One thread per output element, coalesced along N. +// For the exact packed pair (k, n) == (192, 4096) the weight is n-major +// (packed[n][k], column stride k); every other pair keeps the raw logical +// [K, N] row-major layout (column stride n). The int32 summation order is +// identical either way (k-ascending), keeping the fallback bit-identical. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_scalar_fallback_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k, + int b_col_stride) { + const int linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int total = m * n; + if (linear >= total) { + return; + } + const int row = linear / n; + const int col = linear - row * n; + const int8_t* a_row = x_q + row * k; + // b_col_stride is 1 for the n-major packed (192, 4096) pair (each output + // column has a contiguous K run) and n for the raw [K, N] layout. + const int8_t* b_col = weight + col * b_col_stride; + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[kk * b_col_stride]); + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Pack kernels for launch_pack_w8a8_weight (run once per weight tensor +// outside the timed region and outside CUDA/HIP Graph capture, so plain +// kernels are fine): +// * w8a8_transpose_i8_kernel: n-major transpose pack for the exact +// assigned pair (k, n) == (192, 4096): packed[col * k + row] = +// raw[row * n + col] gives every output column a contiguous K run, +// which the timed GEMM stages with two vectorized ds_write_b64 per +// thread instead of byte-granular scatters. +// * w8a8_identity_copy_i8_kernel: identity device-to-device copy for every +// unmatched (k, n) pair (raw [K, N] layout preserved). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_transpose_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int k, + int n) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(k) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + dst[static_cast(col) * k + row] = src[linear]; +} + +__global__ __launch_bounds__(256) void w8a8_identity_copy_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t numel) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < numel) { + dst[linear] = src[linear]; + } +} + +__global__ __launch_bounds__(256) void w8a8_identity_copy_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t numel) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < numel) { + dst[linear] = src[linear]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Host launch symbols (extern "C", consumed by csrc/bindings.cpp). +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + auto* out_bf16 = reinterpret_cast<__hip_bfloat16*>(out); + + // Native INT8 DUMMA m16n16k32 tiled path for the assigned shape family. + // The guard is shape-exact (k == kAssignedK == 192, the only assigned K), + // keeps M >= 128 (small-M API cases fall through to the scalar path), and + // requires N to be a multiple of the active macro-tile N (4096 always + // is). The active macro-tile is fixed at compile time, so the timed + // region never tunes. All three tile instantiations share the same guard. + if (m >= 128 && k == kAssignedK && (n % kActiveBlockN) == 0) { + const dim3 grid( + static_cast(n / kActiveBlockN), + static_cast((m + kActiveBlockM - 1) / kActiveBlockM)); + const dim3 block(static_cast( + kActiveBlockM * kActiveBlockN / 1024 * kWaveSize)); + if (kActiveTileIndex == 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 64, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else if (kActiveTileIndex == 1) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 128, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<128, 64, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + return; + } + + // Generic scalar fallback for every unmatched (m, n, k) and M < 128. + // The packed weight is n-major only for the exact assigned pair + // (k, n) == (192, 4096); every other pair is identity-packed, so the + // fallback decodes b_col with column stride k for that pair and with + // column stride n otherwise (the paired M=2 shape with the same (N, K) + // is covered by the k-stride branch). + const bool packed_nmajor = (k == kAssignedK && n == kAssignedN); + const int b_col_stride = packed_nmajor ? 1 : n; + const int total = m * n; + constexpr int kFallbackThreads = 256; + const dim3 grid(static_cast( + (total + kFallbackThreads - 1) / kFallbackThreads)); + const dim3 block(static_cast(kFallbackThreads)); + hipLaunchKernelGGL( + w8a8_scalar_fallback_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k, b_col_stride); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kCopyThreads = 256; + const int64_t weight_numel = static_cast(k) * n; + const dim3 weight_grid(static_cast( + (weight_numel + kCopyThreads - 1) / kCopyThreads)); + + // Round 1: n-major transpose pack (packed[n][k], K-contiguous per output + // column) for the exact assigned pair (k, n) == (192, 4096) only. The API + // contract (int8_w8a8_gemm_api.py) allows any contiguous opaque packed + // layout; the timed GEMM and the (192, 4096) fallback consume exactly + // this layout, so the timed B staging drops from byte-granular scatters + // to two 8-byte ds_write_b64 per thread. Every other (k, n) pair keeps + // the identity device-to-device copy (raw [K, N] layout), so unmatched + // shapes and paired small-M shapes on other pairs stay on the layout + // their fallback decodes. + if (k == kAssignedK && n == kAssignedN) { + hipLaunchKernelGGL( + w8a8_transpose_i8_kernel, + weight_grid, dim3(kCopyThreads), 0, stream, + raw_weight, packed_weight, k, n); + } else { + hipLaunchKernelGGL( + w8a8_identity_copy_i8_kernel, + weight_grid, dim3(kCopyThreads), 0, stream, + raw_weight, packed_weight, weight_numel); + } + + // The scale copy is unchanged (N-length, order-independent). + const dim3 scale_grid(static_cast( + (n + kCopyThreads - 1) / kCopyThreads)); + hipLaunchKernelGGL( + w8a8_identity_copy_f32_kernel, + scale_grid, dim3(kCopyThreads), 0, stream, + weight_scale, packed_weight_scale, static_cast(n)); +} +// @@end +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_gate_up_proj.hip new file mode 100644 index 00000000..eefe8afc --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_gate_up_proj.hip @@ -0,0 +1,1075 @@ +// @@variant shape=hy3_tp8_shared_gate_up_proj_m4096 commit=95895c8eadee15332d1799fbfd36789ff57dfee7 added=2026-08-27 +// median_us=116 p90_us=116.2 speedup=144.1 baseline_us=1.671e+04 +// source=hy3-dsh-tp8-m4096-1-0ebb994d +// INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. +// +// Worker 2 assigned shape: hy3_tp8_shared_gate_up_proj_m4096 +// (M=4096, N=384, K=4096). +// +// Iteration 2 (architecture round): the bootstrap (iteration 1) was a +// single-buffered 64x128 macro-tile with identity [K, N] packed B. Its ISA +// (PMC bundle) shows the dominant steady-state costs are the B-fragment LDS +// path (with any 16-byte-aligned [K, N] row stride the int8 DUMMA row-major +// B loader reads 8 strided k-rows per lane -> eight ds_read_u8, 16-way bank +// aliasing) plus the per-dword mask/OR fragment reassembly VALU, and the +// single-buffered K loop serializes one global-load round trip in front of +// every stage. This round ports the sibling-validated architecture family +// (hy3 TP4 M4096 shared_gate_up_proj, median 220.7 us at 2x this shape's +// FLOPs, 116.8 TOPS) for the exact (K, N) = (4096, 384) pair: +// * launch_pack_w8a8_weight stores the (K=4096, N=384) weight transposed +// [N, K] (packing runs outside the timed/CUDA-Graph region; all other +// (k, n) pairs keep the identity copy). +// * The prefill kernel gains a kBuffers==2 software-pipelined K loop with +// K stage 64 (one barrier per stage): stage s+1's global loads are +// issued into registers before the stage-s MMA compute and committed to +// the second LDS buffer after it, overlapping the global round trip with +// the LDS-wait-bound compute. Buffer footprint 2 x (64x80 A + 128x80 B) +// = 30,720 B/block -> 2 blocks/CU -> 16 waves/CU (occupancy preserved). +// * B is staged [N, K] in LDS from the transposed packed weight and loaded +// with the col_major fragment loader; each lane's 8 elements are 8 +// consecutive k-values at an 8-byte-aligned LDS address (row stride 80), +// so each B fragment is ONE 8-byte LDS read (ds_read2_b32, ~4-way bank +// floor) instead of eight ds_read_u8 at a 16-way aliasing phase. +// * load_fragment8 writes the same 8 bytes directly into the fragment +// storage (identical element-to-slot mapping and byte order), so the +// v_mmac operand bit patterns and the exact int32 accumulation are +// unchanged while the per-dword mask/OR reassembly VALU disappears. +// * The generic single-buffer 64x128/128x64/64x64 arms, identity packing, +// the scalar fallback (now transposed-aware for the exact (k, n) pair so +// the paired M=2/M=16 decode shapes stay correct), and the A staging / +// epilogue are unchanged in behavior. +// +// Iteration 9 (wave-structure round): the exact assigned shape (M=4096, +// N=384, K=4096) dispatches to w8a8_gemm_prefill_tiled_kernel_w4<64,128,64> +// (see that kernel's comment): identical 64x128 macro-tile, K stage 64 +// double buffer, [N,K] packed B, 30,720 B LDS, 192-block grid, and +// bit-identical int32 accumulation, with 256 threads (4 waves) owning 32x64 +// quadrants and eight independent accumulators per wave instead of 512 +// threads owning 32x32 quadrants with four accumulators. Every other +// large-M shape with the packed (k, n) = (4096, 384) pair keeps the +// accepted 8-wave <64,128,64,2> path. +// +// Iteration 11 (address-math consolidation, HIP only): the accepted w4 +// kernel's exact-source ISA recomputes the whole global/LDS address chain +// every K stage (~34 loop-invariant VALU/SALU per wave per stage: the row*K +// products, the quadrant/lane slot constants, and the stage-count bound), +// although only the 64-byte stage offset and the double-buffer select vary. +// This round hoists every loop-invariant address into per-thread staging +// pointers / LDS commit slots / per-wave fragment-read bases computed once +// before the K loop and walks the three global pointers by kStageK per +// stage. Same bytes, same element-to-slot mapping, same load/store order, +// same barrier structure, grid, LDS, and occupancy -> bit-identical int32 +// accumulation (expected mismatch 0). +// +// Iteration 13 (burst-scheduling probe, HIP only): the w4 steady state +// issues each stage's sixteen v_mmac_i32_16x16x32_i8 as eight kk=0 MMACs +// followed by eight kk=kTileK MMACs into the same eight accumulators. This +// round pins the fragment-load schedule: all twelve load_fragment8 calls +// (six per kk slice) are issued EXPLICITLY before the whole sixteen-MMAC +// burst, so the second slice's six ds_read2_b64 no longer compete with the +// first slice's burst for a liveness-driven placement (the compiler keeps +// both slices' fragments live across the burst, ~+12 VGPR, still <= 128 for +// 2 blocks/CU). No load, byte, element-to-slot mapping, LDS commit, barrier +// (65/block), grid (192 blocks), LDS (30,720 B), threads (256), or MMAC +// emission order changes -> the int32 accumulation is bit-identical +// (expected mismatch 0). +// +// Iteration 18 (wave-count probe, ISA-guided round 2 of 2): the accepted w4 +// cycle accounting (48 single-block CUs at 1 wave/SIMD take 164.7 cycles per +// v_mmac, 72 dual-block CUs at 2 waves/SIMD take 82.4, both landing exactly +// on the 168,654-cycle GRBM wall) isolates the per-SIMD v_mmac issue rate as +// the wall, and rounds 13-15 proved the burst dependency structure and the +// operand-register alternation do not move it. The untested axis is the +// WAVEFRONT COUNT per SIMD: the wall is invariant to occupancy while each +// wavefront still issues 1,024 v_mmac (165 cyc x 1,024 = 168,960 cyc), so +// the exact shape now dispatches to w8a8_gemm_prefill_tiled_kernel_w8 +// <64,128,64> -- an 8-wave (512-thread) clone of the accepted w4 steady +// state (iteration-11 address hoist, iteration-13 burst pinning, same +// 64x128 tile, same 192-block grid, same 30,720-B LDS, same 65 barriers/ +// block, same staging bytes and element-to-slot mapping -> bit-identical +// int32): 32x32 quadrant per wave, four accumulators, 8 MMACs/stage, +// 512 MMACs/wavefront, target <= 64 VGPR -> 2 blocks/CU -> 4 waves/SIMD. +// Falsifiable: 4-deep MMAC pipe -> wall 84,480 cyc (~58 us, ~2x); 2-deep +// pipe (throughput cap 1/82.4 per SIMD) -> wall ~168,755 (flat). +// +// Mathematical contract (exact int32 dot before float scaling): +// out[m, n] = bf16(int32_dot(a[m, :], b[:, n]) * x_scale[m] * weight_scale[n]) +// +// Header order is fixed by the control plane for this DTK: +// hip_runtime.h -> hip_bfloat16.h -> du_mma.h + +#include +#include +#include + +#include + +namespace { + +// gfx928 INT8 DUMMA primitive: m16n16k32, int8 x int8 -> int32. +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; + +// gfx928 wavefront is 64 lanes; every blockDim must be a multiple of 64. +constexpr int kWaveSize = 64; + +// Large-prefill dispatch geometry. Each macro-tile config uses a 32x32 +// quadrant per wave, so waves per block = (BlockM/32) * (BlockN/32). +constexpr int kPrefillMinM = 128; +constexpr int kPrefillStageK = 128; +constexpr int kDefaultBlockM = 64; +constexpr int kDefaultBlockN = 64; + +constexpr int kScalarThreads = 256; + +using namespace du::dumma; + +// Verified gfx928 INT8 DUMMA accumulator ownership: +// row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> columns +// col_mod4 + 4*i. Scale by x_scale[row] and weight_scale[col] and store +// bf16 directly from the fragment (no accumulator LDS round trip). +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Direct 8-byte LDS fragment load for the double-buffered exact-shape path. +// du_load_matrix_sync's int8 loaders assign x[0..7] = 8 consecutive bytes at +// (lane & 15) * ldm + ((lane >> 4) << 3) for both matrix_a row_major and +// matrix_b col_major, and du_mma_sync passes the fragment to the v_mmac +// builtin as one packed 8-byte operand, so the compiler emits a per-byte +// mask/OR reassembly chain (v_and 0xff00/0xff0000/0xff000000 + v_or_b32_sdwa +// + v_or3_b32) for every loaded dword. Writing the same 8 bytes directly into +// the fragment storage keeps the operand bit pattern identical (exact int32 +// accumulation unchanged) and lets the compiler feed the ds_read2_b32 pair +// straight to the v_mmac (still one ds_read2_b32 per fragment, 4-way floor). +__device__ __forceinline__ void load_fragment8( + DUFragmentBase& frag, + const signed char* __restrict__ base, + int stride, + int lane) { + const int off = (lane & 15) * stride + ((lane >> 4) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// Large-prefill path: one block computes a kBlockM x kBlockN output tile. +// kWaveRows x kWaveCols wavefronts each own a 32x32 quadrant (four 16x16 +// DUMMA accumulators), while the block cooperatively stages A[kBlockM, +// kStageK] and B[kStageK, kBlockN] in bank-skewed LDS. Tail M rows are +// zero-filled on load and masked on store, so any M >= 128 is supported. +// +// kBuffers == 1: single-buffered K loop (K stage 128, two barriers per stage): +// stage s+1's global loads cannot start until the stage-s barrier, so each +// stage serializes one global-load round trip (global_load -> vmcnt(0) -> +// ds_write -> lgkmcnt(0) -> barrier) in front of the compute. +// kBuffers == 2: software-pipelined K loop (K stage 64, one barrier per stage): +// stage s+1's loads are issued into registers before the stage-s MMA compute +// and committed to the other LDS buffer after it, overlapping the global +// round trip with the LDS-wait-bound compute. Requires at most one int4 per +// thread per operand (static_asserted), which holds for the exact shape. +// The (K, N) = (4096, 384) packed weight is transposed [N, K], so B is +// staged [N, K] and loaded with the col_major fragment loader. +template +__global__ __launch_bounds__((kBlockM / 32) * (kBlockN / 32) * kWaveSize) void +w8a8_gemm_prefill_tiled_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWaveRows = kBlockM / 32; + constexpr int kWaveCols = kBlockN / 32; + constexpr int kThreads = kWaveRows * kWaveCols * kWaveSize; + // 16-byte bank skew: keeps int4 staging stores aligned while spreading the + // fragment rows across distinct LDS banks (~4-way floor for 8-byte/lane + // reads instead of the 16-way aliasing of power-of-two strides). + constexpr int kAStride = kStageK + sizeof(int4); + // Double-buffered path (exact shape) stages B transposed [N, K] so the + // DUMMA B fragments read 8 consecutive k-values per lane (one ds_read2_b32, + // ~4-way) instead of 8 strided k-rows (eight ds_read_u8, 16-way: with any + // 16-byte-aligned row stride S, B k-rows 8 apart alias onto one bank group + // because (S/4)*8*g == 0 mod 32 for all g). The generic single-buffer path + // keeps [K, N] B and the row-major B loader. + constexpr int kBStride = (kBuffers == 2) ? (kStageK + sizeof(int4)) + : (kBlockN + sizeof(int4)); + constexpr int kBRows = (kBuffers == 2) ? kBlockN : kStageK; + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = + (kBuffers == 2) ? (kStageK / sizeof(int4)) : (kBlockN / sizeof(int4)); + constexpr int kBVectors = kBRows * kBVectorsPerRow; + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWaveCols; + const int wave_col = wave - wave_row * kWaveCols; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBuffers][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBuffers][kBRows * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + if constexpr (kBuffers == 2) { + static_assert(kAVectors <= kThreads && kBVectors <= kThreads, + "double-buffered path stages at most one int4 per thread"); + // The exact shape's packed weight is [N, K] (transposed by + // launch_pack_w8a8_weight), so B fragments load 8 consecutive k-values + // per lane from the [N, K] LDS tile via the col_major loader; the + // element-to-slot mapping (slot i = B[k = kk + 8*g + i][n = col + row]) + // is identical to the row-major loader on a [K, N] tile, so the v_mmac + // operand values and the exact int32 accumulation are unchanged. + DUFragment + b_frag_t0, b_frag_t1; + // Prologue: stage K tile 0 into buffer 0, then one barrier. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + // B is packed [N, K] for the exact shape: each thread stages 16 + // consecutive k-values of one n row (n-stride in global is k). + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + n_local * kBStride)[v] = + *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + const int n_stages = k / kStageK; + for (int s = 0; s < n_stages; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + const bool has_next = (s + 1) < n_stages; + // Issue the stage s+1 global loads now; the data is consumed by the + // ds_write after the compute, so the vmcnt wait lands after the MMA + // loop instead of stalling the front of the stage. + int4 a_reg{0, 0, 0, 0}; + int4 b_reg{0, 0, 0, 0}; + if (has_next) { + const int k1 = (s + 1) * kStageK; + if (tid < kAVectors) { + const int local_row = tid / kAVectorsPerRow; + const int v = tid - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + a_reg = global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k1 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + const int n_local = tid / kBVectorsPerRow; + const int v = tid - n_local * kBVectorsPerRow; + b_reg = *reinterpret_cast( + weight + (n0 + n_local) * k + k1 + + v * static_cast(sizeof(int4))); + } + + // Compute stage s from the buffer staged last iteration. + // Load the int8 fragments as raw 8-byte LDS reads into the fragment + // storage (same bytes, same element-to-slot mapping), so the compiler + // feeds the ds_read2_b32 pair straight to the v_mmac and the per-dword + // mask/OR reassembly VALU disappears from the steady state. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + load_fragment8( + a_frag0, a_tile[cur] + local_row * kAStride + kk, kAStride, + lane); + load_fragment8( + a_frag1, a_tile[cur] + (local_row + kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8( + b_frag_t0, b_tile[cur] + local_col * kBStride + kk, kBStride, + lane); + load_fragment8( + b_frag_t1, b_tile[cur] + (local_col + kTileN) * kBStride + kk, + kBStride, lane); + du_mma_sync(acc00, a_frag0, b_frag_t0, acc00); + du_mma_sync(acc01, a_frag0, b_frag_t1, acc01); + du_mma_sync(acc10, a_frag1, b_frag_t0, acc10); + du_mma_sync(acc11, a_frag1, b_frag_t1, acc11); + } + + // Commit the prefetched stage s+1 into the other LDS buffer, then one + // barrier: it orders both this iteration's compute reads of buffer cur + // (against the next-next prefetch, which reuses cur) and the prefetch + // writes of buffer nxt (against the next iteration's compute reads). + if (has_next) { + if (tid < kAVectors) { + const int local_row = tid / kAVectorsPerRow; + const int v = tid - local_row * kAVectorsPerRow; + reinterpret_cast(a_tile[nxt] + local_row * kAStride)[v] = + a_reg; + } + const int n_local = tid / kBVectorsPerRow; + const int v = tid - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[nxt] + n_local * kBStride)[v] = b_reg; + } + __syncthreads(); + } + } else { + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[kBlockM, kStageK] into LDS (zero-filled tail M rows). + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + // Stage B[kStageK, kBlockN] into LDS from the packed weight. The + // generic path keeps the identity [K, N] layout (pack_weight is an + // identity copy for every (k, n) pair other than the exact shape). + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int kk = vec / kBVectorsPerRow; + const int v = vec - kk * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + kk * kBStride)[v] = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + du_load_matrix_sync( + a_frag0, a_tile[0] + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile[0] + (local_row + kTileM) * kAStride + kk, + kAStride); + du_load_matrix_sync( + b_frag0, b_tile[0] + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile[0] + kk * kBStride + local_col + kTileN, + kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Iteration 9: 4-wave / 8-accumulator variant of the exact-shape +// double-buffered 64x128 kernel (w8a8_gemm_prefill_tiled_kernel_w4, only +// instantiated as <64,128,64>). The macro-tile, K stage 64 double buffer +// (one barrier/stage), [N,K] packed B, bank-skewed LDS layout (30,720 +// B/block), element-to-slot fragment mapping, byte order, and the +// k0-outer/kk-inner int32 accumulation order are IDENTICAL to the accepted +// 8-wave kernel, so the result is bit-identical; the only deltas are the +// wave partition and the per-thread staging split: +// * 512 -> 256 threads (4 wavefronts); each wave owns a 32x64 quadrant +// (two 16-row bands x four 16-col bands) = EIGHT independent m16n16k32 +// accumulators (8-deep MMAC ILP per kk slice) instead of a 32x32 +// quadrant with four accumulators (4-deep ILP). +// * A staging is one int4/thread (256 vectors), B staging two int4/thread +// (512 vectors) at invariant staged bytes (12 KiB/stage/block). +// Per block per stage the LDS fragment-read instructions drop 32 -> 24 (A is +// read once per 64-col wave group instead of once per 32-col group) and the +// LDS write instructions drop 16 -> 12, at invariant total MMA count (64 +// v_mmac per block stage) and invariant bank-conflict floors (8-cycle 2-way +// reads, 16-cycle 4-way writes). Grid (N/128) x (M/64) = 3 x 64 = 192 blocks; +// 30,720 B x 2 = 61,440 B <= 64 KiB LDS and ~80-96 VGPR x 256 x 2 <= 512 KiB +// -> 2 blocks/CU -> 8 waves/CU. +// +// Iteration 11 (address-math consolidation): the steady-state K loop's +// address expressions are strength-reduced to loop-invariant bases computed +// once before the loop -- the three global staging pointers (A/B0/B1) walk +// +kStageK per stage, the three LDS commit slots add only the buffer offset +// nxt * kBufInt4, and the six fragment reads add only the buffer offset +// cur * kBufBytes and the unrolled kk slice. The exact-source ISA of the +// accepted kernel recomputes the full row*K product chains and slot +// constants every stage (~34 VALU/SALU per wave per stage at 0x4A0C-0x4AEC, +// 0x4ABC-0x4AE8, 0x4BFC-0x4C1C of the iteration-9 gfx928.co: v_mul_lo_u32, +// v_mad_u32_u24, v_add3_u32, v_lshlrev_b32, s_ashr_i32 chains whose operands +// are all stage-invariant); this is the concrete compiler limitation recorded +// for the ISA-guided gate. LDS contents, fragment operands, element-to-slot +// mapping, k0-outer/kk-inner int32 accumulation order, barrier count +// (65/block), grid (192 blocks), LDS (30,720 B), threads (256), and VGPR +// class are unchanged -> the result is bit-identical to iteration 9. +template +__global__ __launch_bounds__(4 * kWaveSize) void +w8a8_gemm_prefill_tiled_kernel_w4( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kThreads = 4 * kWaveSize; + constexpr int kAStride = kStageK + sizeof(int4); + constexpr int kBStride = kStageK + sizeof(int4); + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kStageK / sizeof(int4); + constexpr int kBVectors = kBlockN * kBVectorsPerRow; + static_assert(kAVectors == kThreads, + "w4 path stages exactly one A int4 per thread"); + static_assert(kBVectors == 2 * kThreads, + "w4 path stages exactly two B int4 per thread"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / 2; + const int wave_col = wave - wave_row * 2; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBStride]; + + DUFragment acc[2][4]; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 4; ++c) { + du_fill_fragment(acc[r][c], 0); + } + } + + // Prologue: stage K tile 0 into buffer 0 (same vectors and addresses as + // the 512-thread kernel, distributed over 256 threads), then one barrier. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + n_local * kBStride)[v] = + *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + const int n_stages = k / kStageK; + // Iteration 11 (address-math consolidation): hoist every loop-invariant + // address of the steady state out of the K loop. The three global staging + // pointers (A/B0/B1) are precomputed at the stage-1 offset and walk + // +kStageK per stage; the three LDS commit slots are precomputed in buffer + // 0 and add only the buffer offset nxt * kBufInt4; the six fragment reads + // use the per-wave buffer-0 bases and add only the buffer offset + // cur * kBufBytes and the unrolled kk slice. The bytes, the + // element-to-slot mapping, and the load/store order are identical to the + // accepted kernel, so the int32 accumulation is bit-identical. + const int local_row_a = tid / kAVectorsPerRow; + const int v_a = tid - local_row_a * kAVectorsPerRow; + const int global_row_a = m0 + local_row_a; + const bool a_row_valid = global_row_a < m; + const int8_t* __restrict__ a_g = + x_q + static_cast(global_row_a) * k + + v_a * static_cast(sizeof(int4)) + kStageK; + const int n_local0 = tid / kBVectorsPerRow; + const int v_b0 = tid - n_local0 * kBVectorsPerRow; + const int8_t* __restrict__ b_g0 = + weight + static_cast(n0 + n_local0) * k + + v_b0 * static_cast(sizeof(int4)) + kStageK; + const int vec1 = tid + kThreads; + const int n_local1 = vec1 / kBVectorsPerRow; + const int v_b1 = vec1 - n_local1 * kBVectorsPerRow; + const int8_t* __restrict__ b_g1 = + weight + static_cast(n0 + n_local1) * k + + v_b1 * static_cast(sizeof(int4)) + kStageK; + // LDS commit slots as int4 indexes into buffer 0 (the loop adds the + // buffer-1 offset nxt * kBufInt4; both buffers are 16-byte aligned so the + // int4 stores are exactly the accepted ds_write_b128 pattern). + int4* const a_wslot = + reinterpret_cast(a_tile[0] + local_row_a * kAStride) + v_a; + int4* const b_wslot0 = + reinterpret_cast(b_tile[0] + n_local0 * kBStride) + v_b0; + int4* const b_wslot1 = + reinterpret_cast(b_tile[0] + n_local1 * kBStride) + v_b1; + constexpr int kABufInt4 = + (kBlockM * kAStride) / static_cast(sizeof(int4)); + constexpr int kBBufInt4 = + (kBlockN * kBStride) / static_cast(sizeof(int4)); + // Per-wave fragment-read bases in buffer 0; the loop adds the buffer + // offset cur * kBufBytes and the kk slice (same addresses as + // a_tile[cur] + (wave_row * 32) * kAStride + kk, etc.). + const int8_t* const a_rd = a_tile[0] + (wave_row * 32) * kAStride; + const int8_t* const b_rd = b_tile[0] + (wave_col * 64) * kBStride; + constexpr int kABufBytes = kBlockM * kAStride; + constexpr int kBBufBytes = kBlockN * kBStride; + + for (int s = 0; s < n_stages; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + const bool has_next = (s + 1) < n_stages; + // Issue the stage s+1 global loads from the precomputed pointers (A: one + // int4/thread, B: two int4/thread); the vmcnt wait lands after the MMA + // burst, at the commit. + int4 a_reg{0, 0, 0, 0}; + int4 b_reg0{0, 0, 0, 0}; + int4 b_reg1{0, 0, 0, 0}; + if (has_next) { + a_reg = a_row_valid ? *reinterpret_cast(a_g) + : int4{0, 0, 0, 0}; + b_reg0 = *reinterpret_cast(b_g0); + b_reg1 = *reinterpret_cast(b_g1); + a_g += kStageK; + b_g0 += kStageK; + b_g1 += kStageK; + } + + // Compute stage s: each wave owns a 32x64 quadrant (two 16-row bands x + // four 16-col bands), eight independent m16n16k32 accumulators. The + // k0-outer/kk-inner order per output element is unchanged, so the int32 + // accumulation is bit-identical to the accepted kernel. + // + // Iteration 13 (burst-scheduling probe): the twelve fragment loads of + // the stage (six per kk slice) are all issued here, ahead of the whole + // sixteen-MMAC burst, in the same address order as the accepted kernel. + // This pins the second slice's ds_reads above the first slice's MMAC + // burst (both slices' fragments live across the burst); the MMAC + // emission order (kk=0 group, then kk=kTileK group, per accumulator) is + // unchanged, so the result is bit-identical. + static_assert(kStageK == 2 * kTileK, + "w4 burst schedule expects exactly two kk slices per stage"); + DUFragment + a_frag00, a_frag01, a_frag10, a_frag11; + DUFragment + b_frag00, b_frag01, b_frag02, b_frag03, b_frag10, b_frag11, b_frag12, + b_frag13; + const int8_t* const a_cur = a_rd + cur * kABufBytes; + const int8_t* const b_cur = b_rd + cur * kBBufBytes; + // kk = 0 slice loads (row bands x0/x1, col bands 0..3). + load_fragment8(a_frag00, a_cur, kAStride, lane); + load_fragment8(a_frag01, a_cur + kTileM * kAStride, kAStride, lane); + load_fragment8(b_frag00, b_cur, kBStride, lane); + load_fragment8( + b_frag01, b_cur + kTileN * kBStride, kBStride, lane); + load_fragment8( + b_frag02, b_cur + 2 * kTileN * kBStride, kBStride, lane); + load_fragment8( + b_frag03, b_cur + 3 * kTileN * kBStride, kBStride, lane); + // kk = kTileK slice loads (same rows/cols, k offset kTileK). + load_fragment8(a_frag10, a_cur + kTileK, kAStride, lane); + load_fragment8( + a_frag11, a_cur + kTileM * kAStride + kTileK, kAStride, lane); + load_fragment8(b_frag10, b_cur + kTileK, kBStride, lane); + load_fragment8( + b_frag11, b_cur + kTileN * kBStride + kTileK, kBStride, lane); + load_fragment8( + b_frag12, b_cur + 2 * kTileN * kBStride + kTileK, kBStride, lane); + load_fragment8( + b_frag13, b_cur + 3 * kTileN * kBStride + kTileK, kBStride, lane); + // kk = 0 MMAC group. + du_mma_sync(acc[0][0], a_frag00, b_frag00, acc[0][0]); + du_mma_sync(acc[0][1], a_frag00, b_frag01, acc[0][1]); + du_mma_sync(acc[0][2], a_frag00, b_frag02, acc[0][2]); + du_mma_sync(acc[0][3], a_frag00, b_frag03, acc[0][3]); + du_mma_sync(acc[1][0], a_frag01, b_frag00, acc[1][0]); + du_mma_sync(acc[1][1], a_frag01, b_frag01, acc[1][1]); + du_mma_sync(acc[1][2], a_frag01, b_frag02, acc[1][2]); + du_mma_sync(acc[1][3], a_frag01, b_frag03, acc[1][3]); + // kk = kTileK MMAC group (same accumulators, same order). + du_mma_sync(acc[0][0], a_frag10, b_frag10, acc[0][0]); + du_mma_sync(acc[0][1], a_frag10, b_frag11, acc[0][1]); + du_mma_sync(acc[0][2], a_frag10, b_frag12, acc[0][2]); + du_mma_sync(acc[0][3], a_frag10, b_frag13, acc[0][3]); + du_mma_sync(acc[1][0], a_frag11, b_frag10, acc[1][0]); + du_mma_sync(acc[1][1], a_frag11, b_frag11, acc[1][1]); + du_mma_sync(acc[1][2], a_frag11, b_frag12, acc[1][2]); + du_mma_sync(acc[1][3], a_frag11, b_frag13, acc[1][3]); + + // Commit the prefetched stage s+1 into the other LDS buffer, then one + // barrier (same ordering semantics as the accepted kernel). + if (has_next) { + *(a_wslot + nxt * kABufInt4) = a_reg; + *(b_wslot0 + nxt * kBBufInt4) = b_reg0; + *(b_wslot1 + nxt * kBBufInt4) = b_reg1; + } + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 64; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 4; ++c) { + store_prefill_fragment( + acc[r][c], x_scale, weight_scale, out, m, n, + base_row + r * kTileM, base_col + c * kTileN, lane); + } + } +} + +// Iteration 18: 8-wave / 4-accumulator variant of the exact-shape +// double-buffered 64x128 kernel (w8a8_gemm_prefill_tiled_kernel_w8, only +// instantiated as <64,128,64>). The macro-tile, K stage 64 double buffer +// (one barrier/stage), [N,K] packed B, bank-skewed LDS layout (30,720 +// B/block), staging bytes, element-to-slot fragment mapping, byte order, +// and the k0-outer/kk-inner int32 accumulation order are IDENTICAL to the +// accepted w4 kernel, so the result is bit-identical; the only deltas are +// the wave partition and the per-thread staging split: +// * 256 -> 512 threads (8 wavefronts); each wave owns a 32x32 quadrant +// (two 16-row bands x two 16-col bands) = FOUR independent m16n16k32 +// accumulators (4-deep MMAC ILP per kk slice) instead of a 32x64 +// quadrant with eight accumulators. Per-wavefront MMACs per stage drop +// 16 -> 8 (1,024 -> 512 MMACs per wavefront over the K loop). +// * A staging is one int4/thread for the first 256 threads, B staging one +// int4/thread for all 512 threads (same 12,288 B/stage/block as w4). +// Grid (N/128) x (M/64) = 192 blocks; LDS 30,720 B x 2 = 61,440 B <= 64 KiB; +// target <= 64 VGPR x 512 x 2 <= 512 KiB -> 2 blocks/CU -> 4 waves/SIMD +// (must be confirmed in the code object; > 64 VGPR degenerates the probe to +// 1 block/CU = 2 waves/SIMD). Same address hoist (iteration 11) and +// burst-pinned fragment schedule (iteration 13) as the accepted w4. +template +__global__ __launch_bounds__(8 * kWaveSize) void +w8a8_gemm_prefill_tiled_kernel_w8( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kThreads = 8 * kWaveSize; + constexpr int kAStride = kStageK + sizeof(int4); + constexpr int kBStride = kStageK + sizeof(int4); + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kStageK / sizeof(int4); + constexpr int kBVectors = kBlockN * kBVectorsPerRow; + static_assert(kAVectors <= kThreads, + "w8 path stages A with at most one int4 per thread"); + static_assert(kBVectors == kThreads, + "w8 path stages exactly one B int4 per thread"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / 4; + const int wave_col = wave - wave_row * 4; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBStride]; + + DUFragment acc[2][2]; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 2; ++c) { + du_fill_fragment(acc[r][c], 0); + } + } + + // Prologue: stage K tile 0 into buffer 0 (A: first 256 threads, B: all 512 + // threads; same vectors and addresses as the w4 kernel distributed over + // 512 threads), then one barrier. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + n_local * kBStride)[v] = + *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + const int n_stages = k / kStageK; + // Iteration 11 (address-math consolidation), same as w4: hoist every + // loop-invariant address of the steady state out of the K loop. The A + // staging pointer is only dereferenced by threads tid < kAVectors. + const bool a_active = tid < kAVectors; + const int local_row_a = tid / kAVectorsPerRow; + const int v_a = tid - local_row_a * kAVectorsPerRow; + const int global_row_a = m0 + local_row_a; + const bool a_row_valid = global_row_a < m; + const int8_t* __restrict__ a_g = + x_q + static_cast(global_row_a) * k + + v_a * static_cast(sizeof(int4)) + kStageK; + const int n_local = tid / kBVectorsPerRow; + const int v_b = tid - n_local * kBVectorsPerRow; + const int8_t* __restrict__ b_g = + weight + static_cast(n0 + n_local) * k + + v_b * static_cast(sizeof(int4)) + kStageK; + int4* const a_wslot = + reinterpret_cast(a_tile[0] + local_row_a * kAStride) + v_a; + int4* const b_wslot = + reinterpret_cast(b_tile[0] + n_local * kBStride) + v_b; + constexpr int kABufInt4 = + (kBlockM * kAStride) / static_cast(sizeof(int4)); + constexpr int kBBufInt4 = + (kBlockN * kBStride) / static_cast(sizeof(int4)); + // Per-wave fragment-read bases in buffer 0 (32x32 quadrant per wave); the + // loop adds the buffer offset cur * kBufBytes and the kk slice. + const int8_t* const a_rd = a_tile[0] + (wave_row * 32) * kAStride; + const int8_t* const b_rd = b_tile[0] + (wave_col * 32) * kBStride; + constexpr int kABufBytes = kBlockM * kAStride; + constexpr int kBBufBytes = kBlockN * kBStride; + + for (int s = 0; s < n_stages; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + const bool has_next = (s + 1) < n_stages; + // Issue the stage s+1 global loads from the precomputed pointers (A: one + // int4/thread for tid < kAVectors, B: one int4/thread); the vmcnt wait + // lands after the MMA burst, at the commit. + int4 a_reg{0, 0, 0, 0}; + int4 b_reg{0, 0, 0, 0}; + if (has_next) { + if (a_active) { + a_reg = a_row_valid ? *reinterpret_cast(a_g) + : int4{0, 0, 0, 0}; + } + b_reg = *reinterpret_cast(b_g); + a_g += kStageK; + b_g += kStageK; + } + + // Compute stage s: each wave owns a 32x32 quadrant (two 16-row bands x + // two 16-col bands), four independent m16n16k32 accumulators. The + // k0-outer/kk-inner order per output element is unchanged, so the int32 + // accumulation is bit-identical to the accepted w4 kernel. + // + // Iteration 13 (burst-scheduling probe), same as w4: the eight fragment + // loads of the stage (four per kk slice) are all issued here, ahead of + // the whole eight-MMAC burst, in the same address order as the accepted + // w4 kernel. + static_assert(kStageK == 2 * kTileK, + "w8 burst schedule expects exactly two kk slices per stage"); + DUFragment + a_frag00, a_frag01, a_frag10, a_frag11; + DUFragment + b_frag00, b_frag01, b_frag10, b_frag11; + const int8_t* const a_cur = a_rd + cur * kABufBytes; + const int8_t* const b_cur = b_rd + cur * kBBufBytes; + // kk = 0 slice loads (row bands x0/x1, col bands 0..1). + load_fragment8(a_frag00, a_cur, kAStride, lane); + load_fragment8(a_frag01, a_cur + kTileM * kAStride, kAStride, lane); + load_fragment8(b_frag00, b_cur, kBStride, lane); + load_fragment8(b_frag01, b_cur + kTileN * kBStride, kBStride, lane); + // kk = kTileK slice loads (same rows/cols, k offset kTileK). + load_fragment8(a_frag10, a_cur + kTileK, kAStride, lane); + load_fragment8( + a_frag11, a_cur + kTileM * kAStride + kTileK, kAStride, lane); + load_fragment8(b_frag10, b_cur + kTileK, kBStride, lane); + load_fragment8( + b_frag11, b_cur + kTileN * kBStride + kTileK, kBStride, lane); + // kk = 0 MMAC group. + du_mma_sync(acc[0][0], a_frag00, b_frag00, acc[0][0]); + du_mma_sync(acc[0][1], a_frag00, b_frag01, acc[0][1]); + du_mma_sync(acc[1][0], a_frag01, b_frag00, acc[1][0]); + du_mma_sync(acc[1][1], a_frag01, b_frag01, acc[1][1]); + // kk = kTileK MMAC group (same accumulators, same order). + du_mma_sync(acc[0][0], a_frag10, b_frag10, acc[0][0]); + du_mma_sync(acc[0][1], a_frag10, b_frag11, acc[0][1]); + du_mma_sync(acc[1][0], a_frag11, b_frag10, acc[1][0]); + du_mma_sync(acc[1][1], a_frag11, b_frag11, acc[1][1]); + + // Commit the prefetched stage s+1 into the other LDS buffer, then one + // barrier (same ordering semantics as the accepted w4 kernel). + if (has_next) { + if (a_active) { + *(a_wslot + nxt * kABufInt4) = a_reg; + } + *(b_wslot + nxt * kBBufInt4) = b_reg; + } + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 2; ++c) { + store_prefill_fragment( + acc[r][c], x_scale, weight_scale, out, m, n, + base_row + r * kTileM, base_col + c * kTileN, lane); + } + } +} + +// Generic scalar fallback: one output element per grid-stride iteration with +// exact int32 accumulation. Correct for any (m, n, k), including the small-M +// API shapes (M=2/16) and any unmatched geometry. The exact (k, n) = +// (4096, 384) pair reads the transposed [N, K] packed weight (the pack +// specialization is keyed on (k, n), so every consumer of that packed pair, +// including the paired M=2/M=16 decode shapes that reach this fallback, must +// interpret it transposed). +__global__ __launch_bounds__(kScalarThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const bool packed_transposed = (k == 4096 && n == 384); + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + int32_t acc = 0; + const int8_t* a_row = x_q + static_cast(row) * k; + if (packed_transposed) { + // packed[n * k + kk] = raw[kk * n + n]. + const int8_t* b_col = weight + static_cast(col) * k; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[kk]); + } + } else { + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast( + weight[static_cast(kk) * n + col]); + } + } + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// pack_weight. The packed layout is opaque per the API contract; for the +// exact (K=4096, N=384) pair the weight is stored transposed [N, K] +// (n-major) so the double-buffered GEMM stages B in [N, K] LDS order and the +// DUMMA B fragments read 8 consecutive k-values per lane. Every other (k, n) +// pair keeps the identity [K, N] copy (generic fallback unchanged). +__global__ __launch_bounds__(kScalarThreads) void +w8a8_pack_kernel( + const int8_t* __restrict__ raw_weight, + const float* __restrict__ weight_scale, + int8_t* __restrict__ packed_weight, + float* __restrict__ packed_weight_scale, + int64_t weight_elems, + int n, + int k) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (k == 4096 && n == 384) { + // Transpose [K, N] -> [N, K]: packed[n * k + kk] = raw[kk * n + n]. + if (idx < weight_elems) { + const int64_t n_idx = idx / static_cast(k); + const int64_t k_idx = idx - n_idx * static_cast(k); + packed_weight[idx] = raw_weight[k_idx * n + n_idx]; + } + } else if (idx < weight_elems) { + packed_weight[idx] = raw_weight[idx]; + } + if (idx < n) { + packed_weight_scale[idx] = weight_scale[idx]; + } +} + +} // namespace + +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // No split-K: the provided workspace is not touched. The timed operator + // performs no allocation, synchronization, packing, or default-stream + // launch; only the caller-provided out is written. + (void)workspace; + (void)workspace_bytes; + auto* out_bf16 = static_cast(out); + + if (m >= kPrefillMinM && (n % kDefaultBlockN == 0) && + (k % kPrefillStageK == 0)) { + if (n == 384 && k == 4096) { + // Packed (k, n) = (4096, 384) weight is transposed [N, K], so every + // large-M shape with this pair uses the double-buffered transposed-B + // kernel (tail M rows zero-filled on load, masked on store). For the + // exact assigned shape (M=4096, N=384, K=4096): 64x128 macro-tile, + // grid 3x64 = 192 blocks, K stage 64 double-buffered, 30,720 B + // LDS/block -> 2 blocks/CU (see the iteration-18 branch below for the + // current 8-wave dispatch; the accepted 4-wave w4 path is the + // bit-identical 117.50 us best). + const dim3 grid(n / 128, (m + 63) / 64); + if (m == 4096) { + // Iteration 18 (wave-count probe, ISA-guided round 2 of 2): the + // exact shape dispatches to the 8-wave variant w8<64,128,64> (512 + // threads, 32x32 quadrant per wave, 4 accumulators, 512 + // MMACs/wavefront) at the same tile, grid (192 blocks), LDS (30,720 + // B), and barrier count (65/block) as the accepted w4 path. + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel_w8<64, 128, 64>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + // Other large-M shapes with the packed (k, n) pair keep the accepted + // 8-wave <64,128,64,2> kernel (bit-identical path, 169.46 us median + // for M=4096). + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel<64, 128, 64, 2>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (n % 128 == 0) { + // 64x128 tile, 8 waves. + const dim3 grid(n / 128, (m + 63) / 64); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel<64, 128, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (m % 128 == 0) { + // 128x64 tile, 8 waves. + const dim3 grid(n / 64, m / 128); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel<128, 64, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + { + // 64x64 tile, 4 waves (generic default). + const dim3 grid(n / 64, (m + 63) / 64); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel<64, 64, 128>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + } + + const int64_t total = static_cast(m) * n; + const unsigned blocks = static_cast( + (total + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Exact (K, N) = (4096, 384) is stored transposed [N, K]; every other pair + // keeps the identity [K, N] copy. Runs outside the timed/CUDA-Graph region. + const int64_t weight_elems = static_cast(k) * n; + const int64_t total = weight_elems > n ? weight_elems : n; + const unsigned blocks = static_cast( + (total + kScalarThreads - 1) / kScalarThreads); + hipLaunchKernelGGL( + w8a8_pack_kernel, + dim3(blocks), dim3(kScalarThreads), 0, stream, + raw_weight, weight_scale, packed_weight, packed_weight_scale, + weight_elems, n, k); +} +// @@end +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/o_proj.hip new file mode 100644 index 00000000..2acf6d0d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/o_proj.hip @@ -0,0 +1,756 @@ +// @@variant shape=minimax_tp4_o_proj_m4096 commit=ab9604c75a47252917e8d6d9c8dc7a35266de6ac added=2026-08-24 +// median_us=765.9 p90_us=766.8 speedup=38.78 baseline_us=2.97e+04 +// source=minimaxm3-dsh-tp4-m4096-1-0c2f84a9 +// @@variant shape=minimax_tp4_o_proj_m4096 added=2026-08-23 iteration=1 +// INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. +// +// Worker: worker_2 (physical GPU 2), assigned shape: +// minimax_tp4_o_proj_m4096 : M=4096, N=6144, K=2048 +// +// Iteration 20 - swizzled 8-byte-group LDS layout (bank-conflict cut). +// * The accepted kernel's PMC (source digest 33c82f15, the official best +// 789.11 us) shows lds_bank_conflicts 28,704,768 vs lds_instructions +// 4,531,200 (~6.33 conflict cycles per fragment ds_read2 - the largest +// remaining structural inefficiency in the accepted architecture; the +// repaired iteration-8 128x32 class is closed on speed at iterations +// 9/10/11/15/19 ~1079-1139 us, occupancy-4 falsified, so the iteration +// 19 redirect names exactly this axis: a swizzled 64-byte B LDS layout +// with the fragment loader reading the swizzle back). +// * Mechanism: both LDS tiles use the padded 80-byte row stride (20 dwords); +// bank(row, col) = (20*row + col/4) mod 32 has period 8 over the 16 rows +// of every fragment, so lanes r and r+8 (plus the cross-k-chunk overlaps) +// hit the same bank - 4 lanes per bank on every 8-byte fragment read. +// Any 16-byte-aligned pad stride (dword count divisible by 4) keeps the +// same period-8 collision, so padding alone cannot fix it. This round +// re-lays BOTH LDS tiles at dense 64-byte rows (no pad) with the +// bijective 8-byte-group map +// group(row, c) = ((row & 15) + 4*(c & 3)) & 15 + 16*c + 128*(row >> 4) +// (row = tile row, c = k 8-byte group). For any fragment load (16 rows, +// fixed k phase) (row & 15) -> ((row & 15) + 4*(c & 3)) & 15 is a +// permutation, so every 16-lane group reads 16 distinct LDS bank-pairs +// and the old collision pair (r, r+8) lands 8 bank-pairs apart; the +// worst-case lane grouping is 2-way instead of 4-way. The loader and +// both stage writers use the same map (the writers emit two 8-byte LDS +// stores per 16-byte global chunk at group(row, 2j) / group(row, 2j+1)). +// Fragment values, k-byte order inside each 8-byte group, MMAC sequence, +// int32 accumulation order (k0-outer 64, kk-inner 32), barrier count +// (2/stage x 32), grid dim3(96,32) = 3072, element-to-lane mapping and +// the fused scale->bf16 epilogue are bit-identical. LDS drops from +// 16,128 B to 13,056 B/block (a_tile 8,192 + b_tile 4,096 + scales 768); +// the iteration-6 bf16 staging buffer (128 rows x 72 B = 9,216 B) now +// spills into the dead b_tile region [8,192, 9,216) - b_tile is never +// read after the k-loop's final barrier, so the epilogue is unchanged. +// VGPR stays ~72 (a few extra swizzle ALU), occupancy stays 3 blocks/CU +// (VGPR-limited), 4 blocks/CU is still impossible (4 x 256 x 73 > 65,536). +// * Falsifiable prediction: local offline hipcc -O3 --offload-arch=gfx928 +// build reports group_segment_fixed_size 13,056 B, vgpr_count <= 85, +// 0 spills, wavefront 64, max_flat_workgroup_size 256; the profiled +// launch stays 3 blocks/CU and lds_bank_conflicts drops from 28,704,768 +// toward the ~2-way floor (~9-10M) at unchanged lds_instructions. If the +// LDS read bank conflicts are a real per-stage serialization, median +// drops below the official best 789.11 us (>130.63 TOPS, p90 below +// 791.05 us) with correctness passing; if median stays >= 789.11 us with +// conflicts near the floor, the accepted kernel is not LDS-read-bound +// and the redirect is the epilogue/global-load path or the barrier +// structure, NOT more LDS layout work. +// +// Iteration 7 - explicit 64-bit fragment loader (DUMMA issue grouping). +// * The trusted gfx928 ISA of the accepted iteration-6 kernel (compile +// cache key f364bbf5...) shows the steady-state 64-K k-loop body +// (two unrolled m16n16k32 kk steps) holding 16 v_mmac, 12 ds_read2_b32 +// fragment loads, 7 s_waitcnt lgkmcnt gates - and 72 v_and_b32 / +// v_or_b32_sdwa / v_or3_b32 identity byte-reassembly VALU ops +// (~62% of the body's issue slots; 0 shifts, so bit-preserving ORs). +// du_load_matrix_sync for int8 m16n16k32 (both matrix_a row_major and +// matrix_b col_major) loads eight CONSECUTIVE bytes per lane into +// x[0..7], and du_mma_sync reinterprets a.x/b.x as one 64-bit long in +// natural byte order; the backend materializes that 64-bit MMAC +// operand from the library's byte-wise loads by emitting the identity +// shuffle chains, one per fragment, each gated by lgkmcnt(1) waits +// before the consuming MMAC group (lds_wait_instructions 4.18M vs +// lds_instructions 6.89M; valu_active = 2x valu_instructions). +// * This round holds the accepted k-loop structure byte-for-byte +// (single-buffered 64-K stages, n-major [N,K] B col_major fragments, +// 128x64 tile, 256 threads / 4 waves, 8 int32 accumulators, two +// __syncthreads per stage, grid dim3(96,32) = 3072 blocks, 16,128 B +// LDS at 3 blocks/CU, no split-K, no raw asm) and replaces only the +// six per-kk du_load_matrix_sync calls with an explicit loader that +// reads each lane's eight bytes as ONE 64-bit LDS read (8-byte aligned: +// 80-byte row strides and kk in {0,32} keep every fragment address +// 8-byte aligned) and stores them into the fragment's x[0..7] - the +// exact value du_mma_sync reads. The per-lane address must reproduce +// du_load_matrix_sync exactly: row = lane & 15, col = (lane >> 4) << 3, +// eight consecutive bytes at p[row*ldm + col .. +7] (verified in +// /opt/dtk/include/du_mma.hpp); the repair-1 bug was an earlier loader +// that ignored the lane/ldm offset and gave every lane the same eight +// base bytes (near-total mismatch, mismatch_count 25,148,019). +// Fragment values, MMAC sequence, accumulation order and the fused +// epilogue are bit-identical; the only change is which instructions +// sit between the ds_read and each v_mmac (0 reassembly VALU, one +// collapsed wait stream). No new registers are introduced (the freed +// shuffle temps stay available to the scheduler), so VGPR stays 72 +// and occupancy stays 3 blocks/CU. +// * Falsifiable prediction: if the identity reassembly VALU + per-group +// lgkmcnt gates are a real issue-pipe cost, median drops below the +// current best 1031.996 us (>99.88 TOPS, p90 below 1046.02 us) with +// valu_instructions falling from 45.58M toward ~17-20M (72 -> ~0 +// shuffle ops per body per wave, 32 stages x 4 waves x 3072 blocks) +// and lds_wait_instructions falling from 4.18M, at unchanged +// lds_instructions (~6.89M), lds_bank_conflicts (~28.7M) and +// 3 blocks/CU; if the kernel is instead bound by the LDS load issue +// itself or the per-stage barrier/global re-stage latency, median +// stays >= 1031.996 us, falsifying the shuffle mechanism and +// redirecting the next round to register-split kk prefetch or +// stage-level pipelining rather than more loader surgery. +// +// Iteration 6 - fused epilogue: LDS-resident scales + coalesced store. +// * The iteration-5 k-loop (single-buffered 64-K stages, n-major B, +// 128x64 tile, 3 blocks/CU) is held byte-for-byte; only the epilogue +// changes. +// * x_scale[128] and weight_scale[64] are loaded once per block into LDS +// in the prologue (768 B; total LDS 16,128 B/block, still 3 blocks/CU). +// The trusted ISA of iteration 5 showed the old epilogue issued 40 +// global_load_dword scale loads per thread (491K wavefront-instructions +// kernel-wide, ~30% of all vmem_read) with s_waitcnt vmcnt(1)/vmcnt(0) +// serialization before every fragment's multiply chain, and 33 static +// scattered global_store_short_d16_hi 2-byte stores per thread (lane&15 +// selects the row, lane>>4 selects col%4 -> one 2-byte store per output +// element, 64 separate sectors per store instruction, 393K store +// instructions kernel-wide). +// * New epilogue: after the k-loop the dead a_tile (10,240 B) is reused as +// a two-phase bf16 staging buffer (128 rows x 72 B; 32 useful bf16 + 8 B +// pad per row; 72 is 8 B-aligned for vector reads and its 4-byte bank +// offset 18r is a full permutation over 16 consecutive rows, so the +// per-lane 2-byte fragment writes are bank-conflict-free). Each wave +// scales its eight int32 accumulator fragments (x_scale_l[row] * +// w_scale_l[col], fused, then __float2bfloat16) while writing bf16 into +// the staging buffer, and the block then stores each staged 128x32 +// half-tile with fully coalesced 8-byte int2 global stores (consecutive +// lanes cover consecutive 8 B of one output row; every 32 B sector is +// fully written). Global store instructions drop from 128 to 32 per +// block and scale loads from 160 to 3 wavefront-instructions per block; +// the per-element 64-bit v_mad_i64_i32 address arithmetic disappears. +// Bit-identical output is preserved: the int32 accumulation order, the +// fused fragment -> x_scale*weight_scale -> bf16 epilogue math, and the +// element-to-lane mapping are unchanged (only the store path is +// re-routed through LDS); LDS stays 16,128 B/block so occupancy remains +// 3 blocks/CU. The generic scalar fallback and the exact-shape guard +// are untouched. +// +// Iteration 5 - B n-major pack + col_major fragment loads (reassembly cut). +// * Large-prefill path (exact (m, n, k) == (4096, 6144, 2048)): +// native INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 +// output tile per block; 256 threads = 4 wavefronts of 64 lanes; each +// wave owns a 64x32 quadrant built from eight m16n16k32 fragments; the +// block cooperatively vector-loads A[128,64] and B[64,64] into 15 KiB +// of LDS (padded 80-byte row strides, five bank phases); 64-deep K +// single buffering with two block barriers per stage (no split-K, no +// raw asm, no speculative deep pipeline); fused +// dot * x_scale[m] * weight_scale[n] epilogue stored directly as bf16 +// from the accumulator fragments. +// The weight for this shape is packed once (outside the timed region) +// into the n-major [N, K] layout packed[n*K + k] = raw[k*N + n], so B +// stages are [64 n, 64 k] n-major LDS tiles: each m16n16k32 B +// fragment's eight bytes per lane are eight consecutive k bytes and +// the col_major loader emits plain 8-byte LDS reads instead of the +// identity-layout per-byte ds_read_u8 gather + v_and/v_or3 +// byte-reassembly VALU (the gate_up iter-13 / qkv iter-18 lesson: +// accepted M=4096 kernels cut fragment reassembly once B is n-major). +// A stays row-major identity [M, K] (activations cannot be repacked). +// Iteration 2's double-buffered 64-K staging (same identity B) was +// already falsified at 1362 us, so occupancy-2 deeper staging is off +// the table; this round holds the accepted single-buffered pipeline +// and its 3 blocks/CU and only re-lays B. +// * launch_pack_w8a8_weight packs (2048, 6144) n-major once; every other +// (k, n) keeps the bootstrap identity device-to-device copy; packing +// never happens inside the timed GEMM or the CUDA/HIP Graph region. +// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), +// including small-M API cases (M=2/M=16 with the same (K, N) still +// reach the fallback through the exact-shape guard): (k, n) == +// (2048, 6144) decodes the n-major packed layout, every other (K, N) +// pair keeps the identity [K, N] layout. +// +// Graph safety: gemm_out performs no allocation, compilation, autotuning, +// packing, host/device synchronization, or default-stream launch; it runs on +// the caller-provided PyTorch HIP stream and returns the caller's `out`. + +#include +#include +#include + +#include + +namespace { + +using namespace du::dumma; + +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 +constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kTargetM = 4096; +constexpr int kTargetN = 6144; +constexpr int kTargetK = 2048; +constexpr int kBlockM = 128; +constexpr int kBlockN = 64; +constexpr int kStageK = 64; +constexpr int kBlockThreads = 4 * kWaveSize; // 256 + +// Iteration 6 epilogue staging (reuses the dead A/B LDS tiles, 12,288 B after +// the iteration-20 dense swizzle): each staged half-tile is 128 rows x 72 B = +// 9,216 B <= 12,288 B, holding 32 useful bf16 (64 B) plus 8 B pad per row. +// The staging writes spill into the dead b_tile region [8,192, 9,216); b_tile +// is never read after the k-loop's final barrier. 72 is 8 B-aligned (vector +// reads) and its 4-byte-bank offset 18*r is a full permutation over 16 +// consecutive rows, so the per-lane 2-byte fragment staging writes are +// bank-conflict-free. +constexpr int kStageStride = 72; // staging row stride in bytes +constexpr int kStageHalfBf16 = kBlockN / 2; // 32 bf16 per staged row +constexpr int kStageChunks = (kBlockM * kStageHalfBf16) / 4; // 1024 x 8 B + +using bf16_t = hip_bfloat16; + +// --------------------------------------------------------------------------- +// Iteration 6 fused epilogue, stage side: scale one m16n16k32 accumulator +// fragment (verified gfx928 int8 ownership: lane & 15 selects the row, +// lane >> 4 selects col % 4, x[i] maps to columns col%4 + 4*i) with the +// LDS-resident per-row x_scale and per-column weight_scale, convert to bf16, +// and write into the reused a_tile staging buffer (72-byte rows, 32 useful +// bf16 per row). The scale multiply order is identical to iteration 5 +// (acc * x_scale[row] * weight_scale[col]), so the output is bit-identical. +// --------------------------------------------------------------------------- +template +__device__ __forceinline__ void stage_prefill_fragment( + const AccFragment& frag, + int8_t* __restrict__ staging, + const float* __restrict__ x_scale_l, + const float* __restrict__ w_scale_l, + int row0, + int col0, + int wcol0, + int m0, + int m, + int lane) { + const int row = row0 + (lane & 15); + if (m0 + row >= m) { + return; // tail-M masking: padded rows never stage + } + const int col_mod4 = lane >> 4; + const float xs = x_scale_l[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = col0 + col_mod4 + 4 * i; + const float scaled = + static_cast(frag.x[i]) * xs * w_scale_l[wcol0 + col]; + *reinterpret_cast(staging + row * kStageStride + 2 * col) = + __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// Iteration 6 fused epilogue, store side: cooperative fully-coalesced store +// of one staged 128x32 half-tile (8,192 B). Each thread writes four 8-byte +// int2 chunks (q = tid, tid+256, ...); consecutive threads cover consecutive +// 8 B of one output row, so every 32 B sector is fully written and no +// per-element 64-bit address arithmetic is needed. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void store_staged_half_tile( + const int8_t* __restrict__ staging, + bf16_t* __restrict__ out, + int m0, + int n0, + int col_base, + int m, + int n, + int tid) { +#pragma unroll + for (int q = tid; q < kStageChunks; q += kBlockThreads) { + const int row = q >> 3; // 8 chunks (8 B each) per 64 B staging row + const int sub = q & 7; + if (m0 + row >= m) { + continue; + } + *reinterpret_cast(out + static_cast(m0 + row) * n + n0 + + col_base + 4 * sub) = + *reinterpret_cast(staging + row * kStageStride + 8 * sub); + } +} + +// --------------------------------------------------------------------------- +// Iteration 20 swizzled 8-byte-group LDS layout. The accepted kernel's PMC +// shows 28,704,768 lds_bank_conflicts vs 4,531,200 lds_instructions (~6.33 +// conflict cycles per fragment ds_read2). The cause is the padded 80-byte row +// stride (20 dwords): bank(row, col) = (20*row + col/4) mod 32 has period 8 +// over the 16 rows of every fragment, so lanes r and r+8 (plus the +// cross-k-chunk overlaps) hit the same bank - 4 lanes per bank on every 8-byte +// fragment read. Any 16-byte-aligned pad stride keeps the same period-8 +// collision, so padding alone cannot fix it. This round re-lays both LDS +// tiles at dense 64-byte rows with the bijective 8-byte-group map +// group(row, c) = ((row & 15) + 4 * (c & 3)) & 15 + 16 * c + 128 * (row >> 4) +// where row is the tile row and c is the k 8-byte group. For any fragment +// load (16 rows, fixed c phase) (row & 15) -> ((row & 15) + 4 * (c & 3)) & 15 +// is a permutation, so every 16-lane group reads 16 distinct LDS bank-pairs +// and the old collision pair (r, r+8) lands 8 bank-pairs apart; the worst +// case under any 16-lane grouping is 2-way instead of 4-way. The fragment +// loader and both stage writers must use the same map ("reading the swizzle +// back"): each 16-byte global chunk is written as two 8-byte LDS stores at +// group(row, 2j) and group(row, 2j+1). Fragment values, k-byte order within +// each 8-byte group, MMAC sequence, accumulation order and the fused +// epilogue are bit-identical. Iteration 7's explicit 64-bit loader (instead +// of du_load_matrix_sync's per-fragment identity v_and/v_or3 reassembly +// chains) is kept; only the address map changes. +// --------------------------------------------------------------------------- +__device__ __forceinline__ int swizzle_group(int row, int c) { + return (((row & 15) + 4 * (c & 3)) & 15) + 16 * c + 128 * (row >> 4); +} + +// Per-lane (row, k-group) position reproduces du_load_matrix_sync exactly: +// row = row0 + (lane & 15), c = (lane >> 4) + (kk >> 3), eight consecutive k +// bytes in natural byte order (row0 is always a multiple of 16, so +// row & 15 == lane & 15 and row >> 4 == row0 >> 4). +template +__device__ __forceinline__ void load_frag8_swizzled( + Frag& f, const int8_t* __restrict__ p, int lane, int row0, int kk) { + const int row = row0 + (lane & 15); + const int c = (lane >> 4) + (kk >> 3); + *reinterpret_cast(f.x) = + *reinterpret_cast(p + swizzle_group(row, c) * 8); +} + +__device__ __forceinline__ void store_chunk16_swizzled( + int8_t* __restrict__ p, int row, int j, const int4& v) { + *reinterpret_cast(p + swizzle_group(row, 2 * j) * 8) = + *reinterpret_cast(&v); + *reinterpret_cast(p + swizzle_group(row, 2 * j + 1) * 8) = + *reinterpret_cast( + reinterpret_cast(&v) + 8); +} + +// --------------------------------------------------------------------------- +// Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. +// Each wave owns a 64x32 quadrant (eight m16n16k32 int32 accumulators); the +// block cooperatively stages A[128,64] (identity [M, K] row-major) and +// B[64,64] (one-time n-major [N, K] packed layout, n-major LDS tile) into +// single-buffered LDS in 64-deep K stages. B fragments are loaded col_major: +// each lane's eight bytes are eight consecutive k bytes, so the loader emits +// plain 8-byte LDS reads with no per-byte gather/reassembly VALU. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kBlockThreads) void +w8a8_dumma_prefill_128x64x64_identity_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row * 64; + const int local_col = wave_col * 32; + + // Iteration 20: dense 64-byte-row tiles (8,192 + 4,096 B) with the + // swizzled 8-byte-group map; the fragment loader reads the swizzle back. + __shared__ __align__(16) int8_t a_tile[kBlockM * kStageK]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kStageK]; + // Iteration 6: LDS-resident per-row / per-column scales (768 B total; the + // kernel is at 13,056 B LDS per block = 3 blocks/CU, VGPR-limited). Loaded + // once in the prologue, read only by the fused epilogue. + __shared__ __align__(16) float x_scale_l[kBlockM]; + __shared__ __align__(16) float w_scale_l[kBlockN]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + + // Prologue: cooperative vectorized staging of the first 64-K stage. Each + // thread owns one int4 in A rows [0,64) and one in A rows [64,128), plus + // one int4 of B[64,64] (consecutive threads cover consecutive 16-byte + // chunks of each row, so every global load is coalesced). + // Iteration 6: the per-row/per-column scales are also staged into LDS here + // (threads 0..191, one float each); the first __syncthreads below makes + // them visible to the fused epilogue. + if (tid < kBlockM) { + x_scale_l[tid] = x_scale[m0 + tid]; + } else if (tid < kBlockM + kBlockN) { + w_scale_l[tid - kBlockM] = weight_scale[n0 + (tid - kBlockM)]; + } + const int vector_byte_offset = tid * static_cast(sizeof(int4)); + const int stage_row = vector_byte_offset / kStageK; // 0..63 + const int stage_col = vector_byte_offset - stage_row * kStageK; // 0,16,32,48 + store_chunk16_swizzled( + a_tile, stage_row, stage_col >> 4, + *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + stage_col)); + store_chunk16_swizzled( + a_tile, stage_row + kBlockM / 2, stage_col >> 4, + *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + + stage_col)); + // B stage is an n-major [64 n, 64 k] tile of the one-time packed layout + // packed[n*K + k] = raw[k*N + n]: thread tid covers n row (tid>>2) and the + // 16-byte k chunk (tid&3)*16, so the global load is coalesced along k; the + // swizzled LDS store uses the same 8-byte-group map the fragment loader + // reads back. + const int b_n_row = tid >> 2; // 0..63 n rows of the 64x64 B stage + const int b_k_off = (tid & 3) * 16; // 0,16,32,48 within the 64-byte k row + store_chunk16_swizzled( + b_tile, b_n_row, b_k_off >> 4, + *reinterpret_cast( + weight + static_cast(n0 + b_n_row) * k + b_k_off)); + __syncthreads(); + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + // Iteration 7: explicit 64-bit fragment loads (bit-identical values, + // zero reassembly VALU) instead of the library's byte-wise loads; the + // per-lane (row, k-group) position reproduces du_load_matrix_sync. + // Iteration 20: the swizzled loader reads each lane's eight bytes from + // swizzle_group(row, c) instead of the padded 80-byte stride (same + // values, 16 distinct bank-pairs per 16-lane group). + load_frag8_swizzled(b_frag0, b_tile, lane, local_col, kk); + load_frag8_swizzled(b_frag1, b_tile, lane, local_col + kTileN, kk); + load_frag8_swizzled(a_frag0, a_tile, lane, local_row, kk); + load_frag8_swizzled(a_frag1, a_tile, lane, local_row + kTileM, kk); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + load_frag8_swizzled(a_frag0, a_tile, lane, local_row + 2 * kTileM, kk); + load_frag8_swizzled(a_frag1, a_tile, lane, local_row + 3 * kTileM, kk); + du_mma_sync(acc20, a_frag0, b_frag0, acc20); + du_mma_sync(acc21, a_frag0, b_frag1, acc21); + du_mma_sync(acc30, a_frag1, b_frag0, acc30); + du_mma_sync(acc31, a_frag1, b_frag1, acc31); + } + + // Single-buffered re-stage of the next 64-K slice: barrier (every thread + // finished reading the current stage), coalesced vector loads into the + // same buffers, then barrier (all loads visible before the next stage). + const int next_k = k0 + kStageK; + if (next_k < k) { + __syncthreads(); + store_chunk16_swizzled( + a_tile, stage_row, stage_col >> 4, + *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + next_k + + stage_col)); + store_chunk16_swizzled( + a_tile, stage_row + kBlockM / 2, stage_col >> 4, + *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + + next_k + stage_col)); + store_chunk16_swizzled( + b_tile, b_n_row, b_k_off >> 4, + *reinterpret_cast( + weight + static_cast(n0 + b_n_row) * k + next_k + + b_k_off)); + __syncthreads(); + } + } + + // ------------------------------------------------------------------------- + // Iteration 6 fused epilogue: the dead A/B LDS tiles (12,288 B after the + // iteration-20 dense swizzle) are reused as a two-phase bf16 staging buffer + // (128 rows x 72 B, 32 useful bf16 per row; the top 1,024 B spill into the + // dead b_tile region, never read after the k-loop's final barrier). + // Phase A: waves 0 and 2 (wave_col == 0) stage the left 128x32 half-tile + // (cols [0,32)); every thread then stores it coalesced (8-byte chunks). + // Phase B: waves 1 and 3 (wave_col == 1) stage the right half-tile + // (cols [32,64)) and it is stored the same way. LDS is 13,056 B/block so + // occupancy remains 3 blocks/CU (VGPR-limited); the k-loop pipeline is + // untouched. + // ------------------------------------------------------------------------- + __syncthreads(); // every thread is done reading a_tile / b_tile + + const int quad_row = wave_row * kBlockM / 2; // 0 or 64 (staging row base) + const int wcol0 = wave_col * kStageHalfBf16; // 0 or 32 (weight_scale base) + + if (wave_col == 0) { + // Phase A: left half-tile, cols [0,32). + stage_prefill_fragment(acc00, a_tile, x_scale_l, w_scale_l, quad_row, + 0, wcol0, m0, m, lane); + stage_prefill_fragment(acc10, a_tile, x_scale_l, w_scale_l, + quad_row + kTileM, 0, wcol0, m0, m, lane); + stage_prefill_fragment(acc20, a_tile, x_scale_l, w_scale_l, + quad_row + 2 * kTileM, 0, wcol0, m0, m, lane); + stage_prefill_fragment(acc30, a_tile, x_scale_l, w_scale_l, + quad_row + 3 * kTileM, 0, wcol0, m0, m, lane); + stage_prefill_fragment(acc01, a_tile, x_scale_l, w_scale_l, quad_row, + kTileN, wcol0, m0, m, lane); + stage_prefill_fragment(acc11, a_tile, x_scale_l, w_scale_l, + quad_row + kTileM, kTileN, wcol0, m0, m, lane); + stage_prefill_fragment(acc21, a_tile, x_scale_l, w_scale_l, + quad_row + 2 * kTileM, kTileN, wcol0, m0, m, lane); + stage_prefill_fragment(acc31, a_tile, x_scale_l, w_scale_l, + quad_row + 3 * kTileM, kTileN, wcol0, m0, m, lane); + } + __syncthreads(); // half-tile staged + store_staged_half_tile(a_tile, out, m0, n0, 0, m, n, tid); + __syncthreads(); // staging reads done before phase B overwrites a_tile + + if (wave_col == 1) { + // Phase B: right half-tile, cols [32,64). + stage_prefill_fragment(acc00, a_tile, x_scale_l, w_scale_l, quad_row, + 0, wcol0, m0, m, lane); + stage_prefill_fragment(acc10, a_tile, x_scale_l, w_scale_l, + quad_row + kTileM, 0, wcol0, m0, m, lane); + stage_prefill_fragment(acc20, a_tile, x_scale_l, w_scale_l, + quad_row + 2 * kTileM, 0, wcol0, m0, m, lane); + stage_prefill_fragment(acc30, a_tile, x_scale_l, w_scale_l, + quad_row + 3 * kTileM, 0, wcol0, m0, m, lane); + stage_prefill_fragment(acc01, a_tile, x_scale_l, w_scale_l, quad_row, + kTileN, wcol0, m0, m, lane); + stage_prefill_fragment(acc11, a_tile, x_scale_l, w_scale_l, + quad_row + kTileM, kTileN, wcol0, m0, m, lane); + stage_prefill_fragment(acc21, a_tile, x_scale_l, w_scale_l, + quad_row + 2 * kTileM, kTileN, wcol0, m0, m, lane); + stage_prefill_fragment(acc31, a_tile, x_scale_l, w_scale_l, + quad_row + 3 * kTileM, kTileN, wcol0, m0, m, lane); + } + __syncthreads(); + store_staged_half_tile(a_tile, out, m0, n0, kStageHalfBf16, m, n, tid); +} + +// --------------------------------------------------------------------------- +// Generic scalar fallback: one thread per output element. Exact int32 dot +// over K, then fused x_scale * weight_scale, then bf16 store. Handles every +// unmatched (m, n, k), including all small-M API cases. Reads the one-time +// n-major packed [N, K] layout when (k, n) == (2048, 6144), and the +// bootstrap identity [K, N] row-major layout for every other (k, n). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + int32_t acc = 0; + const int8_t* a_row = x_q + static_cast(row) * k; + const bool nmajor_b = (k == kTargetK && n == kTargetN); + const int8_t* b_col = + weight + (nmajor_b ? static_cast(col) * k : col); + const int64_t b_stride = nmajor_b ? 1 : n; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * b_stride]); + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Bootstrap weight packing for every (k, n) except the assigned shape: +// identity device-to-device copy (int4 vectors plus a scalar tail for byte +// counts not a multiple of 16). The assigned (2048, 6144) uses the one-time +// n-major transpose kernel below instead. Scales are always copied identity. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_identity_vec_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int64_t vec_count) { + const int64_t vec = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (vec < vec_count) { + reinterpret_cast(packed)[vec] = + reinterpret_cast(raw)[vec]; + } +} + +__global__ __launch_bounds__(256) void w8a8_pack_identity_tail_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int64_t start, + int64_t count) { + const int64_t idx = + start + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < count) { + packed[idx] = raw[idx]; + } +} + +// One-time n-major transpose for the assigned shape (k, n) == (2048, 6144): +// packed[n*K + k] = raw[k*N + n]. Each thread owns one 16x16 (k-chunk, n-chunk) +// block and writes its 16 transposed 16-byte rows: +// packed[(nc*16 + j2)*K + kc*16 + j1] = raw[(kc*16 + j1)*N + nc*16 + j2], +// for j1, j2 in [0, 16). This is the exact block transpose with full K*N +// coverage (the previous one-int4-per-block version left 15/16 of the packed +// rows untouched and copied each chunk untransposed). Consecutive threads read +// consecutive n chunks of one k row (coalesced 16-byte global reads); writes +// are strided by 16*K (one-time prep, never timed). Requires k % 16 == 0 and +// n % 16 == 0, true for the exact shape. +__global__ __launch_bounds__(256) void w8a8_pack_nmajor_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int n_chunks = n / 16; + const int k_chunks = k / 16; + const int total = n_chunks * k_chunks; + const int tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (tid >= total) { + return; + } + const int kc = tid / n_chunks; + const int nc = tid - kc * n_chunks; + const int64_t src_base = + static_cast(kc) * 16 * n + static_cast(nc) * 16; + const int64_t dst_base = + static_cast(nc) * 16 * k + static_cast(kc) * 16; +#pragma unroll + for (int j2 = 0; j2 < 16; ++j2) { + __align__(16) int8_t chunk[16]; +#pragma unroll + for (int j1 = 0; j1 < 16; ++j1) { + chunk[j1] = raw[src_base + static_cast(j1) * n + j2]; + } + *reinterpret_cast(packed + dst_base + static_cast(j2) * k) = + *reinterpret_cast(chunk); + } +} + +__global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( + const float* __restrict__ raw, + float* __restrict__ packed, + int count) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < count) { + packed[idx] = raw[idx]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Stable host launch symbols consumed by csrc/bindings.cpp. +// Both launchers are pure dispatch: no allocation, no packing, no +// synchronization, no default-stream launch; they run on the caller-provided +// PyTorch HIP stream and are CUDA/HIP-Graph safe. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; // no split-K: the GEMM does not use the workspace + (void)workspace_bytes; + auto* out_bf16 = reinterpret_cast(out); + + // Explicit dispatch. The assigned shape (M=4096, N=6144, K=2048) takes the + // n-major-packed-B DUMMA 128x64 path; every other (m, n, k) - including + // small-M API cases and the paired M=2/M=16 shapes with the same (K, N) - + // takes the scalar fallback, which decodes the n-major packed layout for + // (k, n) == (2048, 6144) and the identity layout otherwise. + if (m == kTargetM && n == kTargetN && k == kTargetK) { + const dim3 grid(kTargetN / kBlockN, kTargetM / kBlockM); + const dim3 block(kBlockThreads); + hipLaunchKernelGGL(w8a8_dumma_prefill_128x64x64_identity_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + constexpr int kBlock = 256; + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + (void)hipGetLastError(); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kBlock = 256; + + // The assigned shape (2048, 6144) is packed once into the n-major + // [N, K] layout (packed[n*K + k] = raw[k*N + n]) so the DUMMA kernel can + // load col_major B fragments as plain 8-byte LDS reads. Every other (k, n) + // keeps the bootstrap identity [K, N] device-to-device copy. + if (k == kTargetK && n == kTargetN) { + const int64_t chunks = (static_cast(k) / 16) * + (static_cast(n) / 16); + const dim3 nmajor_grid( + static_cast((chunks + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_pack_nmajor_kernel, + nmajor_grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else { + // Bootstrap identity weight copy (int4 vectors plus scalar tail). + const int64_t weight_count = static_cast(k) * n; + const int64_t vec_count = weight_count / static_cast(sizeof(int4)); + if (vec_count > 0) { + const dim3 weight_grid( + static_cast((vec_count + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_pack_identity_vec_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, vec_count); + } + const int64_t copied_bytes = vec_count * static_cast(sizeof(int4)); + if (copied_bytes < weight_count) { + const int64_t tail_count = weight_count - copied_bytes; + const dim3 tail_grid( + static_cast((tail_count + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_pack_identity_tail_kernel, + tail_grid, block, 0, stream, + raw_weight, packed_weight, copied_bytes, weight_count); + } + } + + // Identity scale copy for every (k, n). + const dim3 scale_grid(static_cast((n + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, + scale_grid, block, 0, stream, + weight_scale, packed_weight_scale, n); + (void)hipGetLastError(); +} +// @@end +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj.hip new file mode 100644 index 00000000..81b98bb3 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj.hip @@ -0,0 +1,425 @@ +// @@variant shape=minimax_tp4_qkv_proj_m4096 commit=8e400dcec14217e9b03ccc7512c4d99149820b06 added=2026-08-24 +// median_us=1560 p90_us=1575 +// source=minimaxm3-dsh-tp4-m4096-1-0c2f84a9 +// MetaInfer W8A8 INT8 GEMM bootstrap for gfx928 (K500SM_AI). +// +// Worker: worker_0 (physical GPU 0) +// Assigned shape: minimax_tp4_qkv_proj_m4096 (M=4096, N=2304, K=6144) +// +// Operator contract (fixed by the control plane): +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// +// Bootstrap (iteration 1) strategy, kept unchanged by the iteration 2 repair: +// correctness-first but usable as a profiling baseline. A single native INT8 +// DUMMA m16n16k32 tiled kernel +// (64x64 output tile per block, 256 threads = 4 wavefronts, each wave owns a +// 32x32 quadrant = four m16n16k32 int32 accumulator fragments) handles every +// large-M shape with compatible geometry (M >= 128, N % 64 == 0, K % 128 == +// 0) -- this covers the assigned (4096, 2304, 6144): grid = (36, 64) = 2304 +// blocks against 120 CUs. K is staged cooperatively in a SINGLE LDS buffer at +// 128-K granularity (A[64,128] + B[128,64], padded row strides 144 B / 72 B +// to break the int8 fragment-load bank conflicts, 18,432 B/block) with two +// __syncthreads per stage; no split-K, no double buffering, no raw asm, no +// speculative pipelining. A is zero-filled past M and output rows are masked +// in the epilogue for M tails. The epilogue is the verified gfx928 direct +// fragment store (row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> column +// col_mod4 + 4*i) fused with float(dot) * x_scale[row] * weight_scale[col] +// and bf16 conversion. +// +// launch_pack_w8a8_weight is the bootstrap identity device-to-device copy: +// the packed weight is the logical [K, N] row-major int8 weight and the +// packed scales are the logical [N, 1] fp32 scales, valid for every (K, N). +// Later Parallel explore rounds may change only the HIP implementation and +// the matching GEMM interpretation. +// +// Every unmatched (m, n, k) -- including the paired small-M API shapes +// (M=2 / M=16) -- reaches the generic scalar int8/int32 fallback, which +// reads the same logical [K, N] row-major layout. +// +// The timed operator (launch_w8a8_gemm) performs no allocation, compilation, +// autotuning, packing, host/device synchronization, or default-stream launch: +// it only dispatches kernels on the caller-provided HIP stream. +// +// Include order is fixed by the control plane: hip_runtime, hip_bfloat16, +// then du_mma (this DTK's du_mma.h is not self-contained otherwise). + +#include +#include +#include + +#include + +namespace { + +constexpr int kWaveSize = 64; +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; +constexpr int kBlockM = 64; +constexpr int kBlockN = 64; +constexpr int kStageK = 128; +// Padded LDS row strides (bytes) for the int8 fragment loads. With the +// natural strides (A 128, B 64) every du_load_matrix_sync ds_read is +// 8-way (A) / 16-way (B) bank-conflicted: A lane l reads (l&15)*ldm + +// 8*(l>>4) and the row term 32*(l&15) mod 64 collapses to {0,32} (8 lanes +// per bank); B lane l reads 8*(l>>4)*ldm + (l&15) and the k-chunk term +// 8*64 = 512 B is 0 mod 256 B, stacking all four (l>>4) k-groups on the +// same 16 banks. Padding A rows to 144 B (36 dwords) makes 36*(l&15) mod +// 64 take 16 distinct dword-banks; padding B rows to 72 B moves the four +// k-groups 16 banks apart (16*(l>>4)), leaving at most a 4-way residue on +// the byte-granular (l&15)+i diagonal. Pure layout remap: every staged +// value and every fragment element is unchanged, so the int32 accumulation +// order and results stay bit-identical. LDS grows 16,384 -> 18,432 B/block; +// residency drops 4 -> 3 blocks/CU (3 x 18,432 = 55,296 <= 65,536 B LDS, +// 3 x 56 x 256 = 43,008 <= 65,536 VGPR). +constexpr int kAStride = 144; +constexpr int kBStride = 72; +// 4 wavefronts; must remain a multiple of the gfx928 wavefront size (64). +constexpr int kThreadsPerBlock = 256; + +using namespace du::dumma; + +// Direct accumulator epilogue for gfx928 INT8 m16n16k32 (verified against +// du_store_matrix_sync): row = lane & 15, col_mod4 = lane >> 4, +// frag.x[i] maps to columns col_mod4 + 4*i. Stores +// bf16(float(dot) * x_scale[row] * weight_scale[col]); out-of-range rows are +// masked (tail-M handling). +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Large-M prefill kernel: 64x64 output tile per block, K staged in LDS +// (single buffer at 128-K granularity), four waves each owning a 32x32 +// quadrant = four m16n16k32 int8->int32 DUMMA accumulators. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[64,128] (row-major, zero-filled past M) into the padded + // 144-B-row layout (row stride 144, still 16-B vector stores: 144/16 + // = 9 int4 slots per row). + constexpr int kAVectors = kBlockM * kStageK / static_cast(sizeof(int4)); +#pragma unroll(2) + for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int local_row = byte_offset / kStageK; + const int kk = byte_offset - local_row * kStageK; + const int global_row = m0 + local_row; + const int4 v = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + kk) + : int4{0, 0, 0, 0}; + reinterpret_cast(a_tile)[local_row * (kAStride / 16) + kk / 16] = + v; + } + + // Stage B[128,64] (row-major [K, N] logical weight layout) into the + // padded 72-B-row layout. 72 is not a multiple of 16, so each 16-B + // global vector is stored as two 8-B int2 halves (72*kk + col is + // always 8-B aligned). + constexpr int kBVectors = kStageK * kBlockN / static_cast(sizeof(int4)); +#pragma unroll(2) + for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + const int4 v = *reinterpret_cast( + weight + (k0 + kk) * n + n0 + col); + int2* dst = reinterpret_cast(b_tile + kk * kBStride + col); + dst[0] = *reinterpret_cast(&v); + dst[1] = *reinterpret_cast( + reinterpret_cast(&v) + 8); + } + __syncthreads(); + + // Consume the stage: each wave does four m16n16k32 MMAs per 32-K step. + // Four 32-K tiles per stage: pin the full unroll so the int32 + // accumulation order (k0-outer, kk-inner) is deterministic codegen, not + // compiler-heuristic dependent. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll(4) + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + // Single buffer: make sure every wave finished reading LDS before the + // next stage overwrites it. + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Generic scalar int8/int32 fallback for unmatched shapes and small-M API +// cases. One output element per grid-stride step; exact int32 accumulation +// (the assigned K keeps the int8 dot well within int32 range), then the fused +// float scale and bf16 store. Reads the logical [K, N] row-major weight +// layout (bootstrap identity pack). +__global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; + idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = a + static_cast(row) * k; + const int8_t* b_col = b + col; + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + const float scaled = static_cast(acc) * + x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// Identity device-to-device byte copy (pack_weight bootstrap, valid for any +// (K, N)). +__global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +// Identity device-to-device float copy for the weight scales. +__global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +} // namespace + +// Stable host launch symbol consumed by csrc/bindings.cpp. Dispatches on the +// caller-provided stream only; never allocates, synchronizes, or touches the +// default stream. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + auto* out_ptr = static_cast(out); + + // Native INT8 DUMMA m16n16k32 prefill path (2-D macro-tile 64x64, + // single-buffered 128-K K stage) for large-M shapes with compatible + // geometry. Covers the assigned shape (M=4096, N=2304, K=6144): + // 2304 % 64 == 0, 6144 % 128 == 0. The M >= 128 guard keeps the paired + // small-M API shapes (M=2 / M=16) on the scalar fallback. + if (m >= 128 && (n % kBlockN) == 0 && (k % kStageK) == 0) { + const dim3 grid( + static_cast(n / kBlockN), + static_cast((m + kBlockM - 1) / kBlockM)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x64_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Scalar int8/int32 fallback for unmatched (m, n, k) and small-M API cases. + const int64_t total = static_cast(m) * n; + constexpr int kFallbackThreads = 256; + int64_t blocks = (total + kFallbackThreads - 1) / kFallbackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + const dim3 grid(static_cast(blocks)); + const dim3 block(kFallbackThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); +} + +// Stable host launch symbol consumed by csrc/bindings.cpp. Bootstrap: +// identity device-to-device packing, valid for every (K, N) -- the packed +// weight keeps the logical [K, N] row-major layout and the packed scales are +// copied unchanged. Runs only from the optional pack_weight op, outside the +// timed region and outside Graph capture. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kPackThreads = 256; + const int64_t weight_count = static_cast(k) * n; + if (weight_count > 0) { + int64_t blocks = (weight_count + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + hipLaunchKernelGGL( + w8a8_pack_identity_i8_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count); + } + if (n > 0) { + int64_t blocks = (static_cast(n) + kPackThreads - 1) / + kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + hipLaunchKernelGGL( + w8a8_pack_identity_f32_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj_and_indexer_qk.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj_and_indexer_qk.hip new file mode 100644 index 00000000..eca651cf --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj_and_indexer_qk.hip @@ -0,0 +1,1601 @@ +// @@variant shape=minimax_tp4_qkv_proj_and_indexer_qk_m4096 commit=d6cbef3829c738228f5dab5255b89c384c66c6ef added=2026-08-24 +// median_us=1082 p90_us=1089 +// source=minimaxm3-dsh-tp4-m4096-1-0c2f84a9 +// MetaInfer W8A8 INT8 GEMM bootstrap for gfx928 (K500SM_AI). +// +// Worker: worker_1 (physical GPU 1) +// Assigned shape: minimax_tp4_qkv_proj_and_indexer_qk_m4096 +// (M=4096, N=2560, K=6144) +// +// Operator contract (fixed by the control plane): +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// The packed weight layout is opaque to the API: the exact assigned shape +// (K=6144, N=2560) is packed once by launch_pack_w8a8_weight outside the +// timed region into the iteration-5 fragment-interleaved swizzle (element +// (n, k) at ((k>>5)*160 + (n>>4))*512 + (((k>>3)&3)*16 + (n&15))*8 + (k&7); +// same byte count k*n as the identity copy, so allocations and graph-stable +// addresses are unchanged); every other (K, N) keeps the logical [K, N] +// row-major identity copy. +// +// Iteration 1 (DUMMA 2-D macro-tile baseline sweep, mandatory decision): +// * The bootstrap DUMMA 64x128 baseline (this file, prior state) measured +// 1442.6 us median / 89.3 TOPS on the assigned shape (M=4096, N=2560, +// K=6144) with PMC: 36.37M LDS instructions, 32% LDS bank conflicts, +// 40.3M LDS waits (1.11x LDS instructions, dominant stall class), 1.67M +// VMEM reads (1.51 GB through L2), 73.5% L2 hit, 2 blocks/CU. Bottleneck +// hypothesis: the row-major [K, N] B operand is staged into LDS but its +// m16n16k32 B fragments are re-assembled by the generic loader as EIGHT +// byte-granular ds_read_u8 per fragment (16-way bank aliasing with the +// 132-byte stride), which inflates the LDS instruction stream and its +// waits. The trusted hy3 M=4096 lineage (same M and N, K=4096; accepted +// 756.7 us / 113.5 TOPS) fixed exactly this by packing the weight +// n-major ONCE outside the timed region (packed[n][k]) and loading B +// fragments col_major as one 8-byte contiguous ds_read2_b64 per fragment +// (load_b_frag8), hoisting all twelve fragment loads of a 64-K stage +// before the first v_mmac (no lgkmcnt waits inside the 16-MMAC burst), +// and storing the epilogue as one coalesced 8-byte store per lane per +// fragment (100% store sector efficiency). This round ports that accepted +// lineage to the exact assigned shape (K=6144) and installs the complete +// 2-D macro-tile family (64x64, 64x128, 128x64; each 4 wavefronts x 64 +// lanes = 256 threads, one 32x32/32x64/64x32 quadrant per wave) so the +// mandated tile sweep can flip the exact-shape dispatch in the next +// rounds. Round 1 dispatches the exact shape to the 64x128 tile (the +// strongest N=2560 prior evidence: qkv_proj lineage 113.5 TOPS). +// * 64x128 packed-B kernel (iteration-1 layout; the iteration-5 swizzle +// replaces the B tile -- see the top-of-file pack note): A staged +// row-major into 68-byte LDS rows (17 words, odd-word bank spread, +// 4 ds_write_b32 per 16-byte vector), B staged n-major into 80-byte LDS +// rows (20 words, five bank phases, one ds_write_b128 per 16-byte +// vector), B fragments via load_b_frag8 (one ds_read2_b64), +// single-buffered 64-K K stage with two barriers per stage (the accepted +// iteration-3 pipeline), all twelve fragment loads hoisted before the +// 16-MMAC burst, coalesced epilogue. LDS/block = 64x68 + 128x80 = +// 14,592 B -> 3 resident blocks/CU (VGPR-bound at arch_vgpr <= 85, per +// the accepted lineage). Grid (N/128)x(M/64) = 20x64 = 1280 blocks >> +// 120 CUs, no split-K, workspace unused. +// * 64x64 and 128x64 packed-B kernels are compiled family members for the +// mandated tile sweep (grid 40x64 = 2560 and 40x32 = 1280 blocks; LDS +// 9,472 B and 15,360 B); they still consume the n-major 80-byte-row tile +// and are NOT dispatched for the exact shape, which uses 64x128 this +// round and every round since. +// * launch_pack_w8a8_weight packs exactly (k, n) == (6144, 2560) once, +// outside the timed region and Graph capture. Iteration 1 used the +// n-major transpose packed[n][k]; iteration 5 (packing round) replaces +// it with the fragment-interleaved swizzle described at the top of the +// file (same byte count k*n, so allocations and graph-stable addresses +// are unchanged); every other (k, n) keeps the identity copy. +// * The scalar int8/int32 fallback decodes the iteration-5 swizzled layout +// when (k, n) == (6144, 2560) (so the paired M=2 API shape with the same +// (N, K) stays correct) and the identity layout otherwise. +// * Iteration 6 (epilogue round): the timed operator is already a single +// fused kernel -- per-row x_scale, per-column weight_scale, bf16 +// conversion and the coalesced 8-byte-per-lane store all live inside +// w8a8_dumma_prefill_64x128_packedb_kernel, with no workspace/combine +// pass (workspace is unused; launch_w8a8_gemm dispatches exactly one +// kernel). The round restructures the dispatched kernel's epilogue from +// eight independent per-fragment calls into two fused 4-fragment +// row-group calls (store_prefill_rowgroup_coalesced + transpose_frag4 + +// pack_bf16x4): x_scale[row] and the four float4 weight_scale groups are +// loaded once per lane before the first ds_bpermute, the four transposes +// issue as one software-pipelined block (all xor-16 ds_bpermutes before +// xor-32), and the four 8-byte stores issue back-to-back (still 40,960 +// vmem_write instructions, 100% sector efficiency). The int32 transpose +// routing, per-element multiply order and bf16 rounding are byte- +// identical to iteration 5, so stored bits stay exact. The generic +// 64x128/64x64 kernels and the compiled-but-undispatched 64x64/128x64 +// packed-B family members keep the iteration-5 epilogue untouched. +// * Generic 64x128/64x64 DUMMA kernels (identity-packed B) are preserved +// unchanged for every other large-M shape. +// +// The timed operator (launch_w8a8_gemm) performs no allocation, compilation, +// autotuning, packing, host/device synchronization, or default-stream +// launch: it only dispatches kernels on the caller-provided HIP stream and +// uses only the caller-provided out/workspace tensors (workspace is unused +// in this no-split-K kernel). +// +// Include order is fixed by the control plane: hip_runtime, hip_bfloat16, +// then du_mma (this DTK's du_mma.h is not self-contained otherwise). + +#include +#include +#include + +#include + +namespace { + +constexpr int kWaveSize = 64; +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; +// 4 wavefronts; must remain a multiple of the gfx928 wavefront size (64). +constexpr int kThreadsPerBlock = 256; + +// 64x64 generic kernel geometry: K stage of 128 (unpadded rows). +constexpr int kBlockM = 64; +constexpr int kBlockN = 64; +constexpr int kStageK = 128; + +// 64x128 large-prefill kernel geometry: single-buffered 64-K K stage with +// odd-word padded row strides (A: 68 B = 17 words, B: 132 B = 33 words). +constexpr int kBlockM128 = 64; +constexpr int kBlockN128 = 128; +constexpr int kStageK128 = 64; +constexpr int kAStride128 = kStageK128 + 4; // 68 bytes per A row (17 words) +constexpr int kBStride128 = kBlockN128 + 4; // 132 bytes per B row (33 words) + +// Iteration 1 (packed-B family): the exact assigned shape packs the weight +// once (outside timing) and stages B into an LDS tile. The compiled 64x64 +// and 128x64 family members (not dispatched) still use the round-1 n-major +// layout: packed[n][k] staged into an n-major LDS tile with a 16-byte-aligned +// 80-byte row stride (20 words, five bank phases); every col_major +// m16n16k32 B fragment then reads its lane's 8 elements as one contiguous +// 8-byte chunk (one ds_read_b64 per fragment). 128x64 uses the same 80-byte +// stride for the A rows (16-byte-aligned staging stores, five bank phases). +// The dispatched 64x128 kernel switched to the iteration-5 fragment- +// interleaved swizzle (see kPackedBGroupBytes below) and no longer uses +// kPackedBStride. The 64x128/64x64 packed kernels keep the 68-byte A stride +// (4 ds_write_b32 per staging vector). +constexpr int kPackedBStride = 80; // 64 data + 16 pad bytes (20 words) +constexpr int kPackedAStride128x64 = 80; // 64 data + 16 pad (128x64 A tile) +constexpr int kPackedBK = 6144; // exact K of the packed assigned shape +constexpr int kPackedBN = 2560; // exact N of the packed assigned shape + +// Iteration 5 (packing round): the exact-shape pack is swizzled from the +// plain n-major transpose to a fragment-interleaved layout so the dispatched +// 64x128 packed-B kernel's B tile is vector-loadable AND LDS-bank-safe end +// to end. For each 32-k x 16-n DUMMA B fragment block the pack stores the 4 +// x 8-byte k-groups of the 16 n-rows n-interleaved ([k8][n][8], 128 B per +// k8, 512 B per fragment block, k32-major inside the 64-K stage). Every +// 16-byte staging vector is then 2 n-rows x 8 B of one k8 (int4 global load +// + one ds_write_b128 whose 8-lane LDS cycles each cover all 32 banks once), +// and every col_major B fragment load is one ds_read_b64 whose 16-lane group +// spans exactly one 128-B k8 block (all 32 banks exactly once per cycle -> +// conflict-free; the n-major 80-byte-row tile read the same bytes with 2-way +// conflicts). Same byte count k*n -> graph-stable packed layout. +constexpr int kPackedBGroupBytes = 512; // 16 n x 4 k8 x 8 B per 32-k fragment block +constexpr int kPackedB32Bytes = 8 * kPackedBGroupBytes; // 4096 B per 32-k block of the 64-K stage +constexpr int kPackedBN16 = 160; // 2560 / 16 n-groups per 32-k block (global pack stride) + +using namespace du::dumma; + +// Direct accumulator epilogue for gfx928 INT8 m16n16k32 (verified against +// du_store_matrix_sync): row = lane & 15, col_mod4 = lane >> 4, +// frag.x[i] maps to columns col_mod4 + 4*i. Stores +// bf16(float(dot) * x_scale[row] * weight_scale[col]); out-of-range rows are +// masked (tail-M handling). The scale multiply order +// float(dot) * x_scale[row] * weight_scale[col] matches the reference +// (dot.float() * a_scale * b_scale.T), so bf16 bits are exact. +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Coalesced fragment store (packed-B kernels). The m16n16k32 accumulator +// lane mapping (row = lane & 15, column group c4 = lane >> 4, frag.x[i] -> +// column c4 + 4*i) gives each lane four elements strided by 4 columns, so a +// direct store is four 2-byte scalar stores per lane touching every 32-B +// sector at 25% utilization. This epilogue transposes the 4-element groups +// within each 4-lane column group (lanes r, r+16, r+32, r+48; two 2x2 +// shuffle steps) so lane (r, c4) owns the four CONTIGUOUS columns +// 4*c4 .. 4*c4+3, converts to bf16, packs 4 bf16 (8 B), and issues ONE +// 8-byte store per lane per fragment (100% store sector efficiency; the +// assigned shape's vmem_write_instructions drop 163,840 -> 40,960). Only the +// int32 values are re-routed between lanes; the per-element float scale +// multiply order (float(dot) * x_scale[row] * weight_scale[col]) and bf16 +// rounding are unchanged, so the stored bits are identical. The row >= m +// guard is wavefront-uniform (lane & 15 cycles the same 16 rows in every +// 16-lane group) and base_col is a multiple of 4 (multiple of 32 here), so +// the float4 weight_scale load (col0 % 4 == 0) and the 8-byte store (n even) +// are aligned. This DTK lowers __shfl_xor to ds_bpermute (LDS permute at the +// block tail, where the LDS pipe is otherwise idle); there is no global +// round trip and no staging tile. +template +__device__ __forceinline__ void store_prefill_fragment_coalesced( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 4, n is even). + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + row * n + col0) = packed; +} + +// Iteration 6 (epilogue round): 4-element 4x4 lane-group transpose for one +// m16n16k32 accumulator fragment. Identical masks (xor 16, xor 32), select +// structure and lane mapping as store_prefill_fragment_coalesced (step 1 +// swaps the 2x2 in-block pairs, step 2 swaps the 2x2 blocks), so the four +// returned ints are the SAME per-lane contiguous-column values; only the +// surrounding scheduling is restructured. Extracted so the fused +// store_prefill_rowgroup_coalesced can issue the transposes of a whole +// 16-row x 64-col group as one software-pipelined block. +template +__device__ __forceinline__ void transpose_frag4( + const AccFragment& frag, + int lane, + int& f0, + int& f1, + int& f2, + int& f3) { + const int c4 = lane >> 4; + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + f0 = hi ? a0 : t2; + f1 = hi ? a1 : t3; + f2 = hi ? t0 : a2; + f3 = hi ? t1 : a3; +} + +// Scale x 4 contiguous bf16 and pack into one 8-byte store word. The +// per-element multiply order float(dot) * x_scale[row] * weight_scale[col] +// (left to right) and the bf16 rounding are identical to +// store_prefill_fragment_coalesced, so the stored bits are unchanged. +__device__ __forceinline__ uint64_t pack_bf16x4( + int a0, + int a1, + int a2, + int a3, + float xs, + float4 ws) { + const float v0 = static_cast(a0) * xs * ws.x; + const float v1 = static_cast(a1) * xs * ws.y; + const float v2 = static_cast(a2) * xs * ws.z; + const float v3 = static_cast(a3) * xs * ws.w; + return static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); +} + +// Iteration 6 (epilogue round): fused 4-fragment epilogue for ONE 16-row +// group of a wave (the four fragments cover the group's 64 contiguous +// columns). This is the dispatched 64x128 packed-B kernel's epilogue, +// restructured from eight independent per-fragment calls into two +// row-group calls so the tail becomes one explicitly-scheduled block: +// (1) x_scale[row] and the FOUR float4 weight_scale groups are loaded into +// registers ONCE per lane BEFORE the first ds_bpermute (one L1 round trip +// for the whole group instead of relying on cross-call CSE/load sinking), +// overlapping the scale-load latency with the transpose chain; (2) the four +// independent 4x4 transposes are issued as one software-pipelined block +// (step-1 xor-16 ds_bpermutes of all fragments precede step-2 xor-32, so +// the tail's dependent LDS-latency chain is ~2 round trips, not up to 8 +// serialized per-fragment chains); (3) the four 8-byte coalesced stores +// (offsets 0/32/64/96 B from obase) are issued back-to-back at the end, +// keeping 100% store-sector efficiency (same 8-byte stores per lane per +// fragment; vmem_write_instructions unchanged at 40,960). The int32 +// transpose routing, per-element float multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and bf16 rounding are +// byte-identical to the iteration-5 epilogue, so output bits are exact. +template +__device__ __forceinline__ void store_prefill_rowgroup_coalesced( + const F0& f0, + const F1& f1, + const F2& f2, + const F3& f3, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int c4 = lane >> 4; + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 w0 = *reinterpret_cast(weight_scale + col0); + const float4 w1 = *reinterpret_cast(weight_scale + col0 + 16); + const float4 w2 = *reinterpret_cast(weight_scale + col0 + 32); + const float4 w3 = *reinterpret_cast(weight_scale + col0 + 48); + hip_bfloat16* const obase = out + row * n + col0; + // Four independent transposes, issued as one pipelined block (the + // compiler sees all 32 ds_bpermute before the first store). + int u0, u1, u2, u3, v0, v1, v2, v3, q0, q1, q2, q3, r0, r1, r2, r3; + transpose_frag4(f0, lane, u0, u1, u2, u3); + transpose_frag4(f1, lane, v0, v1, v2, v3); + transpose_frag4(f2, lane, q0, q1, q2, q3); + transpose_frag4(f3, lane, r0, r1, r2, r3); + // Four 8-byte coalesced stores, back-to-back (16 rows x 32 B per store + // instruction; every 32-B sector covered by 4 lanes -> 100% efficiency). + // Fragment j's slot is 16 columns (32 B) further along the row: obase + + // 16*j elements. + *reinterpret_cast(obase) = pack_bf16x4(u0, u1, u2, u3, xs, w0); + *reinterpret_cast(obase + 16) = + pack_bf16x4(v0, v1, v2, v3, xs, w1); + *reinterpret_cast(obase + 32) = + pack_bf16x4(q0, q1, q2, q3, xs, w2); + *reinterpret_cast(obase + 48) = + pack_bf16x4(r0, r1, r2, r3, xs, w3); +} + +// Explicit 8-byte loader for the col_major m16n16k32 B fragment (accepted +// hy3 lineage, validated on this DTK). The col_major fragment mapping is +// n = lane & 15, k = 8*(lane >> 4) + i; with the n-major LDS tile at row +// stride 80 the lane's 8 elements are contiguous +// (p[row*ldm + col .. +7], 8-byte aligned because ldm=80 and col are +// multiples of 8), so this compiles to ONE ds_read2_b64 instead of 8 +// byte-granular ds_read_u8 + mask/OR reassembly. The byte placement is +// identical to du_load_matrix_sync, so the v_mmac operand +// registers receive the same values. +__device__ __forceinline__ void load_b_frag8( + DUFragment& f, + const int8_t* __restrict__ p, + int ldm) { + const unsigned row = static_cast(__lane_id()) & 0xfu; + const unsigned col = (static_cast(__lane_id()) >> 4) << 3; + const int64_t v = *reinterpret_cast(p + row * ldm + col); + *reinterpret_cast(&f.x[0]) = v; +} + +// Iteration 5 (packing round): 8-byte loader for the col_major B fragment on +// the swizzled exact-shape tile. The fragment block layout is +// [k8 0..3][n 0..15][8 B] with n = lane&15 and k8 = lane>>4, so lane (n, k8) +// reads p + k8*128 + n*8: one ds_read_b64 per fragment, and each 16-lane +// group of the wavefront spans exactly one 128-B k8 block, so every LDS +// cycle touches all 32 banks exactly once (conflict-free; the n-major +// 80-byte-row tile delivered the same bytes with 2-way conflicts). The byte +// placement is identical to du_load_matrix_sync / load_b_frag8, +// so the v_mmac operand registers and the int32 accumulation order are +// unchanged. +__device__ __forceinline__ void load_b_frag8_swz( + DUFragment& f, + const int8_t* __restrict__ p) { + const unsigned n = static_cast(__lane_id()) & 0xfu; + const unsigned k8 = static_cast(__lane_id()) >> 4; + const int64_t v = + *reinterpret_cast(p + (k8 << 7) + (n << 3)); + *reinterpret_cast(&f.x[0]) = v; +} + +// Large-M prefill kernel: 64x128 output tile per block, 256 threads +// (4 wavefronts), single-buffered 64-K K stage, identity-packed B. Each wave +// owns a 32x64 quadrant = eight m16n16k32 int8->int32 DUMMA accumulators. +// Serves large-M shapes with N % 128 == 0 except the exact assigned shape +// (which routes to w8a8_dumma_prefill_64x128_packedb_kernel above). +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x128_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM128; + const int n0 = static_cast(blockIdx.x) * kBlockN128; + + __shared__ __align__(16) int8_t a_tile[kBlockM128 * kAStride128]; + __shared__ __align__(16) int8_t b_tile[kStageK128 * kBStride128]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; + DUFragment + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); + + const int num_stages = k / kStageK128; + + // Cooperative staging mapping (fixed per thread, reused every stage): + // A[64,64] -> 256 int4 vectors; thread `tid` owns vector `tid` + // (row = tid/4, 16-B column group = (tid%4)*16), so each + // wavefront covers 16 rows x 64 B contiguous per row. + // B[64,128] -> 512 int4 vectors; thread `tid` owns vectors `tid` and + // `tid+256` (kk = tid/8, column group = (tid%8)*16), so + // each wavefront covers 8 kk rows x 128 B contiguous. + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + const int b_kk0 = tid >> 3; + const int b_col16 = (tid & 7) << 4; + + // VGPR payload for one full 64-K stage (3 x int4 = 12 VGPR), loaded just + // before it is committed (no one-stage-ahead overlap in this control). + int4 vA, vB0, vB1; + + // Prologue: load stage 0 into registers, commit it to the single buffer, + // and make it visible before the first burst. + { + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + vB0 = *reinterpret_cast( + weight + b_kk0 * n + n0 + b_col16); + vB1 = *reinterpret_cast( + weight + (b_kk0 + kStageK128 / 2) * n + n0 + b_col16); + // Commit as int32 stores (skips the 4-byte pad column; 4 x ds_write_b32 + // per 16-byte vector because the odd strides are 4 B mod 16). + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + int32_t* bdst0 = reinterpret_cast( + b_tile + b_kk0 * kBStride128 + b_col16); + bdst0[0] = vB0.x; + bdst0[1] = vB0.y; + bdst0[2] = vB0.z; + bdst0[3] = vB0.w; + int32_t* bdst1 = bdst0 + (kStageK128 / 2) * (kBStride128 / 4); + bdst1[0] = vB1.x; + bdst1[1] = vB1.y; + bdst1[2] = vB1.z; + bdst1[3] = vB1.w; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + // Consume the single buffer: each wave does eight m16n16k32 MMAs per + // 32-K step over its 32x64 quadrant. + const int local_row = wave_row * 32; + const int local_col = wave_col * 64; +#pragma unroll + for (int kk = 0; kk < kStageK128; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride128 + kk, kAStride128); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride128 + kk, + kAStride128); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride128 + local_col, kBStride128); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride128 + local_col + kTileN, + kBStride128); + du_load_matrix_sync( + b_frag2, b_tile + kk * kBStride128 + local_col + 2 * kTileN, + kBStride128); + du_load_matrix_sync( + b_frag3, b_tile + kk * kBStride128 + local_col + 3 * kTileN, + kBStride128); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + } + // All consumption of the single buffer is complete only after every + // wavefront passes this barrier; the stage s+1 stores below overwrite + // that same buffer, so they must wait for it. + __syncthreads(); + + // Load stage s+1 and commit it into the single buffer. The global-load + // latency is fully exposed here: the loads are issued after the burst + // (nothing to overlap) and the compiler's vmcnt wait before the DS + // stores is on the critical path. `s + 1 < num_stages` is block-uniform, + // so the branch and its barrier are free of divergence. + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vB0 = *reinterpret_cast( + weight + (s1 + b_kk0) * n + n0 + b_col16); + vB1 = *reinterpret_cast( + weight + (s1 + b_kk0 + kStageK128 / 2) * n + n0 + b_col16); + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + int32_t* bdst0 = reinterpret_cast( + b_tile + b_kk0 * kBStride128 + b_col16); + bdst0[0] = vB0.x; + bdst0[1] = vB0.y; + bdst0[2] = vB0.z; + bdst0[3] = vB0.w; + int32_t* bdst1 = bdst0 + (kStageK128 / 2) * (kBStride128 / 4); + bdst1[0] = vB1.x; + bdst1[1] = vB1.y; + bdst1[2] = vB1.z; + bdst1[3] = vB1.w; + // Make the stage s+1 stores visible to every wavefront before the + // next burst. + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 64; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc02, x_scale, weight_scale, out, m, n, + base_row, base_col + 2 * kTileN, lane); + store_prefill_fragment( + acc03, x_scale, weight_scale, out, m, n, + base_row, base_col + 3 * kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); + store_prefill_fragment( + acc12, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + 2 * kTileN, lane); + store_prefill_fragment( + acc13, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + 3 * kTileN, lane); +} + +// --------------------------------------------------------------------------- +// Packed-B family for the exact assigned shape (K=6144, N=2560). +// launch_pack_w8a8_weight stores the logical [K, N] weight once, outside the +// timed region, in the iteration-5 fragment-interleaved swizzle (the 64x64 +// and 128x64 compiled-but-undispatched family members still assume the +// round-1 n-major packed[n][k] / 80-byte-row tile and are never launched by +// launch_w8a8_gemm). The dispatched 64x128 kernel stages B into the swizzled +// [k32][n16][k8][n][8] tile and consumes col_major B fragments with the +// 8-byte loader load_b_frag8_swz (one conflict-free ds_read_b64 per +// fragment). A stays row-major (68-byte LDS rows). All three kernels share +// the accepted pipeline: single-buffered 64-K K stage with two +// __syncthreads per stage, all twelve fragment loads hoisted before the +// 16-MMAC burst, fused coalesced scale/bf16 epilogue. The int32 accumulation +// order is unchanged from the generic kernels, so results are bit-identical +// to the logical [K, N] layout. +// --------------------------------------------------------------------------- + +// 64x128 tile: 4 waves, 32x64 quadrant per wave, eight accumulators. This is +// the exact-shape dispatch (strongest N=2560 prior evidence: the accepted +// hy3 qkv_proj lineage at 113.5 TOPS, K=4096). Grid (20, 64) = 1280 blocks. +// LDS/block = 64x68 + 64x128 = 12,544 B (iteration 5: the swizzled B tile +// needs no padding) -> 3 blocks/CU (VGPR-bound at arch_vgpr <= 85). +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x128_packedb_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM128; + const int n0 = static_cast(blockIdx.x) * kBlockN128; + + __shared__ __align__(16) int8_t a_tile[kBlockM128 * kAStride128]; + // Swizzled B tile (iteration 5): the exact-shape pack stores the weight + // fragment-interleaved [k32 0..1][n16 0..7][k8 0..3][n 0..15][8 B] + // (64-K stage = 2 x 4096 B), so staging is int4 -> ds_write_b128 and the + // col_major B fragment loads are one conflict-free ds_read_b64 each. + __shared__ __align__(16) int8_t b_tile[kStageK128 * kBlockN128]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; + DUFragment + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); + + const int num_stages = k / kStageK128; + + // A staging (row-major x_q -> 68-byte LDS rows): thread `tid` owns the + // 16-byte k-run (a_col16..a_col16+15) of row a_row. + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + // B staging on the swizzled pack: thread `tid` owns the 16-byte vectors + // (k32=0, n16, k8, n_pair) = v = tid and (k32=1, n16, k8, n_pair) = v = + // tid+256, i.e. bytes + // ((k32*160 + n0/16 + n16)*512 + k8*128 + n_pair*16) of packed_b, which + // are the 8-byte k-runs of n-rows 2*n_pair and 2*n_pair+1 in one k8 group. + const int b_n16 = tid >> 5; // 0..7 (n16 group inside the 128-n tile) + const int b_k8 = (tid >> 3) & 3; // 0..3 (8-byte k8 group inside the 32-k block) + const int b_np = tid & 7; // 0..7 (n pair inside the 16-n group) + // Tile-local offset of the 16-byte chunk (k32 block 0, n16 group b_n16, + // k8 group, n pair) inside b_tile. + const int b_off = + b_n16 * kPackedBGroupBytes + (b_k8 << 7) + (b_np << 4); + // Global packed_b offset of the same 16-byte chunk inside ONE 32-k block + // (the n16-group term is already carried by the (k32g*kPackedBN16 + + // n0_16 + b_n16) index, so it must NOT be repeated here). + const int b_pack_off = (b_k8 << 7) + (b_np << 4); + + // VGPR payload for one full 64-K stage (3 x int4 = 12 VGPR). + int4 vA, vB0, vB1; + + // Prologue: load stage 0 into registers, commit it to the single buffers, + // and make it visible before the first burst. + { + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + const int n0_16 = n0 >> 4; + vB0 = *reinterpret_cast( + packed_b + (n0_16 + b_n16) * kPackedBGroupBytes + b_pack_off); + vB1 = *reinterpret_cast( + packed_b + + (kPackedBN16 + n0_16 + b_n16) * kPackedBGroupBytes + b_pack_off); + // Commit A as int32 stores (68-byte rows are 4 B mod 16 -> no b128). + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + // Commit B as 16-byte vector stores (all offsets are 0 mod 16): each + // 8-lane LDS cycle covers all 32 banks exactly once (conflict-free). + *reinterpret_cast(b_tile + b_off) = vB0; + *reinterpret_cast(b_tile + kPackedB32Bytes + b_off) = vB1; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + // Consume the single buffer: each wave does eight m16n16k32 MMAs per + // 32-K step over its 32x64 quadrant. All twelve fragment loads of both + // 32-K steps are hoisted before the first v_mmac, so the LDS latency + // concentrates in one wait and the 16-MMAC stream has no lgkmcnt waits. + const int local_row = wave_row * 32; + const int8_t* abase = a_tile + local_row * kAStride128; + // B fragment j covers n16 group wave_col*4 + j of the swizzled tile + // (512 B per 16-n fragment block); the kk=32 burst lives in k32 block 1. + const int8_t* bbase = b_tile + (wave_col << 2) * kPackedBGroupBytes; + du_load_matrix_sync(a_frag0, abase, kAStride128); + du_load_matrix_sync( + a_frag1, abase + kTileM * kAStride128, kAStride128); + load_b_frag8_swz(b_frag0, bbase); + load_b_frag8_swz(b_frag1, bbase + kPackedBGroupBytes); + load_b_frag8_swz(b_frag2, bbase + 2 * kPackedBGroupBytes); + load_b_frag8_swz(b_frag3, bbase + 3 * kPackedBGroupBytes); + DUFragment + a2_frag0, a2_frag1; + DUFragment + b2_frag0, b2_frag1, b2_frag2, b2_frag3; + du_load_matrix_sync(a2_frag0, abase + kTileK, kAStride128); + du_load_matrix_sync( + a2_frag1, abase + kTileM * kAStride128 + kTileK, kAStride128); + load_b_frag8_swz(b2_frag0, bbase + kPackedB32Bytes); + load_b_frag8_swz( + b2_frag1, bbase + kPackedB32Bytes + kPackedBGroupBytes); + load_b_frag8_swz( + b2_frag2, bbase + kPackedB32Bytes + 2 * kPackedBGroupBytes); + load_b_frag8_swz( + b2_frag3, bbase + kPackedB32Bytes + 3 * kPackedBGroupBytes); + // kk = 0 burst (eight m16n16k32 MMAs). + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + // kk = 32 burst (same accumulators, same int32 accumulation order). + du_mma_sync(acc00, a2_frag0, b2_frag0, acc00); + du_mma_sync(acc01, a2_frag0, b2_frag1, acc01); + du_mma_sync(acc02, a2_frag0, b2_frag2, acc02); + du_mma_sync(acc03, a2_frag0, b2_frag3, acc03); + du_mma_sync(acc10, a2_frag1, b2_frag0, acc10); + du_mma_sync(acc11, a2_frag1, b2_frag1, acc11); + du_mma_sync(acc12, a2_frag1, b2_frag2, acc12); + du_mma_sync(acc13, a2_frag1, b2_frag3, acc13); + // All consumption of the single buffer is complete only after every + // wavefront passes this barrier; the stage s+1 stores below overwrite + // that same buffer. + __syncthreads(); + + // Load stage s+1 and commit it into the single buffer. + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int k32g = s1 >> 5; // global 32-k block of the stage s+1 payload + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + const int n0_16 = n0 >> 4; + vB0 = *reinterpret_cast( + packed_b + + (k32g * kPackedBN16 + n0_16 + b_n16) * kPackedBGroupBytes + + b_pack_off); + vB1 = *reinterpret_cast( + packed_b + + ((k32g + 1) * kPackedBN16 + n0_16 + b_n16) * kPackedBGroupBytes + + b_pack_off); + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + *reinterpret_cast(b_tile + b_off) = vB0; + *reinterpret_cast(b_tile + kPackedB32Bytes + b_off) = vB1; + // Make the stage s+1 stores visible to every wavefront before the + // next burst. + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 64; + // Iteration 6 (epilogue round): fused 4-fragment row-group epilogues. + // acc00..acc03 cover rows base_row..base_row+15, columns base_col..+63; + // acc10..acc13 cover the same columns one 16-row group lower. Each call + // loads x_scale[row] and the four float4 weight_scale groups once per + // lane before any ds_bpermute, issues the four transposes as one + // software-pipelined block, then stores back-to-back. + store_prefill_rowgroup_coalesced( + acc00, acc01, acc02, acc03, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_rowgroup_coalesced( + acc10, acc11, acc12, acc13, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); +} + +// 64x64 packed-B tile: 4 waves, 32x32 quadrant per wave, four accumulators. +// Compiled family member for the mandated tile sweep (grid 40x64 = 2560 +// blocks for the assigned shape; LDS/block = 64x68 + 64x80 = 9,472 B). Not +// dispatched this round; the exact shape uses the 64x128 packed-B kernel. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x64_packedb_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride128]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kPackedBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + const int num_stages = k / kStageK128; + + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + const int b_n = tid >> 2; + const int b_k16 = (tid & 3) << 4; + + int4 vA, vB; + + // Prologue: stage 0. + { + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + vB = *reinterpret_cast( + packed_b + (n0 + b_n) * k + b_k16); + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + const int8_t* abase = a_tile + local_row * kAStride128; + const int8_t* bbase = b_tile + local_col * kPackedBStride; + du_load_matrix_sync(a_frag0, abase, kAStride128); + du_load_matrix_sync( + a_frag1, abase + kTileM * kAStride128, kAStride128); + load_b_frag8(b_frag0, bbase, kPackedBStride); + load_b_frag8( + b_frag1, bbase + kTileN * kPackedBStride, kPackedBStride); + DUFragment + a2_frag0, a2_frag1; + DUFragment + b2_frag0, b2_frag1; + du_load_matrix_sync(a2_frag0, abase + kTileK, kAStride128); + du_load_matrix_sync( + a2_frag1, abase + kTileM * kAStride128 + kTileK, kAStride128); + load_b_frag8(b2_frag0, bbase + kTileK, kPackedBStride); + load_b_frag8( + b2_frag1, bbase + kTileN * kPackedBStride + kTileK, kPackedBStride); + // kk = 0 burst. + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + // kk = 32 burst. + du_mma_sync(acc00, a2_frag0, b2_frag0, acc00); + du_mma_sync(acc01, a2_frag0, b2_frag1, acc01); + du_mma_sync(acc10, a2_frag1, b2_frag0, acc10); + du_mma_sync(acc11, a2_frag1, b2_frag1, acc11); + __syncthreads(); + + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vB = *reinterpret_cast( + packed_b + (n0 + b_n) * k + s1 + b_k16); + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB; + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment_coalesced( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment_coalesced( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment_coalesced( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment_coalesced( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// 128x64 packed-B tile: 4 waves, 64x32 quadrant per wave, eight +// accumulators (the o_proj lineage geometry). Compiled family member for the +// mandated tile sweep (grid 40x32 = 1280 blocks for the assigned shape; +// LDS/block = 128x80 + 64x80 = 15,360 B). Not dispatched this round. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_128x64_packedb_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM128 * 2; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM128 * 2 * kPackedAStride128x64]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kPackedBStride]; + + DUFragment + a_frag0, a_frag1, a_frag2, a_frag3; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + + const int num_stages = k / kStageK128; + + // A staging: thread owns the 16-byte k-run (a_col16..+15) of A rows + // a_row and a_row + 64 (A[128,64] = 512 int4 = 2 per thread). + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + // B staging: thread owns the 16-byte k-run (b_k16..+15) of n-row b_n. + const int b_n = tid >> 2; + const int b_k16 = (tid & 3) << 4; + + int4 vA0, vA1, vB; + + // Prologue: stage 0. + { + const int g_row = m0 + a_row; + vA0 = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + vA1 = (g_row + kBlockM128 < m) + ? *reinterpret_cast( + x_q + (g_row + kBlockM128) * k + a_col16) + : int4{0, 0, 0, 0}; + vB = *reinterpret_cast( + packed_b + (n0 + b_n) * k + b_k16); + // 80-byte A rows are 0 mod 16: one b128 store per vector. + *reinterpret_cast( + a_tile + a_row * kPackedAStride128x64 + a_col16) = vA0; + *reinterpret_cast( + a_tile + (a_row + kBlockM128) * kPackedAStride128x64 + a_col16) = vA1; + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMAs per 32-K + // step (4 A fragments x 2 B fragments), all fragment loads hoisted. + const int local_row = wave_row * 64; + const int local_col = wave_col * 32; + const int8_t* abase = a_tile + local_row * kPackedAStride128x64; + const int8_t* bbase = b_tile + local_col * kPackedBStride; + du_load_matrix_sync(a_frag0, abase, kPackedAStride128x64); + du_load_matrix_sync( + a_frag1, abase + kTileM * kPackedAStride128x64, + kPackedAStride128x64); + du_load_matrix_sync( + a_frag2, abase + 2 * kTileM * kPackedAStride128x64, + kPackedAStride128x64); + du_load_matrix_sync( + a_frag3, abase + 3 * kTileM * kPackedAStride128x64, + kPackedAStride128x64); + load_b_frag8(b_frag0, bbase, kPackedBStride); + load_b_frag8( + b_frag1, bbase + kTileN * kPackedBStride, kPackedBStride); + DUFragment + a2_frag0, a2_frag1, a2_frag2, a2_frag3; + DUFragment + b2_frag0, b2_frag1; + du_load_matrix_sync(a2_frag0, abase + kTileK, kPackedAStride128x64); + du_load_matrix_sync( + a2_frag1, abase + kTileM * kPackedAStride128x64 + kTileK, + kPackedAStride128x64); + du_load_matrix_sync( + a2_frag2, abase + 2 * kTileM * kPackedAStride128x64 + kTileK, + kPackedAStride128x64); + du_load_matrix_sync( + a2_frag3, abase + 3 * kTileM * kPackedAStride128x64 + kTileK, + kPackedAStride128x64); + load_b_frag8(b2_frag0, bbase + kTileK, kPackedBStride); + load_b_frag8( + b2_frag1, bbase + kTileN * kPackedBStride + kTileK, kPackedBStride); + // kk = 0 burst. + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc20, a_frag2, b_frag0, acc20); + du_mma_sync(acc21, a_frag2, b_frag1, acc21); + du_mma_sync(acc30, a_frag3, b_frag0, acc30); + du_mma_sync(acc31, a_frag3, b_frag1, acc31); + // kk = 32 burst. + du_mma_sync(acc00, a2_frag0, b2_frag0, acc00); + du_mma_sync(acc01, a2_frag0, b2_frag1, acc01); + du_mma_sync(acc10, a2_frag1, b2_frag0, acc10); + du_mma_sync(acc11, a2_frag1, b2_frag1, acc11); + du_mma_sync(acc20, a2_frag2, b2_frag0, acc20); + du_mma_sync(acc21, a2_frag2, b2_frag1, acc21); + du_mma_sync(acc30, a2_frag3, b2_frag0, acc30); + du_mma_sync(acc31, a2_frag3, b2_frag1, acc31); + __syncthreads(); + + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int g_row = m0 + a_row; + vA0 = (g_row < m) + ? *reinterpret_cast( + x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vA1 = (g_row + kBlockM128 < m) + ? *reinterpret_cast( + x_q + (g_row + kBlockM128) * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vB = *reinterpret_cast( + packed_b + (n0 + b_n) * k + s1 + b_k16); + *reinterpret_cast( + a_tile + a_row * kPackedAStride128x64 + a_col16) = vA0; + *reinterpret_cast( + a_tile + (a_row + kBlockM128) * kPackedAStride128x64 + a_col16) = + vA1; + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB; + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 64; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment_coalesced( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment_coalesced( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment_coalesced( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment_coalesced( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); + store_prefill_fragment_coalesced( + acc20, x_scale, weight_scale, out, m, n, + base_row + 2 * kTileM, base_col, lane); + store_prefill_fragment_coalesced( + acc21, x_scale, weight_scale, out, m, n, + base_row + 2 * kTileM, base_col + kTileN, lane); + store_prefill_fragment_coalesced( + acc30, x_scale, weight_scale, out, m, n, + base_row + 3 * kTileM, base_col, lane); + store_prefill_fragment_coalesced( + acc31, x_scale, weight_scale, out, m, n, + base_row + 3 * kTileM, base_col + kTileN, lane); +} + +// Generic large-M prefill kernel: 64x64 output tile per block, K staged in +// LDS (single buffer), four waves each owning a 32x32 quadrant = four +// m16n16k32 int8->int32 DUMMA accumulators. Covers large-M shapes whose N is +// a multiple of 64 but not 128 (and K a multiple of 128). +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM * kStageK]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBlockN]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[64,128] (row-major, zero-filled past M). + constexpr int kAVectors = kBlockM * kStageK / static_cast(sizeof(int4)); + for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int local_row = byte_offset / kStageK; + const int kk = byte_offset - local_row * kStageK; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile)[vec] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + kk) + : int4{0, 0, 0, 0}; + } + + // Stage B[128,64] (row-major [K, N] logical weight layout). + constexpr int kBVectors = kStageK * kBlockN / static_cast(sizeof(int4)); + for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + reinterpret_cast(b_tile)[vec] = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + col); + } + __syncthreads(); + + // Consume the stage: each wave does four m16n16k32 MMAs per 32-K step. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kStageK + kk, kStageK); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kStageK + kk, kStageK); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBlockN + local_col, kBlockN); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBlockN + local_col + kTileN, kBlockN); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + // Single buffer: make sure every wave finished reading LDS before the + // next stage overwrites it. + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Generic scalar int8/int32 fallback for unmatched shapes and small-M API +// cases. One output element per grid-stride step; exact int32 accumulation +// (K <= 8192 keeps the int8 dot well within int32 range: 8192*128*128 = +// 134,217,728 < 2^31), then the fused float scale and bf16 store. The exact +// assigned shape (K=6144, N=2560) uses the iteration-5 swizzled pack (set up +// once by launch_pack_w8a8_weight; decoded per element below); every other +// (k, n) keeps the logical [K, N] row-major layout. The branch is +// grid-uniform per launch, so the paired M=2 API shape with the same (N, K) +// stays correct. +__global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + const bool packed_b = (k == kPackedBK && n == kPackedBN); + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; + idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = a + static_cast(row) * k; + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + int32_t bw; + if (packed_b) { + // Iteration-5 swizzled exact-shape layout: element (n=col, k=kk) at + // ((kk>>5)*160 + (col>>4))*512 + (((kk>>3)&3)*16 + (col&15))*8 + + // (kk&7) (fragment-interleaved [k32][n16][k8][n][8]). + const int64_t off = + ((static_cast(kk >> 5) * kPackedBN16 + (col >> 4)) * + kPackedBGroupBytes) + + ((((kk >> 3) & 3) * 16 + (col & 15)) << 3) + (kk & 7); + bw = static_cast(b[off]); + } else { + bw = static_cast( + b[static_cast(kk) * n + col]); + } + acc += static_cast(a_row[kk]) * bw; + } + const float scaled = static_cast(acc) * + x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// Identity device-to-device byte copy (pack_weight fallback, valid for any +// (K, N)). +__global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +// One-time swizzled pack of the logical [K, N] weight for the exact assigned +// shape (K=6144, N=2560), iteration 5 (packing round). Runs only from +// launch_pack_w8a8_weight, outside the timed region and outside Graph +// capture. The packed buffer keeps the same byte count (k*n) as the identity +// pack, so allocations and graph-stable addresses are unchanged. Layout: +// element (n, k) lands at +// ((k>>5)*n16g + (n>>4))*512 + (((k>>3)&3)*16 + (n&15))*8 + (k&7) +// with n16g = n>>4 (160 for the exact shape): for each 32-k x 16-n DUMMA B +// fragment block the 4 k8 groups x 16 n-rows x 8 B are stored n-interleaved +// (k8-major within the block). Each thread writes one contiguous 16-byte +// chunk (two n-rows of one k8) with a strided 2x8 gather on the read side +// (packing is outside timing, so the scalar gather is off the critical +// path). The guard guarantees n % 16 == 0. +__global__ __launch_bounds__(256) void w8a8_pack_b_swz_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int k, + int n) { + const int64_t chunks = (static_cast(k) * n) >> 4; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + const int n16g = n >> 4; + for (int64_t c = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + c < chunks; + c += stride) { + const int k32 = static_cast(c / (static_cast(n16g) * 32)); + const int rem = + static_cast(c % (static_cast(n16g) * 32)); + const int n16 = rem >> 5; + const int w = rem & 31; + const int k8 = w >> 3; + const int np = w & 7; + const int n0 = (n16 << 4) + (np << 1); + const int k0 = (k32 << 5) + (k8 << 3); + int8_t* dst16 = dst + (c << 4); +#pragma unroll + for (int i = 0; i < 8; ++i) { + // 16-byte chunk = [n0 row: k0..k0+7][n0+1 row: k0..k0+7] (n-major + // within the k8 group, matching the [k8][n][8] tile layout). + dst16[i] = src[static_cast(k0 + i) * n + n0]; + dst16[8 + i] = src[static_cast(k0 + i) * n + n0 + 1]; + } + } +} + +// Identity device-to-device float copy for the weight scales. +__global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +} // namespace + +// Stable host launch symbol consumed by csrc/bindings.cpp. Dispatches on the +// caller-provided stream only; never allocates, synchronizes, or touches the +// default stream. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + auto* out_ptr = static_cast(out); + + // Packed-B family: the exact assigned shape (K=6144, N=2560) routes to + // the swizzled-packed-B 64x128 kernel (iteration 5 layout). The guard is + // exact (m >= 128, n == 2560, k == 6144) and sits BEFORE the generic + // 64x128 path; every other shape keeps its existing path (generic 64x128 + // for n % 128 == 0, 64x64 for n % 64 == 0, scalar fallback otherwise) and + // the identity-packed layout. Grid (N/128)x(M/64) = 20x64 = 1280 blocks. + if (m >= 128 && n == kPackedBN && k == kPackedBK) { + const dim3 grid( + static_cast(n / kBlockN128), + static_cast((m + kBlockM128 - 1) / kBlockM128)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x128_packedb_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Native INT8 DUMMA m16n16k32 prefill path (2-D macro-tile 64x128, + // single-buffered 64-K K stage, identity-packed B) for every other + // large-M shape with compatible geometry. + if (m >= 128 && (n % kBlockN128) == 0 && (k % 128) == 0) { + const dim3 grid( + static_cast(n / kBlockN128), + static_cast((m + kBlockM128 - 1) / kBlockM128)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x128_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Native INT8 DUMMA m16n16k32 prefill path (64x64 tile) for large-M shapes + // whose N is not a multiple of 128 but is a multiple of 64. + if (m >= 128 && (n % kBlockN) == 0 && (k % kStageK) == 0) { + const dim3 grid( + static_cast(n / kBlockN), + static_cast((m + kBlockM - 1) / kBlockM)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x64_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Scalar int8/int32 fallback for unmatched (m, n, k) and small-M API cases + // (M < 128, including the paired M=2 API shape with the same (N,K)). + const int64_t total = static_cast(m) * n; + constexpr int kFallbackThreads = 256; + int64_t blocks = (total + kFallbackThreads - 1) / kFallbackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + const dim3 grid(static_cast(blocks)); + const dim3 block(kFallbackThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); +} + +// Stable host launch symbol consumed by csrc/bindings.cpp. For the exact +// assigned shape (K=6144, N=2560) it performs the one-time swizzled B pack +// (iteration 5: logical [K, N] -> fragment-interleaved +// [k32][n16][k8][n][8]) outside the timed region and outside Graph capture; +// every other (K, N) keeps the identity device-to-device copy, valid for +// every (K, N). The packed buffer size is k*n in both cases, so the +// allocation and graph-stable packed layout are unchanged. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kPackThreads = 256; + const int64_t weight_count = static_cast(k) * n; + if (weight_count > 0) { + if (k == kPackedBK && n == kPackedBN) { + const int64_t chunks = weight_count >> 4; + int64_t blocks = (chunks + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_b_swz_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + int64_t blocks = (weight_count + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_identity_i8_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count); + } + } + if (n > 0) { + int64_t blocks = (static_cast(n) + kPackThreads - 1) / + kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_identity_f32_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/shared_gate_up_proj.hip new file mode 100644 index 00000000..f9b6e0ad --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/shared_gate_up_proj.hip @@ -0,0 +1,629 @@ +// @@variant shape=minimax_tp4_shared_gate_up_proj_m4096 commit=5c8bfa17fe45a29749afdfc9124782e9068e415c added=2026-08-24 +// median_us=986.5 p90_us=987.7 +// source=minimaxm3-dsh-tp4-m4096-1-0c2f84a9 +// MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). +// worker_3 bootstrap (iteration 1): correctness-first native INT8 DUMMA +// m16n16k32 tiled prefill kernels for the assigned MiniMax TP4 M=4096 +// shapes: +// * minimax_tp4_shared_gate_up_proj_m4096: (M=4096, N=1536, K=6144) +// * minimax_tp4_shared_down_proj_m4096: (M=4096, N=6144, K=768) +// +// This file provides the two stable host launch symbols consumed by +// csrc/bindings.cpp: +// * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's +// HIP stream; no allocation, packing, +// autotuning, or synchronization. +// * launch_pack_w8a8_weight(...) - bootstrap identity device-to-device +// copy (opaque packed layout = raw +// [K, N] layout; later rounds may change +// only this and the GEMM interpretation). +// +// Strategy (correctness-first bootstrap, then measured tuning): +// * Large-M prefill (M >= 128): a single templated native INT8 DUMMA +// m16n16k32 kernel computes a (kBlockM x kBlockN) output tile from +// (kBlockM/32) x (kBlockN/32) wavefronts, each wave owning one 32x32 +// quadrant (four 16x16 int32 accumulators) resident across the whole K +// loop. A and B are cooperatively staged into bank-skewed LDS with +// 16-byte vectorized coalesced global loads. The K pipeline is chosen by +// the kDoubleBuffer template flag: +// - false: single-buffered LDS, kStageK=128 chunks, two barriers per +// stage (stage in, compute, buffer-protect). Used by the 128x64 and +// 64x64 arms. +// - true: double-buffered LDS, kStageK=32 chunks (iteration 6 halved +// the stage depth 64 -> 32 so the two 32-K LDS buffers fit 4 +// blocks/CU), one barrier per stage, loop-carried int4 register +// payload prefetching one stage ahead (global loads issued before the +// MMAC burst, published into the idle buffer before the next stage's +// MMAs). Used by the 64x128 arm (both assigned shapes). No split-K, +// no raw asm. +// * LDS row strides are bank-phased for the DUMMA fragment loads: A rows +// keep a 16-byte pad (stride kStageK+16: 144 for 128-K stages, 80 for +// 64-K stages, 48 for 32-K stages) so the A-fragment ds_read2_b32 spread +// over all 32 banks +// (2-way), while B rows use a 4-byte pad (stride kBlockN+4 = 132) +// because any 16-byte-aligned B stride collapses the four k-groups of +// every B-fragment byte load onto the same bank phase (4 banks, 16-way; +// stride 132 moves the groups to phases {0,8,16,24}, 16 banks, 4-way). +// * Three macro-tiles are instantiated (64x64, 64x128, 128x64) and the +// active one is selected at launch time by explicit geometry dispatch: +// - N % 128 == 0 -> 64x128 tile (8 wavefronts, 512 threads) +// - else M % 128 == 0 -> 128x64 tile (8 wavefronts, 512 threads) +// - else -> 64x64 tile (4 wavefronts, 256 threads) +// Both assigned shapes route to the 64x128 tile (double-buffered 64-K +// stages since iteration 3, halved to 32-K stages in iteration 6): +// - gate_up : grid (12, 64) = 768 blocks, 192 K stages, 6144 waves +// - down_proj: grid (48, 64) = 3072 blocks, 24 K stages, 24576 waves +// The M x N output-tile grid already dwarfs the 120 CUs, so no split-K. +// * Tail M rows are zero-filled on A load and masked on store, so any +// M >= 128 is correct; M < 128 (decode API shapes, e.g. M=2/M=16) +// reaches the scalar fallback below. +// * Everything else (M < 128 or unmatched geometry): scalar int8/int32 +// fallback with exact int32 accumulation in k order. +// +// Mathematical contract (exact int32 dot before float scaling): +// out[m, n] = bf16(int32_dot(a[m, :], b[:, n]) * x_scale[m] * weight_scale[n]) +// +// Header order is fixed by the control plane for this DTK: +// hip_runtime.h -> hip_bfloat16.h -> du_mma.h + +#include +#include +#include + +#include + +namespace { + +// gfx928 INT8 DUMMA primitive: m16n16k32, int8 x int8 -> int32. +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; + +// gfx928 wavefront is 64 lanes; every blockDim must be a multiple of 64. +constexpr int kWaveSize = 64; + +// Large-prefill dispatch geometry. +constexpr int kPrefillMinM = 128; +constexpr int kPrefillStageK = 128; +// Stage depth of the double-buffered arm (iteration 6: halved 64 -> 32 so +// the two 32-K LDS buffers (a_tile[2][64x48] + b_tile[2][32x132] = 14,592 +// B/block) fit 4 blocks/CU (58,368 B <= 65,536 B) -> up to 32 resident +// waves/CU vs 16 for the 64-K depth (27,136 B/block -> 2 blocks/CU). Staged +// bytes per MAC, (1/kBlockM + 1/kBlockN), are independent of the stage depth, +// so total global loads and LDS traffic stay bit-identical to the 64-K +// champion; the change is purely an occupancy/latency-hiding experiment +// against the measured LDS-latency-bound profile (1.65 lds_wait per LDS +// instruction). A naive 128-K double buffer (52,224 B/block -> 1 block/CU) +// remains rejected. +constexpr int kPrefillDoubleStageK = 32; +constexpr int kDefaultBlockM = 64; +constexpr int kDefaultBlockN = 64; + +// LDS padding in bytes added to each staged row. A keeps 16 so every A row +// start stays 16-byte aligned for int4 staging while shifting fragment rows +// off the same 32-bank phase (power-of-two strides alias every row onto +// identical banks). B deliberately uses only 4 bytes of pad: with the gfx928 +// INT8 DUMMA B-fragment layout (lane (c, g) reads bytes at k-rows 8g..8g+7, +// one byte per row, ldm apart), any 16-byte-aligned row stride makes +// 2*stride*g == 0 (mod 32), so all four k-groups of every ds_read_u8 land on +// the same bank phase -> 4 banks x 16 lanes (16-way). Stride 132 (== 4 mod +// 16) moves the k-groups to bank phases {0,8,16,24} -> 16 banks x 4 lanes. +constexpr int kLdsPad = 16; +constexpr int kBLdsPad = 4; + +constexpr int kScalarThreads = 256; + +using namespace du::dumma; + +// Verified gfx928 INT8 DUMMA accumulator ownership: +// row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> columns +// col_mod4 + 4*i. Scale by x_scale[row] and weight_scale[col] and store +// bf16 directly from the fragment (no accumulator LDS round trip). +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Load one flat staging vector (index v over the A-then-B region) of the K +// stage starting at k0 into a register slot. Tail-M A rows are zero-filled so +// the DUMMA math stays defined; B rows are contiguous kBlockN chunks at global +// row stride k (bootstrap identity [K, N] packed layout). Used only by the +// double-buffered pipeline's loop-carried payload. +template +__device__ __forceinline__ int4 prefill_load_stage_vector( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int v, + int k0, + int m, + int k, + int n, + int m0, + int n0) { + constexpr int kAVectorsPerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kBlockN / static_cast(sizeof(int4)); + if (v < kAVectors) { + const int local_row = v / kAVectorsPerRow; + const int vv = v - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + return global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + + vv * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + const int bv = v - kAVectors; + const int kk = bv / kBVectorsPerRow; + const int vv = bv - kk * kBVectorsPerRow; + return *reinterpret_cast( + weight + (k0 + kk) * n + n0 + + vv * static_cast(sizeof(int4))); +} + +// Publish one flat staging vector into the given LDS buffers. A rows keep +// 16-byte-aligned int4 stores at row stride kAStride; B rows are committed as +// four int32 stores at row stride kBStride (4-byte pad, see kBLdsPad). +template +__device__ __forceinline__ void prefill_publish_stage_vector( + int8_t* __restrict__ a_tile, + int8_t* __restrict__ b_tile, + int v, + int4 val) { + constexpr int kAVectorsPerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kBlockN / static_cast(sizeof(int4)); + if (v < kAVectors) { + const int local_row = v / kAVectorsPerRow; + const int vv = v - local_row * kAVectorsPerRow; + reinterpret_cast(a_tile + local_row * kAStride)[vv] = val; + } else { + const int bv = v - kAVectors; + const int kk = bv / kBVectorsPerRow; + const int vv = bv - kk * kBVectorsPerRow; + int32_t* b_row = reinterpret_cast( + b_tile + kk * kBStride + vv * static_cast(sizeof(int4))); + b_row[0] = val.x; + b_row[1] = val.y; + b_row[2] = val.z; + b_row[3] = val.w; + } +} + +// Large-M prefill path: one block computes a kBlockM x kBlockN output tile. +// kWaveRows x kWaveCols wavefronts each own a 32x32 quadrant (four 16x16 +// DUMMA accumulators), while the block cooperatively stages +// A[kBlockM, kStageK] and B[kStageK, kBlockN] in bank-skewed LDS. With +// kDoubleBuffer=false the LDS is single-buffered with two barriers per stage +// (stage in, compute, protect); with kDoubleBuffer=true two LDS buffers +// alternate across the stage with a loop-carried int4 register payload and +// one barrier per stage (see the branch for the pipeline). Tail M rows are +// zero-filled on load and masked on store, so any M >= 128 is supported. The +// packed weight is the bootstrap identity [K, N] layout, so B rows are staged +// n-major contiguous chunks. +template +__global__ __launch_bounds__((kBlockM / 32) * (kBlockN / 32) * kWaveSize) void +w8a8_dumma_prefill_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWaveRows = kBlockM / 32; + constexpr int kWaveCols = kBlockN / 32; + constexpr int kThreads = kWaveRows * kWaveCols * kWaveSize; + constexpr int kAStride = kStageK + kLdsPad; + constexpr int kBStride = kBlockN + kBLdsPad; + static_assert(kAStride % sizeof(int4) == 0, + "A LDS row stride must stay 16-byte aligned"); + static_assert(kBStride % sizeof(int32_t) == 0, + "B LDS row stride must stay 4-byte aligned (int32 stores)"); + static_assert(kStageK % kTileK == 0, + "K stage must be a multiple of the DUMMA K unit"); + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kBlockN / sizeof(int4); + constexpr int kBVectors = kStageK * kBVectorsPerRow; + constexpr int kStageVectors = kAVectors + kBVectors; + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWaveCols; + const int wave_col = wave - wave_row * kWaveCols; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + if constexpr (kDoubleBuffer) { + // Double-buffered K pipeline (one barrier per stage): + // * two LDS buffers; stage s lives in buffer s & 1; + // * a loop-carried int4 payload prefetches one stage ahead: the global + // loads for stage s+2 are issued before stage s's MMAC burst, and + // stage s+1's payload is published into the idle buffer (vmcnt wait + // implicit) before stage s+1's MMAs. The stage-top vmem wait and the + // staging LDS-write burst no longer sit on the barrier's critical + // path the way the single-buffer stage-in barrier forces. + static_assert(kStageVectors <= 2 * kThreads, + "double-buffered staging needs at most 2 payload slots"); + constexpr int kPayloadSlots = (kStageVectors + kThreads - 1) / kThreads; + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[2][kStageK * kBStride]; + + const int nstages = k / kStageK; + int4 payload[kPayloadSlots]; + + // Prologue: stage 0 is staged straight into LDS buffer 0, then the + // global prefetch of stage 1 is issued into the payload so its latency + // overlaps the prologue barrier and stage 0's compute. +#pragma unroll + for (int p = 0; p < kPayloadSlots; ++p) { + const int v = tid + p * kThreads; + if (v < kStageVectors) { + prefill_publish_stage_vector( + a_tile[0], b_tile[0], v, + prefill_load_stage_vector( + x_q, weight, v, 0, m, k, n, m0, n0)); + } + } + if (nstages > 1) { +#pragma unroll + for (int p = 0; p < kPayloadSlots; ++p) { + const int v = tid + p * kThreads; + if (v < kStageVectors) { + payload[p] = prefill_load_stage_vector( + x_q, weight, v, kStageK, m, k, n, m0, n0); + } + } + } + __syncthreads(); + + for (int s = 0; s < nstages; ++s) { + const int buf = s & 1; + // Publish the prefetched stage s+1 into the idle buffer before this + // stage's MMAC burst. Buffer (s+1)&1 was last read by stage s-1's MMAs, + // which the previous barrier ordered ahead of this publish. The barrier + // at the bottom of the iteration then (a) makes the published stage s+1 + // visible to every thread before iteration s+1's MMAs, and (b) keeps all + // reads of buffer buf ahead of the stage s+2 publish (same-parity + // buffer) in iteration s+1. + if (s + 1 < nstages) { + const int nbuf = buf ^ 1; +#pragma unroll + for (int p = 0; p < kPayloadSlots; ++p) { + const int v = tid + p * kThreads; + if (v < kStageVectors) { + prefill_publish_stage_vector(a_tile[nbuf], b_tile[nbuf], + v, payload[p]); + } + } + // Issue the global prefetch of stage s+2 (latency hidden behind this + // stage's MMAC burst and the barrier). + if (s + 2 < nstages) { +#pragma unroll + for (int p = 0; p < kPayloadSlots; ++p) { + const int v = tid + p * kThreads; + if (v < kStageVectors) { + payload[p] = prefill_load_stage_vector( + x_q, weight, v, (s + 2) * kStageK, m, k, n, m0, n0); + } + } + } + } + + // MMAC burst on buffer buf. Each wave owns a 32x32 quadrant: two A + // fragments (rows 0-15 / 16-31) and two B fragments (cols 0-15 / 16-31) + // per kTileK step, four du_mma_sync per step. The per-kk sequence + // (k0-outer, kk-inner) is identical to the single-buffered arm, so the + // int32 accumulation order stays bit-identical. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile[buf] + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile[buf] + (local_row + kTileM) * kAStride + kk, + kAStride); + du_load_matrix_sync( + b_frag0, b_tile[buf] + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile[buf] + kk * kBStride + local_col + kTileN, + kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + __syncthreads(); + } + } else { + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Cooperatively stage A[kBlockM, kStageK] rows into LDS. Out-of-range + // (tail-M) rows are zero-filled so the DUMMA math stays defined. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + // Cooperatively stage B[kStageK, kBlockN] rows into LDS. The packed + // weight is the bootstrap identity [K, N] layout, so each B row is a + // contiguous chunk of kBlockN columns at row stride n. Global loads stay + // 16-byte vectorized and fully coalesced (one int4 per element); the LDS + // row stride kBStride = kBlockN + 4 is not 16-byte aligned by design (see + // kBLdsPad), so ds_write_b128 is unavailable and each int4 is committed + // as four int32 stores (identical total bank traffic to one b128). + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int kk = vec / kBVectorsPerRow; + const int v = vec - kk * kBVectorsPerRow; + const int4 b = *reinterpret_cast( + weight + (k0 + kk) * n + n0 + + v * static_cast(sizeof(int4))); + int32_t* b_row = reinterpret_cast( + b_tile + kk * kBStride + v * static_cast(sizeof(int4))); + b_row[0] = b.x; + b_row[1] = b.y; + b_row[2] = b.z; + b_row[3] = b.w; + } + __syncthreads(); + + // Compute the kTileK steps of this stage. Each wave owns a 32x32 + // quadrant: two A fragments (rows 0-15 / 16-31) and two B fragments + // (cols 0-15 / 16-31) per step, four du_mma_sync per step. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + + // Protect the single LDS buffer from being overwritten by the next + // stage's staging while this stage's fragment loads are still in flight. + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Generic scalar fallback: one output element per thread with exact int32 +// accumulation in k order. Correct for any (m, n, k), including the small-M +// API shapes (M=2/M=16) and any unmatched geometry. The packed weight is the +// bootstrap identity [K, N] layout (column stride n). +__global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = x_q + static_cast(row) * k; + const int8_t* b_col = weight + static_cast(col); + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); +} + +// Bootstrap pack: identity device-to-device copies. The packed layout is +// opaque per the API contract; until a later round changes the pack layout +// (and the matching GEMM interpretation together), packed == raw [K, N]. +__global__ __launch_bounds__(kScalarThreads) void w8a8_identity_copy_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t numel) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < numel) { + dst[idx] = src[idx]; + } +} + +__global__ __launch_bounds__(kScalarThreads) void +w8a8_identity_copy_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t numel) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < numel) { + dst[idx] = src[idx]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Host launch symbols (extern "C", consumed by csrc/bindings.cpp). +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // Bootstrap uses no split-K, so the workspace is not touched. The timed + // operator performs no allocation, synchronization, packing, or + // default-stream launch; it writes only the caller-provided out. + (void)workspace; + (void)workspace_bytes; + auto* out_bf16 = static_cast(out); + + // Native INT8 DUMMA m16n16k32 tiled path for large-M prefill. All three + // macro-tile instantiations share the exact guard; the launch geometry is + // selected explicitly so every assigned shape is covered (no + // single-shape specialization). The guard divisibility (k % 128 == 0, + // n % 64 == 0) holds for all supported catalog shapes. + if (m >= kPrefillMinM && (n % kDefaultBlockN == 0) && + (k % kPrefillStageK == 0)) { + if (n % 128 == 0) { + // 64x128 tile, 8 wavefronts / 512 threads, double-buffered 32-K stage. + // Both assigned shapes: gate_up grid (12, 64) = 768 blocks x 192 stages, + // down_proj grid (48, 64) = 3072 blocks x 24 stages. + const dim3 grid( + static_cast(n / 128), + static_cast((m + 63) / 64)); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 128, + kPrefillDoubleStageK, true>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (m % 128 == 0) { + // 128x64 tile, 8 wavefronts / 512 threads (single-buffered arm). + const dim3 grid( + static_cast(n / 64), + static_cast(m / 128)); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_kernel<128, 64, kPrefillStageK, false>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + { + // 64x64 tile, 4 wavefronts / 256 threads (generic default, + // single-buffered arm). + const dim3 grid( + static_cast(n / 64), + static_cast((m + 63) / 64)); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_kernel<64, 64, kPrefillStageK, false>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + } + + // Generic scalar fallback for every unmatched (m, n, k) and M < 128 + // (including the paired M=2/M=16 API shapes). + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast( + (total + kScalarThreads - 1) / kScalarThreads)); + const dim3 block(static_cast(kScalarThreads)); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Bootstrap: identity device-to-device copy for every (k, n). The scale + // copy is unchanged in later rounds (N-length, order-independent). + const int64_t weight_numel = static_cast(k) * n; + const dim3 weight_grid(static_cast( + (weight_numel + kScalarThreads - 1) / kScalarThreads)); + hipLaunchKernelGGL( + w8a8_identity_copy_i8_kernel, + weight_grid, dim3(kScalarThreads), 0, stream, + raw_weight, packed_weight, weight_numel); + + const dim3 scale_grid(static_cast( + (n + kScalarThreads - 1) / kScalarThreads)); + hipLaunchKernelGGL( + w8a8_identity_copy_f32_kernel, + scale_grid, dim3(kScalarThreads), 0, stream, + weight_scale, packed_weight_scale, static_cast(n)); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/o_proj.hip new file mode 100644 index 00000000..f771caa3 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/o_proj.hip @@ -0,0 +1,1028 @@ +// @@variant shape=minimax_tp8_o_proj_m16 commit=35384b7737eacb9fb8946dc1c5f984e8bf364aae added=2026-08-28 +// median_us=11.37 p90_us=11.39 +// source=minimax-dsh-tp8-m16-1-78402260 +// W8A8 INT8 GEMM - gfx928 (K500SM_AI) - worker_2 bootstrap (iteration 1). +// +// Operator contract (logical): +// out[m, n] = bf16( int32_dot(x_q[m, :], packed_weight[:, n]) +// * x_scale[m, 0] * packed_weight_scale[n, 0] ) +// +// Bootstrap strategy (correctness first, no DUMMA/split-K machinery): +// * one thread computes exactly one output element; +// * 2D grid: blockIdx.y = row (M), blockIdx.x = 256-wide N tile; +// 256 threads/block = 4 wavefronts of 64; +// * complete K loop accumulated exactly in int32 (|acc| <= 127*127*K fits +// trivially for every supported K); +// * float32 scaling in the exact reference order +// (float(dot) * x_scale[row]) * weight_scale[col], then bf16 RN store; +// * pack_weight is an identity device-to-device copy (raw layout), so the +// GEMM reads packed_weight as plain [K, N] row-major int8. +// +// This scalar kernel is the generic fallback for every (m, n, k), including +// the paired M=2 API shapes. Later optimization rounds may add exact-shape +// guards BEFORE this fallback inside launch_w8a8_gemm, and may change the +// pack_weight HIP implementation together with the matching GEMM +// interpretation. The host launcher signatures below are stable and consumed +// by csrc/bindings.cpp. +// +// Iteration 5 (HIP-only packed-weight round): the exact (M=16, N=6144, +// K=1024) shape now runs the packed-B DUMMA kernel below (packed +// [n_tile][k_step][lane][8] B fragment-slot layout, dwordx2 register +// transport, A kept staged in LDS), and launch_pack_w8a8_weight permutes +// that shape's weight into the packed layout (one-time, untimed, outside +// Graph capture). The scalar fallback decodes the packed layout for the +// exact (n, k) == (6144, 1024) pair so the paired M=2 API shape stays exact; +// every other (m, n, k) keeps the identity pack and direct row-major reads. +// +// Iteration 6 (staging-family round): the staging decision stays A-only LDS +// (A staged once per block, B direct register transport -- B is the cold +// once-read stream with zero cross-block reuse, so an LDS hop for B would add +// 6 MiB of LDS writes + reads without removing any global request). The +// focused change is the B transport width: the packed layout becomes +// [n_tile][k_step_pair][lane][16] -- each lane's 16 B slot holds the +// fragments of TWO consecutive k_steps (2 x 8 B, byte-for-byte the +// du_mma.hpp matrix_b row_major lane order) -- so each lane fetches two +// k-steps' whole B fragments with ONE aligned 16-byte global_load_dwordx4 +// (1 KiB contiguous per wave per instruction = 8 sectors) instead of two +// dwordx2. Same-size permutation (6 MiB), same k-ascending int32 +// accumulation, so output stays bit-identical to iterations 1/4/5. +// +// Iteration 7 (pipeline round: single vs double buffering): the exact ISA of +// the iteration-6 kernel shows the depth-1 "prefetch" is shallower than its +// source comment claims -- per K=128 stage the 2 B global_load_dwordx4 are +// issued before the 4-MMAC burst but the single s_waitcnt vmcnt(0) lands +// immediately AFTER the burst (before the 16-mov fragment rotate), so each B +// HBM round trip has only ~1 MMAC-burst (~120-200 cycles) of in-flight time +// against a ~1,700+ cycle contended round trip; 17.2 us = 8 stages x ~3,100 +// cycles, of which ~3,000 is the exposed B wait. All three pipeline gates +// hold (K=1024 >= 1024; L2 36.6% < 70%; doubled LDS 2 x 16,640 = 33,280 B < +// 48 KiB). The round resolves as: register buffering depth 1 (single) -> 2 +// (double) via a lagged 3-set rotation (burst stage s from set 0, rotate +// sets 0<-1<-2, refill set 2 with stage s+3 AFTER the rotate, so the +// compiler's wait for the refill loads lands at the NEXT iteration's rotate +// -- one full loop iteration of in-flight time per B dwordx4). LDS-ring +// double buffering is rejected by evidence (adds 2 ds_write + 2 ds_read + a +// self-barrier per stage while removing zero B round trips). Barriers per +// k-step stay 0 (the single-wave __syncthreads compiles to s_waitcnt only). +// Same fragments, k-ascending int32 order -> output bit-identical to +// iterations 1/4/5/6. +// +// Iteration 8 (HIP-only occupancy round, mandate: tune ONE occupancy limiter +// with PMC evidence -- waves per block / VGPR live range / LDS footprint / +// spills): the exact ISA + PMC of the accepted iteration-7 kernel show the +// limiter is the WAVE COUNT, not registers or LDS. arch_vgpr = 88 allows +// 2 waves/SIMD (256/88) and LDS 16,640 B/block leaves the measured 3.2 +// blocks/CU resident, but grid = 384 one-wave blocks on 120 CUs = 3.2 +// waves/CU = 0.8 waves/SIMD ACHIEVED: no SIMD in the whole GPU runs a second +// wave, so every B HBM round trip (~3,000 cycles per K=128 stage, the +// iteration-7 measured limiter) is fully exposed with nothing to switch to. +// The round resolves the trusted occupancy-probe split set {2,3,4,5} as an +// IN-BLOCK split-K = 2 ("waves per block" is the tuned limiter): workgroup +// 64 -> 128 (two 64-lane wavefronts), grid stays 384, each wave computes the +// SAME 16x16 tile over a disjoint K half (512 K = 4 stages, stage-aligned), +// so waves/CU doubles to ~6.4 (1.6 waves/SIMD) and the two co-resident waves +// of a block hide each other's B waits. A/B HBM bytes are UNCHANGED (each +// byte still read exactly once -- disjoint K slices, no repeated reads), the +// packed [n_tile][k_step_pair][lane][16] layout is untouched, and the int32 +// partials are combined in LDS (1,024 B) with one real 2-wave s_barrier: +// integer addition is associative, so the int32 total (and therefore the bf16 +// output) stays BIT-IDENTICAL to iterations 1/4/5/6/7. LDS per block grows +// 16,640 -> 17,664 B (fits the measured >= 4-block/CU capacity with margin). +// The single-wave iteration-7 kernel remains defined below as the recorded +// accepted best (fallback reference only). +// +// Falsifiable prediction (iteration 8): if the B round-trip wait was the +// limiter and 2 co-resident waves per block hide it (1.6 waves/SIMD +// achieved), the per-wave K halves and median should fall to ~9-12 us with +// P90 tracking median (1.4-1.9x vs the 16.847 us official best); if the +// hardware cannot co-schedule the second wave on the same SIMD slot, median +// stays ~16-17 us within noise and the next round must attack the 4-way LDS +// bank conflicts (49,152 events) or the 4-MMAC dependency chain instead. + +#include +#include +#include + +#include + +namespace { + +// 256 threads = 4 gfx928 wavefronts of 64 lanes. +constexpr int kScalarBlockThreads = 256; +constexpr int kCopyBlockThreads = 256; + +// --------------------------------------------------------------------------- +// Exact-shape constants for the packed-B DUMMA specialization (iteration 5). +// kTargetN/kTargetK identify the minimax_tp8_o_proj_m16 weight (K=1024, +// N=6144) that gets the non-identity pack; the paired M=2 API shape shares +// this (N, K) weight and decodes the packed layout in the scalar fallback. +// --------------------------------------------------------------------------- +constexpr int kTargetN = 6144; +constexpr int kTargetK = 1024; +constexpr int kDummaBlockThreads = 64; // one gfx928 wavefront per block + +// gfx928 DUMMA INT8 tile geometry (m16n16k32). +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; + +// Iteration-5/6 packed-B layout constants (depend on the tile constants +// above). Bytes of one (n_tile, k_step) B fragment: 16 cols x 32 k rows = +// 512 B, stored as 64 lanes x 8 B. Iteration 6 pairs two consecutive +// k_steps per lane: [n_tile][k_step_pair][lane][16] with 16 B per lane (= +// the 8-B fragment slots of k_step 2p and 2p+1), so one aligned 16-byte +// global_load_dwordx4 per lane per pair fetches two whole fragments +// (1 KiB contiguous per wave per instruction = 8 sectors). +constexpr int kPackedFragBytes = kDummaTileN * kDummaTileK; // 512 +constexpr int kPackedPairBytes = 2 * kPackedFragBytes; // 1024: one pair slot +constexpr int kPackedLaneBytes = kPackedPairBytes / kDummaBlockThreads; // 16 +// K=128 stage = 4 DUMMA steps = 2 k_step pairs: 2 B global dwordx4 (2 KiB) +// + 4 A ds_read2 from LDS per stage refill, 4 v_mmac per stage. Iteration 7: +// the loop keeps THREE stage register sets (current / next / next-next); one +// refill (2 KiB of B) is in flight per wave in steady state with the prologue +// issuing 3 refills (6 KiB) back-to-back. +constexpr int kStageK = 128; +constexpr int kStageSteps = kStageK / kDummaTileK; // 4 + +// Iteration 8 (occupancy round): in-block split-K = 2. Each block runs TWO +// 64-lane wavefronts (workgroup 128) that split the K=1024 range in half +// (512 K = 4 stages per wave); int32 partials are combined in LDS. The pack +// layout and kDummaBlockThreads (= 64 lanes per wavefront) are unchanged. +constexpr int kDummaSplitK = 2; // in-block K split (waves per block) +constexpr int kDummaBlockThreads2 = kDummaBlockThreads * kDummaSplitK; // 128 +constexpr int kSplitWaveK = kTargetK / kDummaSplitK; // 512 +constexpr int kSplitStages = kSplitWaveK / kStageK; // 4 + +// --------------------------------------------------------------------------- +// Packed-B decode for the generic scalar fallback: returns logical +// weight[kk, col] from the packed [n_tile][k_step_pair][lane][16] layout +// produced by w8a8_pack_o_proj_m16_kernel (see the pack kernel for the exact +// permutation). Only consulted when (n, k) == (kTargetN, kTargetK), i.e. the +// paired M=2 API shape shares the packed weight buffer with the M=16 path. +// --------------------------------------------------------------------------- +__device__ __forceinline__ int8_t +packed_b_element(const int8_t* __restrict__ packed, int col, int kk) { + const int n_tile = col >> 4; + const int nn = col & 15; + const int k_step = kk >> 5; + const int kk8 = kk & 31; + const int lane = (kk8 >> 3) * 16 + nn; + const int i = kk8 & 7; + const int pair = k_step >> 1; + const int half = k_step & 1; // 0: low 8 B of the lane slot, 1: high 8 B + const int k_pairs = (kTargetK / kDummaTileK) >> 1; // 16 + return packed[((static_cast(n_tile) * k_pairs + pair) * + kPackedPairBytes + + static_cast(lane) * kPackedLaneBytes + half * 8 + + i)]; +} + +// --------------------------------------------------------------------------- +// Scalar INT8 GEMM kernel (generic fallback for every (m, n, k)). +// +// One thread per output element. Adjacent lanes map to adjacent N columns +// (fastest-changing dimension), giving coalesced B and out accesses. For the +// exact (n, k) == (6144, 1024) pair the weight tensor is the packed +// [n_tile][k_step_pair][lane][16] layout (o_proj M=16 pack), decoded +// elementwise. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarBlockThreads) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, // [M, K] row-major int8 + const int8_t* __restrict__ b, // [K, N] (packed for 6144x1024) + const float* __restrict__ x_scale, // [M, 1] + const float* __restrict__ weight_scale, // [N, 1] + hip_bfloat16* __restrict__ out, // [M, N] bf16 + int n, + int k) { + const int row = static_cast(blockIdx.y); + const int col = static_cast(blockIdx.x) * kScalarBlockThreads + + static_cast(threadIdx.x); + if (col >= n) { + return; + } + + const int8_t* __restrict__ a_row = a + static_cast(row) * k; + const bool packed_b = (n == kTargetN && k == kTargetK); + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + const int8_t b_val = packed_b ? packed_b_element(b, col, kk) + : b[static_cast(kk) * n + col]; + acc += static_cast(a_row[kk]) * static_cast(b_val); + } + + // Exact reference order: (float(dot) * x_scale[row]) * weight_scale[col]. + // For every supported K the int32 dot is exact and |dot| < 2^24, so the + // float32 conversion is exact and both multiplications round identically + // to the trusted int64 reference (0 mismatch expected). + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// M=16 INT8 DUMMA m16n16k32 one-wave kernel (accepted best, iteration 1). +// +// One 64-lane wavefront per block computes one 16x16 output tile with +// explicit int32 accumulation: +// * du_fill_fragment(acc, 0), then k/32 du_mma_sync steps; +// * A/B fragments loaded directly from the row-major global tensors +// (lda = k, ldb = n); identity packed weight = raw [K, N] layout; +// * no LDS, no __syncthreads (single wave, no cross-wave barrier); +// * direct accumulator-fragment epilogue using the verified gfx928 INT8 +// m16n16k32 ownership: row = lane & 15, col_mod4 = lane >> 4, +// acc.x[i] maps to column col_mod4 + 4*i; +// * exact reference scaling order float(acc) * x_scale[row] * +// weight_scale[col], bf16 RN store. +// +// Superseded for the exact (16,6144,1024) shape by the A-only-staged kernel +// below (iteration 4); kept unchanged as the recorded accepted best. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(64) void w8a8_dumma_m16n16k32_onewave_kernel( + const int8_t* __restrict__ a, // [M, K] row-major int8 + const int8_t* __restrict__ b, // [K, N] row-major int8 (identity) + const float* __restrict__ x_scale, // [M, 1] + const float* __restrict__ weight_scale, // [N, 1] + hip_bfloat16* __restrict__ out, // [M, N] bf16 + int n, + int k) { + const int lane = static_cast(threadIdx.x); + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + for (int k0 = 0; k0 < k; k0 += kDummaTileK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); + du::dumma::du_load_matrix_sync(b_frag, b + static_cast(k0) * n + n0, + n); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + const int row = lane & 15; + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int out_col = n0 + col_mod4 + 4 * i; + const float scaled = static_cast(acc_frag.x[i]) * x_scale[row] * + weight_scale[out_col]; + out[static_cast(row) * n + out_col] = + __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// M=16 INT8 DUMMA m16n16k32 A-only-staged kernel (iteration 4). +// +// Identical geometry and math to the accepted iteration-1 one-wave kernel: +// grid = 384 blocks x 64 lanes (>= 120 CUs), one 16x16 output tile per wave, +// no split-K, direct accumulator-fragment epilogue, B loaded directly from +// the row-major identity-packed global tensor. The single pipeline change is +// A-only staging: the whole A tile (16 x 1024 = 16 KiB) is written ONCE per +// block into LDS by the prologue with 16 fully-coalesced dwordx4 loads per +// lane (64 lanes x 16 B = 1 KiB contiguous per load instruction), stored at +// a padded row stride of 1040 B (1024 + 16) for bank skew and 8-B alignment; +// the K loop then reads the A fragment from LDS via du_load_matrix_sync +// (same (row, k) -> byte mapping as the proven global path, stride 1040). +// Per k-step this removes 8 scattered global_load_ubyte (A), ~8 incremental +// s_waitcnt vmcnt waits, ~32 fragment-reassembly VALU and ~20 A-address VALU +// (all A addresses become loop-invariant LDS offsets), i.e. roughly half the +// loop-body instructions and half the wait chain, and it removes the L1/L2 +// sector pressure of 8 loads that each touch 16 rows x 4 k-chunks. A staged +// bytes are reused 16x per MMA across the tile's 16 output columns and are +// L2-shared across all 384 blocks; B bytes remain the cold once-read stream +// (each B[k][n] read from global once, reused 16x inside the tensor core). +// Single wave -> the one __syncthreads is a self-barrier that only orders +// this wave's own LDS writes before its LDS reads (no cross-wave barrier). +// Exact int32 accumulation order is unchanged (single accumulator, k +// ascending), so the bf16 output is bit-identical to iteration 1. +// +// Superseded for the exact (16,6144,1024) shape by the packed-B kernel below +// (iteration 5); kept unchanged as the recorded accepted best. +// --------------------------------------------------------------------------- +constexpr int kAStageStride = 1040; // K = 1024 + 16 pad: 8-B aligned + skew +// Iteration 16 (repair 1): the in-block split-K=2 kernel +// (w8a8_dumma_m16n16k32_packedb_sk2_kernel) stages each wave's K-HALF of A +// only (wave w consumes k in [w*512, (w+1)*512)) into its OWN half-size A LDS +// tile (s_a4[wave]): 16 rows x 33 int4 = 512 cols + 16 B pad = 528 B row +// stride (8-B aligned; 33 dwords mod 32 = 1 -> the 16 rows and 4 col groups +// still land on distinct LDS bank phases, same skew property as 1040). Each +// tile is 8,448 B; two tiles (one per wave) are 16,896 B/block, so A LDS per +// block is 16,896 B (vs the accepted 16,640-B single full tile) and total LDS +// 18,048 B -> still 3 blocks/CU. Repair note: the iteration-16 draft used ONE +// shared 8,448-B tile for both waves, which cannot hold both K-halves +// concurrently (its staging loop also covered only 8 of 16 rows per wave) -> +// wrong A data for 8 rows per wave; the repair gives each wave its own tile. +constexpr int kSk2AStageStride = 528; // K-half = 512 + 16 pad + +__global__ __launch_bounds__(64) void w8a8_dumma_m16n16k32_astage_kernel( + const int8_t* __restrict__ a, // [M, K] row-major int8 + const int8_t* __restrict__ b, // [K, N] row-major int8 (identity) + const float* __restrict__ x_scale, // [M, 1] + const float* __restrict__ weight_scale, // [N, 1] + hip_bfloat16* __restrict__ out, // [M, N] bf16 + int n, + int k) { + // 16,640 B LDS per block: 16 rows x 1040 B padded stride. + __shared__ int4 s_a4[(kDummaTileM * kAStageStride) / 16]; + + const int lane = static_cast(threadIdx.x); + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + // Prologue: stage the full A tile (16,384 B = 1024 x int4 chunks). Lane l + // owns chunk g = i*64 + l, so each of the 16 load instructions covers 1 KiB + // of contiguous bytes (fully coalesced) and every 16-B chunk maps to its + // padded LDS slot exactly once. Rows land 65 int4 (1040 B) apart, which is + // 8-B aligned and spreads the 16 rows across distinct LDS bank phases. + { + const int4* __restrict__ a4 = reinterpret_cast(a); +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int g = i * 64 + lane; // 16-B chunk index in [0, 1024) + s_a4[(g >> 6) * 65 + (g & 63)] = a4[g]; + } + } + // Single wave: self-barrier ordering this wave's own LDS writes before its + // LDS reads; no cross-wave dependency, no grid sync. + __syncthreads(); + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + const int8_t* __restrict__ s_a = + reinterpret_cast(s_a4); // padded row-major A in LDS + for (int k0 = 0; k0 < k; k0 += kDummaTileK) { + du::dumma::du_load_matrix_sync(a_frag, s_a + k0, kAStageStride); + du::dumma::du_load_matrix_sync(b_frag, b + static_cast(k0) * n + n0, + n); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + const int row = lane & 15; + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int out_col = n0 + col_mod4 + 4 * i; + const float scaled = static_cast(acc_frag.x[i]) * x_scale[row] * + weight_scale[out_col]; + out[static_cast(row) * n + out_col] = + __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// M=16 INT8 DUMMA m16n16k32 packed-B kernel (iterations 5/6). +// +// HIP-only packed-weight round: the SAME accepted geometry as iteration 4 +// (grid = 384 blocks x 64 lanes >= 120 CUs, one 16x16 output tile per wave, +// no split-K, A staged once into LDS at padded stride 1040, direct +// accumulator-fragment epilogue) but the B operand now comes from a packed +// weight layout instead of 8 scattered global ubyte loads per k-step. +// +// Iteration-6 packed layout [n_tile][k_step_pair][lane][16] (see +// w8a8_pack_o_proj_m16_kernel): the 16x32 B fragments of the two consecutive +// k_steps (2p, 2p+1) of a pair are stored as 64 lanes x 16 B with lane's 16 B +// = fragment bytes of k_step 2p (low 8 B) followed by k_step 2p+1 (high 8 B), +// each fragment byte-for-byte the du_mma.hpp matrix_b row_major loader order +// (lane's 8 B = logical B[k_step*32 + col + i, n_tile*16 + row] for i = +// 0..7, row = lane & 15, col = (lane >> 4) << 3), so one aligned 16-byte +// (dwordx4) load per lane per pair fills TWO consecutive b_frag slots +// exactly (1 KiB contiguous per wave per instruction = 8 sectors; 2 pairs = +// 2 KiB contiguous per stage). The A fragment is read from the staged LDS +// tile with the same u64 transport (matrix_a row_major: 8 contiguous bytes +// at row*1040 + col, identical to the ds_read2_b32 the iteration-4 library +// loader compiled to). +// +// K-loop transport (iteration 7, pipeline round): the loop is a lagged +// THREE-set register pipeline with 0 barriers per k-step. Prologue fills +// stages 0/1/2 back-to-back (3 cold-start DRAM latencies, 6 KiB in flight, +// overlapped with the A-staging prologue); each loop iteration bursts stage +// s from the current set, rotates current<-next<-next-next (consuming the +// stage loaded a FULL iteration earlier, so the compiler's s_waitcnt +// vmcnt(0)/lgkmcnt(0) lands after the burst that precedes the rotate), then +// refills the freed third set with stage s+3's six loads (2 B global +// dwordx4 + 4 A ds_read2 = 2 KiB) whose first read is the NEXT iteration's +// rotate. Every B dwordx4 therefore stays in flight for one whole loop +// iteration (~4x the in-flight time of iteration 6's compiled depth-1, where +// the wait landed right after the same-iteration burst), and B vmem +// instructions per replay drop from 6,912 (iteration 6) to 6,144 (384 x 16 +// dwordx4 = 3 prologue + 5 in-loop refills). A staged bytes are +// unchanged (16 KiB per block, read once from HBM, 16x reused per MMA), B +// HBM bytes unchanged (6 MiB, every byte read exactly once). Single wave -> +// the one __syncthreads is a self-barrier (compiles to s_waitcnt, zero +// s_barrier instructions in the kernel). Exact int32 accumulation is +// unchanged (single accumulator, k ascending), so the bf16 output is +// bit-identical to iterations 1, 4, 5 and 6. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void load_stage_fragments_packedb( + const int8_t* __restrict__ s_a, // staged A in LDS (wave-local K-half + // stride for the sk2 kernel, full-K + // stride 1040 for the one-wave kernel) + const int8_t* __restrict__ bpacked, // packed [n_tile][pair][lane][16] + int lane, + int a_off, + int64_t b_base, + int s, // GLOBAL stage index (0..7): selects the wave's disjoint + // K-half for B (pairs 2s, 2s+1) + int a_stage, // A stage index within the wave's staged K-half (0..3 for + // the sk2 kernel; == s for the full-K one-wave kernel) + uint64_t& a0, uint64_t& a1, uint64_t& a2, uint64_t& a3, + uint64_t& b0, uint64_t& b1, uint64_t& b2, uint64_t& b3) { + const int64_t k0 = static_cast(s) * kStageK; + const int64_t ak0 = static_cast(a_stage) * kStageK; + a0 = *reinterpret_cast(s_a + a_off + ak0); + a1 = *reinterpret_cast(s_a + a_off + ak0 + kDummaTileK); + a2 = *reinterpret_cast(s_a + a_off + ak0 + 2 * kDummaTileK); + a3 = *reinterpret_cast(s_a + a_off + ak0 + 3 * kDummaTileK); + // B: one aligned 16-B (dwordx4) cooperative load per k_step pair; the + // pair index covers steps 4s..4s+3 as pairs (2s, 2s+1). q.x = the 8-B + // fragment slot of the even step, q.y = the odd step (little-endian u64, + // same byte order du_mma_sync consumes). + const int64_t pair = (k0 / kDummaTileK) >> 1; + const int64_t b_off = b_base + pair * kPackedPairBytes + + static_cast(lane) * kPackedLaneBytes; + const ulonglong2 q0 = *reinterpret_cast(bpacked + b_off); + const ulonglong2 q1 = + *reinterpret_cast(bpacked + b_off + kPackedPairBytes); + b0 = q0.x; // k_step 4s + b1 = q0.y; // k_step 4s+1 + b2 = q1.x; // k_step 4s+2 + b3 = q1.y; // k_step 4s+3 +} + +__global__ __launch_bounds__(kDummaBlockThreads) void +w8a8_dumma_m16n16k32_packedb_kernel( + const int8_t* __restrict__ a, // [M, K] row-major int8 + const int8_t* __restrict__ b, // packed [n_tile][pair][lane][16] + const float* __restrict__ x_scale, // [M, 1] + const float* __restrict__ weight_scale, // [N, 1] + hip_bfloat16* __restrict__ out, // [M, N] bf16 + int n, + int k) { + // 16,640 B LDS per block: 16 rows x 1040 B padded stride (as iteration 4). + __shared__ int4 s_a4[(kDummaTileM * kAStageStride) / 16]; + + const int lane = static_cast(threadIdx.x); + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + // Prologue: stage the full A tile (16,384 B = 1024 x int4 chunks). Lane l + // owns chunk g = i*64 + l, so each of the 16 load instructions covers 1 KiB + // of contiguous bytes (fully coalesced); rows land 65 int4 (1040 B) apart. + { + const int4* __restrict__ a4 = reinterpret_cast(a); +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int g = i * 64 + lane; // 16-B chunk index in [0, 1024) + s_a4[(g >> 6) * 65 + (g & 63)] = a4[g]; + } + } + // Single wave: self-barrier ordering this wave's own LDS writes before its + // LDS reads; no cross-wave dependency, no grid sync. + __syncthreads(); + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Per-lane A ownership is loop-invariant: row = lane & 15, + // col = (lane >> 4) << 3, x[i] = s_a[row*1040 + col + i] (du_mma.hpp + // matrix_a row_major). Per-lane B slot offset is loop-invariant too: + // b_base + pair*1024 + lane*16. + const int8_t* __restrict__ s_a = reinterpret_cast(s_a4); + const int a_row = lane & 15; + const int a_col = (lane >> 4) << 3; + const int a_off = a_row * kAStageStride + a_col; + const int64_t b_base = static_cast(blockIdx.x) * + (k / kDummaTileK / 2) * kPackedPairBytes; + + // Iteration-7 lagged 3-set pipeline. Prologue: stages 0, 1 and 2 land + // directly in registers (three cold-start DRAM latencies per block, issued + // back-to-back = 6 KiB in flight, overlapped with the A-staging prologue). + const int n_stages = k / kStageK; // 8 for K = 1024 + uint64_t ca0, ca1, ca2, ca3, cb0, cb1, cb2, cb3; // stage s + uint64_t na0, na1, na2, na3, nb0, nb1, nb2, nb3; // stage s+1 + uint64_t ma0, ma1, ma2, ma3, mb0, mb1, mb2, mb3; // stage s+2 + load_stage_fragments_packedb(s_a, b, lane, a_off, b_base, 0, 0, ca0, ca1, + ca2, ca3, cb0, cb1, cb2, cb3); + load_stage_fragments_packedb(s_a, b, lane, a_off, b_base, 1, 1, na0, na1, + na2, na3, nb0, nb1, nb2, nb3); + load_stage_fragments_packedb(s_a, b, lane, a_off, b_base, 2, 2, ma0, ma1, + ma2, ma3, mb0, mb1, mb2, mb3); + + for (int s = 0; s < n_stages; ++s) { + // Stage-s MMAC burst straight from registers, in the exact a_frag.x / + // b_frag.x byte order du_mma_sync consumes (little-endian u64 slot = + // bytes 0..7 of the fragment). +#pragma unroll + for (int t = 0; t < kStageSteps; ++t) { + const uint64_t av = + (t == 0) ? ca0 : (t == 1) ? ca1 : (t == 2) ? ca2 : ca3; + const uint64_t bv = + (t == 0) ? cb0 : (t == 1) ? cb1 : (t == 2) ? cb2 : cb3; + __builtin_memcpy(a_frag.x, &av, 8); + __builtin_memcpy(b_frag.x, &bv, 8); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + // Lagged rotation: stage s+1 -> current, stage s+2 -> next. The ma/mb + // set (stage s+2) was loaded a FULL iteration earlier, so the compiler's + // s_waitcnt vmcnt(0)/lgkmcnt(0) for it lands at THIS rotate -- after the + // burst, one whole loop iteration of in-flight time per B dwordx4 (vs + // ~1 MMAC burst in iteration 6's compiled depth-1). + ca0 = na0; + ca1 = na1; + ca2 = na2; + ca3 = na3; + cb0 = nb0; + cb1 = nb1; + cb2 = nb2; + cb3 = nb3; + na0 = ma0; + na1 = ma1; + na2 = ma2; + na3 = ma3; + nb0 = mb0; + nb1 = mb1; + nb2 = mb2; + nb3 = mb3; + // Refill the just-freed third set with stage s+3's six loads (2 B global + // dwordx4 + 4 A ds_read2 = 2 KiB in flight). The rotate above reads + // ma/mb (WAR), so the compiler cannot hoist these loads above it; their + // first read is the NEXT iteration's rotate, one full iteration after + // issue. The guard keeps the tail stages from reading past the packed B + // buffer (stage 8+ never loads; stale ma/mb at s=6/7 are never consumed). + if (s + 3 < n_stages) { + load_stage_fragments_packedb(s_a, b, lane, a_off, b_base, s + 3, s + 3, + ma0, ma1, ma2, ma3, mb0, mb1, mb2, mb3); + } + } + + // Direct accumulator-fragment epilogue, unchanged from iterations 1 and 4 + // (gfx928 INT8 m16n16k32 ownership: row = lane & 15, col_mod4 = lane >> 4, + // acc.x[i] maps to column col_mod4 + 4*i). + const int row = lane & 15; + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int out_col = n0 + col_mod4 + 4 * i; + const float scaled = static_cast(acc_frag.x[i]) * x_scale[row] * + weight_scale[out_col]; + out[static_cast(row) * n + out_col] = + __float2bfloat16(scaled); + } +} + +// --------------------------------------------------------------------------- +// M=16 INT8 DUMMA m16n16k32 packed-B IN-BLOCK split-K=2 kernel (iteration 8, +// occupancy round). Same packed [n_tile][k_step_pair][lane][16] layout, same +// 16x16 tile geometry, same iteration-7 lagged 3-set pipeline, same A-only +// LDS staging (16 x 1040 B) -- but the workgroup is now TWO 64-lane +// wavefronts (128 threads) that split K in half: wave w computes the same +// tile over k in [w*512, w*512+512) (4 K=128 stages each, disjoint B pairs +// 8w..8w+7 via the global stage index s = wave*4 + sw). Waves per block 1->2 +// is the tuned occupancy limiter: grid stays 384 blocks, so waves/CU goes +// 3.2 -> ~6.4 (1.6 waves/SIMD vs the 88-VGPR 2-waves/SIMD capacity), giving +// the wave scheduler a second co-resident wavefront to switch to during each +// wave's exposed B HBM round trip. A/B HBM bytes are unchanged (disjoint K +// slices, every byte read exactly once -- no repeated reads for occupancy). +// Combine: wave 1 publishes its int32 accumulator (4 ints per lane = 1,024 B) +// to LDS, one real 2-wave s_barrier orders it, wave 0 adds it in int32 and +// runs the unchanged direct bf16 epilogue (wave 1 exits). Integer addition is +// associative, so the int32 total and the bf16 output stay bit-identical to +// iterations 1/4/5/6/7. LDS per block: 16,640 (A) + 1,024 (partials) = +// 17,664 B (fits the measured >= 4-block/CU capacity with margin). 2 real +// s_barrier per block (post-staging + pre-combine) -> 0.0625 per k-step. +// +// Iteration 13 (final conditional HIP-only round; control plane keeps +// plateau=false, raw_inline_asm_allowed=false, so ONE HIP-only consolidation +// change on this accepted kernel -- the file is restored to the accepted +// iteration-8 digest f8b708... after the iteration-12 barrier-removal variant +// regressed to 12.716 us). The accepted ISA's epilogue (after the pre-combine +// s_barrier) issues FIVE per-lane global_load_dword (1 x_scale + 4 +// weight_scale) with FOUR serialized s_waitcnt vmcnt(1)/vmcnt(0) windows -- +// each column's weight_scale load is issued only after the previous column's +// store, then waited: ~4 x L2 latency of serial global round trips on the +// block's drain tail. Iteration 13 prefetches the scales into LDS: wave 0 +// stores its 16-row x_scale slice + this block's 16-column weight_scale slice +// (n0..n0+15, 128 B) right after the A-staging loop, before the staging +// barrier; the two scale loads issue at kernel top (waits resolve during the +// barrier/B-wait) and the epilogue reads conflict-free LDS broadcasts with +// zero global latency. Same float values, same scaling order, +128 B LDS +// (17,792 B/block; 3 blocks/CU = 53,376 B <= 64 KiB unchanged), no new +// barrier, VGPR unchanged (~80) -> bf16 output BIT-IDENTICAL. +// +// Iteration 16, repair 1 (final conditional HIP-only round; control plane +// keeps plateau=false, raw_inline_asm_allowed=false: recent valid candidates +// are +6.80% it13 accepted, -4.05% it14 balanced-epilogue REGRESSED +// (reverted), -5.63% it15 epilogue-LDS-hoist REGRESSED (reverted); the +// epilogue axis is closed, so this round makes ONE HIP-only consolidation +// change on the untouched structural axis). The round proposes wave-private +// A-side staging: each wave stages ONLY its own K-half into a half-size +// 528-B-stride LDS tile instead of the full 16,640-B tile staged by both +// waves. Repair 1 corrects the draft's staging: the single shared 8,448-B +// tile cannot hold both waves' K-halves concurrently, and the draft loop +// covered only 8 of 16 rows per wave (per-thread loads 8 -> 4 halves the +// staged bytes, not the tile) -> wrong A fragments for half the rows of each +// wave's partial. The repaired kernel gives EACH wave its own 8,448-B half +// tile (s_a4[wave][528], 16,896 B/block A LDS; total LDS 16,896 + 1,024 + +// 128 = 18,048 B -> 3 blocks/CU unchanged) and stages each wave's FULL +// 512-chunk K-half (8 dwordx4 per thread, same per-thread count and +// coalescing as the accepted cooperative staging; every A byte read exactly +// once). The two s_barrier per block are kept, A fragment values, per-wave K +// halves, int32 accumulation order and the scaling math are UNCHANGED -> +// bf16 output stays BIT-IDENTICAL to iterations 8/13. Grid 384 x workgroup +// 128 unchanged; A/B HBM bytes unchanged; packed +// [n_tile][k_step_pair][lane][16] layout untouched; shape guards and scalar +// fallback untouched. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDummaBlockThreads2) void +w8a8_dumma_m16n16k32_packedb_sk2_kernel( + const int8_t* __restrict__ a, // [M, K] row-major int8 + const int8_t* __restrict__ b, // packed [n_tile][pair][lane][16] + const float* __restrict__ x_scale, // [M, 1] + const float* __restrict__ weight_scale, // [N, 1] + hip_bfloat16* __restrict__ out, // [M, N] bf16 + int n, + int k) { + // 18,048 B LDS per block: TWO per-wave 16 x 528 B padded A-stride K-half + // tiles (iteration 16 repair 1: each wave stages ONLY its own K-half into + // its own 8,448-B tile, s_a4[0]/s_a4[1]) + 1,024 B int32 partial buffer + // (16-B aligned int4 x 64 lanes, wave 1 -> wave 0) + 128 B epilogue scale + // cache (iteration 13: wave 0's 16-row x_scale slice + this block's + // 16-column weight_scale slice). A LDS per block is 16,896 B; total LDS + // 18,048 B keeps co-residency at 3 blocks/CU (18,048 x 3 = 54,144 B <= 64 + // KiB; 4 x 18,048 = 72,192 B > 64 KiB), same as the accepted 17,792-B + // layout. A single shared 8,448-B tile (iteration-16 draft) could not hold + // both waves' K-halves concurrently, so the repair allocates one tile per + // wave. + __shared__ int4 s_a4[2][(kDummaTileM * kSk2AStageStride) / 16]; + __shared__ int4 s_partial4[kDummaBlockThreads]; + __shared__ float s_xscale[kDummaTileM]; // x_scale[0..15] (64 B) + __shared__ float s_wscale[kDummaTileN]; // weight_scale[n0..n0+15] (64 B) + + const int tid = static_cast(threadIdx.x); // 0..127 + const int wave = tid >> 6; // 0..1: in-block K split + const int lane = tid & 63; + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + // Prologue (iteration 16, repair 1): each wave stages ONLY its own K-half + // into ITS OWN half-size tile s_a4[wave]. Wave w owns global chunk + // g = r*64 + (w*32) + c32 for r in 0..15, c32 in 0..31 (k in + // [w*512, (w+1)*512) -- all 512 chunks of the wave's half), mapped to its + // tile's slot r*33+c32. Per-thread staging loads stay 8 dwordx4 (one full + // K-half per wave; the draft's "8 -> 4" halved the staged bytes, not the + // tile, and left 8 rows per wave unstaged); each instruction covers the + // wave's 512-B segment of 2 rows (1,024 B contiguous, fully coalesced -- + // the two waves together cover the accepted loop's 2-KiB contiguous + // segment) and every 16-B chunk maps to its padded LDS slot exactly once. + // Rows land 33 int4 (528 B) apart (8-B aligned, 33 mod 32 = 1 -> the 16 + // rows and 4 col groups sit on distinct LDS bank phases, same skew class + // as the 1040-B stride). Wave w later reads ONLY slots it wrote in its own + // tile, so the staging barrier still orders each wave's own writes before + // its own reads (barrier kept: the accepted ISA's waitcnt placement stays + // unchanged). + { + const int4* __restrict__ a4 = reinterpret_cast(a); +#pragma unroll + for (int i = 0; i < kSplitWaveK / 64; ++i) { // 8 iterations + const int flat = i * 64 + lane; // 16-B chunk in [0, 512) of this half + const int r = flat >> 5; // M row 0..15 (32 chunks per row half) + const int c32 = flat & 31; // 16-B chunk within the row's K-half + const int g = r * 64 + (wave << 5) + c32; // global chunk in [0, 1024) + s_a4[wave][r * 33 + c32] = a4[g]; + } + } + // Iteration 13 (epilogue scale prefetch): wave 0 caches its 16-row x_scale + // slice and this block's 16-column weight_scale slice (n0..n0+15) in LDS + // right after the A staging, before the staging barrier. The accepted ISA + // shows the epilogue issuing FIVE per-lane global_load_dword (1 x_scale + + // 4 weight_scale) AFTER the pre-combine s_barrier with FOUR serialized + // s_waitcnt vmcnt(1)/vmcnt(0) windows (each column's weight_scale load is + // issued only after the previous column's store, then waited) -- ~4 x L2 + // latency of serial global round trips on the block's drain tail. Anchored + // by these LDS stores, the two scale loads issue here at kernel top (their + // vmcnt waits resolve during the barrier/B-wait) and the epilogue reads + // LDS broadcast slots (4-way broadcast, conflict-free) with zero global + // latency. Same float values, same scaling order + // (float(dot) * x_scale[row]) * weight_scale[col] -> bf16 output stays + // BIT-IDENTICAL. Wave 0 both writes and reads these slots (same-wave LDS + // ordering is the compiler's lgkmcnt waits; no new barrier). + if (wave == 0) { + if (lane < kDummaTileM) { + s_xscale[lane] = x_scale[lane]; + s_wscale[lane] = weight_scale[n0 + lane]; + } + } + // Real 2-wave barrier: both waves' staging writes are visible before any + // fragment read (first of two s_barrier per block). + __syncthreads(); + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + const int8_t* __restrict__ s_a = reinterpret_cast(s_a4[wave]); + const int a_row = lane & 15; + const int a_col = (lane >> 4) << 3; + const int a_off = a_row * kSk2AStageStride + a_col; + const int64_t b_base = static_cast(blockIdx.x) * + (k / kDummaTileK / 2) * kPackedPairBytes; + + // Iteration-7 lagged 3-set pipeline, kSplitStages = 4 stages per wave. The + // GLOBAL stage index s = wave*4 + sw picks the wave's disjoint K half: + // B pairs s*2/s*2+1 via s, and the A slice is read from the wave's OWN + // half-size LDS tile (s_a4[wave], 8,448 B) via the LOCAL stage index sw + // (a_stage) -- each wave reads only slots it staged itself. + const int n_stages = kSplitStages; // 4 + uint64_t ca0, ca1, ca2, ca3, cb0, cb1, cb2, cb3; // stage s + uint64_t na0, na1, na2, na3, nb0, nb1, nb2, nb3; // stage s+1 + uint64_t ma0, ma1, ma2, ma3, mb0, mb1, mb2, mb3; // stage s+2 + const int s0 = wave * n_stages; // global stage of the wave's first K-half + load_stage_fragments_packedb(s_a, b, lane, a_off, b_base, s0 + 0, 0, ca0, + ca1, ca2, ca3, cb0, cb1, cb2, cb3); + load_stage_fragments_packedb(s_a, b, lane, a_off, b_base, s0 + 1, 1, na0, + na1, na2, na3, nb0, nb1, nb2, nb3); + load_stage_fragments_packedb(s_a, b, lane, a_off, b_base, s0 + 2, 2, ma0, + ma1, ma2, ma3, mb0, mb1, mb2, mb3); + + for (int sw = 0; sw < n_stages; ++sw) { + const int s = wave * n_stages + sw; + // Stage-s MMAC burst straight from registers, in the exact a_frag.x / + // b_frag.x byte order du_mma_sync consumes (little-endian u64 slot = + // bytes 0..7 of the fragment). +#pragma unroll + for (int t = 0; t < kStageSteps; ++t) { + const uint64_t av = + (t == 0) ? ca0 : (t == 1) ? ca1 : (t == 2) ? ca2 : ca3; + const uint64_t bv = + (t == 0) ? cb0 : (t == 1) ? cb1 : (t == 2) ? cb2 : cb3; + __builtin_memcpy(a_frag.x, &av, 8); + __builtin_memcpy(b_frag.x, &bv, 8); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + // Lagged rotation: stage s+1 -> current, stage s+2 -> next. The ma/mb + // set (stage s+2) was loaded a FULL iteration earlier, so the compiler's + // s_waitcnt vmcnt(0)/lgkmcnt(0) for it lands at THIS rotate -- one whole + // loop iteration of in-flight time per B dwordx4 (unchanged from + // iteration 7). + ca0 = na0; + ca1 = na1; + ca2 = na2; + ca3 = na3; + cb0 = nb0; + cb1 = nb1; + cb2 = nb2; + cb3 = nb3; + na0 = ma0; + na1 = ma1; + na2 = ma2; + na3 = ma3; + nb0 = mb0; + nb1 = mb1; + nb2 = mb2; + nb3 = mb3; + // Refill the freed third set with stage s+3's six loads (2 B global + // dwordx4 + 4 A ds_read2 = 2 KiB in flight). Only sw = 0 refills within + // the wave's 4 stages (global stage wave*4 + 3); the guard keeps tail + // stages from reading past the wave's packed-B half. + if (sw + 3 < n_stages) { + // a_stage = sw + 3 stays within the wave's own 4-stage K-half (the + // guard bounds it to 3), so the A reads never leave the 8,448-B tile. + load_stage_fragments_packedb(s_a, b, lane, a_off, b_base, s + 3, sw + 3, + ma0, ma1, ma2, ma3, mb0, mb1, mb2, mb3); + } + } + + // In-block split-K combine. Wave 1 publishes its int32 accumulator (one + // aligned 16-B int4 per lane); the real 2-wave barrier orders the publish + // before wave 0's read; wave 0 adds in int32 (exact: integer addition is + // associative, so the total is bit-identical to the single-pass k-ascending + // order) and runs the unchanged direct epilogue; wave 1 exits. + if (wave == 1) { + s_partial4[lane] = + make_int4(acc_frag.x[0], acc_frag.x[1], acc_frag.x[2], acc_frag.x[3]); + } + __syncthreads(); // second (and last) real s_barrier per block + if (wave == 0) { + const int4 p = s_partial4[lane]; + const int total[4] = {acc_frag.x[0] + p.x, acc_frag.x[1] + p.y, + acc_frag.x[2] + p.z, acc_frag.x[3] + p.w}; + const int row = lane & 15; + const int col_mod4 = lane >> 4; + // Iteration 13: scales come from the LDS cache staged at kernel top + // (s_xscale / s_wscale) instead of five serialized global loads; the + // scaling math and order are unchanged. + const float xs = s_xscale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int out_col = n0 + col_mod4 + 4 * i; + const float scaled = static_cast(total[i]) * xs * + s_wscale[col_mod4 + 4 * i]; + out[static_cast(row) * n + out_col] = + __float2bfloat16(scaled); + } + } +} + +// --------------------------------------------------------------------------- +// One-time pack for the exact (k, n) == (1024, 6144) o_proj weight: permute +// the logical [K, N] int8 layout into packed[n_tile][k_step_pair][lane][16] +// B fragment slots (iteration 6: two consecutive k_steps per lane, see the +// packed-B kernel comment for the exact du_mma.hpp matrix_b row_major lane +// mapping: row = lane & 0xf, col = (lane >> 4) << 3, x[i] = p[(col + i) * +// ldm + row]). One thread per (n_tile, k_step_pair, lane): 16 pairs x 384 +// n_tiles x 64 lanes = 393,216 threads, each writing one aligned 16-byte +// slot. Runs once per weight, outside the timed GEMM and outside Graph +// capture; byte count is unchanged (6 MiB), so the packed tensor is a +// same-size permutation. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kCopyBlockThreads) void +w8a8_pack_o_proj_m16_kernel(const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + const int k_pairs = (static_cast(k) / kDummaTileK) >> 1; // 16 + const int64_t slots_per_tile = + static_cast(k_pairs) * kDummaBlockThreads; // 1024 + const int64_t total = slots_per_tile * (n / kDummaTileN); + if (idx >= total) { + return; + } + const int n_tile = static_cast(idx / slots_per_tile); + const int64_t rem = idx % slots_per_tile; + const int pair = static_cast(rem / kDummaBlockThreads); + const int lane = static_cast(rem % kDummaBlockThreads); + const int row = lane & 15; + const int col = (lane >> 4) << 3; + // v0 = 8-B fragment slot of k_step 2p (low 8 B of the 16-B lane slot), + // v1 = fragment slot of k_step 2p+1 (high 8 B). + uint64_t v0 = 0; + uint64_t v1 = 0; +#pragma unroll + for (int s = 0; s < 2; ++s) { + const int k_step = 2 * pair + s; + const int8_t* __restrict__ src = + raw + (static_cast(k_step) * kDummaTileK + col) * n + + static_cast(n_tile) * kDummaTileN + row; + uint64_t v = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) { + v |= (static_cast(static_cast(src[i * n]))) << (8 * i); + } + if (s == 0) { + v0 = v; + } else { + v1 = v; + } + } + ulonglong2* __restrict__ dst = reinterpret_cast( + packed + ((static_cast(n_tile) * k_pairs + pair) * + kPackedPairBytes + + static_cast(lane) * kPackedLaneBytes)); + ulonglong2 u; + u.x = v0; + u.y = v1; + *dst = u; +} + +// --------------------------------------------------------------------------- +// Identity device-to-device pack (generic fallback for every (K, N) except +// the exact (1024, 6144) o_proj weight, which uses the pack kernel above). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kCopyBlockThreads) void +w8a8_identity_copy_i8_kernel(const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t i = static_cast(blockIdx.x) * kCopyBlockThreads + + static_cast(threadIdx.x); + if (i < count) { + dst[i] = src[i]; + } +} + +__global__ __launch_bounds__(kCopyBlockThreads) void +w8a8_identity_copy_f32_kernel(const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t i = static_cast(blockIdx.x) * kCopyBlockThreads + + static_cast(threadIdx.x); + if (i < count) { + dst[i] = src[i]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Stable host launcher symbols consumed by csrc/bindings.cpp. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // The scalar bootstrap needs no workspace; the buffer is still required by + // the API contract so later split-K rounds can use it without ABI changes. + (void)workspace; + (void)workspace_bytes; + + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + + // Exact-shape specialization: minimax_tp8_o_proj_m16 (M=16, N=6144, K=1024) + // -> packed-B DUMMA m16n16k32 in-block split-K=2 (iteration 8): 384 blocks + // x 128 lanes (TWO 64-lane wavefronts per block), each wave computes the + // same 16x16 output tile over a disjoint K half (512 K = 4 stages), int32 + // partials combined in LDS, A tile staged once into LDS, B read from the + // packed [n_tile][k_step_pair][lane][16] fragment-slot layout (produced by + // launch_pack_w8a8_weight for this (k, n)). Waves/CU doubles from 3.2 to + // ~6.4 (1.6 waves/SIMD) with no repeated A/B HBM reads. All other + // (m, n, k) -- including the paired M=2 TP8 o_proj shape with the same + // (N, K) -- keep the scalar generic fallback below (which decodes the + // packed layout for that exact (n, k)). + if (m == 16 && n == kTargetN && k == kTargetK) { + const dim3 block(kDummaBlockThreads2); // 128 = 2 x 64-lane wavefronts + const dim3 grid(kTargetN / kDummaTileN); + hipLaunchKernelGGL(w8a8_dumma_m16n16k32_packedb_sk2_kernel, grid, block, 0, + stream, a, b, x_scale, weight_scale, + reinterpret_cast(out), n, k); + return; + } + + const dim3 block(kScalarBlockThreads); + const dim3 grid( + static_cast((n + kScalarBlockThreads - 1) / + kScalarBlockThreads), + static_cast(m)); + hipLaunchKernelGGL(w8a8_scalar_gemm_kernel, grid, block, 0, stream, a, b, + x_scale, weight_scale, + reinterpret_cast(out), n, k); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_count = static_cast(k) * n; + const dim3 block(kCopyBlockThreads); + if (k == kTargetK && n == kTargetN) { + // Exact o_proj M=16 weight: permute into [n_tile][k_step_pair][lane][16] + // B fragment slots (one-time, untimed, outside Graph capture). + const int64_t total = (static_cast(k) / kDummaTileK / 2) * + (n / kDummaTileN) * kDummaBlockThreads; + const dim3 grid(static_cast( + (total + kCopyBlockThreads - 1) / kCopyBlockThreads)); + hipLaunchKernelGGL(w8a8_pack_o_proj_m16_kernel, grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else if (weight_count > 0) { + const dim3 grid(static_cast( + (weight_count + kCopyBlockThreads - 1) / kCopyBlockThreads)); + hipLaunchKernelGGL(w8a8_identity_copy_i8_kernel, grid, block, 0, stream, + raw_weight, packed_weight, weight_count); + } + if (n > 0) { + const dim3 grid(static_cast( + (n + kCopyBlockThreads - 1) / kCopyBlockThreads)); + hipLaunchKernelGGL(w8a8_identity_copy_f32_kernel, grid, block, 0, stream, + weight_scale, packed_weight_scale, n); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj.hip new file mode 100644 index 00000000..bf0cbf03 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj.hip @@ -0,0 +1,585 @@ +// @@variant shape=minimax_tp8_qkv_proj_m16 commit=b86aaad82e2a39632cd692ec7fe88dcaeeadaaed added=2026-08-28 +// median_us=33.2 p90_us=33.31 +// source=minimax-dsh-tp8-m16-1-78402260 +// MetaInfer W8A8 INT8 GEMM — gfx928 (K500SM_AI) bootstrap. +// +// worker_0 owns physical GPU 0 and the exact shape +// minimax_tp8_qkv_proj_m16 : M=16, N=1280, K=6144 +// +// Iteration 0 (bootstrap, correctness-first): one simple scalar int8 +// dot-product kernel is used for every shape, plus an identity device-to- +// device `pack_weight`. No DUMMA, no split-K, no LDS staging, no inline +// assembly: the goal is a kernel that compiles on gfx928, is CUDA/HIP-Graph +// capture safe, and reproduces the exact int64 CPU reference bit-for-bit. +// +// Iteration 1 (DUMMA bootstrap): the exact assigned shape +// minimax_tp8_qkv_proj_m16 (M=16, N=1280, K=6144) ran the minimal +// DUMMA m16n16k32 INT8 kernel: one 64-lane wavefront per 16x16 output +// tile (80 one-wave blocks, no cross-wave barrier, no LDS staging, no +// split-K), direct global fragment loads and explicit int32 accumulation, +// then x_scale[m] * packed_weight_scale[n] scaling and bf16 store. +// Measured 206.9 us median (logical 1.216 TOPs, 38.7 GB/s): latency/ +// occupancy-bound, not bandwidth-bound. +// +// Iteration 2 (architecture round, grid parallelism): same minimal direct +// DUMMA m16n16k32 tile and one-wave zero-barrier blocks, grid raised from +// 80 blocks to 240 with split-K=3 (exactly 2 one-wave blocks per CU on all +// 120 CUs). Each block computes its K/3=2048 (64 k32 steps) int32 partial +// over a uniform K-slice with the same direct global fragment loads, stores +// it to a workspace plane, and a separate 80-block combine kernel sums the +// 3 planes ascending (exact int32), applies the scales and stores bf16. +// Measured floor improved (min 162 us) but the median exploded +// (416.7 us, p90 781.4): the direct byte-load path's per-step latency is +// already inflated ~1.27x at 2 waves/CU, so pure split-K occupancy past +// 2 blocks/CU has to be probed before any load-path micro-optimization. +// +// Iteration 3 (architecture round, split-K occupancy sweep completion): +// keep the identical direct-load DUMMA m16n16k32 partial + combine family, +// but make the split-K count a compile-time template over the two legal +// probes that remain unmeasured for this shape: +// SPLIT_K=2 (default, measured): grid = 160 one-wave zero-barrier +// blocks = a finer one-wave grid (80 -> 160) with the mildest +// occupancy step (40 CUs carry 2 waves, 80 CUs carry 1 wave) and +// the per-wave K-chain halved from 192 to 96 k32 steps; predicted +// to stay in the stable latency regime of iteration 1 while +// shortening every wave's serial load->mma chain. +// SPLIT_K=6 (compiled alternative, one constant away): grid = 480 +// one-wave blocks = exactly 4 blocks per CU (CU-aligned, +// non-power-of-two), per-wave K-chain = 32 k32 steps; the strongest +// CU-aligned occupancy probe in the trusted set {2,3,4,5,6,8,9,12} +// whose K=6144 chunk (1024) stays k32-aligned. +// Both configurations write int32 partials into the caller workspace plane +// (tile, s) and run the ascending-sum combine+scale+bf16 kernel in the +// timed Graph replay, so split-K combine cost is part of the operator wall. +// All other (m, n, k) keep the generic scalar fallback. +// +// Iteration 10 (final conditional inline-asm round; HIP-only consolidation +// per the control-plane policy: plateau=false, raw_inline_asm_allowed=false): +// the measured geometry is flipped from SPLIT_K=2 to SPLIT_K=6, completing +// the split-K occupancy sweep the control plane keeps mandating before any +// micro-optimization ("benchmark a finer one-wave zero-barrier grid or +// multiple legal split-K candidates including combine cost"). Iteration 3 +// (SPLIT_K=2: 160 blocks, 96 k32 steps/wave) is the accepted champion at +// 106.8 us; iteration 2 (SPLIT_K=3: 240 blocks, 64 steps/wave) was +// pathological (median 416.7 us); iteration 9's LDS staging on the S=2 +// geometry regressed to 115.5 us and thereby falsified 'byte loads are the +// per-step poison' (per its own acceptance criterion), pointing at the +// one-wave grid / per-CU occupancy as the remaining lever. SPLIT_K=6 is the +// strongest CU-aligned candidate in the trusted set {2,3,4,5,6,8,9,12} for +// this shape: grid = (N/16)*6 = 480 one-wave zero-barrier blocks = exactly +// 4 blocks per CU on all 120 CUs (non-power-of-two, CU-aligned), K-chunk = +// 1024 = 32 k32 steps per wave (k32-aligned), 6 int32 partial planes +// (491,520 B) inside the 16-plane / 1,310,720 B workspace capacity, and the +// 80-block ascending-sum combine stays in the timed Graph replay, so the +// full combine cost is measured. Load path, layout, accumulation order and +// kernels are otherwise byte-identical to the accepted S=2 geometry, so +// correctness stays exact int32/bf16 (0 mismatches) and graph capture is +// unchanged. Prediction: if per-step latency does NOT inflate with waves/CU +// (S=2 measured ~0.56 us/step), 4 waves x 32 steps = 128 serial step-times +// vs 2 x 96 = 192 -> median ~72-80 us including combine (1.35-1.5x vs +// 106.8); if it does inflate like the S=3 probe (~1.27x at 2 waves/CU), the +// median lands well above 106.8 and falsifies further occupancy probing. +// +// Iteration 11 (HIP-only, rejected 66.34 us): SPLIT_K=12 (960 blocks = 8/CU +// = 2 waves/SIMD) sat just above the S=6 champion (64.46 us) with the same +// 491,520 L2 sector requests -> the L2 request-throughput ceiling is +// ~7.4-7.6 G sector/s and wave count >= 480 is sufficient, so pure split-K +// occupancy is closed (the remaining trusted candidates S=4/S=8 have uneven +// 2.67/5.33 blocks-per-CU tails and cannot beat the champion's 128 step +// path). Iteration 12 (HIP-only, rejected 83.65 us): depth-1 register +// prefetch (double-buffered fragments) regressed, falsifying HIP-level load +// scheduling. The untouched axis is the request COUNT itself: A is re-read +// once per n16 tile (80x), i.e. A contributes 245,760 of the 491,520 L2 +// sectors. Iteration 13: one wave owns TWO adjacent n16 tiles (n32 block) +// sharing a single A fragment; with SPLIT_K=12 the grid stays 480 one-wave +// zero-barrier blocks (4/CU = 1/SIMD, the champion's MLP level), steps per +// wave halve to 16 (2 mmacs per step), and L2 sectors drop to 368,640 +// (-25%). See the kernel comment for the full prediction. +// +// Iteration 14 (HIP-only consolidation, champion 60.374 us): the partial +// kernel's per-wave 384-load chain (16 steps x 24 byte loads) and the +// combine kernel's per-lane 48-load chain (12 planes x 4 strided dword +// loads) are both serialized-load latency bound. The remaining HIP lever at +// the champion's 480-block / 1-wave-per-SIMD geometry is the combine: the +// plane element order is permuted to lane-major (lane l owns plane[l*4..l*4+3]) +// in BOTH kernels, so the partial stores one dwordx4 per lane per plane +// (vmem_write 3,840 -> 960) and the combine reads one dwordx4 per lane per +// split (48 scalar loads -> 12 dwordx4 loads per lane). Values, ascending +// int32 order, workspace bytes, grid, graph structure and shape guards are +// unchanged; sub-480-block probes (n64 S=12 -> 240 blocks, n32 S=6 -> 240 +// blocks) are closed by the session's own sector model (S=2 160 blocks 4.60 +// G/s, S=3 240 blocks ~1.2-3 G/s vs 7.4-7.6 G/s at >= 480 waves). +// +// Iteration 15 (final conditional inline-asm round; HIP-only consolidation +// per the control-plane policy: plateau=false, raw_inline_asm_allowed=false, +// skill_allowed=false): apply iteration 14's "vector loads beat serialized +// scalar chains" medicine to the partial kernel's A fragment. The accepted +// code object emits 24 scalar global_load_ubyte per lane per k32 step in six +// 4-load groups, each group followed by a full count-limited s_waitcnt +// vmcnt(3..0) chain and a 4-instruction byte-pack (v_and/v_lshlrev/v_or3), +// i.e. six serialized wait+pack segments per step. du_load_matrix_sync +// (matrix_a, row_major, ldm = k) maps lane l to the EIGHT CONTIGUOUS K +// elements a[row = l & 15][col = ((l >> 4) << 3) .. +7], and du_mma_sync +// consumes the fragment as one 64-bit little-endian value (8 bytes), so one +// 8-byte-aligned global_load_dwordx2 replaces A's 8 scalar byte loads + 2 +// wait chains + 8 pack VALU per step with identical register values (0 +// mismatches). B fragments stay scalar: matrix_b row_major bytes are strided +// by the K-major row stride (n) per lane, i.e. contiguous only in a repacked +// n-major layout, the axis iteration 5 already falsified. This cuts +// vmem_read 184,320 -> 130,560 (-29%, requests only; the 368,640-sector L2 +// footprint is unchanged), so the round discriminates the session's two +// candidate ceilings: if per-wait-chain serialization is on the critical +// path, per-step ~3.3 us -> ~2.2-2.6 us (partial ~35-42 us); if the L2 +// sector-rate ceiling (~7 G sector/s) is binding instead, the median stays +// flat ~52-55 us partial. +// +// Contract (see int8_w8a8_gemm_api.py): +// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) +// * x_scale[m, 0] * weight_scale[n, 0] ) +// x_q : int8 [M, K] row-major +// weight : int8 [K, N] row-major (identity packed layout this round; +// later rounds may change pack_weight + GEMM together) +// x_scale : fp32 [M, 1] +// weight_scale : fp32 [N, 1] +// out : bf16 [M, N] row-major, caller-provided +// workspace : uint8, caller-allocated before Graph capture; the exact +// shape path uses SPLIT_K x 80 x 256 int32 +// (SPLIT_K=12: 983,040 B default this round; +// SPLIT_K=6: 491,520 B; SPLIT_K=2: 163,840 B) for +// partial planes, all inside the 16-plane / +// 1,310,720 B capacity for this shape (see +// workspace_split_k_capacity). +// +// The timed operator (`launch_w8a8_gemm`) performs no allocation, +// compilation, autotuning, packing, host/device synchronization, or +// default-stream launch: it only validates the exact shape guard and +// launches the split-K partial + combine kernels on the caller-provided +// stream. + +#include +#include +#include + +#include + +namespace { + +// gfx928 wavefront = 64 lanes; block sizes must be multiples of 64. +constexpr int kScalarBlockThreads = 128; +constexpr int kPackBlockThreads = 256; + +// --------------------------------------------------------------------------- +// Generic scalar INT8 dot-product GEMM. +// +// One thread computes one output element out[m, n]. The full K loop runs in +// int32 (exact for every assigned K; max K=6144 -> |dot| <= 6144*127*127 +// ~= 9.9e7 << 2^31), then the two fp32 scales are applied and the result is +// stored as bf16. +// +// Thread t maps to (row = t / n, col = t % n): adjacent lanes land on +// adjacent addresses in the fastest-changing N dimension, so the per-K-row +// B reads (b[k*n + col]) stay lane-coalesced and the A reads are broadcast +// within a row. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_scalar_gemm_kernel(const int8_t* __restrict__ a, // [M, K] + const int8_t* __restrict__ b, // [K, N] + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [M, N] + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t t = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (t >= total) { + return; + } + const int row = static_cast(t / n); + const int col = static_cast(t - static_cast(row) * n); + + const int8_t* a_row = a + static_cast(row) * k; + const int8_t* b_col = b + col; + + int32_t acc = 0; + const int8_t* b_ptr = b_col; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_ptr[0]); + b_ptr += n; + } + + // Same arithmetic order as the harness reference + // (dot.to(fp32) * x_scale * weight_scale.T): fp32 multiply left-to-right, + // then round-to-nearest-even bf16. + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[t] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Identity device-to-device weight packing (bootstrap). +// raw_weight [K, N] -> packed_weight [K, N], weight_scale [N] -> +// packed_weight_scale [N]. Later optimization rounds may replace this with a +// real layout transform and a matching GEMM interpretation. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kPackBlockThreads) void +w8a8_pack_weight_identity_kernel(const int8_t* __restrict__ raw_weight, + int8_t* __restrict__ packed_weight, + int64_t count) { + const int64_t t = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (t < count) { + packed_weight[t] = raw_weight[t]; + } +} + +__global__ __launch_bounds__(kPackBlockThreads) void +w8a8_pack_scale_identity_kernel(const float* __restrict__ raw_scale, + float* __restrict__ packed_scale, + int count) { + const int t = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (t < count) { + packed_scale[t] = raw_scale[t]; + } +} + +// Shared scalar launch helper: works for any (m, n, k) satisfying the API +// contract, so it doubles as the generic fallback. +void launch_scalar_gemm(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + int m, + int n, + int k, + hipStream_t stream) { + const int64_t total = static_cast(m) * n; + const int64_t blocks = + (total + kScalarBlockThreads - 1) / kScalarBlockThreads; + hipLaunchKernelGGL(w8a8_scalar_gemm_kernel, + dim3(static_cast(blocks)), + dim3(kScalarBlockThreads), 0, stream, a, b, x_scale, + weight_scale, reinterpret_cast(out), m, + n, k); +} + +// --------------------------------------------------------------------------- +// DUMMA INT8 m16n16k32 split-K partial kernel (exact assigned shape M=16). +// +// Occupancy-probe family (iterations 1-12): one 64-lane wavefront owned one +// (16x16) partial tile over one uniform K-slice of the template SPLIT_K with +// direct global byte-load fragments; grid = (N/16)*SPLIT_K. Measured +// geometry ladder: S=1 80 blocks 206.9 us, S=3 240 blocks pathological +// 416.7 us, S=2 160 blocks 106.8 us (iter-3 champion), S=6 480 blocks +// 64.46 us (iter-10 champion), S=12 960 blocks 66.34 us. S=6/S=12 both move +// the same 491,520 L2 sectors at ~7.4-7.6 G sector/s, so pure split-K +// occupancy is closed; load width (iter 5) and LDS staging (iter 9) and +// register prefetch (iter 12) all regressed. +// +// Iteration 13 (n32 block, A-sharing): each wavefront now owns TWO adjacent +// n16 tiles (tiles 2*tile2 and 2*tile2+1) over the same K-slice. Both mmacs +// consume the SAME A fragment (A is loaded once per k32 step for both +// tiles), which halves A's L2 sector traffic (A was re-read once per n16 +// tile, 80x total). B is two separate n16 fragments at n0 and n0+16. The +// partials layout is unchanged ([tile][s][16][16] with tile the n16-tile +// index, so the combine kernel is untouched) and every output element still +// accumulates the identical k-ascending int32 sequence, so correctness +// stays bit-exact. Grid = (N/32)*SPLIT_K: with SPLIT_K=12 -> 40*12 = 480 +// one-wave zero-barrier blocks = exactly 4 blocks per CU (1 wave per SIMD, +// the iteration-10 champion's MLP level), 16 k32 steps x 2 mmacs per wave. +// --------------------------------------------------------------------------- +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; +constexpr int kDummaBlockThreads = 64; +// Iteration 13: adjacent n16 tiles processed per wave (shared A fragment). +constexpr int kDummaTilesPerBlock = 2; + +template +__global__ __launch_bounds__(kDummaBlockThreads) void +w8a8_dumma_m16n16k32_splitk_partial_kernel( + const int8_t* __restrict__ a, // [16, K] row-major + const int8_t* __restrict__ b, // [K, N] row-major + int32_t* __restrict__ partials, // [N/16][SPLIT_K][16][16] + int n, + int k) { + const int tile2 = static_cast(blockIdx.x) / SPLIT_K; + const int s = static_cast(blockIdx.x) % SPLIT_K; + const int n0 = tile2 * kDummaTilesPerBlock * kDummaTileN; + + // Uniform split; the exact shape (K=6144) divides evenly by every + // SPLIT_K in {2, 6, 12} (3072 / 1024 / 512), each k32-aligned. + const int k_chunk = k / SPLIT_K; + const int k_begin = k_chunk * s; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag0; + du::dumma::DUFragment + b_frag1; + du::dumma::DUFragment + acc_frag0; + du::dumma::DUFragment + acc_frag1; + du::dumma::du_fill_fragment(acc_frag0, 0); + du::dumma::du_fill_fragment(acc_frag1, 0); + + const int lane = static_cast(threadIdx.x); + for (int k0 = k_begin; k0 < k_begin + k_chunk; k0 += kDummaTileK) { + // One A fragment feeds both n16 tiles (halves A's L2 sector traffic). + // Iteration 15: du_load_matrix_sync (matrix_a row_major, ldm = k) maps + // lane l to the eight CONTIGUOUS K elements a[row = l & 15][col = + // ((l >> 4) << 3) .. +7], and du_mma_sync consumes the fragment as one + // 64-bit little-endian value, so one 8-byte-aligned dwordx2 load + // replaces A's 8 scalar global_load_ubyte + 2 count-limited s_waitcnt + // vmcnt chains + 8 byte-pack VALU ops per k32 step with identical + // register values (bit-identical int32 accumulation; 0 mismatches). + // B stays scalar: matrix_b row_major bytes are strided by the K-major + // row stride (n), not contiguous, so they cannot be vector-loaded from + // the identity layout (that axis was falsified at iteration 5). + const int a_row = lane & 15; + const int a_col = (lane >> 4) << 3; + const int8_t* a_ptr = static_cast( + __builtin_assume_aligned(a + k0 + a_row * k + a_col, 8)); + *reinterpret_cast(a_frag.x) = + *reinterpret_cast(a_ptr); + du::dumma::du_load_matrix_sync(b_frag0, b + k0 * n + n0, n); + du::dumma::du_load_matrix_sync(b_frag1, b + k0 * n + n0 + kDummaTileN, + n); + du::dumma::du_mma_sync(acc_frag0, a_frag, b_frag0, acc_frag0); + du::dumma::du_mma_sync(acc_frag1, a_frag, b_frag1, acc_frag1); + } + + // Single wave per block: no barrier needed before the direct global + // stores. Iteration 14: each lane's four int32 accumulator elements are + // stored CONTIGUOUSLY (lane-major plane: lane l occupies plane[l*4 .. l*4+3], + // one dwordx4 store per lane per plane instead of four strided dword + // stores), so the combine kernel reads every plane with one dwordx4 load + // per lane per split instead of four strided scalar loads. The raw + // fragment values are unchanged: element i of lane l is the partial of + // (row = l & 15, col = (l >> 4) + 4*i) (gfx928 m16n16k32 int32 accumulator + // layout, verified against du_store_matrix_sync and the accepted pre-edit + // ISA; see the combine kernel for the matching interpretation). The plane + // layout stays [tile][s][16][16] with tile the n16-tile index (2*tile2 and + // 2*tile2+1); only the intra-plane element order is permuted, consistently + // in both kernels, so every int32 partial value and the ascending-sum + // combine order are bit-identical. + int32_t* plane0 = + partials + + ((tile2 * kDummaTilesPerBlock) * SPLIT_K + s) * + (kDummaTileM * kDummaTileN); + *reinterpret_cast(plane0 + lane * 4) = + make_int4(acc_frag0.x[0], acc_frag0.x[1], acc_frag0.x[2], + acc_frag0.x[3]); + int32_t* plane1 = + partials + + ((tile2 * kDummaTilesPerBlock + 1) * SPLIT_K + s) * + (kDummaTileM * kDummaTileN); + *reinterpret_cast(plane1 + lane * 4) = + make_int4(acc_frag1.x[0], acc_frag1.x[1], acc_frag1.x[2], + acc_frag1.x[3]); +} + +// --------------------------------------------------------------------------- +// Split-K combine kernel (same exact shape). One 64-lane wavefront per 16x16 +// output tile: 4 elements per lane, each summed ascending over the SPLIT_K +// planes (exact int32, split-k-ascending + k-ascending inside each split is +// bit-identical to the reference int64 dot), then x_scale[m] * +// packed_weight_scale[n] and bf16 store. Part of the timed operator wall. +// Iteration 14: reads the lane-major plane permutation written by the +// partial kernel (lane l owns plane[l*4 .. l*4+3]) with one dwordx4 load per +// lane per split instead of four strided dword loads. +// --------------------------------------------------------------------------- +template +__global__ __launch_bounds__(kDummaBlockThreads) void +w8a8_dumma_m16_splitk_combine_kernel( + const int32_t* __restrict__ partials, // [N/16][SPLIT_K][16][16] + const float* __restrict__ x_scale, // [16] + const float* __restrict__ weight_scale, // [N] + hip_bfloat16* __restrict__ out, // [16, N] row-major + int n) { + const int tile = static_cast(blockIdx.x); + const int n0 = tile * kDummaTileN; + const int lane = static_cast(threadIdx.x); + // Iteration 14 lane-major plane layout: lane l's four elements are stored + // contiguously (one dwordx4 per lane per split), so the combine reads + // plane[l*4 .. l*4+3] with one dwordx4 load per lane per split. Repair + // (iteration 14 round 1): the gfx928 m16n16k32 int32 accumulator fragment + // layout (verified against du_store_matrix_sync and the accepted pre-edit + // ISA) is row = lane & 15, col = (lane >> 4) + 4*i for element i, i.e. + // lane l's i-th element is the partial of (row = l & 15, col = (l >> 4) + + // 4*i). Element i of lane l is stored at plane[l*4 + i], so acc0..acc3 + // (v.x..v.w) map to row = lane % 16 and columns (lane / 16) + {0, 4, 8, + // 12}. The split-ascending int32 sum per output element is unchanged + // (bit-identical values; only the (lane, i) -> (row, col) interpretation + // was wrong and is fixed here). + const int row = lane % kDummaTileN; // 0..15: fragment row = l & 15 + const int col_lo = lane / kDummaTileN; // 0..3: fragment col base = l >> 4 + const int plane_base = tile * SPLIT_K * (kDummaTileM * kDummaTileN); + + int32_t acc0 = 0, acc1 = 0, acc2 = 0, acc3 = 0; +#pragma unroll + for (int s = 0; s < SPLIT_K; ++s) { + const int4 v = *reinterpret_cast( + partials + plane_base + s * (kDummaTileM * kDummaTileN) + + lane * 4); + acc0 += v.x; + acc1 += v.y; + acc2 += v.z; + acc3 += v.w; + } + const float s0 = static_cast(acc0) * x_scale[row] * + weight_scale[n0 + col_lo]; + out[row * n + n0 + col_lo] = __float2bfloat16(s0); + const float s1 = static_cast(acc1) * x_scale[row] * + weight_scale[n0 + col_lo + 4]; + out[row * n + n0 + col_lo + 4] = __float2bfloat16(s1); + const float s2 = static_cast(acc2) * x_scale[row] * + weight_scale[n0 + col_lo + 8]; + out[row * n + n0 + col_lo + 8] = __float2bfloat16(s2); + const float s3 = static_cast(acc3) * x_scale[row] * + weight_scale[n0 + col_lo + 12]; + out[row * n + n0 + col_lo + 12] = __float2bfloat16(s3); +} + +// Split-K DUMMA launch for the exact M=16 shape with N % 16 == 0, +// K % 32 == 0 and K % SPLIT_K == 0. Partial + combine both run inside the +// timed Graph replay on the caller stream. The partial kernel processes +// kDummaTilesPerBlock adjacent n16 tiles per wave (grid = N/(16*2) * +// SPLIT_K); the combine still launches one block per n16 tile over the +// unchanged [tile][s][16][16] plane layout. +template +void launch_dumma_m16_splitk(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int n, + int k, + hipStream_t stream) { + const int tiles = n / kDummaTileN; + const int64_t required_bytes = + static_cast(SPLIT_K) * tiles * kDummaTileM * kDummaTileN * + static_cast(sizeof(int32_t)); + if (workspace_bytes < required_bytes) { + // Defensive: workspace capacity for this shape is 16 planes, so this + // never triggers for SPLIT_K in {2, 6, 12}; keep the scalar path as the + // generic fallback anyway. + launch_scalar_gemm(a, b, x_scale, weight_scale, out, kDummaTileM, n, k, + stream); + return; + } + int32_t* partials = reinterpret_cast(workspace); + const int partial_blocks = + tiles / kDummaTilesPerBlock * SPLIT_K; + hipLaunchKernelGGL(w8a8_dumma_m16n16k32_splitk_partial_kernel, + dim3(static_cast(partial_blocks)), + dim3(kDummaBlockThreads), 0, stream, a, b, partials, n, + k); + hipLaunchKernelGGL(w8a8_dumma_m16_splitk_combine_kernel, + dim3(static_cast(tiles)), + dim3(kDummaBlockThreads), 0, stream, partials, x_scale, + weight_scale, reinterpret_cast(out), n); +} + +} // namespace + +// --------------------------------------------------------------------------- +// Host launch symbols consumed by csrc/bindings.cpp (TORCH_LIBRARY zth_w8a8). +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // Exact assigned shape: minimax_tp8_qkv_proj_m16 (M=16, N=1280, K=6144). + // Iteration 13: default SPLIT_K=12 on the n32 (2 tiles/wave, shared A) + // partial kernel -> 40*12 = 480 one-wave zero-barrier blocks = exactly + // 4 blocks per CU on all 120 CUs (1 wave per SIMD, the iteration-10 + // champion's MLP level), 16 k32 steps x 2 mmacs per wave, 12 int32 + // partial planes (983,040 B) inside the 16-plane / 1,310,720 B capacity. + // Iteration 14: plane element order is lane-major (both kernels changed + // together, values/order bit-identical), so the combine's per-lane + // serialized load chain drops from 48 strided dword loads to 12 dwordx4 + // loads and the partial's stores drop from 8 to 2 dwordx4 per block. + // Iteration 15: the partial kernel's A fragment load is vectorized to one + // 8-byte dwordx2 per lane per k32 step (the 8 bytes per lane are + // contiguous K elements in the identity row-major A layout; du_mma_sync + // consumes the fragment as one 64-bit value), cutting A's 8 scalar + // global_load_ubyte + 2 wait chains + 8 byte-pack VALU ops per step with + // bit-identical values; B stays scalar (strided by n). vmem_read drops + // 184,320 -> 130,560 at unchanged grid (480) / planes (12) / sectors. + // SPLIT_K=6 (240 blocks) and SPLIT_K=2 (80 blocks) stay instantiated as + // measured alternatives; flip kDefaultSplitK to switch the measured + // geometry without any other source change. The guard is exact, so paired + // M=2 shapes with the same (N, K) keep taking the generic scalar fallback + // below. + constexpr int kDefaultSplitK = 12; + if (m == 16 && n == 1280 && k == 6144) { + switch (kDefaultSplitK) { + case 12: + launch_dumma_m16_splitk<12>(a, b, x_scale, weight_scale, out, + workspace, workspace_bytes, n, k, stream); + break; + case 6: + launch_dumma_m16_splitk<6>(a, b, x_scale, weight_scale, out, + workspace, workspace_bytes, n, k, stream); + break; + case 2: + launch_dumma_m16_splitk<2>(a, b, x_scale, weight_scale, out, + workspace, workspace_bytes, n, k, stream); + break; + default: + launch_scalar_gemm(a, b, x_scale, weight_scale, out, kDummaTileM, n, + k, stream); + break; + } + return; + } + + // Generic scalar fallback for every unmatched (m, n, k). + launch_scalar_gemm(a, b, x_scale, weight_scale, out, m, n, k, stream); +} + +extern "C" void launch_pack_w8a8_weight(const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Identity device-to-device copy, valid for every (K, N). Out of the timed + // region (pack_weight op), so multiple launches are fine here. + const int64_t weight_count = static_cast(k) * n; + const int64_t weight_blocks = + (weight_count + kPackBlockThreads - 1) / kPackBlockThreads; + hipLaunchKernelGGL(w8a8_pack_weight_identity_kernel, + dim3(static_cast(weight_blocks)), + dim3(kPackBlockThreads), 0, stream, raw_weight, + packed_weight, weight_count); + + const int64_t scale_blocks = + (static_cast(n) + kPackBlockThreads - 1) / kPackBlockThreads; + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, + dim3(static_cast(scale_blocks)), + dim3(kPackBlockThreads), 0, stream, weight_scale, + packed_weight_scale, n); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj_and_indexer_qk.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj_and_indexer_qk.hip new file mode 100644 index 00000000..3f4a48e3 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj_and_indexer_qk.hip @@ -0,0 +1,714 @@ +// @@variant shape=minimax_tp8_qkv_proj_and_indexer_qk_m16 commit=e72025dfcfa50ff3bf3fcb00b87663e65f4e0fc4 added=2026-08-28 +// median_us=39.73 p90_us=39.81 +// source=minimax-dsh-tp8-m16-1-78402260 +// csrc/w8a8_gemm_hip.hip +// +// Worker 1 (physical GPU 1) — assigned shape: +// minimax_tp8_qkv_proj_and_indexer_qk_m16 : M=16, N=1536, K=6144 +// +// Iteration 0: correctness-first scalar bootstrap (one thread per output +// element, blockDim 128). Preserved below as the generic fallback for every +// unmatched (m, n, k) shape, including the paired M=2 shape with the same +// (N, K). +// +// Iteration 1 (DUMMA bootstrap, accepted): minimal native gfx928 DUMMA INT8 +// m16n16k32 tile, one 64-lane wavefront per block, one 16x16 output tile per +// block, direct global fragment loads, explicit int32 accumulation, direct +// scaled-bf16 fragment epilogue. Measured 203.2 us median vs the 111.5 us +// fixed Triton baseline; grid = N/16 = 96 blocks on 120 CUs (0.8 blocks/CU, +// below the two-blocks-per-CU latency-hiding target). +// +// Iteration 2 (architecture round: measure grid parallelism before +// polishing): replace the exact (m==16 && n==1536 && k==6144) arm with a +// split-K=4 partial GEMM + separate combine: +// * partial kernel: grid = 96 N-tiles x 4 K-splits = 384 one-wave blocks +// (64 lanes each, one 16x16 output tile per block); uniform K=1536 +// slice per split (48 m16n16k32 steps, k ascending within the slice, +// exact int32 order); direct global-to-fragment loads (same load path +// as iteration 1 — vectorized loads/LDS staging are deliberately +// deferred to a later polishing round); each block stores its 16x16 +// int32 partial to its exclusive plane partials[split_id][16][N]; +// * combine kernel: grid = 96 one-wave blocks; sums the 4 planes in +// ascending split order (split s covers k in [s*1536,(s+1)*1536) → +// bit-identical int32 accumulation order), applies +// x_scale[row] * weight_scale[col], stores scaled bf16. +// Zero LDS, zero barriers, zero atomics in both kernels. The two launches +// are ordered by the caller-provided stream, so the combine kernel always +// observes the partial kernel's writes and no workspace clear is needed +// (every launch overwrites every plane element before combine reads it). +// 384 blocks = 3.2 blocks/CU on the 120-CU device: all CUs are covered and +// co-resident one-wave blocks can overlap the global-load latency chains — +// the falsifiable grid-parallelism claim of this round. +// +// launch_pack_w8a8_weight stays an identity device-to-device copy (packed +// layout == logical [K, N] row-major), unchanged from iteration 0. +// +// Iteration 6 (pipeline round: B-only LDS staging with coalesced 8/16-byte +// cooperative loads): keep the accepted iteration-2 geometry (split-K=4, +// 384 one-wave blocks), workspace, combine kernel, identity pack, scalar +// fallback, and bit-exact int32 order; replace ONLY the partial kernel's +// fragment transport. L2/VMEM evidence picks B as the operand to stage: B +// is the only compulsory-HBM operand (9,437,184 B read exactly once, zero +// reuse) while A is 96 KiB unique, re-read 96x, and L2-hot (measured 76.5% +// L2 hit rate). Each block stages its [1536][16] B K-slice once into a +// compact 24,704 B n-major LDS tile s_b[16][1544] (stride 1544 = 1536 + 8: +// rows are 8-B aligned for ds_read_b64 and the stride spreads the 64 lanes' +// bank pairs to the 4-cycle conflict floor; 2 blocks/CU co-residency, vs +// iteration 4's whole-slice A+B tile at 1 block/CU) via 24 coalesced +// dwordx4 (16 B/lane) cooperative loads + conflict-free ds_write_b8 + one +// __syncthreads; the 48-step K loop then contains no global memory op for B +// at all: A is one aligned global_load_dwordx2 per step (verified 8-byte +// loader) and B is one ds_read_b64 per lane per step, with a depth-1 +// register pipeline issuing step s+1's loads before step s's v_mmac. +// Whether the issued-ahead loads actually overlap the MMAC is measured, +// not assumed. HBM bytes are unchanged (B is still fetched exactly once, +// now in 24 wide loads per block instead of 384 byte-loads); LDS bytes go +// 0 -> 49,152 B/block (24,576 staged + 24,576 read) = 18.9 MB over the +// launch; vmem_read instructions drop 294,912 -> 27,648 (9,216 staging + +// 18,432 A). +// +// Iteration 7 (pipeline round decision: SINGLE buffering kept — the +// double-buffering comparison gate fails on both criteria: L2 hit rate is +// 77.8% (not below 70%, because A is L2-hot and B is already staged) and +// doubled LDS would be 2 x 24,704 B = 49,408 B (not below 48 KiB), so +// double buffering is not indicated; barriers per K step are counted at 0 +// both before and after — the staging __syncthreads compiles out for +// one-wave blocks (0 s_barrier in the iteration-6 ISA)). Deepen the +// single-buffered pipeline instead: the iteration-6 depth-1 ping-pong +// compiled to ONE 8-step unrolled group iterated 6x whose group ENTRY +// issues the group's first A global load and then fully drains it with +// s_waitcnt vmcnt(0) lgkmcnt(0) before the first MMAC — 6 full L2 round +// trips per block on the critical path (steps 1..7 of each group get +// staggered-wait coverage, step 0 of each group does not). Replace it +// with a depth-8 single-buffered register ring for A: every A global load +// is issued exactly 8 steps before its MMAC so no MMAC ever waits on an +// L2 round trip; B stays on the 24,704 B LDS tile (one ds_read_b64 per +// step, hoisted/batched by the scheduler); the K loop still has zero +// barriers; LDS stays 24,704 B/block (2 blocks/CU). Everything else +// (split-K=4 grid, workspace, combine, identity pack, scalar fallback, +// exact int32 order) is unchanged. +// +// Iteration 8 (HIP-only resource round: LDS-footprint occupancy limiter): +// the fresh exact-source PMC launch records (iteration8/pmc.csv) show the +// iteration-7 kernel at arch_vgpr=168, lds=24,704 B, scratch=0, 1 wave per +// block, 384 blocks — the BINDING occupancy limiter is the LDS footprint: +// 64 KiB/CU / 24,704 B = 2.65 -> 2 blocks/CU co-residency (3 x 24,704 B = +// 74,112 B > 64 KiB), i.e. only 2 of the CU's 4 SIMDs carry a wave (VGPRs +// at 168 allow 1 wave/SIMD = 4 waves/CU, not binding at 2 blocks). The +// exact-source ISA (iteration8/current-best-isa/isa.txt) also shows the +// depth-8 ring compiled to a FULL 48-step unroll: all 48 A +// global_load_dwordx2 issued in one burst (v[69:70]..v[163:164]), 16 +// ds_read2_b64 for the B fragments, ONE s_waitcnt vmcnt(47) lgkmcnt(14) at +// 0x2D64, then 48 dependent v_mmac on the single accumulator with no +// inter-MMAC waits — the per-wave critical path is one A round trip + a +// 48-deep dependent MMAC chain hidden by only 2 resident waves/CU. +// Falsifiable claim: halve the per-block staged B slice by moving the +// exact (16,1536,6144) arm from split-K=4 to split-K=8 — grid 96 N-tiles x +// 8 splits = 768 one-wave blocks (6.4 blocks/CU), B tile s_b[16][776] +// (12,288 B payload + 8 B/row pad, 776/4 = 194 == 2 mod 32 keeps the +// 4-cycle bank-conflict floor, 12,416 B/block -> LDS co-residency 5 +// blocks/CU, VGPR-capped at 4-5), K loop 24 steps (24 A loads, 12 +// ds_read2_b64, one wait, 24 dependent MMACs — half the serial chain per +// wave), workspace 8 x 98,304 B = 786,432 B (<= the 16-plane 1,572,864 B +// capacity allocate_workspace guarantees for this shape), combine sums the +// 8 planes in ascending split order (int32 addition is exact mod 2^32, so +// the split grouping does not change the accumulated value), identity pack +// and scalar fallback unchanged. HBM byte accounting is unchanged: B is +// still read exactly once (768 blocks x 12,288 B = 9,437,184 B), A read +// volume is unchanged (768 x 12,288 B = 9,437,184 B, still 96x re-read of +// the 96 KiB L2-hot slice), vmem_read instructions stay 27,648 (9,216 +// staging + 18,432 A) and LDS bytes stay 18,874,368 B/launch (12,288 +// staged + 12,288 read per block x 768). No repeated HBM operand reads are +// traded for the occupancy gain. If median does not improve below 74.018 +// us / p90 below 74.777 us, the LDS-footprint/co-residency lever is +// falsified and the next round must attack the 24-48-deep dependent MMAC +// accumulation chain itself (e.g. split accumulators) at a fixed split-K. +// +// Iteration 9 (rejected, NOT in this source): the staging-burst candidate +// (median 66.626 us, -26.07% vs iteration 8) tried to collapse the three +// serial staging rounds into one 12-load burst. Its exact code object +// (iterations/.../iteration9/isa/isa.txt, digest 1290fe7b) shows the +// compiler limitation: it refused to emit a single 12-load burst — it +// unrolled the staging into straight-line 8+2+2 global_load_dwordx4 with +// partial vmcnt drains (vmcnt(5) after 8 loads, vmcnt(6)->vmcnt(2) after +// the 9th/10th, vmcnt(1)->vmcnt(0) after the 11th/12th) — and the 48 live +// staging VGPRs lifted arch_vgpr 104 -> 137, dropping co-residency 5 -> 4 +// blocks/CU (1 wave/SIMD). Both the unchanged round structure AND the +// register budget lost; the round is not in the source (restored to the +// accepted iteration-8 digest 4a6e82d before this round). +// +// Iteration 10 (this round, HIP-only occupancy probe): the accepted +// iteration-8 exact-source ISA (profiles/.../iteration9/current-best-isa/ +// isa.txt) shows the staging #pragma unroll 4 loop compiled to a 3-round +// loop whose body issues 4 global_load_dwordx4, drains them progressively +// to vmcnt(0), writes 64 ds_write_b8, and branches back — 3 serial HBM +// round trips per wave on the critical path (600 concurrent waves x +// 12,288 B = 7.2 MB in flight per round), i.e. the per-wave serial staging +// RTT chain and co-residency, not the MMAC chain, bind the partial kernel +// (measured co-res 2 -> 5 moved 74.018 -> 49.259 us; co-res 5 -> 4 moved +// it back to 66.626 us). Falsifiable claim: move ONLY the split count +// 8 -> 12 inside the exact (16,1536,6144) arm — keeping the m16n16k32 +// DUMMA tile, the one-wave/one-tile block structure, the B-only LDS +// staging loop recipe (coalesced dwordx4 + conflict-free ds_write_b8 + one +// __syncthreads), the depth-8 A register ring (kSliceSteps 16 % 8 == 0, +// ring closes exactly on step 15), the combine kernel (now summing 12 +// planes in ascending split order, same code path), the identity pack, the +// scalar fallback, and the bit-exact int32 order — with: grid 96 N-tiles x +// 12 splits = 1,152 one-wave blocks (9.6 blocks/CU); per-block B slice +// 512 k -> LDS tile s_b[16][520] = 8,320 B (8,192 B payload + 8 B/row +// pad; 520 % 8 == 0 keeps rows 8-B aligned and 520/4 = 130 == 2 mod 32 +// keeps the 4-cycle bank-conflict floor) -> co-residency 7 blocks/CU by +// LDS (65,536 / 8,320 = 7.88; VGPR-capped at 8, arch_vgpr must stay +// <= 128 or co-res drops to 4 and the probe is falsified); K loop 16 steps +// (16 A loads, 8 ds_read2_b64, one wait, 16 dependent MMACs); staging 8 +// dwordx4 chunks -> the unroll-4 loop runs 2 serial rounds instead of 3 +// (one fewer HBM RTT per wave); workspace 12 x 98,304 B = 1,179,648 B <= +// the 16-plane 1,572,864 B capacity (the workspace_bytes guard passes); +// combine sums the 12 planes in ascending split order. HBM and LDS byte +// accounting is unchanged — no repeated HBM reads are traded: B is still +// read exactly once (1,152 blocks x 8,192 B = 9,437,184 B), A read volume +// is unchanged (1,152 x 8,192 B = 9,437,184 B, still the 96x re-read of +// the 96 KiB L2-hot slice), vmem_read instructions stay 27,648 (9,216 +// staging + 18,432 A), LDS bytes stay 18,874,368 B/launch (8,192 staged + +// 8,192 read per block x 1,152). Predicted deltas: (a) partial-kernel +// staging rounds 3 -> 2, (b) co-residency 5 -> 7 (840 concurrent waves, +// 1.37 rounds vs 1.28), (c) MMAC chain 24 -> 16, (d) combine reads 12 +// planes (1,179,648 B) instead of 8 (786,432 B) — combine grows ~5.28 -> +// ~6.5-7 us while the partial kernel is predicted to drop ~38.9 -> +// ~27-33 us. If median does not improve below 49.259 us / p90 below +// 49.415 us, the occupancy probe at co-res 7 is falsified and the next +// round must either probe split-K=16 (co-res 8, but 1.6 rounds and 16 +// combine planes) or attack the remaining serial staging rounds +// (e.g. wider per-round loads) at a fixed split-K. +// +// Iteration 14 (accepted, 43.186 us median / 43.361 us p90): the +// staging-burst (iteration 11: 85.049 us), fused-combine (iteration 12: +// 45.115 us) and split-K=16 occupancy (iteration 13: 45.345 us median / +// 202.93 us p90) probes all falsified their levers at fixed split-K=12 +// geometry; the remaining serial wall component was the STANDALONE combine +// kernel's load pattern: each lane read its 4 elements x 12 planes as 48 +// scalar int32 global loads grouped in four serial load-drain rounds (ISA +// staircase vmcnt drains), and the combine runs stream-serial after the +// 1,152-block partial kernel, so its HBM-latency-bound drains (L2 hit +// 16.288%, measured at split-K=8) add directly to the operator wall. +// Change was ONLY the combine kernel: each lane owns 4 consecutive n +// columns of one row (row = lane>>2, tcol4 = (lane&3)<<2) so the 12 plane +// reads per lane become 16-B-aligned int4 loads (one drain group instead +// of four) and the bf16 epilogue becomes one 8-B uint64 store; per-element +// ascending split sum s = 0..11 is unchanged -> bit-identical int32 order. +// Accepted at 43.186 us median / 43.361 us p90 (below the 44.283/44.502 +// falsifier), so the combine's load-drain structure WAS a wall component. +// +// Iteration 15 (this round, HIP-only consolidation): the accepted +// iteration-14 exact-source ISA (iteration14/isa/isa.txt, source digest +// d67fd4d5) records the partial kernel at arch_vgpr=72, sgpr=25, lds=8,320 +// B, scratch=0, grid 1,152 and shows the 16-step K loop compiled as one +// pre-loop 16-load A burst + 8 ds_read2_b64 batch + 16 v_mmac chained on a +// SINGLE accumulator v[1:4] with only a vmcnt(15)..vmcnt(0)/ +// lgkmcnt(7)..lgkmcnt(0) staircase between them — a 16-deep dependent MMAC +// chain per wave (the iteration-14 falsifier names this as the remaining +// probe: 'the next round must attack the dependent 16-MMAC accumulator +// chain with split accumulators (value-identical mod 2^32) at fixed +// split-K=12, or stop at co-res 7'). Change ONLY the partial kernel's +// accumulator structure: two independent accumulator fragments acc0/acc1 — +// steps 0..7 accumulate into acc0, steps 8..15 into acc1 (k ascending +// within each chain), then acc0 += acc1 before the plane store; int32 +// addition is exact mod 2^32 so the grouped ascending sum is bit-identical +// (0 mismatches expected). The dependent MMAC chain per wave halves 16 -> +// 8; MMAC issue count (16), staging recipe (2-round unroll-4, 8,320 B LDS +// tile), depth-8 A ring (refill loads stay unconditional, ring closes +// exactly on step 15), plane store, combine kernel (vectorized int4, +// grid 96), identity pack, scalar fallback, paired M=2 fallback, and +// workspace 1,179,648 B are untouched. Predicted deltas: (a) per-wave +// MMAC critical path 16 x MMAC-latency -> 8 x MMAC-latency + 4 int32 adds +// (the adds are VALU on a different pipe and overlap other waves), exposed +// mainly in the 1.37-round tail where fewer than 7 waves/CU remain to hide +// the chain; (b) arch_vgpr expected ~72-80 (acc1 adds 4 VGPRs; 65,536 / +// (7 x 64) = 146.3, so co-res 7 by LDS is unaffected below arch_vgpr +// 146); scratch_bytes > 0 in the next PMC would falsify the no-spill +// premise. Falsifiers: median not below 43.186 us / p90 not below 43.361 +// us (the accumulator chain is not the remaining wall component at co-res +// 7; the family is at its measured floor and the next round must stop or +// re-probe with fresh PMC); any correctness mismatch falsifies the +// bit-exact premise. + +#include +#include +#include + +#include + +namespace { + +// gfx928 wavefront is 64 lanes; blockDim must be a multiple of 64. +constexpr int kScalarBlockThreads = 128; + +// Minimal DUMMA tile constants (gfx928 INT8 support is m16n16k32). +constexpr int kDummaM = 16; +constexpr int kDummaN = 16; +constexpr int kDummaK = 32; +constexpr int kDummaThreads = 64; // one wavefront (64 lanes) per block + +// Iteration 2 architecture: split-K = 4 for the assigned shape. 96 N-tiles +// x 4 splits = 384 blocks = 3.2 blocks/CU on the 120-CU device (>= the +// two-blocks-per-CU latency-hiding target; every CU covered). Uniform +// k32-aligned slices: K=6144 / 4 = 1536 = 48 m16n16k32 steps. Workspace: +// 4 x [16][1536] int32 planes = 393,216 B, inside the 16-plane capacity +// (16 x 98,304 B = 1,572,864 B) that allocate_workspace guarantees for this +// shape. Trusted occupancy-probe split list [2,3,4,5,7,8,10,15]: only +// {2,3,4,8,12} divide 6144 into k32-aligned slices. Iteration 8 moved the +// arm from 4 to 8 (96 N-tiles x 8 splits = 768 one-wave blocks, LDS tile +// 24,704 B -> 12,416 B, co-residency 2 -> 5, MMAC chain 48 -> 24 steps; +// accepted 74.018 -> 49.259 us). Iteration 9 (staging burst at split-K=8) +// regressed to 66.626 us and is NOT in this source: its exact code object +// shows the compiler refused to emit one 12-load burst (it scheduled 8+2+2 +// dwordx4 loads with partial vmcnt drains) and arch_vgpr 104 -> 137 dropped +// co-residency 5 -> 4. Iteration 10 (this round) probes split-K = 12: +// 96 N-tiles x 12 splits = 1,152 one-wave blocks, LDS tile 12,416 B -> +// 8,320 B -> co-residency 7 by LDS (VGPR-capped at 8), staging 12 -> 8 +// dwordx4 chunks (the unroll-4 staging loop runs 2 serial rounds instead of +// 3), MMAC chain 24 -> 16 steps, workspace 12 x 98,304 B = 1,179,648 B +// inside the 16-plane 1,572,864 B capacity. +constexpr int kSplitK = 12; + +// Iteration 6 pipeline constants, valid only under the exact-shape guard +// (m == 16 && n == 1536 && k == 6144) that is the sole caller of this +// kernel: uniform K slice = k / kSplitK = 512 at split-K=12 (was 768 at +// split-K=8, 1536 at split-K=4) = 16 m16n16k32 steps. +constexpr int kSliceK = 512; // K per split (6144 / 12) +constexpr int kSliceKStride = 520; // LDS row stride: 512 + 8 pad +constexpr int kSliceSteps = kSliceK / kDummaK; // 16 +constexpr int kStageInstrs = (kSliceK * kDummaN) / (kDummaThreads * 16); // 8 +// Iteration 7 pipeline depth: A register-ring depth (must be a power of two +// and divide kSliceSteps so the ring closes exactly: 16 % 8 == 0). +constexpr int kPrefetchDepth = 8; + +// Scalar int8 GEMM: out[m,n] = bf16(int32_dot(a[m,:], b[:,n]) * +// x_scale[m] * weight_scale[n]). +// a : [M, K] int8, row-major +// b : [K, N] int8, row-major +// x_scale : [M, 1] float +// weight_scale: [N, 1] float +// out : [M, N] bf16, row-major +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int total = m * n; + const int linear = blockIdx.x * kScalarBlockThreads + threadIdx.x; + if (linear >= total) { + return; + } + + // Adjacent lanes land on adjacent addresses in the fastest-changing N + // dimension (linear = row * n + col). + const int row = linear / n; + const int col = linear - row * n; + + const int8_t* __restrict__ a_row = a + row * k; + const int8_t* __restrict__ b_col = b + col; + + int32_t dot = 0; + for (int kk = 0; kk < k; ++kk) { + dot += static_cast(a_row[kk]) * + static_cast(b_col[kk * n]); + } + + const float scaled = + static_cast(dot) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// Iteration 6 fragment transport helpers (byte placement identical to the +// library loaders, verified in the iteration-5 round): +// du_load_matrix_sync(matrix_a, row_major) and (matrix_b, col_major) +// both place lane l's 8 elements at p[(l & 15) * ldm + ((l >> 4) << 3) + +// i], i = 0..7 (confirmed against du_mma.hpp). With ldm and offsets that +// are multiples of 8, each lane's fragment is ONE 8-B-aligned vector +// load: A = global_load_dwordx2 from the logical [16][K] row-major +// activation (L2-hot); B = ds_read_b64 from the n-major LDS tile +// s_b[16][kSliceKStride] (stride 1544 at split-K=4, 776 at split-K=8, +// 520 at split-K=12; stride % 8 == 0 keeps every row 8-B aligned and +// stride/4 == 2 mod 32 spreads the 64 lanes' bank pairs to the 4-cycle +// conflict floor for a full-wave 512-B ds_read_b64). +__device__ __forceinline__ void load_a_frag8( + du::dumma::DUFragment& f, + const int8_t* __restrict__ p, + int ldm) { + const unsigned row = static_cast(threadIdx.x) & 0xfu; + const unsigned col = (static_cast(threadIdx.x) >> 4) << 3; + const int64_t v = *reinterpret_cast(p + row * ldm + col); + *reinterpret_cast(&f.x[0]) = v; +} + +__device__ __forceinline__ void load_b_lds8( + du::dumma::DUFragment& f, + const int8_t* __restrict__ s_b, + int k_local) { + const unsigned row = static_cast(threadIdx.x) & 0xfu; + const unsigned col = (static_cast(threadIdx.x) >> 4) << 3; + const int64_t v = *reinterpret_cast( + s_b + row * kSliceKStride + k_local + col); + *reinterpret_cast(&f.x[0]) = v; +} + +// Split-K partial GEMM kernel (exact (16, 1536, 6144) arm, S = kSplitK): +// block b -> tile_id = b / S (0..95), split_id = b % S (0..11); +// K slice [split_id * 512, (split_id + 1) * 512), k ascending in +// m16n16k32 steps; stores the 16x16 int32 partial into its exclusive +// plane partials[split_id][16][N] (no atomic, no zeroing needed: each +// split block owns its plane and every launch rewrites all elements). +// One 64-lane wavefront per block; iteration-6 B-only LDS staging +// (iteration 8 halves the staged slice with split-K=8; iteration 10 halves +// it again with split-K=12): +// * stage: 8 coalesced dwordx4 global loads (16 B/lane, one B row per +// lane) transpose the block's [512][16] B slice into the n-major LDS +// tile s_b[n][kk] = raw[(k_base + kk) * n + n0 + n] via 16 +// conflict-free ds_write_b8 per lane per instruction (64 lanes hit 16 +// distinct LDS words), then one __syncthreads; the #pragma unroll 4 +// staging loop compiles to 2 serial rounds (4 loads drained to +// vmcnt(0) each) instead of the 3 rounds at split-K=8; iteration 17 +// splits the staging by K-half: rows 0..255 are staged and barriered +// as before, rows 256..511 loads are ISSUED right after the barrier +// (before MMAC 0) and their ds_writes are DEFERRED to between the two +// K half-loops with one extra __syncthreads, so the second staging +// round trip overlaps MMACs 0..7 instead of sitting serially behind +// round 1's drain (same 4-load in-flight quantum — no burst widening); +// * K loop (16 steps, no global op for B, no barrier): A = one +// global_load_dwordx2 (load_a_frag8), B = one ds_read_b64 +// (load_b_lds8); iteration 7 deepens the single-buffered pipeline to a +// depth-8 register ring for A: load_a_frag8(a_ring[slot], step s+8) is +// issued in iteration s, consumed by the MMAC of iteration s+8, so no +// MMAC ever waits on an L2 round trip (the iteration-6 ISA drained +// vmcnt(0)+lgkmcnt(0) at every 8-step group entry, 6 full L2 round +// trips per block on the critical path — removed here). B stays LDS +// and is loaded per step (scheduler hoists/batches the ds_reads). +// With split-K=12 the per-wave dependent MMAC chain is 16 steps instead +// of 24, and the 8,320 B LDS tile (16 x 520) raises co-residency to +// 7 blocks/CU by LDS (VGPR-capped at 8; was 5 at 12,416 B) — the +// iteration-10 occupancy probe. +// a : [16, K] int8 row-major (ldm = k) +// b : [K, N] int8 row-major (ldm = n), raw logical layout (no pack) +// partials : [S][16][N] int32 planes (workspace, preallocated pre-Graph) +__global__ __launch_bounds__(kDummaThreads) void +w8a8_dumma_m16n16k32_sk4_partial_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int* __restrict__ partials, + int n, + int k) { + __shared__ int8_t s_b[kDummaN][kSliceKStride]; + + const int lane = static_cast(threadIdx.x); + const int tile_id = static_cast(blockIdx.x) / kSplitK; + const int split_id = static_cast(blockIdx.x) % kSplitK; + const int n0 = tile_id * kDummaN; + const int k_base = split_id * (k / kSplitK); + + // Stage the block's B K-slice once (8,192 B read exactly once from + // global, same bytes as the previous per-step loads; 8 dwordx4 + // cooperative loads instead of 128 byte loads). Iteration 17 splits the + // staging by K-half: rows 0..255 are staged and barriered as before, but + // the rows 256..511 loads are ISSUED right after the barrier (before any + // MMAC) and their LDS writes are DEFERRED to between the two K + // half-loops (one extra __syncthreads), so round 2's L2/DRAM round trip + // overlaps MMACs 0..7 instead of sitting serially behind round 1's drain + // (the accepted iteration-15 ISA drains the two 4-load rounds strictly + // one after the other). The in-flight VMEM quantum stays at 4 staging + // dwordx4 + the 16 A dwordx2 burst — no wider burst (the iteration-9/11 + // falsification) — and the deferred write + second barrier preserve the + // exact same LDS contents and k-ascending MMAC order (bit-identical). + const int8_t* __restrict__ b_slice = + b + static_cast(k_base) * n + n0; +#pragma unroll + for (int i = 0; i < kStageInstrs / 2; ++i) { + const int row = i * kDummaThreads + lane; // 0..255 + const int4 v = *reinterpret_cast(b_slice + row * n); + const int8_t* src = reinterpret_cast(&v); +#pragma unroll + for (int c = 0; c < kDummaN; ++c) { + s_b[c][row] = src[c]; + } + } + __syncthreads(); + // Round 2's loads are issued AFTER the A burst below (not before it): the + // waitcnt staircase counts ALL outstanding VMEM ops in issue order, so the + // MMAC 0..7 waits must not be gated by the cold B loads. Issued last, + // they complete in the background while MMACs 0..7 drain the L2-hot A + // burst, and the deferred writes below wait only vmcnt(3)..vmcnt(0) for + // them after the first half-loop. + int4 b_stage2[kStageInstrs / 2]; + + du::dumma::DUFragment + a_ring[kPrefetchDepth]; + // Iteration 15: two split accumulators. Steps 0..7 accumulate into acc0, + // steps 8..15 into acc1 (each chain k-ascending), then acc0 += acc1 — + // int32 addition is exact mod 2^32, so the grouped sum is bit-identical + // to the single-accumulator ascending sum (and to the reference order). + // The dependent MMAC chain per wave halves 16 -> 8 while the MMAC issue + // count (16), the A/B/plane byte traffic, the LDS tile, and the plane + // layout stay unchanged. + du::dumma::DUFragment + acc0; + du::dumma::DUFragment + acc1; + du::dumma::du_fill_fragment(acc0, 0); + du::dumma::du_fill_fragment(acc1, 0); + + // Iteration 7: depth-8 single-buffered register ring for A. The + // iteration-6 depth-1 ping-pong compiled to one 8-step unrolled group + // iterated 6x, and each group entry issued its first A load and then + // fully drained it (s_waitcnt vmcnt(0) lgkmcnt(0)) before the first MMAC + // — 6 full L2 round trips per block on the critical path. Here every A + // load for step s+8 is issued in iteration s (8 MMAC issue slots before + // its use), so no MMAC sits on an L2 round trip. kSliceSteps (16 at + // split-K=12) is a multiple of kPrefetchDepth (8), so the ring closes + // exactly on step 15; the modulo index is a compile-time constant under + // full unroll, keeping the ring in registers. B is read from LDS per + // step; the scheduler hoists and batches the ds_reads as before. MMAC + // order stays k ascending 0..15 -> bit-identical int32 accumulation. +#pragma unroll + for (int i = 0; i < kPrefetchDepth; ++i) { + load_a_frag8(a_ring[i], a + k_base + i * kDummaK, k); + } + // Round-2 (rows 256..511) B loads issued after the A burst: newest in the + // VMEM issue order, so the MMAC 0..7 waitcnt staircase tracks the A loads + // only, and these DRAM/L2 round trips overlap the first K half-loop. +#pragma unroll + for (int i = 0; i < kStageInstrs / 2; ++i) { + const int row = (kStageInstrs / 2 + i) * kDummaThreads + lane; // 256..511 + b_stage2[i] = *reinterpret_cast(b_slice + row * n); + } +#pragma unroll + for (int s = 0; s < kSliceSteps / 2; ++s) { + const int slot = s & (kPrefetchDepth - 1); + du::dumma::DUFragment + b_frag; + load_b_lds8(b_frag, s_b[0], s * kDummaK); + du::dumma::du_mma_sync(acc0, a_ring[slot], b_frag, acc0); + // Unconditional refill (s + kPrefetchDepth <= 15 < kSliceSteps): the + // depth-8 ring closes exactly on step 15; the compiler keeps all 16 A + // fragments in the accepted iteration-14 pre-loop burst (one + // global_load_dwordx2 burst + one s_waitcnt vmcnt(15) lgkmcnt(7)). + load_a_frag8(a_ring[slot], a + k_base + (s + kPrefetchDepth) * kDummaK, k); + } + // Drain round 2's loads (rows 256..511) and write them to LDS. These + // ds_writes target s_b rows 256..511 only, disjoint from the rows 0..255 + // that MMACs 0..7 read, so they do not race the first half-loop; the + // __syncthreads makes them visible to every lane before MMAC 8 reads + // them. Same bytes, same LDS layout, same k-ascending MMAC order — the + // result is bit-identical to the single-staged iteration-15 kernel. +#pragma unroll + for (int i = 0; i < kStageInstrs / 2; ++i) { + const int row = (kStageInstrs / 2 + i) * kDummaThreads + lane; + const int8_t* src = reinterpret_cast(&b_stage2[i]); +#pragma unroll + for (int c = 0; c < kDummaN; ++c) { + s_b[c][row] = src[c]; + } + } + __syncthreads(); + // Steps 8..15 -> acc1. No refill needed: s + kPrefetchDepth >= + // kSliceSteps, all remaining A fragments were already loaded above. +#pragma unroll + for (int s = kSliceSteps / 2; s < kSliceSteps; ++s) { + const int slot = s & (kPrefetchDepth - 1); + du::dumma::DUFragment + b_frag; + load_b_lds8(b_frag, s_b[0], s * kDummaK); + du::dumma::du_mma_sync(acc1, a_ring[slot], b_frag, acc1); + } + // Combine the two chains: acc0 (steps 0..7) += acc1 (steps 8..15) — + // ascending split sum, bit-identical mod 2^32. +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc0.x[i] += acc1.x[i]; + } + + // Store the 16x16 int32 partial to the split plane (verified gfx928 + // accumulator ownership: row = lane & 15, col_mod4 = lane >> 4, + // frag.x[i] holds column col_mod4 + 4*i). + const int row = lane & 15; + const int col_mod4 = lane >> 4; + int* __restrict__ plane = partials + split_id * (kDummaM * n) + row * n + n0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + plane[col_mod4 + 4 * i] = acc0.x[i]; + } +} + +// Split-K combine kernel (exact (16, 1536, 6144) arm, S = kSplitK): 96 +// one-wave blocks; each block sums the S int32 partial planes for its 16x16 +// tile in ascending split order (split 0 covers k [0,512), ..., split 11 +// covers [5632,6144)) — bit-identical int32 accumulation order to the +// reference (int32 addition is exact mod 2^32, so any ascending split +// grouping of the same products is value-identical) — then applies +// x_scale[row] * weight_scale[col] and stores scaled bf16. Same-stream +// ordering guarantees the partial kernel's writes are visible; no barrier, +// no atomic, no LDS. +// +// Iteration 14 (vectorized plane reads, fixed split-K=12): each lane now +// owns 4 CONSECUTIVE n columns of one row (row = lane >> 2, tcol4 = +// (lane & 3) << 2) instead of 4 column-aligned elements in different rows +// (e = lane + 64*i), so all S plane reads per lane group into 16-B-aligned +// int4 loads: S int4 global_load_dwordx4 per lane (was 4*S scalar +// global_load_dword spread over four serial load-drain groups per lane, see +// the iteration-13 ISA staircase vmcnt drains) and the bf16 epilogue +// becomes one 8-B-aligned uint64 store (was four 2-B stores). The +// per-element split sum is unchanged (ascending s = 0..S-1 over the same +// planes) so the output is bit-identical; the partial kernel, grid, LDS +// tile, workspace, identity pack, scalar fallback, and exact int32 order +// are untouched. +// partials : [S][16][N] int32 planes +// x_scale : [16, 1] float +// weight_scale : [N, 1] float +// out : [16, N] bf16 row-major +__global__ __launch_bounds__(kDummaThreads) void +w8a8_dumma_m16n16k32_sk4_combine_kernel( + const int* __restrict__ partials, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n) { + const int n0 = static_cast(blockIdx.x) * kDummaN; + const int lane = static_cast(threadIdx.x); + const int plane_stride = kDummaM * n; // int32 per split plane + + // 16-B-aligned int32 base: row * n and n0 are multiples of 4 int32 and + // tcol4 is a multiple of 4 int32 (16 B). + const int row = lane >> 2; + const int tcol4 = (lane & 3) << 2; + const int idx = row * n + n0 + tcol4; + + int4 acc = *reinterpret_cast(partials + idx); +#pragma unroll + for (int s = 1; s < kSplitK; ++s) { + const int4 v = *reinterpret_cast( + partials + s * plane_stride + idx); + acc.x += v.x; + acc.y += v.y; + acc.z += v.z; + acc.w += v.w; + } + + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + n0 + tcol4); + // Repair 2 (compile fix, no design change): __float2bfloat16 returns the + // INTERNAL __hip_bfloat16, which exposes ~11 implicit integer/float + // conversion operators (amd_hip_bf16.h) and has no implicit conversion to + // the public hip_bfloat16 (whose float ctor is explicit, so copy-init + // `const hip_bfloat16 bN = __float2bfloat16(...)` fails overload + // resolution with "no viable conversion" — the 4 compile errors). Each + // bf16 is therefore direct-initialized from the float scalar via the + // public explicit hip_bfloat16(float) ctor, whose float_to_bfloat16 RNE + // is byte-identical to __float2bfloat16's float_2_bfloatraw on gfx928 + // (same rounding and NaN handling), so every output bit is unchanged. + const hip_bfloat16 b0(static_cast(acc.x) * xs * ws.x); + const hip_bfloat16 b1(static_cast(acc.y) * xs * ws.y); + const hip_bfloat16 b2(static_cast(acc.z) * xs * ws.z); + const hip_bfloat16 b3(static_cast(acc.w) * xs * ws.w); + // hip_bfloat16 (public type, amd_hip_bfloat16.h) exposes its raw 16-bit + // bf16 pattern as the public `data` member; it has no conversion to + // __hip_bfloat16_raw (that operator exists only on the internal + // __hip_bfloat16 class), so pack the four 16-bit patterns directly. + const uint64_t packed = + (static_cast(b0.data)) | + (static_cast(b1.data) << 16) | + (static_cast(b2.data) << 32) | + (static_cast(b3.data) << 48); + *reinterpret_cast(out + idx) = packed; +} + +} // namespace + +// Host launch symbol consumed by csrc/bindings.cpp (extern "C"). +// Graph-safe: launches only on the caller-provided stream; performs no +// allocation, compilation, autotuning, weight packing, host/device +// synchronization, or default-stream launch. `workspace` holds the +// split-K int32 partial planes for the assigned shape (preallocated before +// Graph capture by the API contract). +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + hip_bfloat16* out_ptr = reinterpret_cast(out); + const int total = m * n; + const int blocks = (total + kScalarBlockThreads - 1) / kScalarBlockThreads; + + if (m == 16 && n == 1536 && k == 6144) { + // Assigned shape minimax_tp8_qkv_proj_and_indexer_qk_m16. + // Iteration 10 (HIP-only occupancy probe): split-K=12 partial GEMM + // (grid 1,152 = 9.6 blocks/CU; B-only LDS staging at 8,320 B/block -> + // 7 blocks/CU co-residency by LDS, up from 5 at split-K=8; 16-step K + // loop; one barrier; combine included in the operator wall). The guard + // keeps the paired M=2 (and any other) shape on the generic scalar + // fallback below. Workspace 12 x 98,304 B = 1,179,648 B <= the + // 1,572,864 B (16-plane) capacity allocate_workspace guarantees for + // this shape. + const int64_t sk_need = + static_cast(kSplitK) * kDummaM * n * sizeof(int32_t); + if (workspace != nullptr && workspace_bytes >= sk_need) { + int* partials = reinterpret_cast(workspace); + const int tiles = n / kDummaN; // 96 + hipLaunchKernelGGL(w8a8_dumma_m16n16k32_sk4_partial_kernel, + dim3(tiles * kSplitK), dim3(kDummaThreads), 0, + stream, a, b, partials, n, k); + hipLaunchKernelGGL(w8a8_dumma_m16n16k32_sk4_combine_kernel, + dim3(tiles), dim3(kDummaThreads), 0, stream, + partials, x_scale, weight_scale, out_ptr, n); + return; + } + // Workspace too small for kSplitK partial planes: correctness-first + // scalar fallback for the exact shape (never writes out of bounds). + } + + // Generic scalar fallback for every unmatched (m, n, k), including the + // paired M=2 API shape with the same (N, K) and the M=3072 prefill shape. + hipLaunchKernelGGL(w8a8_scalar_gemm_kernel, dim3(blocks), + dim3(kScalarBlockThreads), 0, stream, a, b, x_scale, + weight_scale, out_ptr, m, n, k); +} + +// Host launch symbol consumed by csrc/bindings.cpp (extern "C"). +// Identity device-to-device copy (packed layout == logical [K, N] row-major +// layout). Runs through the optional zth_w8a8::pack_weight op, outside the +// timed region and outside Graph capture. For unmatched (K, N) this +// identity copy remains the generic fallback. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + hipMemcpyAsync(packed_weight, raw_weight, + static_cast(k) * static_cast(n) * + sizeof(int8_t), + hipMemcpyDeviceToDevice, stream); + hipMemcpyAsync(packed_weight_scale, weight_scale, + static_cast(n) * sizeof(float), + hipMemcpyDeviceToDevice, stream); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_down_proj.hip new file mode 100644 index 00000000..f5930cb7 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_down_proj.hip @@ -0,0 +1,1409 @@ +// @@variant shape=minimax_tp8_shared_down_proj_m16 commit=2d8b403746ea59e59c41eabce07da2e6b2447072 added=2026-08-28 +// median_us=8.953 p90_us=9.069 +// source=minimax-dsh-tp8-m16-1-78402260 +// MetaInfer W8A8 INT8 GEMM - HIP implementation for Hygon K500SM_AI / gfx928. +// +// Iteration 1 (minimal DUMMA bootstrap): the assigned M=16 MiniMax TP8 +// shared_gate_up_proj shape (16x768x6144) ran on a minimal 16x16x32 DUMMA +// tile with one 64-lane wavefront per output tile, explicit int32 +// accumulation, direct row-major global fragment loads, and no LDS staging / +// no cross-wave barriers. Measured 182.1 us median Graph replay (0.574x vs +// the 104.527 us Triton baseline): the one-wave grid of 48 blocks fills only +// 48 of 120 CUs (0.4 blocks/CU) and each block serially streams its full +// K=6144 loop, so the kernel is launch-geometry/latency bound, not ALU bound. +// +// Iteration 2 (architecture round, grid parallelism): keep the exact same +// zero-barrier one-wave direct-fragment tile, but split K into kSplitK = 5 +// independent 32-aligned slices per output tile: +// * grid = 48 N tiles x 5 K slices = 240 one-wave blocks = exactly 2 +// blocks per device CU (120 CUs) -- the two-blocks-per-CU latency-hiding +// target of the control plane; +// * each partial block publishes its int32 accumulator fragment to one +// workspace plane (every (tile, slice) pair is written exactly once per +// launch, so no workspace clear is needed); +// * a 48-block 64-thread combine kernel sums the five planes in ascending +// slice order (exact int32 order, bit-identical to the scalar reference), +// applies x_scale[m] * packed_weight_scale[n], and stores bf16; both +// kernels are enqueued on the caller stream so the combine cost is inside +// the timed operator / Graph replay; +// * every other shape, including the paired M=2 API shape with the same +// (N, K) and the paired 16x6144x384 down_proj shape, keeps the scalar +// generic fallback below; +// * launch_pack_w8a8_weight stays an identity device-to-device copy +// (generic fallback for every (K, N)). +// +// Iteration 3 (architecture round, occupancy probe): keep the exact +// iteration-2 architecture (zero-barrier one-wave direct-fragment tile, +// split-K partial planes + 48-block combine in the timed Graph) but raise +// kSplitK from 5 to 10, so the grid doubles from 240 to 480 one-wave blocks +// = exactly 4.0 blocks per device CU (120 CUs). K = 6144 = 192 chunks of 32 +// = 10*19 + 2, so slices 0..1 carry 20 chunks (640 elements) and slices 2..9 +// carry 19 chunks (608 elements): every boundary stays a multiple of the +// DUMMA K step and the imbalance is <= 6%. The workspace then holds 10 +// int32 partial planes (491520 B), well inside the 786432 B (16-plane) the +// API allocates for this shape. Hypothesis: at 2.0 blocks/CU the partial +// kernel is still global-load-latency bound (PMC: ~6% VALU issue +// utilization, 45.8 us perturbed partial vs 50.6 us total median), so +// doubling the independent resident one-wave streams per CU should hide more +// of that latency; the combine kernel grows from 5 to 10 plane-sums but +// stays a tiny 48-block 64-thread launch over L2-hot plane traffic. +// Measured: 45.47 us median Graph replay (2.30x), so occupancy scaling +// helped (~10%) but the direct-fragment load path remains the binding cost. +// +// Iteration 5 (load-path round): ISA on the accepted iteration-3 code object +// shows the library fragment loaders expand each k32 chunk into 16 +// global_load_ubyte instructions per lane -- 8 contiguous bytes for A that +// the compiler never fuses plus 8 bytes strided by ldm for B -- followed by +// ~15 s_waitcnt vmcnt before a single v_mmac_i32_16x16x32_i8: 147456 +// vmem_read_instructions over the 480-block partial launch. Fix: keep the +// exact zero-barrier one-wave m16n16k32 tile, the 480-block split-K=10 grid, +// and the 48-block combine, but replace the byte-gather fragment loads with +// one aligned 8-byte (global_load_dwordx2) load per lane per fragment: +// * A stays the logical row-major activation [16, K]: lane (row = lane&15, +// col_group = (lane>>4)*8) owns A[row][k0+col_group .. k0+col_group+7], +// which is 8 contiguous bytes (ldm = K = 6144); +// * B is packed once, outside the timed region and outside Graph capture, +// into the n-major transpose packed[n*K+kk] = raw[kk*n+n] for the exact +// (k,n) == (6144,768) (same byte count and buffer, so Graph addresses +// are unchanged; the identity copy remains the pack fallback for every +// other (k,n)); lane (row = lane&15, col_group = (lane>>4)*8) then owns +// packed[(n0+row)*K + k0+col_group .. +7], also 8 contiguous bytes; +// * du_mma_sync consumes the identical byte pattern the library loader +// would produce, so the int32 accumulation stays bit-identical +// (k-ascending per slice, ascending-slice combine, 0 mismatches); +// * the generic scalar fallback decodes the n-major pack for +// (n,k) == (768,6144) (including the paired M=2 API shape with the same +// (N,K)) and the workspace-too-small direct kernel reads B with ldm = k +// through the same packed pointer. +// Falsifiable: partial-kernel vmem_read_instructions must drop from 147456 +// to ~18432 (2 dwordx2 loads x ~19.2 chunks x 480 blocks) and the median +// must beat the 45.47 us iteration-3 best. +// +// down_proj bootstrap (this round, M=16 N=6144 K=384): the paired MiniMax +// TP8 shared_down_proj shape ran the generic scalar fallback (45.83 us +// median Graph replay vs the 23.152 us Triton baseline). Establish the +// minimal m16n16k32 DUMMA tile for this exact shape with the same proven +// load path as the accepted gate_up iterations 5-16 (n-major packed B + +// one aligned 8-byte global load per lane per fragment + depth-1 register +// prefetch), one 64-lane wavefront per block, one 16-column N tile per +// block, no LDS and no cross-wave barrier. Grid = 6144/16 = 384 one-wave +// blocks = 3.2 blocks/CU over 120 CUs; K = 384 = 12 k32 chunks. B is +// n-major packed for the exact (k, n) == (384, 6144) pair by +// launch_pack_w8a8_weight (outside the timed region), and the scalar +// generic fallback decodes the same n-major pack for the paired M=2 API +// shape with the same (N, K). Measured: 29.947 us median Graph replay +// (0.773x vs the 23.152 us Triton baseline; -53% vs the scalar 45.83 us), +// so the direct-fragment tile is correct and Graph-safe but the one-wave +// grid is still latency/issue bound: 384 tiny blocks each run a strict +// load -> mma -> load -> mma chain with only depth-1 prefetch slack. +// +// Iteration 2 (architecture round, grid parallelism): keep the exact tile, +// the packed-B load path, and the depth-1 prefetch, but regroup the 384 +// one-wave blocks into 192 two-wave blocks: each block runs 2 adjacent +// 16-column N tiles (one wavefront per tile, wave w computes tile +// 2*blockIdx.x + w; grid = 6144/16/2 = 192 blocks x 128 threads = 1.6 +// blocks/CU over 120 CUs, still enough independent blocks to cover every +// device CU). Two guaranteed co-resident wavefronts per block let the CU +// scheduler interleave one wave's global-load latency with the other's +// v_mmac issue, and block dispatch prologue halves vs 384 blocks. No LDS, +// no barrier, no split-K: every output element is still accumulated +// k-ascending by one wavefront, so the int32 result stays bit-identical to +// iteration 1 and the scalar reference. Falsifiable: median Graph replay +// must beat the 29.947 us iteration-1 official best (p90 guard vs 32.091). +// Explored and rejected: 4 waves/block x 1 tile (96 blocks < 120 CUs +// underfills the device) and 1 wave x 2/4 adjacent N tiles (192/96 blocks +// with a sequential per-wave tile loop that adds no intra-CU overlap). +// +// Iteration 4 (architecture/pipeline round, bounded LDS prefetch): the +// iteration-2 code object (this exact shape's ISA artifact) shows each +// wavefront still runs 12 fully serialized wait rounds -- 2 x +// global_load_dwordx2 -> s_waitcnt vmcnt(0) -> 1 x v_mmac per k32 chunk, +// K = 384 = 12 chunks -- with the compiler sinking every register +// prefetch into the loop head, so each v_mmac stalls for nearly the full +// global/L2 latency and the kernel is a pure per-wave load-wait chain. +// Regrouping (iteration 2) moved nothing and splitting the chain +// (iteration 3's split-K=2 + separate combine kernel) only moved the wait +// cost into a second launch (31.04 us, rejected). Fix: bounded LDS +// prefetch -- prefetch the ENTIRE per-tile K=384 window into LDS in one +// coalesced stage per block (one wait group + one __syncthreads), then +// run the K loop with zero global loads: +// * A [16, 384] = 6144 B staged ONCE per block into s_a (padded row +// stride 392 = 384 + 8: every 8-byte ds_read_b64 fragment read stays +// 8-B aligned and conflict-free, see the kernel doc) and SHARED by +// both wavefronts: the same 6144 A bytes are reused by the block's +// two output tiles, halving A global traffic vs the iteration-1/2 +// per-wave re-read (1.15 MB vs 2.3 MB per launch); +// * each wavefront stages its own tile's B slice [16, 384] = 6144 B +// (one contiguous 6-KiB region of the n-major pack) into s_b[wave]; +// every B byte is read from global exactly once (2.36 MB total; no +// cross-tile B reuse is possible -- tiles are disjoint column slices); +// * the stage is 18 aligned 8-byte global_load_dwordx2 per lane (6 A + +// 12 B) + 18 ds_write_b64 + one barrier, so the 12 per-chunk +// global-latency waits collapse into ONE ~L stage wait per wavefront +// and the K loop becomes 12 rounds of ds_read_b64 + v_mmac with only +// ~30-cycle LDS latency (hidden by the loop's register prefetch); +// * grid stays 192 blocks x 128 threads = 1.6 blocks/CU over 120 CUs +// (enough independent blocks to cover every device CU), no workspace, +// LDS = 18816 B/block (2 co-resident blocks = 37.6 KiB < 64 KiB/CU); +// * the K loop consumes the identical k-ascending bytes as iteration +// 1/2 (each output element still accumulated by exactly one +// wavefront), so the int32 result stays bit-identical to the scalar +// reference (0 mismatches expected). +// Falsifiable: median Graph replay must beat the 29.755 us iteration-2 +// shadow median (official best 29.947 us) and p90 must stay within the +// 1.05x guard of 31.552 us, with unchanged Graph capture/replay semantics +// (one launch on the caller stream, no allocation/sync/compile). +// +// Iteration 16 (final conditional inline-asm round -> HIP-only +// consolidation): raw asm stays forbidden (isa_policy.phase=hip_only, +// plateau=false: recent_valid_improvements_percent = [+3.44, +2.02, -4.52] +// after the rejected split-K=15 three-wave probe, so the required three +// recent valid HIP candidates within [-2%, +2%) of best do not exist), so +// this round consolidates the two-kernel operator into ONE launch with the +// qkv-proven monotonic-arrival fused combine tail (the exact pattern +// accepted on this stack for hy3_tp4_qkv_proj_m16, round 23): +// * the partial kernel becomes template ; FUSED=true adds the +// tail: after wave 1 publishes the plane, one __syncthreads, then +// (wave 1, lane 0) does __threadfence() + atomicAdd(&counters[tile], 1) +// + __threadfence(); the block whose atomicAdd returns +// (arrived % kSplitK) == kSplitK - 1 is the LAST arrival for its tile +// and sums the tile's 10 planes in ascending slice order with the exact +// iteration-12 combine epilogue (int4 plane loads, x_scale scalar + +// weight_scale float4 hoisted, hip_bfloat16 conversion, 4 x 2-byte +// stores), then writes the bf16 output; +// * the 48-block combine kernel and its launch gap disappear from the +// timed operator; the per-tile monotonic counters (n_tiles x int32 = +// 192 B) live in the last 192 B of the 786432-B (16-plane) contract +// workspace (plane 15's tail; planes 0..9 are the only planes used), +// zeroed once per workspace by a static-pointer-guarded hipMemsetAsync +// on the caller stream before the first launch (eager warmup runs +// before Graph capture, so the memset is never part of the captured +// graph and every replay is exactly one kernel launch); +// * the two-kernel path (FUSED=false + the combine kernel) remains as the +// fallback for workspaces below the 16-plane contract, byte-identical +// to iteration 14; +// * per output element the int32 accumulation is still the +// ascending-slice sum over the same bytes with the same conversion, so +// the bf16 output is bit-identical to iteration 14 (0 mismatches +// expected) and Graph capture/replay semantics are unchanged (one launch +// on the caller stream, no allocation/sync/compile in the launcher). +// +// Logical operation (see int8_w8a8_gemm_api.py): +// out[m, n] = bf16(int32_dot(x_q[m, :], packed_weight[:, n]) +// * x_scale[m, 0] * packed_weight_scale[n, 0]) +// +// Headers are included in the DTK-known-good order: hip_runtime, hip_bfloat16, +// then du_mma (this DTK's du_mma.h is not self-contained when included before +// the HIP runtime headers). +// +// The timed entry point (launch_w8a8_gemm) performs no allocation, no +// compilation, no autotuning, no weight packing, no host/device +// synchronization, and no default-stream launch: it only computes launch +// geometry and enqueues on the caller-provided stream, so it stays safe under +// torch.cuda.CUDAGraph capture and replay. + +#include +#include +#include + +#include + +namespace { + +// gfx928 native wavefront is 64 lanes; the block size must be a multiple of 64. +constexpr int kScalarBlockThreads = 128; +constexpr int kCopyBlockThreads = 256; + +// Minimal DUMMA tile for the assigned M=16 decode shape: INT8 m16n16k32 with +// int32 accumulation, one 64-lane wavefront owning one independent fragment. +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; +constexpr int kDummaWaveThreads = 64; +// Iteration 14: the split-K partial kernel runs 128 threads = 2 wavefronts +// per block (block size stays a multiple of the 64-lane wavefront). +constexpr int kPartialBlockThreads = 2 * kDummaWaveThreads; + +// Split-K geometry for the assigned MiniMax TP8 gate_up shape (iteration 3 +// occupancy probe): 10 independent K slices per 16x16 output tile -> grid = +// 48 N tiles x 10 slices = 480 one-wave zero-barrier blocks = exactly 4.0 +// blocks per device CU (120 CUs), doubling the iteration-2 split-K = 5 point +// (240 blocks, 2.0 blocks/CU, 50.55 us median). K = 6144 = 192 chunks of 32 +// = 10*19 + 2, so slices 0..1 carry 20 chunks (640 elements) and slices 2..9 +// carry 19 chunks (608 elements); every slice boundary stays a multiple of +// the DUMMA K step and the imbalance is <= 6%. The workspace holds kSplitK +// int32 partial planes of [16][768] (49152 B each); the Python API allocates +// 16 planes (786432 B) for this shape, so kSplitK = 10 (491520 B) always +// fits the workspace and 256-B stage alignment. +constexpr int kSplitK = 10; +constexpr int kM16N768Elems = 16 * 768; // one int32 partial plane (12288) + +// Iteration 16: the fused combine tail needs the full 16-plane contract +// workspace the Python API allocates for this shape (786432 B): the +// per-tile monotonic arrival counters live in the last n_tiles*4 B of +// plane 15's tail (planes 0..9 are the only partial planes used, so there +// is no overlap), exactly like the accepted qkv fused path. +constexpr int64_t kSkWorkspaceContractBytes = + static_cast(16) * kM16N768Elems * sizeof(int32_t); // 786432 B + +// Scalar INT8 dot-product kernel: one thread -> one out[m, n] element. +// +// a: [M, K] int8 row-major (logical activation, stride K) +// b: [K, N] int8 row-major (identity-packed logical weight, stride N) +// x_scale: [M] fp32 +// weight_scale: [N] fp32 +// out: [M, N] bf16 +// +// Linear indexing keeps adjacent lanes on adjacent n, so B-column reads and +// the bf16 stores are coalesced in the fastest-changing N dimension. +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int linear = blockIdx.x * blockDim.x + threadIdx.x; + if (linear >= m * n) { + return; + } + const int row = linear / n; + const int col = linear - row * n; + + const int8_t* a_row = a + row * k; + + // The exact gate_up (n, k) == (768, 6144) and down_proj (n, k) == + // (6144, 384) pairs use the n-major packed weight (packed[n*K+kk] = + // raw[kk*n+n], produced by launch_pack_w8a8_weight); all other shapes + // keep the identity [K, N] row-major layout. Decode both so the generic + // fallback -- including the paired M=2 API shapes with the same (N, K) -- + // stays correct. + const bool nmajor_packed = + (n == 768 && k == 6144) || (n == 6144 && k == 384); + const int8_t* b_col = + nmajor_packed ? (b + static_cast(col) * k) : (b + col); + const int64_t b_stride = nmajor_packed ? 1 : static_cast(n); + + // Exact int32 dot product. Largest assigned K is 6144: + // 6144 * 128 * 128 = 100,663,296 < 2^31, so the accumulation is exact and + // cannot overflow. int8 products are exact in int32, and the int32 sum is + // bit-identical to the float32 dot of (A.float() @ B.float()). + int32_t acc = 0; + for (int i = 0; i < k; ++i) { + acc += static_cast(a_row[i]) * + static_cast(b_col[i * b_stride]); + } + + // Reference: (A.float() @ B.float()) * x_scale * weight_scale.T -> bf16. + // Both scales are applied left-associatively as in the torch reference. + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + // hip_bfloat16's supported float conversion rounds to nearest even, which + // matches the torch float->bfloat16 conversion. + out[linear] = hip_bfloat16(scaled); +} + +// Minimal INT8 DUMMA kernel for M=16 decode: one wave (64 lanes) per block, +// one 16x16 output tile per block, explicit int32 accumulation. Reachable +// only under the exact (16, 768, 6144) shape guard when the workspace cannot +// hold the ten split-K partial planes, so B is the n-major packed weight +// (packed[n*K+kk] = raw[kk*n+n]) and is loaded with ldm = k. +// +// a: [16, K] int8 row-major activation (M == 16 is enforced by the caller +// guard; fragment loads need ldm = k) +// b: [N, K] int8 n-major packed weight for the exact guard that is the +// only caller (fragment loads need ldm = k) +// x_scale: [16] fp32 +// weight_scale: [N] fp32 +// out: [16, N] bf16 +// +// Each K step loads a 16x32 A fragment and a 32x16 B fragment directly from +// global memory with du::dumma::du_load_matrix_sync (row_major fragments, +// 8-byte-aligned base offsets because K and the N tile origin are multiples +// of 32/16 bytes) and accumulates with du::dumma::du_mma_sync into an int32 +// accumulator fragment. No LDS, no __syncthreads, no cross-wave barrier: +// each block is a single wavefront and all stores are fragment-owner direct. +__global__ __launch_bounds__(kDummaWaveThreads) void +w8a8_gemm_dumma_m16_direct_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n, + int k) { + const int lane = static_cast(threadIdx.x); // 0..63 + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Complete K loop accumulated in exact int32 (same values as the scalar + // reference; K = 6144 * 128 * 128 < 2^31 so no overflow). B is the n-major + // packed weight for this shape, so the row_major B fragment loader with + // ldm = k reads packed[(n0 + lane&15)*k + k0 + (lane>>4)*8 + i], exactly + // W[k0 + (lane>>4)*8 + i][n0 + lane&15]. + for (int k0 = 0; k0 < k; k0 += kDummaTileK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); + du::dumma::du_load_matrix_sync( + b_frag, b + static_cast(n0) * k + k0, k); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Direct fragment epilogue. The int32 accumulator ownership for this DTK + // (verified against du_store_matrix_sync with mem_row_major) is: + // row = lane & 15, col_group = lane >> 4, + // acc_frag.x[i] holds output column col_group + 4*i. + // The four 2-byte bf16 values of one row are written by lanes l, l+16, + // l+32, l+48 (same lane&15); the whole 16x16 tile is only 512 bytes of + // output so the scattered 2-byte stores are absorbed by L2 write combining. + const int row = lane & 15; + const int col_group = lane >> 4; + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = n0 + col_group + 4 * i; + const float scaled = + static_cast(acc_frag.x[i]) * xs * weight_scale[col]; + out[row * n + col] = hip_bfloat16(scaled); + } +} + +// Iteration-5 (HIP-only staging round: B-only staging, A register-resident) +// INT8 DUMMA kernel for the assigned MiniMax TP8 down_proj shape M=16, +// N=6144, K=384, retiled in iteration 8 (HIP-only resource round) from two +// 64-lane wavefronts per block (128 threads, two adjacent 16x16 tiles per +// block) to ONE 64-lane wavefront per block (64 threads, one 16x16 tile per +// block), explicit int32 accumulation, the n-major packed-B load path of +// the iteration-1/2/4 lineage. +// +// a: [16, K] int8 row-major activation (M == 16 is enforced by the caller +// guard) +// b: [N, K] int8 n-major packed weight for the exact (k, n) == (384, 6144) +// guard (packed[n*K+kk] = raw[kk*n+n], produced outside the +// timed region by launch_pack_w8a8_weight) +// x_scale: [16] fp32 +// weight_scale: [N] fp32 +// out: [16, N] bf16 +// +// Why B-only staging (iteration-4 ISA, exact code object in the +// iteration-4 artifact): iteration 4 staged BOTH A and B into LDS (18 +// dwordx2 global loads + 18 ds_write_b64 + one __syncthreads per block, +// then a 12-chunk K loop of 2 x ds_read_b64 -> v_mmac) and measured the +// accepted 20.784 us median Graph replay vs the 23.152 us Triton baseline. +// The iteration-4 code object shows two residual costs that this round +// removes: +// * a block-wide s_barrier after the stage forces the two wavefronts to +// rendezvous, so wave 0's K loop cannot start until wave 1's cold B +// loads have arrived (per-block wall time = the SLOWEST of the two +// waves' stages); +// * the A fragment is staged through LDS (6 ds_write + 12 ds_read per +// wave) even though A is a 6-KiB tensor that is L1/L2-hot and +// byte-identical for every one of the 192 blocks. +// The per-chunk wait is B-dominated (A is L2-hot; iteration-4 design note), +// so the load-bearing part of iteration 4 is the ONE-wait B stage; A needs +// no LDS round trip at all -- each lane's entire A fragment for K = 384 is +// exactly 12 x 8 B = 96 B = 24 VGPRs, preloadable once. +// +// Fix (B-only staging, A register-resident, zero barriers): +// * s_a disappears; LDS is only s_b[2][16 * 392] = 12544 B/block (two +// co-resident blocks = 25.1 KiB < 64 KiB/CU); +// * each lane preloads its own A fragment bytes for all 12 k32 chunks +// into registers once (12 aligned 8-byte global_load_dwordx2 into +// a_reg[12]; lane (row = lane & 15, col_group = (lane >> 4) * 8) +// needs a[row*k + j*32 + col_group .. +8) per chunk j; all 12 loads +// are independent and L1/L2-hot, so they join the same one-wait stage +// group as B and never appear in the K loop); +// * each wavefront stages its own tile's B slice [16, 384] = 6144 B (one +// contiguous 6-KiB region of the n-major pack) into s_b[wave] exactly +// as iteration 4: 12 aligned 8-byte dwordx2 per lane + 12 ds_write_b64, +// padded row stride 392 = 384 + 8 (392 % 8 == 0 keeps every 8-byte +// ds_read_b64 aligned; 392/8 = 49 odd keeps the 16 lanes of each bank +// phase on 16 distinct 8-byte LDS slots, conflict-free); every B byte +// is read from global exactly once (2.36 MB per launch); +// * the __syncthreads disappears: each wave writes and reads only its +// own s_b[wave] (A never touches LDS), so the two wavefronts of a +// block are fully independent and the CU scheduler can overlap wave +// 0's K loop with wave 1's stage/loop instead of a barrier rendezvous; +// * the K loop is fully unrolled over the 12 compile-time chunks +// (kChunks = 384 / 32; the exact-shape caller guard fixes k == 384): +// chunk j's a_frag comes from a_reg[j] (a register copy, no wait) and +// chunk j's b_frag from ONE ds_read_b64 (s_b[wave] + row*392 + j*32 + +// col_group); the unrolled straight-line body lets the compiler hoist +// the independent ds_read_b64s ahead of the dependent v_mmac chain +// (the iteration-4 rotated for(;;) loop was compiled with the reads at +// the loop head consumed by a per-chunk s_waitcnt lgkmcnt(0) right +// before each v_mmac -- the same wait-split the iteration-2 ISA showed +// for global loads). +// +// Stage per lane: 24 aligned 8-byte dwordx2 global loads (12 A register +// preload + 12 B) + 12 ds_write_b64, ONE wait group, ZERO barriers. +// +// Iteration 8 (HIP-only resource round, occupancy limiter = waves per +// block): the iteration-5 exact code object (profiles/.../iteration6/ +// current-best-isa, source digest 298b5137) records resources vgpr 81 / +// sgpr 24 / LDS 12544 B / spills 0, and the launch geometry is 192 two-wave +// blocks x 128 threads = 384 wavefronts over 120 CUs x 4 SIMDs = 480 SIMD +// slots = 0.8 waves/SIMD: balanced dispatch puts 2 blocks (4 waves, one per +// SIMD) on 72 CUs and 1 block (2 waves, two SIMDs idle) on 48 CUs, so every +// active SIMD hosts EXACTLY one wavefront and all stage/loop/epilogue +// memory stalls are fully exposed per wave. VGPR 81 allows 3 waves/SIMD +// (243 <= 256) and LDS 12544 B allows 5 blocks/CU (62.7 KiB < 64 KiB), so +// neither register nor LDS capacity binds; the binding occupancy limiter is +// the waves-per-block grouping granularity. Fix: split each two-wave block +// into two one-wave blocks (grid 192 -> 384, block 128 -> 64 threads, LDS +// 12544 -> 6272 B/block): balanced dispatch then puts 3 blocks on 96 CUs +// and 4 on 24 CUs, so no CU runs fewer than 3 resident wavefronts (was 2 on +// 48 CUs) and every CU keeps 3-4 independent stage/K-loop streams to +// overlap the single stage wait and the epilogue's serialized weight_scale +// loads. B is still read from global exactly once (2.36 MB/launch); the +// only traffic increase is A re-reads (1.18 -> 2.36 MB), all L1/L2 hits on +// the byte-identical 6-KiB A tensor (no repeated HBM reads). +// Grid = 6144/16 = 384 one-wave blocks (3.2 blocks/CU over 120 CUs, +// enough independent blocks to cover every device CU). The K loop +// consumes the identical k-ascending bytes of the same fragments (A from +// the logical row-major tensor, B from the n-major pack 1:1 through LDS), +// so each output element is still accumulated k-ascending (chunks 0..11) +// by exactly one wavefront and the int32 result -- and therefore the bf16 +// output -- is bit-identical to iteration 4/5 and the scalar reference (0 +// mismatches expected). No split-K, no atomics. +// +// Iteration 9 (HIP-only epilogue pipeline round on the iteration-8 +// geometry): the accepted iteration-8 exact code object (isa artifact in +// iterations/.../iteration8, source digest f8bd16a4) records the same +// per-wave instruction mix as iteration 5 (29 global_load, 30 s_waitcnt: +// 12 B stage + 12 A preload + 4 epilogue weight_scale dword + 1 x_scale +// dword) and still shows the per-wave serial tail after the last v_mmac: +// three global_load_dword (weight_scale +16/+32/+48 bytes from the v[4:5] +// base) each followed by a serialized s_waitcnt vmcnt(1) before its +// multiply-store -- a strict load -> wait -> mul -> store chain per output +// element for elements 1..3, fully exposed because the K loop is already +// finished and every active SIMD hosts exactly one wavefront (384 one-wave +// blocks over 120 CUs x 4 SIMDs = 0.8 waves/SIMD, so no co-resident +// partner exists to overlap the tail). Fix (the iteration-7 in-round- +// repaired pattern, applied to the accepted iteration-8 geometry): +// preload x_scale[row] (scalar) and the epilogue's four weight_scale +// columns n0+cg+4*i (i = 0..3) as four scalar dword loads in the same +// source block as the stage group -- they are independent of the A/B +// stage, so the compiler issues them with the stage loads and the stage's +// one wait group covers them -- and the epilogue becomes a pure ALU tail +// (4 x v_cvt_f32_i32 + 4 x v_mul_f32 + 4 x 2-byte stores, ZERO global +// loads, ZERO waits). The four columns are strided by 4 (acc_frag.x[i] +// holds column cg + 4*i), so one contiguous float4 at weight_scale[n0+cg] +// would alias columns owned by other lane groups and would be misaligned +// for cg > 0 (iteration-7 repair); four scalar loads carry the same +// values as the iteration-1..8 epilogue's weight_scale[col] reads. Same +// values, same left-associative scaling (acc * x_scale * weight_scale), +// same k-ascending int32 accumulation (K loop untouched), so the bf16 +// output is bit-identical to iteration 8 and the scalar reference (0 +// mismatches expected). Graph capture/replay semantics unchanged: one +// launch on the caller stream, no allocation, no sync, no compile. +// Falsifiable: median Graph replay must beat the 17.5837 us iteration-8 +// official best (official_best_median_us 17.58368492126465) and p90 must +// stay within the 1.05x guard of 19.2542 us (18.337279558181763 x 1.05), +// with exact int32 output (0 mismatches) and graph_capture_passed true. +// +// Iteration 12 (HIP-only epilogue round on the accepted iteration-9 kernel): +// the accepted iteration-9 exact code object (profiles/.../iteration9/ +// current-best-isa/isa.txt, source digest c0cca6e3, the object behind the +// 11.8823 us official best) disproves the "epilogue has ZERO global loads" +// claim of the iteration-9 write-up: the compiler hoisted x_scale[row] and +// weight_scale[n0+cg] into the stage group, but REMATERIALIZED the other +// three columns (ws[1..3]) back into the tail -- the epilogue still runs +// three global_load_dword at weight_scale +16/+32/+48 bytes from the v[4:5] +// base, each followed by a serialized s_waitcnt vmcnt(1) before its +// multiply-store (strict load -> wait -> mul -> store per output element +// for elements 1..3), fully exposed at 0.8 waves/SIMD (384 one-wave blocks +// over 480 SIMD slots; no co-resident partner exists to overlap the tail). +// Each round trip is a hot-L2 hit on the 24-KB weight_scale working set +// (~200-400 cycles), so the tail costs roughly 3 x (200-400) cycles per +// wave, comparable to the B stage wait. Fix (the validated TP4 down_proj +// recipe -- w8a8_dumma_m16n16k32_sk2_kernel: "scales prefetched into LDS +// s_scale[32] at kernel top with stores deferred until after both staging +// loops (removes prologue vmcnt window)"): route the four strided +// weight_scale columns through LDS. Each lane loads ws[0..3] +// (weight_scale[n0+cg+4*i], i = 0..3) in the stage group -- same issue +// batch as the A/B stage, so the stage's B-dominated wait covers their +// arrival -- and stores them into a 16-float s_scale plane indexed by the +// tile-local column (cg + 4*i; the 16 lanes of each column write the same +// value -> LDS broadcast, conflict-free). The four scale VGPRs die at +// those stores, before the K loop, so K-loop register pressure is +// unchanged (82 VGPRs expected) -- which is exactly why this beats the +// iteration-9 register-preload form the compiler rematerialized. Verified +// in the exact local gfx928 object of this round (hipcc -c -O3 +// --offload-arch=gfx928, exit 0): all 29 global loads (12 B dwordx2 + 12 A +// dwordx2 + xs + ws[0..3]) issue in the stage region; the compiler then +// proves each lane's epilogue read s_scale[cg+4*i] equals its own earlier +// store (same lane, same index), eliminates the four LDS reads, and the +// epilogue consumes the scale REGISTERS directly -- a pure ALU tail (4 x +// v_cvt_f32_i32 + 4 x v_mul_f32 + 4 x 2-byte stores) with ZERO global +// loads, ZERO vmcnt waits and ZERO ds_reads after the last v_mmac; the two +// ds_write2_b32 scale stores remain only as dead side effects under one +// pre-satisfied s_waitcnt vmcnt(0). Resources: 79 VGPR / 24 SGPR / LDS +// 6336 B / spills 0 / scratch 0 (iteration 9: 82 VGPR / 21 SGPR / 6272 B). +// Same values (weight_scale[n0+cg+4*i] read from the same global memory), +// same left-associative scaling (acc_f32 * xs * ws[i] == acc_f32 * xs * +// weight_scale[col]), same k-ascending int32 accumulation (K loop +// untouched), so the bf16 output is bit-identical to iteration 9 and the +// scalar reference (0 mismatches expected). LDS grows 6272 -> 6336 +// B/block (4 blocks/CU = 25.3 KiB < 64 KiB; occupancy unchanged; no new +// barriers). Graph capture/replay semantics unchanged: one launch on the +// caller stream, no allocation, no sync, no compile. Falsifiable: median +// Graph replay must beat the 11.882290244102478 us iteration-9 official +// best and p90 must stay within the 1.05x guard of 12.990516364574429 us +// (12.371920347213745 x 1.05), with exact int32 output (0 mismatches) and +// graph_capture_passed true. +// +// Iteration 13 (HIP-only K-loop consolidation on the accepted iteration-12 +// kernel): the accepted iteration-12 exact code object (profiles/.../ +// iteration13/current-best-isa/isa.txt, source digest 4b707c7d, the object +// behind the 11.7602 us official best) shows the K loop accumulating all 12 +// k32 chunks into ONE 4-register fragment (12 x v_mmac_i32_16x16x32_i8 all +// writing v[0:3], addresses 0x5960..0x59FC): a strictly serialized 12-deep +// v_mmac dependency chain on every wave's critical path at 0.8 waves/SIMD +// (384 one-wave blocks over 480 SIMD slots, so no co-resident partner hides +// the dependent latency). Fix: accumulate into TWO independent 16x16 int32 +// fragments -- acc_even consumes chunks 0,2,4,6,8,10 and acc_odd chunks +// 1,3,5,7,9,11 (the existing ds_read2_b64 pairing already delivers the two +// chunks of each round in one LDS read, so each round's two v_mmacs are +// operand-ready together) -- turning one 12-deep chain into two 6-deep +// chains that the scheduler interleaves, then combine with one 4-wide int32 +// add before the epilogue. Integer addition is exact and associative and +// per-element |partial| ~ 6.2e6 << 2^31 (no wrap), so acc_even + acc_odd +// equals the iteration-12 k-ascending single-accumulator sum exactly -- the +// int32 value (and therefore the bf16 output) is bit-identical, 0 mismatches +// expected. Same a_reg / s_b reads, same ds_read_b64 -> v_mmac operands, +// same fragment ownership, same epilogue (untouched), same grid/staging/ +// scales; only the accumulation grouping changes. Verified in the exact +// local gfx928 object of this round (hipcc -c -O3 --offload-arch=gfx928, +// exit 0, only the pre-existing -Wreturn-type warnings): the K loop now runs +// two interleaved 6-deep v_mmac chains (v[4:7] and v[0:3], strictly +// alternating at 0x597C..0x5A18) instead of one 12-deep chain (was all 12 +// into v[0:3]), the four combine v_add_u32s are scheduled into the epilogue +// head, and resources are UNCHANGED from the accepted object: 79 VGPR / 24 +// SGPR / LDS 6336 B / spills 0 / scratch 0 (the two chains fit the same +// allocation); s_waitcnt 26, zero barriers, zero branches, same 29 global +// loads, same 6 ds_read2_b64. Graph capture/replay +// semantics unchanged: one launch on the caller stream, no allocation, no +// sync, no compile; the exact-shape guard (m == 16 && n == 6144 && k == 384) +// and the generic scalar fallback (including the paired M=2 API shape) are +// untouched. Falsifiable: median Graph replay must beat the +// 11.760229468345642 us iteration-12 official best and p90 must stay within +// the 1.05x guard of 12.990516364574429 us (12.371920347213745 x 1.05), with +// exact int32 output (0 mismatches) and graph_capture_passed true. +__global__ __launch_bounds__(kDummaWaveThreads) void +w8a8_gemm_dumma_m16n6144_lds_staged_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n, + int k) { + // k == 384 is enforced by the exact-shape caller guard; the body + // hardcodes the 12 k32 chunks of this shape. + constexpr int kChunks = 384 / kDummaTileK; // 12 + const int lane = static_cast(threadIdx.x); // 0..63 (one wavefront) + // Each block owns one 16-column N tile: n0 = 16 * blockIdx.x (0..6128). + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + // B-only staging: per-wave B slice [16, 384], padded row stride 392. + // One wavefront per block writes and reads only its own s_b, so NO + // __syncthreads is needed anywhere in the kernel. + __shared__ int8_t s_b[kDummaTileN * 392]; + // Iteration 12: 16-float LDS plane holding this block's four strided + // weight_scale columns (tile-local index cg + 4*i), written in the stage + // region so the epilogue has ZERO global loads. LDS grows + // 6272 -> 6336 B/block (4 blocks/CU = 25.3 KiB < 64 KiB). + __shared__ float s_scale[16]; + + // Verified fragment ownership (see kernel doc): lane owns 8 contiguous + // bytes at k offset col_group = (lane >> 4) * 8 within each k32 chunk. + // A's row stride is ldm = k = 384 and the n-major packed B base stride + // is also k = 384; both are multiples of 8, so every global address is + // 8-byte aligned. + const int row = lane & 15; + const int col_group = (lane >> 4) << 3; // {0, 8, 16, 24} + + // ---- Stage: B into LDS (coalesced) + A into registers ---------------- + // B: each wave stages its own tile's 768 8-byte units (12 per lane): + // unit u -> srow = u / 48, kk8 = (u % 48) * 8, source + // b[(n0 + srow)*k + kk8] (the tile is a contiguous 6-KiB pack region). +#pragma unroll + for (int j = 0; j < 12; ++j) { + const int u = lane + 64 * j; + const int srow = u / 48; + const int kk8 = (u - srow * 48) * 8; + const long bv = *reinterpret_cast( + b + (static_cast(n0) + srow) * k + kk8); + *reinterpret_cast(s_b + srow * 392 + kk8) = bv; + } + // A: each lane preloads the 8 fragment bytes of all 12 k32 chunks into + // registers (24 VGPRs): a_reg[j] = a[row*k + j*32 + col_group .. +8). + // All 12 loads are independent and L1/L2-hot (the same 6144 A bytes are + // read by every block), so they sit in the same one-wait stage group + // and the K loop never touches A memory again. + long a_reg[kChunks]; +#pragma unroll + for (int j = 0; j < kChunks; ++j) { + a_reg[j] = *reinterpret_cast( + a + static_cast(row) * k + j * kDummaTileK + col_group); + } + + // Iteration 9 (epilogue scale hoist): x_scale[row] preloads HERE, in the + // stage group -- independent of the A/B stage, covered by the stage's + // B-dominated wait group, and its epilogue wait is pre-satisfied. + // Iteration 12 (LDS-anchored epilogue scales): the accepted iteration-9 + // exact code object shows the compiler rematerialized ws[1..3] back into + // the epilogue tail (three serialized global_load_dword + s_waitcnt + // vmcnt(1) chains after the last v_mmac, ~200-400 hot-L2 cycles each at + // one wave per SIMD). Fix: each lane loads the four strided weight_scale + // columns n0+cg+4*i (i = 0..3) HERE -- same issue batch as the A/B stage, + // so the B-dominated stage wait covers their arrival -- and stores them + // into the 16-float s_scale LDS plane (tile-local column index cg + 4*i; + // the 16 lanes of each column write the same value -> LDS broadcast, no + // bank conflict). The stores anchor the loads in the stage region, so + // the compiler can no longer rematerialize them into the tail; in the + // verified object it then proves each lane's epilogue read equals its own + // store (same lane, same index) and eliminates the LDS reads, leaving the + // scales live in registers through the K loop (verified: 79 VGPR). + // Same values (weight_scale[n0+cg+4*i] from the same global memory) and + // same left-associative scaling as iteration 9, so the bf16 output is + // bit-identical. + const int cg = lane >> 4; + const float xs = x_scale[row]; + const float ws0 = weight_scale[n0 + cg + 0]; + const float ws1 = weight_scale[n0 + cg + 4]; + const float ws2 = weight_scale[n0 + cg + 8]; + const float ws3 = weight_scale[n0 + cg + 12]; + float* s_scale_ptr = reinterpret_cast(s_scale); + s_scale_ptr[cg + 0] = ws0; + s_scale_ptr[cg + 4] = ws1; + s_scale_ptr[cg + 8] = ws2; + s_scale_ptr[cg + 12] = ws3; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + // Iteration 13: two independent accumulator fragments break the accepted + // object's single 12-deep serialized v_mmac chain (12 x v_mmac all writing + // v[0:3]) into two 6-deep chains that the scheduler interleaves; the two + // partial sums are combined with one 4-wide int32 add before the epilogue. + // Exact/associative int32 addition (no overflow, |partial| ~ 6.2e6), so + // the combined sum is bit-identical to the k-ascending accumulation. + du::dumma::DUFragment + acc_even, acc_odd; + du::dumma::du_fill_fragment(acc_even, 0); + du::dumma::du_fill_fragment(acc_odd, 0); + + // Zero-barrier, zero-global K loop over the 12 compile-time chunks, + // k-ascending: chunk j's a_frag comes from a_reg[j] (register copy, no + // wait) and chunk j's b_frag from ONE ds_read_b64 (s_b + row*392 + // + j*32 + col_group). Fully unrolled so the independent ds_read_b64s + // can be hoisted ahead of the dependent v_mmac chain. Iteration 13: each + // round covers TWO chunks (j -> acc_even, j+1 -> acc_odd) -- the two + // v_mmacs are independent, so the serial accumulator dependency is 6-deep + // per fragment instead of 12-deep on one fragment. +#pragma unroll + for (int j = 0; j < kChunks; j += 2) { + *reinterpret_cast(a_frag.x) = a_reg[j]; + *reinterpret_cast(b_frag.x) = + *reinterpret_cast( + s_b + row * 392 + j * kDummaTileK + col_group); + du::dumma::du_mma_sync(acc_even, a_frag, b_frag, acc_even); + *reinterpret_cast(a_frag.x) = a_reg[j + 1]; + *reinterpret_cast(b_frag.x) = + *reinterpret_cast( + s_b + row * 392 + (j + 1) * kDummaTileK + col_group); + du::dumma::du_mma_sync(acc_odd, a_frag, b_frag, acc_odd); + } + + // Iteration 13: combine the two partial sums (one 4-wide int32 add; exact + // and associative, same int32 value as the iteration-12 single-accumulator + // k-ascending sum, so the bf16 output is bit-identical). +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc_frag.x[i] = acc_even.x[i] + acc_odd.x[i]; + } + + // Direct fragment epilogue: row = lane & 15, cg = lane >> 4 (0..3), + // acc_frag.x[i] holds output column cg + 4*i; apply x_scale[row] * + // weight_scale[col] and store bf16. The whole 16x16 tile is 512 B of + // output, so the scattered 2-byte stores are absorbed by L2 write + // combining. + // Iteration 12: xs and the four ws columns were loaded in the stage + // region above (their LDS stores anchor the loads; the verified object + // eliminates the epilogue's s_scale reads because each lane reads exactly + // what it wrote -- same lane, same index -- and consumes the REGISTER + // values), so this tail is a pure ALU sequence -- v_cvt_f32 + v_mul_f32 + + // 2-byte store per element -- with ZERO global loads, ZERO vmcnt waits + // and ZERO ds_reads after the last v_mmac (the accepted iteration-9 + // object still ran three serialized global_load_dword + vmcnt(1) chains + // for elements 1..3). Same values and same left-associative scaling as + // the iteration-1..9 epilogues, so the bf16 output is bit-identical. + const float* s_scale_ro = reinterpret_cast(s_scale); + const float ws[4] = {s_scale_ro[cg + 0], s_scale_ro[cg + 4], + s_scale_ro[cg + 8], s_scale_ro[cg + 12]}; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = n0 + cg + 4 * i; + const float scaled = + static_cast(acc_frag.x[i]) * xs * ws[i]; + out[row * n + col] = hip_bfloat16(scaled); + } +} + +// Split-K partial DUMMA kernel (iteration 2 architecture, iteration 5 load +// path): identical m16n16k32 tile and zero-barrier one-wave grid as the +// iteration-1 kernel, but each block accumulates only one 32-aligned K slice +// (kSplitK total) and publishes its int32 accumulator fragment to workspace +// plane s instead of the final bf16 output. No LDS, no __syncthreads, no +// cross-wave barrier; every (N tile, K slice) pair is written exactly once +// per launch, so no workspace clear is needed; k-ascending accumulation +// within each slice. +// +// Iteration-5 fragment transport: the library du_load_matrix_sync expands a +// chunk into 16 global_load_ubyte + ~15 vmcnt waits per lane (the known +// decode load-path poison: 147456 vmem_read_instructions over the launch). +// Replace it with one aligned 8-byte global_load_dwordx2 per lane per +// fragment, using the verified m16n16k32 int8 fragment ownership: +// row = lane & 15, col_group = (lane >> 4) * 8, +// a_frag.x[i] = A[row][k0 + col_group + i] (logical [16, K], +// 8 contiguous bytes, base a + row*k + k0 + col_group); +// b_frag.x[i] = B[k0 + col_group + i][n0 + row] (n-major packed +// packed[n*K+kk] = raw[kk*n+n], 8 contiguous bytes, base +// b + (n0 + row)*k + k0 + col_group). +// Both addresses are 8-byte aligned (k0 % 32 == 0, col_group % 8 == 0, +// k == 6144) and du_mma_sync consumes the identical byte pattern the library +// loader would produce, so the int32 accumulation stays bit-identical. +// +// Iteration 8: the chunk loop is software-pipelined with depth-1 register +// prefetch (next chunk's two dwordx2 loads issued after the current v_mmac; +// see the in-kernel comment). Same loads, same k-ascending order, same +// bit-identical int32 accumulation; only the load/compute overlap changes. +// +// Iteration 13 (HIP-only consolidation): pair two adjacent k32 chunks into +// one k64 round. The accepted iteration-8/12 steady state is one serialized +// 2 x global_load_dwordx2 -> wait -> 1 x v_mmac chain per k32 chunk +// (19-20 wait rounds per block at ~450-530 cycles of global/L2 latency +// each), so partial-kernel time is proportional to the number of wait +// rounds. Each round now covers TWO k32 sub-chunks with FOUR dwordx2 loads +// (sub-chunk 0 at k0, sub-chunk 1 at k0+32; the lane's 8 bytes of the two +// sub-chunks are 32 bytes apart in k, so they stay two 8-byte loads, never a +// dwordx4) prefetched one round ahead, ONE wait group, and TWO v_mmac. The +// grid stays exactly 480 blocks / 4.0 blocks per CU and every block runs +// exactly 10 rounds (slices 0..1: 10 pairs of 64; slices 2..9: 9 pairs + one +// plain 32-row tail chunk) instead of 19-20, with every prefetch address +// bounded below k_end. Each v_mmac consumes the identical k-ascending byte +// pattern of its k32 range in the same relative order, so the int32 +// accumulation per output element (and therefore the bf16 output) is +// bit-identical to iteration 8/12. +// +// Iteration 14 (HIP-only consolidation): the iteration-13 code object shows +// the compiler wait-splits the 4-load k64 round into vmcnt(2) -> v_mmac -> +// vmcnt(0) -> v_mmac, so the second sub-chunk still exposes a full load +// latency and halving wait rounds gained only +3.4% (18.9963 -> 18.3637 us): +// at the current exactly-one-wavefront-per-SIMD grid (480 blocks x 64 +// threads = 480 wavefronts over 120 CUs x 4 SIMDs) the kernel is ~99% +// memory-latency stall (~306 VALU issued per block vs ~27.7k cycles per +// block), and no HIP change can shorten a single wavefront's serial wait +// chain (counter-split waits are raw-asm-only, and the HIP plateau is not +// proven, so raw asm stays forbidden). Raise SIMD-level concurrency +// instead: each block becomes 128 threads = 2 wavefronts; wave 0 accumulates +// chunks [0, half) of the slice and wave 1 accumulates [half, end) -- +// exactly 5 k64 rounds per wavefront, perfectly balanced for both slice +// sizes (20-chunk slices: 5 pairs each; 19-chunk slices: wave 0 runs 4 pairs +// + 1 tail chunk, wave 1 runs 5 pairs) -- then wave 0 publishes its int32 +// 16x16 tile to a 1 KiB LDS tile, one __syncthreads, and wave 1 adds it +// (ascending slice order: wave 0's k-range is strictly below wave 1's, and +// int32 addition is exact/associative with max |acc| ~ 1e8 < 2^31, so the +// combined sum -- and therefore the bf16 output -- is bit-identical to the +// single-wave ascending accumulation) and writes the plane. Grid stays 480 +// blocks / 4.0 blocks per CU = 960 wavefronts = exactly 2 per SIMD, so each +// SIMD can interleave a second wavefront during the load stalls; the load +// set (vmem_read 18432), plane stores (vmem_write 1920), workspace (10 +// planes), the combine kernel, and every fallback are unchanged. +// +// grid = 48 N tiles x 10 K slices = 480 two-wave blocks (4.0 blocks/CU at +// 120 CUs = 2 wavefronts per SIMD); slice s bounds are computed from kSplitK +// and k so the launcher stays allocation/sync-free. +// +// Iteration 16: the kernel becomes template . FUSED=false is +// byte-identical to the accepted iteration-14 kernel (the tail parameters +// x_scale / weight_scale / out / counters are never dereferenced). +// FUSED=true appends the qkv-proven monotonic-arrival fused combine tail +// (see the file header): wave 1 publishes the plane, one __syncthreads, +// then (wave 1, lane 0) does __threadfence() + atomicAdd(&counters[tile], +// 1) + __threadfence(); the block whose atomicAdd returns +// (arrived % kSplitK) == kSplitK - 1 is the LAST arrival for its tile and +// runs the iteration-12 combine epilogue (one aligned int4 load per plane +// per lane in ascending slice order, x_scale scalar + weight_scale float4 +// hoisted, hip_bfloat16 conversion, 4 x 2-byte stores) so the bf16 output +// is bit-identical to the two-kernel path. +template +__global__ __launch_bounds__(kPartialBlockThreads) void +w8a8_gemm_dumma_m16_direct_sk_partial_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int32_t* __restrict__ partials, // [kSplitK][16][768] int32 + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int32_t* __restrict__ counters, // [n_tiles] monotonic arrival counts + int n, + int k) { + const int wave = static_cast(threadIdx.x) >> 6; // 0 or 1 + const int lane = static_cast(threadIdx.x) & 63; // 0..63 + const int n_tiles = n / kDummaTileN; // 48 for n == 768 + const int tile = static_cast(blockIdx.x) % n_tiles; + const int s = static_cast(blockIdx.x) / n_tiles; + const int n0 = tile * kDummaTileN; + + // 32-aligned K slice bounds for slice s (remainder spread over the first + // slices so every boundary is a multiple of the DUMMA K step). + const int total_chunks = k / kDummaTileK; // 192 for k == 6144 + const int base_chunks = total_chunks / kSplitK; + const int rem_chunks = total_chunks % kSplitK; + const int slice_chunks = base_chunks + (s < rem_chunks ? 1 : 0); + const int k_start = (s * base_chunks + (s < rem_chunks ? s : rem_chunks)) * + kDummaTileK; + const int k_end = k_start + slice_chunks * kDummaTileK; + + // Iteration-14 per-wave chunk split: wave w accumulates chunks + // [w*half, ...) of the slice so both wavefronts are exactly balanced + // (20-chunk slices: 10 per wave; 19-chunk slices: 9 in wave 0, 10 in + // wave 1) and every chunk stays k-ascending within its wave with wave 0's + // range strictly below wave 1's, so the per-block LDS combine below + // preserves the ascending-slice int32 order. + const int half_chunks = slice_chunks >> 1; // 9 or 10 + const int chunk_lo = wave * half_chunks; + const int chunk_hi = (wave == 0) ? half_chunks : slice_chunks; + const int k_lo = k_start + chunk_lo * kDummaTileK; + const int k_hi = k_start + chunk_hi * kDummaTileK; + const int wave_chunks = chunk_hi - chunk_lo; // 9 or 10 + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Iteration-5 load path: one aligned 8-byte vector load per lane per + // fragment (see kernel doc for the ownership mapping). a_frag.x[0..7] and + // b_frag.x[0..7] are 8 contiguous bytes; du_mma_sync itself reinterprets + // each as one 64-bit operand, so the fragment registers are filled exactly + // as the library loader would fill them. + // + // Iteration-8 depth-1 register prefetch: the accepted iteration-5 steady + // state is 2 x global_load_dwordx2 -> s_waitcnt vmcnt(0) -> 1 x v_mmac per + // chunk, so every v_mmac stalls on loads issued in the same iteration + // (latency-bound despite 4.0 blocks/CU). Software-pipeline the K loop + // instead: hold the current chunk's two 8-byte words in plain locals and + // issue chunk (k0 + kDummaTileK)'s two dwordx2 loads AFTER chunk k0's + // v_mmac. Each v_mmac then consumes words loaded a full iteration earlier, + // and at the point the wait must be placed the only outstanding loads are + // that older pair -- no loads younger than the consumed ones are in flight + // -- so a plain vmcnt(0) covers exactly the right data and the prefetch + // pair stays in flight across the next iteration. The fragment registers + // still receive the identical k-ascending byte pattern, so the int32 + // accumulation order (and therefore the bf16 output) is bit-identical to + // iteration 5. + const int arow = lane & 15; + const int acol = (lane >> 4) << 3; // {0, 8, 16, 24}: k offset within a + // 32-wide chunk (lane owns 8 bytes) + const int brow = lane & 15; // n index within the 16-wide tile + const int bcol = (lane >> 4) << 3; // k offset within the 32-wide chunk + // Iteration-13 k64 pairing, applied to each wavefront's own chunk range: + // every wavefront runs exactly 5 wait rounds (wave_chunks 9 or 10 -> 4 + // pairs + 1 tail chunk, or 5 pairs). Four dwordx2 words (two k32 + // sub-chunks) are software-pipelined one round ahead exactly like the + // iteration-8 pair, so the wait before each round covers only the + // previously issued 4-load set; all prefetch addresses stay strictly below + // k_hi (the loop prefetches only pairs that exist, and the tail chunk is + // loaded plainly with no prefetch). + const int pairs = wave_chunks >> 1; // 5 per wavefront + const int tail_chunks = wave_chunks & 1; // wave 0 only on odd slices + const int k_pair_end = k_lo + pairs * 2 * kDummaTileK; + // Prefetch pair 0: sub-chunk 0 at k_lo, sub-chunk 1 at k_lo + 32. + long a_word0 = *reinterpret_cast( + a + static_cast(arow) * k + k_lo + acol); + long b_word0 = *reinterpret_cast( + b + static_cast(n0 + brow) * k + k_lo + bcol); + long a_word1 = *reinterpret_cast( + a + static_cast(arow) * k + k_lo + kDummaTileK + acol); + long b_word1 = *reinterpret_cast( + b + static_cast(n0 + brow) * k + k_lo + kDummaTileK + bcol); + int k0 = k_lo; + for (; k0 + 2 * kDummaTileK < k_pair_end; k0 += 2 * kDummaTileK) { + // Consume sub-chunk 0 (k0 .. k0+31), then sub-chunk 1 (k0+32 .. k0+63): + // same k-ascending v_mmac sequence as two adjacent iteration-8 chunks. + *reinterpret_cast(a_frag.x) = a_word0; + *reinterpret_cast(b_frag.x) = b_word0; + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + *reinterpret_cast(a_frag.x) = a_word1; + *reinterpret_cast(b_frag.x) = b_word1; + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + // Prefetch pair p+1 (four dwordx2 loads) after pair p's two v_mmacs, so + // all four words are in flight one full round before consumption. + a_word0 = *reinterpret_cast( + a + static_cast(arow) * k + (k0 + 2 * kDummaTileK) + acol); + b_word0 = *reinterpret_cast( + b + static_cast(n0 + brow) * k + + (k0 + 2 * kDummaTileK) + bcol); + a_word1 = *reinterpret_cast( + a + static_cast(arow) * k + (k0 + 3 * kDummaTileK) + acol); + b_word1 = *reinterpret_cast( + b + static_cast(n0 + brow) * k + + (k0 + 3 * kDummaTileK) + bcol); + } + // Last pair: consume the words prefetched by the final loop iteration. + *reinterpret_cast(a_frag.x) = a_word0; + *reinterpret_cast(b_frag.x) = b_word0; + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + *reinterpret_cast(a_frag.x) = a_word1; + *reinterpret_cast(b_frag.x) = b_word1; + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + // Odd-slice tail: the remaining k32 chunk (slices 2..9: 608 = 9*64 + 32) + // at k_pair_end, loaded plainly with no prefetch (nothing follows; the + // address is a multiple of 32 and stays below k_end). + if (tail_chunks) { + *reinterpret_cast(a_frag.x) = *reinterpret_cast( + a + static_cast(arow) * k + k_pair_end + acol); + *reinterpret_cast(b_frag.x) = *reinterpret_cast( + b + static_cast(n0 + brow) * k + k_pair_end + bcol); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Iteration-14 per-block int32 combine: wave 0 publishes its partial + // 16x16 accumulator tile to a 1 KiB LDS tile, one barrier, then wave 1 + // adds wave 0's partial (ascending slice order: wave 0 holds chunks + // [0, half), wave 1 holds [half, end); int32 addition is exact and + // associative with max |acc| ~ 1e8 < 2^31, so the combined sum is + // bit-identical to the single-wave ascending accumulation) and publishes + // the plane with the same direct fragment ownership as the iteration-1 + // epilogue: row = lane & 15, col_group = lane >> 4, acc_frag.x[i] -> + // column col_group + 4*i. Plane strides are multiples of 256 B so every + // store stays 4-B aligned (workspace base is 256-B aligned). + __shared__ int32_t s_part[kDummaTileM * kDummaTileN]; // 1 KiB + const int row = lane & 15; + const int col_group = lane >> 4; + if (wave == 0) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + s_part[row * kDummaTileN + col_group + 4 * i] = acc_frag.x[i]; + } + } + __syncthreads(); + if (wave == 1) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc_frag.x[i] += s_part[row * kDummaTileN + col_group + 4 * i]; + } + int32_t* plane = partials + static_cast(s) * kM16N768Elems; +#pragma unroll + for (int i = 0; i < 4; ++i) { + plane[row * n + n0 + col_group + 4 * i] = acc_frag.x[i]; + } + } + + // Iteration-16 fused combine tail (FUSED=true only): the 48-block combine + // kernel and its launch gap are consolidated into the partial launch with + // the exact qkv-proven monotonic-arrival pattern (accepted on this stack + // for hy3_tp4_qkv_proj_m16, round 23): after the plane publish, one + // __syncthreads, then (wave 1, lane 0) does + // __threadfence() + atomicAdd(&counters[tile], 1) + __threadfence(). + // The block whose atomicAdd returns (arrived % kSplitK) == kSplitK - 1 is + // the LAST arrival for its tile (counters are monotonic, zeroed once per + // workspace by the launcher before the first launch; int32 wraparound + // would need ~2^31/10 replays): the release fence orders this block's + // plane store before its arrival atomic, the acquire fence orders its + // plane reads after observing every sibling arrival, and the intervening + // barriers broadcast the flag to the wave, so no spin loop or residency + // assumption is needed. The last block then sums the tile's kSplitK + // planes in ascending slice order with the exact iteration-12 combine + // epilogue (one aligned int4 load per plane per lane, x_scale scalar + + // weight_scale float4 hoisted, hip_bfloat16 conversion, 4 x 2-byte + // stores), so the bf16 output is bit-identical to iteration 14 (same + // planes, same ascending int32 sum, same conversion and store addresses). + if (FUSED) { + __syncthreads(); + __shared__ int s_is_last; + if (wave == 1 && lane == 0) { + __threadfence(); // release: this block's plane stores are visible to + // the observer of this block's arrival atomic + const int arrived = atomicAdd(&counters[tile], 1); + __threadfence(); // acquire: this block's plane reads below (after the + // barrier) see every sibling store released before + // its arrival atomic + s_is_last = ((arrived % kSplitK) == kSplitK - 1); + } + __syncthreads(); + if (s_is_last && wave == 1) { + // Epilogue, byte-identical to w8a8_gemm_m16_n768_sk_combine_kernel: + // lane -> (erow = lane >> 2, eq = lane & 3) owns four consecutive + // columns ecol0 = n0 + 4*eq (16-B-aligned int4 plane reads and + // 8-B-aligned bf16 stores, as in the iteration-12 vectorized combine). + const int erow = lane >> 2; + const int eq = lane & 3; + const int ecol0 = n0 + 4 * eq; + const int eidx = erow * n + ecol0; + const float xs = x_scale[erow]; + const float4 ws = *reinterpret_cast(weight_scale + ecol0); + int32_t eacc0 = 0, eacc1 = 0, eacc2 = 0, eacc3 = 0; +#pragma unroll + for (int s2 = 0; s2 < kSplitK; ++s2) { + const int4 v = *reinterpret_cast( + partials + static_cast(s2) * kM16N768Elems + eidx); + eacc0 += v.x; + eacc1 += v.y; + eacc2 += v.z; + eacc3 += v.w; + } + const float es0 = static_cast(eacc0) * xs * ws.x; + const float es1 = static_cast(eacc1) * xs * ws.y; + const float es2 = static_cast(eacc2) * xs * ws.z; + const float es3 = static_cast(eacc3) * xs * ws.w; + hip_bfloat16* eo = out + eidx; + eo[0] = hip_bfloat16(es0); + eo[1] = hip_bfloat16(es1); + eo[2] = hip_bfloat16(es2); + eo[3] = hip_bfloat16(es3); + } + } +} + +// Split-K combine kernel (iteration 2; iteration 12 vectorized plane reads): +// one 64-lane wavefront per 16x16 output tile (48 blocks). Each lane sums +// the kSplitK partial planes in ascending slice order -- exact int32 order, +// bit-identical to the scalar reference (int32 addition is associative and +// max |acc| ~ 1.0e8 < 2^31) -- then applies x_scale[m] * weight_scale[n] and +// stores bf16. This kernel is part of the timed operator: it is enqueued on +// the caller stream right after the partial kernel, inside the same Graph +// capture. +// +// Iteration 12 (HIP-only consolidation): the accepted iteration-2 scalar +// combine's ISA (profiled 5.28 us over the 48-block launch, 23-25% of the +// operator aggregate) is 45 global_load_dword + 26 s_waitcnt per wavefront: +// each lane reads its four output columns (col_group + 4*i, strided by 4) +// with 4 x 10 scalar dword loads and four partially-serialized wait chains. +// Remap the 64 lanes to (row = lane >> 2, quarter = lane & 3) so each lane +// owns four CONSECUTIVE columns and replaces its 40 dword loads with 10 +// aligned 16-byte global_load_dwordx4 (one int4 per plane: plane stride +// 49152 B, row stride 3072 B, and quarter base 16 B are all multiples of 16, +// so every int4 address is 16-B aligned). The ten plane loads are +// independent, so one wait group covers the set instead of 26 per-wavefront +// waits, and the per-wavefront 64-lane footprint of one load instruction is +// 16 full 64-B lines instead of 64 scattered dwords (the four lanes of a row +// wrote the full line, so every line read is complete and L2-resident). +// Scale loads are hoisted to the kernel top so they issue with the plane +// burst (the prior vectorized object issued the weight_scale dwordx4 late and +// exposed a second serialized wait after the plane chain). Per element the +// accumulation is still the same ascending-slice int32 sum over the same +// bytes, so the bf16 output is bit-identical to iteration 2; only the load +// granularity / lane mapping / load order change. Stores stay 4 x 2-byte +// bf16 per lane (L2 write combining absorbs them, as in the partial-kernel +// epilogue). +__global__ __launch_bounds__(kDummaWaveThreads) void +w8a8_gemm_m16_n768_sk_combine_kernel( + const int32_t* __restrict__ partials, // [kSplitK][16][768] int32 + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n) { + const int lane = static_cast(threadIdx.x); // 0..63 + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + const int row = lane >> 2; // 0..15 + const int q = lane & 3; // 0..3: 4-column quarter within the tile + const int col0 = n0 + 4 * q; // 4 consecutive int32 columns (16-B aligned) + const int idx = row * n + col0; + + // Scale loads hoisted to the kernel top (x_scale scalar, weight_scale as + // one aligned float4) so they issue with the plane burst below instead of + // trailing the plane chain with their own serialized wait. + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + + // One 16-byte int4 load per plane, ascending slice order; all ten loads + // are independent, so a single wait group covers the set (see kernel doc). + int32_t acc0 = 0, acc1 = 0, acc2 = 0, acc3 = 0; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + const int4 v = *reinterpret_cast( + partials + static_cast(s) * kM16N768Elems + idx); + acc0 += v.x; + acc1 += v.y; + acc2 += v.z; + acc3 += v.w; + } + + // Same per-element scaling and 2-byte bf16 stores as iteration 2. + const float scaled0 = static_cast(acc0) * xs * ws.x; + const float scaled1 = static_cast(acc1) * xs * ws.y; + const float scaled2 = static_cast(acc2) * xs * ws.z; + const float scaled3 = static_cast(acc3) * xs * ws.w; + hip_bfloat16* o = out + idx; + o[0] = hip_bfloat16(scaled0); + o[1] = hip_bfloat16(scaled1); + o[2] = hip_bfloat16(scaled2); + o[3] = hip_bfloat16(scaled3); +} + +// Identity device-to-device copy for the bootstrap pack operation (int8 +// weights). Grid-stride loop: generic fallback for every (K, N). +__global__ void w8a8_pack_identity_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t numel) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < numel; idx += stride) { + dst[idx] = src[idx]; + } +} + +// Identity device-to-device copy for the bootstrap pack operation (fp32 +// per-column weight scales). +__global__ void w8a8_pack_identity_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t numel) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < numel; idx += stride) { + dst[idx] = src[idx]; + } +} + +// N-major transpose pack (iteration 5) for the exact MiniMax TP8 gate_up +// shape (k, n) == (6144, 768): dst[n * k + kk] = src[kk * n + n]. Runs once, +// outside the timed region and outside Graph capture, into the same +// caller-owned packed_weight buffer (byte count unchanged, so captured +// addresses stay valid). Each m16n16k32 B-fragment lane then owns 8 +// contiguous bytes (see w8a8_gemm_dumma_m16_direct_sk_partial_kernel), +// enabling one aligned global_load_dwordx2 per fragment instead of 8 +// strided global_load_ubyte. The identity copy remains the fallback for +// every other (K, N). +__global__ void w8a8_pack_nmajor_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t numel, + int n, + int k) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < numel; idx += stride) { + const int64_t n_idx = idx / k; + const int64_t k_idx = idx - n_idx * k; + dst[idx] = src[k_idx * n + n_idx]; + } +} + +} // namespace + +// gemm_out backend entry point (called by csrc/bindings.cpp on PyTorch's +// current HIP stream). This function must stay allocation/sync/compile-free +// and enqueue only on the caller stream: it is captured into a CUDA/HIP Graph +// and replayed with changed tensor contents. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // workspace (caller-owned, allocated before Graph capture) holds the + // kSplitK int32 partial planes used only by the exact-shape guard below + // (plus the per-tile fused-combine arrival counters in the last 192 B of + // the 16-plane contract workspace); its capacity is validated by the + // Python API layer. + + hip_bfloat16* out_bf16 = static_cast(out); + + // Exact-shape guard for the optimized M=16 MiniMax TP8 gate_up shape + // (16x768x6144): iteration-16 fused single-kernel split-K = 10 DUMMA + // m16n16k32 (480 two-wave partial blocks = 4.0 blocks/CU = 2 + // wavefronts/SIMD, one aligned 8-byte vector load per lane per fragment + // against the n-major packed weight, per-block LDS combine, qkv-proven + // monotonic-arrival fused combine tail replacing the 48-block combine + // kernel), on the caller stream. Every other + // (m, n, k) -- including the paired M=2 API shape with the same (N, K) + // and the paired 16x6144x384 down_proj shape -- falls through to the + // scalar generic fallback below (which decodes the n-major pack for + // (n, k) == (768, 6144)). + if (m == 16 && n == 768 && k == 6144) { + const int n_tiles = n / kDummaTileN; // 48 tiles of 16 columns + constexpr int64_t kSkWorkspaceBytes = + static_cast(kSplitK) * kM16N768Elems * sizeof(int32_t); + if (workspace_bytes >= kSkWorkspaceContractBytes) { + // Iteration-16 fused single-kernel path: the combine runs inside the + // partial kernel (last arrival per tile), so the operator is ONE + // launch and the 48-block combine kernel + its launch gap disappear. + // Grid stays 48 tiles x 10 slices = 480 two-wave blocks (2 + // wavefronts per SIMD). The per-tile monotonic arrival counters + // (n_tiles x int32 = 192 B) live in the last 192 B of the 786432-B + // (16-plane) contract workspace (plane 15's tail; planes 0..9 are + // the only partial planes used, so there is no overlap). The + // counters are zeroed once per workspace with an async memset on the + // caller stream (static workspace-pointer guard): the first launch + // happens during the eager warmup before Graph capture, so the + // memset is never part of the captured graph and every replay is + // exactly one kernel launch. + int32_t* partials = static_cast(workspace); + int32_t* counters = reinterpret_cast( + static_cast(workspace) + kSkWorkspaceContractBytes) - + n_tiles; + static const void* s_fused_counters_ws = nullptr; + if (s_fused_counters_ws != workspace) { + hipMemsetAsync(counters, 0, + static_cast(n_tiles) * sizeof(int32_t), + stream); + s_fused_counters_ws = workspace; + } + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_gemm_dumma_m16_direct_sk_partial_kernel), + dim3(n_tiles * kSplitK), dim3(kPartialBlockThreads), 0, stream, + a, b, partials, x_scale, weight_scale, out_bf16, counters, n, k); + return; + } + if (workspace_bytes >= kSkWorkspaceBytes) { + // Two-kernel path for workspaces below the 16-plane contract (cannot + // happen through the validated API path): the iteration-14 partial + // kernel with the fused tail compiled out (FUSED=false) + the + // iteration-12 vectorized combine kernel, byte-identical to the + // accepted iteration-14 behavior. Grid = 48 tiles x 10 slices = 480 + // two-wave blocks = exactly 4.0 blocks per device CU (120 CUs) = 2 + // wavefronts per SIMD; the FUSED=false instantiation never + // dereferences x_scale / weight_scale / out / counters. + int32_t* partials = static_cast(workspace); + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_gemm_dumma_m16_direct_sk_partial_kernel), + dim3(n_tiles * kSplitK), dim3(kPartialBlockThreads), 0, stream, + a, b, partials, x_scale, weight_scale, out_bf16, nullptr, n, k); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_m16_n768_sk_combine_kernel), + dim3(n_tiles), dim3(kDummaWaveThreads), 0, stream, + partials, x_scale, weight_scale, out_bf16, n); + return; + } + // Workspace too small for the ten partial planes (cannot happen through + // the validated API path): keep the iteration-1 direct kernel under the + // same shape guard. + const int blocks = n / kDummaTileN; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_dumma_m16_direct_kernel), + dim3(blocks), dim3(kDummaWaveThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, n, k); + return; + } + + // Exact-shape guard for the assigned MiniMax TP8 down_proj M=16 shape + // (16x6144x384, iteration-5 HIP-only staging round, iteration-8 occupancy + // round, iteration-9 epilogue scale-hoist round, iteration-12 LDS-staged + // epilogue scales round): m16n16k32 DUMMA kernel + // with ONE 64-lane wavefront per block (64 + // threads) and one 16-column N tile per block, B-only LDS staging (each + // wavefront stages its tile's K=384 window once -- 12 aligned dwordx2 + + // 12 ds_write_b64, ONE wait group, ZERO barriers) with A register-resident + // (each lane preloads its 12 chunk-fragment words into registers once) + // and a fully unrolled k-ascending K loop, grid = 6144/16 = 384 one-wave + // blocks = 3.2 blocks/CU over 120 CUs (iteration 8: waves per block 2 -> + // 1, so every CU holds 3-4 resident wavefronts instead of 2 on 48 CUs; + // LDS halves to 6272 B/block), against the n-major packed weight for the + // exact (k, n) == (384, 6144) pair, on the caller stream. Every other + // (m, n, k) -- including the paired M=2 API shape with the same (N, K) -- + // falls through to the scalar generic fallback below (which decodes the + // n-major pack for (n, k) == (6144, 384)). No workspace is needed for + // this path. + if (m == 16 && n == 6144 && k == 384) { + const int blocks = n / kDummaTileN; // 384 one-wave blocks + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_dumma_m16n6144_lds_staged_kernel), + dim3(blocks), dim3(kDummaWaveThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, n, k); + return; + } + + // Generic scalar fallback for every unmatched (m, n, k), including M=2 + // API shapes that share an assigned (N, K). + const int threads = kScalarBlockThreads; + const int total = m * n; + const int blocks = (total + threads - 1) / threads; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_kernel), + dim3(blocks), dim3(threads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); +} + +// pack_weight backend entry point (called by csrc/bindings.cpp on PyTorch's +// current HIP stream, outside the timed region and outside Graph capture). +// +// Bootstrap: identity device-to-device copy of raw_weight[K, N] and +// weight_scale[N, 1]. Later Parallel-explore rounds may change only this +// function's HIP implementation (and the matching GEMM interpretation) to +// test packed layouts; the identity copy remains the generic fallback for +// unmatched (K, N). +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_numel = static_cast(k) * n; + const int weight_blocks = static_cast( + (weight_numel + kCopyBlockThreads - 1) / kCopyBlockThreads); + const int scale_blocks = + (n + kCopyBlockThreads - 1) / kCopyBlockThreads; + + if ((k == 6144 && n == 768) || (k == 384 && n == 6144)) { + // Iteration-5 n-major transpose pack for the exact gate_up (6144, 768) + // and down_proj (384, 6144) pairs (packed[n*K+kk] = raw[kk*n+n]); the + // matching GEMM kernels and the scalar generic fallback decode this + // layout. Identity copy for every other (k, n). + hipLaunchKernelGGL( + w8a8_pack_nmajor_i8_kernel, + dim3(weight_blocks), dim3(kCopyBlockThreads), 0, stream, + raw_weight, packed_weight, weight_numel, n, k); + } else { + hipLaunchKernelGGL( + w8a8_pack_identity_i8_kernel, + dim3(weight_blocks), dim3(kCopyBlockThreads), 0, stream, + raw_weight, packed_weight, weight_numel); + } + hipLaunchKernelGGL( + w8a8_pack_identity_f32_kernel, + dim3(scale_blocks), dim3(kCopyBlockThreads), 0, stream, + weight_scale, packed_weight_scale, static_cast(n)); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_gate_up_proj.hip new file mode 100644 index 00000000..4400c8bd --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_gate_up_proj.hip @@ -0,0 +1,914 @@ +// @@variant shape=minimax_tp8_shared_gate_up_proj_m16 commit=4358dea23da7468d3868ce258b971667a50c139e added=2026-08-28 +// median_us=17.11 p90_us=17.13 +// source=minimax-dsh-tp8-m16-1-78402260 +// MetaInfer W8A8 INT8 GEMM - HIP implementation for Hygon K500SM_AI / gfx928. +// +// Iteration 1 (minimal DUMMA bootstrap): the assigned M=16 MiniMax TP8 +// shared_gate_up_proj shape (16x768x6144) ran on a minimal 16x16x32 DUMMA +// tile with one 64-lane wavefront per output tile, explicit int32 +// accumulation, direct row-major global fragment loads, and no LDS staging / +// no cross-wave barriers. Measured 182.1 us median Graph replay (0.574x vs +// the 104.527 us Triton baseline): the one-wave grid of 48 blocks fills only +// 48 of 120 CUs (0.4 blocks/CU) and each block serially streams its full +// K=6144 loop, so the kernel is launch-geometry/latency bound, not ALU bound. +// +// Iteration 2 (architecture round, grid parallelism): keep the exact same +// zero-barrier one-wave direct-fragment tile, but split K into kSplitK = 5 +// independent 32-aligned slices per output tile: +// * grid = 48 N tiles x 5 K slices = 240 one-wave blocks = exactly 2 +// blocks per device CU (120 CUs) -- the two-blocks-per-CU latency-hiding +// target of the control plane; +// * each partial block publishes its int32 accumulator fragment to one +// workspace plane (every (tile, slice) pair is written exactly once per +// launch, so no workspace clear is needed); +// * a 48-block 64-thread combine kernel sums the five planes in ascending +// slice order (exact int32 order, bit-identical to the scalar reference), +// applies x_scale[m] * packed_weight_scale[n], and stores bf16; both +// kernels are enqueued on the caller stream so the combine cost is inside +// the timed operator / Graph replay; +// * every other shape, including the paired M=2 API shape with the same +// (N, K) and the paired 16x6144x384 down_proj shape, keeps the scalar +// generic fallback below; +// * launch_pack_w8a8_weight stays an identity device-to-device copy +// (generic fallback for every (K, N)). +// +// Iteration 3 (architecture round, occupancy probe): keep the exact +// iteration-2 architecture (zero-barrier one-wave direct-fragment tile, +// split-K partial planes + 48-block combine in the timed Graph) but raise +// kSplitK from 5 to 10, so the grid doubles from 240 to 480 one-wave blocks +// = exactly 4.0 blocks per device CU (120 CUs). K = 6144 = 192 chunks of 32 +// = 10*19 + 2, so slices 0..1 carry 20 chunks (640 elements) and slices 2..9 +// carry 19 chunks (608 elements): every boundary stays a multiple of the +// DUMMA K step and the imbalance is <= 6%. The workspace then holds 10 +// int32 partial planes (491520 B), well inside the 786432 B (16-plane) the +// API allocates for this shape. Hypothesis: at 2.0 blocks/CU the partial +// kernel is still global-load-latency bound (PMC: ~6% VALU issue +// utilization, 45.8 us perturbed partial vs 50.6 us total median), so +// doubling the independent resident one-wave streams per CU should hide more +// of that latency; the combine kernel grows from 5 to 10 plane-sums but +// stays a tiny 48-block 64-thread launch over L2-hot plane traffic. +// Measured: 45.47 us median Graph replay (2.30x), so occupancy scaling +// helped (~10%) but the direct-fragment load path remains the binding cost. +// +// Iteration 5 (load-path round): ISA on the accepted iteration-3 code object +// shows the library fragment loaders expand each k32 chunk into 16 +// global_load_ubyte instructions per lane -- 8 contiguous bytes for A that +// the compiler never fuses plus 8 bytes strided by ldm for B -- followed by +// ~15 s_waitcnt vmcnt before a single v_mmac_i32_16x16x32_i8: 147456 +// vmem_read_instructions over the 480-block partial launch. Fix: keep the +// exact zero-barrier one-wave m16n16k32 tile, the 480-block split-K=10 grid, +// and the 48-block combine, but replace the byte-gather fragment loads with +// one aligned 8-byte (global_load_dwordx2) load per lane per fragment: +// * A stays the logical row-major activation [16, K]: lane (row = lane&15, +// col_group = (lane>>4)*8) owns A[row][k0+col_group .. k0+col_group+7], +// which is 8 contiguous bytes (ldm = K = 6144); +// * B is packed once, outside the timed region and outside Graph capture, +// into the n-major transpose packed[n*K+kk] = raw[kk*n+n] for the exact +// (k,n) == (6144,768) (same byte count and buffer, so Graph addresses +// are unchanged; the identity copy remains the pack fallback for every +// other (k,n)); lane (row = lane&15, col_group = (lane>>4)*8) then owns +// packed[(n0+row)*K + k0+col_group .. +7], also 8 contiguous bytes; +// * du_mma_sync consumes the identical byte pattern the library loader +// would produce, so the int32 accumulation stays bit-identical +// (k-ascending per slice, ascending-slice combine, 0 mismatches); +// * the generic scalar fallback decodes the n-major pack for +// (n,k) == (768,6144) (including the paired M=2 API shape with the same +// (N,K)) and the workspace-too-small direct kernel reads B with ldm = k +// through the same packed pointer. +// Falsifiable: partial-kernel vmem_read_instructions must drop from 147456 +// to ~18432 (2 dwordx2 loads x ~19.2 chunks x 480 blocks) and the median +// must beat the 45.47 us iteration-3 best. +// +// Iteration 16 (final conditional inline-asm round -> HIP-only +// consolidation): raw asm stays forbidden (isa_policy.phase=hip_only, +// plateau=false: recent_valid_improvements_percent = [+3.44, +2.02, -4.52] +// after the rejected split-K=15 three-wave probe, so the required three +// recent valid HIP candidates within [-2%, +2%) of best do not exist), so +// this round consolidates the two-kernel operator into ONE launch with the +// qkv-proven monotonic-arrival fused combine tail (the exact pattern +// accepted on this stack for hy3_tp4_qkv_proj_m16, round 23): +// * the partial kernel becomes template ; FUSED=true adds the +// tail: after wave 1 publishes the plane, one __syncthreads, then +// (wave 1, lane 0) does __threadfence() + atomicAdd(&counters[tile], 1) +// + __threadfence(); the block whose atomicAdd returns +// (arrived % kSplitK) == kSplitK - 1 is the LAST arrival for its tile +// and sums the tile's 10 planes in ascending slice order with the exact +// iteration-12 combine epilogue (int4 plane loads, x_scale scalar + +// weight_scale float4 hoisted, hip_bfloat16 conversion, 4 x 2-byte +// stores), then writes the bf16 output; +// * the 48-block combine kernel and its launch gap disappear from the +// timed operator; the per-tile monotonic counters (n_tiles x int32 = +// 192 B) live in the last 192 B of the 786432-B (16-plane) contract +// workspace (plane 15's tail; planes 0..9 are the only planes used), +// zeroed once per workspace by a static-pointer-guarded hipMemsetAsync +// on the caller stream before the first launch (eager warmup runs +// before Graph capture, so the memset is never part of the captured +// graph and every replay is exactly one kernel launch); +// * the two-kernel path (FUSED=false + the combine kernel) remains as the +// fallback for workspaces below the 16-plane contract, byte-identical +// to iteration 14; +// * per output element the int32 accumulation is still the +// ascending-slice sum over the same bytes with the same conversion, so +// the bf16 output is bit-identical to iteration 14 (0 mismatches +// expected) and Graph capture/replay semantics are unchanged (one launch +// on the caller stream, no allocation/sync/compile in the launcher). +// +// Logical operation (see int8_w8a8_gemm_api.py): +// out[m, n] = bf16(int32_dot(x_q[m, :], packed_weight[:, n]) +// * x_scale[m, 0] * packed_weight_scale[n, 0]) +// +// Headers are included in the DTK-known-good order: hip_runtime, hip_bfloat16, +// then du_mma (this DTK's du_mma.h is not self-contained when included before +// the HIP runtime headers). +// +// The timed entry point (launch_w8a8_gemm) performs no allocation, no +// compilation, no autotuning, no weight packing, no host/device +// synchronization, and no default-stream launch: it only computes launch +// geometry and enqueues on the caller-provided stream, so it stays safe under +// torch.cuda.CUDAGraph capture and replay. + +#include +#include +#include + +#include + +namespace { + +// gfx928 native wavefront is 64 lanes; the block size must be a multiple of 64. +constexpr int kScalarBlockThreads = 128; +constexpr int kCopyBlockThreads = 256; + +// Minimal DUMMA tile for the assigned M=16 decode shape: INT8 m16n16k32 with +// int32 accumulation, one 64-lane wavefront owning one independent fragment. +constexpr int kDummaTileM = 16; +constexpr int kDummaTileN = 16; +constexpr int kDummaTileK = 32; +constexpr int kDummaWaveThreads = 64; +// Iteration 14: the split-K partial kernel runs 128 threads = 2 wavefronts +// per block (block size stays a multiple of the 64-lane wavefront). +constexpr int kPartialBlockThreads = 2 * kDummaWaveThreads; + +// Split-K geometry for the assigned MiniMax TP8 gate_up shape (iteration 3 +// occupancy probe): 10 independent K slices per 16x16 output tile -> grid = +// 48 N tiles x 10 slices = 480 one-wave zero-barrier blocks = exactly 4.0 +// blocks per device CU (120 CUs), doubling the iteration-2 split-K = 5 point +// (240 blocks, 2.0 blocks/CU, 50.55 us median). K = 6144 = 192 chunks of 32 +// = 10*19 + 2, so slices 0..1 carry 20 chunks (640 elements) and slices 2..9 +// carry 19 chunks (608 elements); every slice boundary stays a multiple of +// the DUMMA K step and the imbalance is <= 6%. The workspace holds kSplitK +// int32 partial planes of [16][768] (49152 B each); the Python API allocates +// 16 planes (786432 B) for this shape, so kSplitK = 10 (491520 B) always +// fits the workspace and 256-B stage alignment. +constexpr int kSplitK = 10; +constexpr int kM16N768Elems = 16 * 768; // one int32 partial plane (12288) + +// Iteration 16: the fused combine tail needs the full 16-plane contract +// workspace the Python API allocates for this shape (786432 B): the +// per-tile monotonic arrival counters live in the last n_tiles*4 B of +// plane 15's tail (planes 0..9 are the only partial planes used, so there +// is no overlap), exactly like the accepted qkv fused path. +constexpr int64_t kSkWorkspaceContractBytes = + static_cast(16) * kM16N768Elems * sizeof(int32_t); // 786432 B + +// Scalar INT8 dot-product kernel: one thread -> one out[m, n] element. +// +// a: [M, K] int8 row-major (logical activation, stride K) +// b: [K, N] int8 row-major (identity-packed logical weight, stride N) +// x_scale: [M] fp32 +// weight_scale: [N] fp32 +// out: [M, N] bf16 +// +// Linear indexing keeps adjacent lanes on adjacent n, so B-column reads and +// the bf16 stores are coalesced in the fastest-changing N dimension. +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int linear = blockIdx.x * blockDim.x + threadIdx.x; + if (linear >= m * n) { + return; + } + const int row = linear / n; + const int col = linear - row * n; + + const int8_t* a_row = a + row * k; + + // The exact gate_up (n, k) == (768, 6144) uses the n-major packed weight + // (packed[n*K+kk] = raw[kk*n+n], produced by launch_pack_w8a8_weight); all + // other shapes keep the identity [K, N] row-major layout. Decode both so + // the generic fallback -- including the paired M=2 API shape with the same + // (N, K) -- stays correct. + const bool nmajor_packed = (n == 768 && k == 6144); + const int8_t* b_col = + nmajor_packed ? (b + static_cast(col) * k) : (b + col); + const int64_t b_stride = nmajor_packed ? 1 : static_cast(n); + + // Exact int32 dot product. Largest assigned K is 6144: + // 6144 * 128 * 128 = 100,663,296 < 2^31, so the accumulation is exact and + // cannot overflow. int8 products are exact in int32, and the int32 sum is + // bit-identical to the float32 dot of (A.float() @ B.float()). + int32_t acc = 0; + for (int i = 0; i < k; ++i) { + acc += static_cast(a_row[i]) * + static_cast(b_col[i * b_stride]); + } + + // Reference: (A.float() @ B.float()) * x_scale * weight_scale.T -> bf16. + // Both scales are applied left-associatively as in the torch reference. + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + // hip_bfloat16's supported float conversion rounds to nearest even, which + // matches the torch float->bfloat16 conversion. + out[linear] = hip_bfloat16(scaled); +} + +// Minimal INT8 DUMMA kernel for M=16 decode: one wave (64 lanes) per block, +// one 16x16 output tile per block, explicit int32 accumulation. Reachable +// only under the exact (16, 768, 6144) shape guard when the workspace cannot +// hold the ten split-K partial planes, so B is the n-major packed weight +// (packed[n*K+kk] = raw[kk*n+n]) and is loaded with ldm = k. +// +// a: [16, K] int8 row-major activation (M == 16 is enforced by the caller +// guard; fragment loads need ldm = k) +// b: [N, K] int8 n-major packed weight for the exact guard that is the +// only caller (fragment loads need ldm = k) +// x_scale: [16] fp32 +// weight_scale: [N] fp32 +// out: [16, N] bf16 +// +// Each K step loads a 16x32 A fragment and a 32x16 B fragment directly from +// global memory with du::dumma::du_load_matrix_sync (row_major fragments, +// 8-byte-aligned base offsets because K and the N tile origin are multiples +// of 32/16 bytes) and accumulates with du::dumma::du_mma_sync into an int32 +// accumulator fragment. No LDS, no __syncthreads, no cross-wave barrier: +// each block is a single wavefront and all stores are fragment-owner direct. +__global__ __launch_bounds__(kDummaWaveThreads) void +w8a8_gemm_dumma_m16_direct_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n, + int k) { + const int lane = static_cast(threadIdx.x); // 0..63 + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Complete K loop accumulated in exact int32 (same values as the scalar + // reference; K = 6144 * 128 * 128 < 2^31 so no overflow). B is the n-major + // packed weight for this shape, so the row_major B fragment loader with + // ldm = k reads packed[(n0 + lane&15)*k + k0 + (lane>>4)*8 + i], exactly + // W[k0 + (lane>>4)*8 + i][n0 + lane&15]. + for (int k0 = 0; k0 < k; k0 += kDummaTileK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); + du::dumma::du_load_matrix_sync( + b_frag, b + static_cast(n0) * k + k0, k); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Direct fragment epilogue. The int32 accumulator ownership for this DTK + // (verified against du_store_matrix_sync with mem_row_major) is: + // row = lane & 15, col_group = lane >> 4, + // acc_frag.x[i] holds output column col_group + 4*i. + // The four 2-byte bf16 values of one row are written by lanes l, l+16, + // l+32, l+48 (same lane&15); the whole 16x16 tile is only 512 bytes of + // output so the scattered 2-byte stores are absorbed by L2 write combining. + const int row = lane & 15; + const int col_group = lane >> 4; + const float xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = n0 + col_group + 4 * i; + const float scaled = + static_cast(acc_frag.x[i]) * xs * weight_scale[col]; + out[row * n + col] = hip_bfloat16(scaled); + } +} + +// Split-K partial DUMMA kernel (iteration 2 architecture, iteration 5 load +// path): identical m16n16k32 tile and zero-barrier one-wave grid as the +// iteration-1 kernel, but each block accumulates only one 32-aligned K slice +// (kSplitK total) and publishes its int32 accumulator fragment to workspace +// plane s instead of the final bf16 output. No LDS, no __syncthreads, no +// cross-wave barrier; every (N tile, K slice) pair is written exactly once +// per launch, so no workspace clear is needed; k-ascending accumulation +// within each slice. +// +// Iteration-5 fragment transport: the library du_load_matrix_sync expands a +// chunk into 16 global_load_ubyte + ~15 vmcnt waits per lane (the known +// decode load-path poison: 147456 vmem_read_instructions over the launch). +// Replace it with one aligned 8-byte global_load_dwordx2 per lane per +// fragment, using the verified m16n16k32 int8 fragment ownership: +// row = lane & 15, col_group = (lane >> 4) * 8, +// a_frag.x[i] = A[row][k0 + col_group + i] (logical [16, K], +// 8 contiguous bytes, base a + row*k + k0 + col_group); +// b_frag.x[i] = B[k0 + col_group + i][n0 + row] (n-major packed +// packed[n*K+kk] = raw[kk*n+n], 8 contiguous bytes, base +// b + (n0 + row)*k + k0 + col_group). +// Both addresses are 8-byte aligned (k0 % 32 == 0, col_group % 8 == 0, +// k == 6144) and du_mma_sync consumes the identical byte pattern the library +// loader would produce, so the int32 accumulation stays bit-identical. +// +// Iteration 8: the chunk loop is software-pipelined with depth-1 register +// prefetch (next chunk's two dwordx2 loads issued after the current v_mmac; +// see the in-kernel comment). Same loads, same k-ascending order, same +// bit-identical int32 accumulation; only the load/compute overlap changes. +// +// Iteration 13 (HIP-only consolidation): pair two adjacent k32 chunks into +// one k64 round. The accepted iteration-8/12 steady state is one serialized +// 2 x global_load_dwordx2 -> wait -> 1 x v_mmac chain per k32 chunk +// (19-20 wait rounds per block at ~450-530 cycles of global/L2 latency +// each), so partial-kernel time is proportional to the number of wait +// rounds. Each round now covers TWO k32 sub-chunks with FOUR dwordx2 loads +// (sub-chunk 0 at k0, sub-chunk 1 at k0+32; the lane's 8 bytes of the two +// sub-chunks are 32 bytes apart in k, so they stay two 8-byte loads, never a +// dwordx4) prefetched one round ahead, ONE wait group, and TWO v_mmac. The +// grid stays exactly 480 blocks / 4.0 blocks per CU and every block runs +// exactly 10 rounds (slices 0..1: 10 pairs of 64; slices 2..9: 9 pairs + one +// plain 32-row tail chunk) instead of 19-20, with every prefetch address +// bounded below k_end. Each v_mmac consumes the identical k-ascending byte +// pattern of its k32 range in the same relative order, so the int32 +// accumulation per output element (and therefore the bf16 output) is +// bit-identical to iteration 8/12. +// +// Iteration 14 (HIP-only consolidation): the iteration-13 code object shows +// the compiler wait-splits the 4-load k64 round into vmcnt(2) -> v_mmac -> +// vmcnt(0) -> v_mmac, so the second sub-chunk still exposes a full load +// latency and halving wait rounds gained only +3.4% (18.9963 -> 18.3637 us): +// at the current exactly-one-wavefront-per-SIMD grid (480 blocks x 64 +// threads = 480 wavefronts over 120 CUs x 4 SIMDs) the kernel is ~99% +// memory-latency stall (~306 VALU issued per block vs ~27.7k cycles per +// block), and no HIP change can shorten a single wavefront's serial wait +// chain (counter-split waits are raw-asm-only, and the HIP plateau is not +// proven, so raw asm stays forbidden). Raise SIMD-level concurrency +// instead: each block becomes 128 threads = 2 wavefronts; wave 0 accumulates +// chunks [0, half) of the slice and wave 1 accumulates [half, end) -- +// exactly 5 k64 rounds per wavefront, perfectly balanced for both slice +// sizes (20-chunk slices: 5 pairs each; 19-chunk slices: wave 0 runs 4 pairs +// + 1 tail chunk, wave 1 runs 5 pairs) -- then wave 0 publishes its int32 +// 16x16 tile to a 1 KiB LDS tile, one __syncthreads, and wave 1 adds it +// (ascending slice order: wave 0's k-range is strictly below wave 1's, and +// int32 addition is exact/associative with max |acc| ~ 1e8 < 2^31, so the +// combined sum -- and therefore the bf16 output -- is bit-identical to the +// single-wave ascending accumulation) and writes the plane. Grid stays 480 +// blocks / 4.0 blocks per CU = 960 wavefronts = exactly 2 per SIMD, so each +// SIMD can interleave a second wavefront during the load stalls; the load +// set (vmem_read 18432), plane stores (vmem_write 1920), workspace (10 +// planes), the combine kernel, and every fallback are unchanged. +// +// grid = 48 N tiles x 10 K slices = 480 two-wave blocks (4.0 blocks/CU at +// 120 CUs = 2 wavefronts per SIMD); slice s bounds are computed from kSplitK +// and k so the launcher stays allocation/sync-free. +// +// Iteration 16: the kernel becomes template . FUSED=false is +// byte-identical to the accepted iteration-14 kernel (the tail parameters +// x_scale / weight_scale / out / counters are never dereferenced). +// FUSED=true appends the qkv-proven monotonic-arrival fused combine tail +// (see the file header): wave 1 publishes the plane, one __syncthreads, +// then (wave 1, lane 0) does __threadfence() + atomicAdd(&counters[tile], +// 1) + __threadfence(); the block whose atomicAdd returns +// (arrived % kSplitK) == kSplitK - 1 is the LAST arrival for its tile and +// runs the iteration-12 combine epilogue (one aligned int4 load per plane +// per lane in ascending slice order, x_scale scalar + weight_scale float4 +// hoisted, hip_bfloat16 conversion, 4 x 2-byte stores) so the bf16 output +// is bit-identical to the two-kernel path. +template +__global__ __launch_bounds__(kPartialBlockThreads) void +w8a8_gemm_dumma_m16_direct_sk_partial_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + int32_t* __restrict__ partials, // [kSplitK][16][768] int32 + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int32_t* __restrict__ counters, // [n_tiles] monotonic arrival counts + int n, + int k) { + const int wave = static_cast(threadIdx.x) >> 6; // 0 or 1 + const int lane = static_cast(threadIdx.x) & 63; // 0..63 + const int n_tiles = n / kDummaTileN; // 48 for n == 768 + const int tile = static_cast(blockIdx.x) % n_tiles; + const int s = static_cast(blockIdx.x) / n_tiles; + const int n0 = tile * kDummaTileN; + + // 32-aligned K slice bounds for slice s (remainder spread over the first + // slices so every boundary is a multiple of the DUMMA K step). + const int total_chunks = k / kDummaTileK; // 192 for k == 6144 + const int base_chunks = total_chunks / kSplitK; + const int rem_chunks = total_chunks % kSplitK; + const int slice_chunks = base_chunks + (s < rem_chunks ? 1 : 0); + const int k_start = (s * base_chunks + (s < rem_chunks ? s : rem_chunks)) * + kDummaTileK; + const int k_end = k_start + slice_chunks * kDummaTileK; + + // Iteration-14 per-wave chunk split: wave w accumulates chunks + // [w*half, ...) of the slice so both wavefronts are exactly balanced + // (20-chunk slices: 10 per wave; 19-chunk slices: 9 in wave 0, 10 in + // wave 1) and every chunk stays k-ascending within its wave with wave 0's + // range strictly below wave 1's, so the per-block LDS combine below + // preserves the ascending-slice int32 order. + const int half_chunks = slice_chunks >> 1; // 9 or 10 + const int chunk_lo = wave * half_chunks; + const int chunk_hi = (wave == 0) ? half_chunks : slice_chunks; + const int k_lo = k_start + chunk_lo * kDummaTileK; + const int k_hi = k_start + chunk_hi * kDummaTileK; + const int wave_chunks = chunk_hi - chunk_lo; // 9 or 10 + + du::dumma::DUFragment + a_frag; + du::dumma::DUFragment + b_frag; + du::dumma::DUFragment + acc_frag; + du::dumma::du_fill_fragment(acc_frag, 0); + + // Iteration-5 load path: one aligned 8-byte vector load per lane per + // fragment (see kernel doc for the ownership mapping). a_frag.x[0..7] and + // b_frag.x[0..7] are 8 contiguous bytes; du_mma_sync itself reinterprets + // each as one 64-bit operand, so the fragment registers are filled exactly + // as the library loader would fill them. + // + // Iteration-8 depth-1 register prefetch: the accepted iteration-5 steady + // state is 2 x global_load_dwordx2 -> s_waitcnt vmcnt(0) -> 1 x v_mmac per + // chunk, so every v_mmac stalls on loads issued in the same iteration + // (latency-bound despite 4.0 blocks/CU). Software-pipeline the K loop + // instead: hold the current chunk's two 8-byte words in plain locals and + // issue chunk (k0 + kDummaTileK)'s two dwordx2 loads AFTER chunk k0's + // v_mmac. Each v_mmac then consumes words loaded a full iteration earlier, + // and at the point the wait must be placed the only outstanding loads are + // that older pair -- no loads younger than the consumed ones are in flight + // -- so a plain vmcnt(0) covers exactly the right data and the prefetch + // pair stays in flight across the next iteration. The fragment registers + // still receive the identical k-ascending byte pattern, so the int32 + // accumulation order (and therefore the bf16 output) is bit-identical to + // iteration 5. + const int arow = lane & 15; + const int acol = (lane >> 4) << 3; // {0, 8, 16, 24}: k offset within a + // 32-wide chunk (lane owns 8 bytes) + const int brow = lane & 15; // n index within the 16-wide tile + const int bcol = (lane >> 4) << 3; // k offset within the 32-wide chunk + // Iteration-13 k64 pairing, applied to each wavefront's own chunk range: + // every wavefront runs exactly 5 wait rounds (wave_chunks 9 or 10 -> 4 + // pairs + 1 tail chunk, or 5 pairs). Four dwordx2 words (two k32 + // sub-chunks) are software-pipelined one round ahead exactly like the + // iteration-8 pair, so the wait before each round covers only the + // previously issued 4-load set; all prefetch addresses stay strictly below + // k_hi (the loop prefetches only pairs that exist, and the tail chunk is + // loaded plainly with no prefetch). + const int pairs = wave_chunks >> 1; // 5 per wavefront + const int tail_chunks = wave_chunks & 1; // wave 0 only on odd slices + const int k_pair_end = k_lo + pairs * 2 * kDummaTileK; + // Prefetch pair 0: sub-chunk 0 at k_lo, sub-chunk 1 at k_lo + 32. + long a_word0 = *reinterpret_cast( + a + static_cast(arow) * k + k_lo + acol); + long b_word0 = *reinterpret_cast( + b + static_cast(n0 + brow) * k + k_lo + bcol); + long a_word1 = *reinterpret_cast( + a + static_cast(arow) * k + k_lo + kDummaTileK + acol); + long b_word1 = *reinterpret_cast( + b + static_cast(n0 + brow) * k + k_lo + kDummaTileK + bcol); + int k0 = k_lo; + for (; k0 + 2 * kDummaTileK < k_pair_end; k0 += 2 * kDummaTileK) { + // Consume sub-chunk 0 (k0 .. k0+31), then sub-chunk 1 (k0+32 .. k0+63): + // same k-ascending v_mmac sequence as two adjacent iteration-8 chunks. + *reinterpret_cast(a_frag.x) = a_word0; + *reinterpret_cast(b_frag.x) = b_word0; + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + *reinterpret_cast(a_frag.x) = a_word1; + *reinterpret_cast(b_frag.x) = b_word1; + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + // Prefetch pair p+1 (four dwordx2 loads) after pair p's two v_mmacs, so + // all four words are in flight one full round before consumption. + a_word0 = *reinterpret_cast( + a + static_cast(arow) * k + (k0 + 2 * kDummaTileK) + acol); + b_word0 = *reinterpret_cast( + b + static_cast(n0 + brow) * k + + (k0 + 2 * kDummaTileK) + bcol); + a_word1 = *reinterpret_cast( + a + static_cast(arow) * k + (k0 + 3 * kDummaTileK) + acol); + b_word1 = *reinterpret_cast( + b + static_cast(n0 + brow) * k + + (k0 + 3 * kDummaTileK) + bcol); + } + // Last pair: consume the words prefetched by the final loop iteration. + *reinterpret_cast(a_frag.x) = a_word0; + *reinterpret_cast(b_frag.x) = b_word0; + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + *reinterpret_cast(a_frag.x) = a_word1; + *reinterpret_cast(b_frag.x) = b_word1; + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + // Odd-slice tail: the remaining k32 chunk (slices 2..9: 608 = 9*64 + 32) + // at k_pair_end, loaded plainly with no prefetch (nothing follows; the + // address is a multiple of 32 and stays below k_end). + if (tail_chunks) { + *reinterpret_cast(a_frag.x) = *reinterpret_cast( + a + static_cast(arow) * k + k_pair_end + acol); + *reinterpret_cast(b_frag.x) = *reinterpret_cast( + b + static_cast(n0 + brow) * k + k_pair_end + bcol); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Iteration-14 per-block int32 combine: wave 0 publishes its partial + // 16x16 accumulator tile to a 1 KiB LDS tile, one barrier, then wave 1 + // adds wave 0's partial (ascending slice order: wave 0 holds chunks + // [0, half), wave 1 holds [half, end); int32 addition is exact and + // associative with max |acc| ~ 1e8 < 2^31, so the combined sum is + // bit-identical to the single-wave ascending accumulation) and publishes + // the plane with the same direct fragment ownership as the iteration-1 + // epilogue: row = lane & 15, col_group = lane >> 4, acc_frag.x[i] -> + // column col_group + 4*i. Plane strides are multiples of 256 B so every + // store stays 4-B aligned (workspace base is 256-B aligned). + __shared__ int32_t s_part[kDummaTileM * kDummaTileN]; // 1 KiB + const int row = lane & 15; + const int col_group = lane >> 4; + if (wave == 0) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + s_part[row * kDummaTileN + col_group + 4 * i] = acc_frag.x[i]; + } + } + __syncthreads(); + if (wave == 1) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc_frag.x[i] += s_part[row * kDummaTileN + col_group + 4 * i]; + } + int32_t* plane = partials + static_cast(s) * kM16N768Elems; +#pragma unroll + for (int i = 0; i < 4; ++i) { + plane[row * n + n0 + col_group + 4 * i] = acc_frag.x[i]; + } + } + + // Iteration-16 fused combine tail (FUSED=true only): the 48-block combine + // kernel and its launch gap are consolidated into the partial launch with + // the exact qkv-proven monotonic-arrival pattern (accepted on this stack + // for hy3_tp4_qkv_proj_m16, round 23): after the plane publish, one + // __syncthreads, then (wave 1, lane 0) does + // __threadfence() + atomicAdd(&counters[tile], 1) + __threadfence(). + // The block whose atomicAdd returns (arrived % kSplitK) == kSplitK - 1 is + // the LAST arrival for its tile (counters are monotonic, zeroed once per + // workspace by the launcher before the first launch; int32 wraparound + // would need ~2^31/10 replays): the release fence orders this block's + // plane store before its arrival atomic, the acquire fence orders its + // plane reads after observing every sibling arrival, and the intervening + // barriers broadcast the flag to the wave, so no spin loop or residency + // assumption is needed. The last block then sums the tile's kSplitK + // planes in ascending slice order with the exact iteration-12 combine + // epilogue (one aligned int4 load per plane per lane, x_scale scalar + + // weight_scale float4 hoisted, hip_bfloat16 conversion, 4 x 2-byte + // stores), so the bf16 output is bit-identical to iteration 14 (same + // planes, same ascending int32 sum, same conversion and store addresses). + if (FUSED) { + __syncthreads(); + __shared__ int s_is_last; + if (wave == 1 && lane == 0) { + __threadfence(); // release: this block's plane stores are visible to + // the observer of this block's arrival atomic + const int arrived = atomicAdd(&counters[tile], 1); + __threadfence(); // acquire: this block's plane reads below (after the + // barrier) see every sibling store released before + // its arrival atomic + s_is_last = ((arrived % kSplitK) == kSplitK - 1); + } + __syncthreads(); + if (s_is_last && wave == 1) { + // Epilogue, byte-identical to w8a8_gemm_m16_n768_sk_combine_kernel: + // lane -> (erow = lane >> 2, eq = lane & 3) owns four consecutive + // columns ecol0 = n0 + 4*eq (16-B-aligned int4 plane reads and + // 8-B-aligned bf16 stores, as in the iteration-12 vectorized combine). + const int erow = lane >> 2; + const int eq = lane & 3; + const int ecol0 = n0 + 4 * eq; + const int eidx = erow * n + ecol0; + const float xs = x_scale[erow]; + const float4 ws = *reinterpret_cast(weight_scale + ecol0); + int32_t eacc0 = 0, eacc1 = 0, eacc2 = 0, eacc3 = 0; +#pragma unroll + for (int s2 = 0; s2 < kSplitK; ++s2) { + const int4 v = *reinterpret_cast( + partials + static_cast(s2) * kM16N768Elems + eidx); + eacc0 += v.x; + eacc1 += v.y; + eacc2 += v.z; + eacc3 += v.w; + } + const float es0 = static_cast(eacc0) * xs * ws.x; + const float es1 = static_cast(eacc1) * xs * ws.y; + const float es2 = static_cast(eacc2) * xs * ws.z; + const float es3 = static_cast(eacc3) * xs * ws.w; + hip_bfloat16* eo = out + eidx; + eo[0] = hip_bfloat16(es0); + eo[1] = hip_bfloat16(es1); + eo[2] = hip_bfloat16(es2); + eo[3] = hip_bfloat16(es3); + } + } +} + +// Split-K combine kernel (iteration 2; iteration 12 vectorized plane reads): +// one 64-lane wavefront per 16x16 output tile (48 blocks). Each lane sums +// the kSplitK partial planes in ascending slice order -- exact int32 order, +// bit-identical to the scalar reference (int32 addition is associative and +// max |acc| ~ 1.0e8 < 2^31) -- then applies x_scale[m] * weight_scale[n] and +// stores bf16. This kernel is part of the timed operator: it is enqueued on +// the caller stream right after the partial kernel, inside the same Graph +// capture. +// +// Iteration 12 (HIP-only consolidation): the accepted iteration-2 scalar +// combine's ISA (profiled 5.28 us over the 48-block launch, 23-25% of the +// operator aggregate) is 45 global_load_dword + 26 s_waitcnt per wavefront: +// each lane reads its four output columns (col_group + 4*i, strided by 4) +// with 4 x 10 scalar dword loads and four partially-serialized wait chains. +// Remap the 64 lanes to (row = lane >> 2, quarter = lane & 3) so each lane +// owns four CONSECUTIVE columns and replaces its 40 dword loads with 10 +// aligned 16-byte global_load_dwordx4 (one int4 per plane: plane stride +// 49152 B, row stride 3072 B, and quarter base 16 B are all multiples of 16, +// so every int4 address is 16-B aligned). The ten plane loads are +// independent, so one wait group covers the set instead of 26 per-wavefront +// waits, and the per-wavefront 64-lane footprint of one load instruction is +// 16 full 64-B lines instead of 64 scattered dwords (the four lanes of a row +// wrote the full line, so every line read is complete and L2-resident). +// Scale loads are hoisted to the kernel top so they issue with the plane +// burst (the prior vectorized object issued the weight_scale dwordx4 late and +// exposed a second serialized wait after the plane chain). Per element the +// accumulation is still the same ascending-slice int32 sum over the same +// bytes, so the bf16 output is bit-identical to iteration 2; only the load +// granularity / lane mapping / load order change. Stores stay 4 x 2-byte +// bf16 per lane (L2 write combining absorbs them, as in the partial-kernel +// epilogue). +__global__ __launch_bounds__(kDummaWaveThreads) void +w8a8_gemm_m16_n768_sk_combine_kernel( + const int32_t* __restrict__ partials, // [kSplitK][16][768] int32 + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int n) { + const int lane = static_cast(threadIdx.x); // 0..63 + const int n0 = static_cast(blockIdx.x) * kDummaTileN; + const int row = lane >> 2; // 0..15 + const int q = lane & 3; // 0..3: 4-column quarter within the tile + const int col0 = n0 + 4 * q; // 4 consecutive int32 columns (16-B aligned) + const int idx = row * n + col0; + + // Scale loads hoisted to the kernel top (x_scale scalar, weight_scale as + // one aligned float4) so they issue with the plane burst below instead of + // trailing the plane chain with their own serialized wait. + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + + // One 16-byte int4 load per plane, ascending slice order; all ten loads + // are independent, so a single wait group covers the set (see kernel doc). + int32_t acc0 = 0, acc1 = 0, acc2 = 0, acc3 = 0; +#pragma unroll + for (int s = 0; s < kSplitK; ++s) { + const int4 v = *reinterpret_cast( + partials + static_cast(s) * kM16N768Elems + idx); + acc0 += v.x; + acc1 += v.y; + acc2 += v.z; + acc3 += v.w; + } + + // Same per-element scaling and 2-byte bf16 stores as iteration 2. + const float scaled0 = static_cast(acc0) * xs * ws.x; + const float scaled1 = static_cast(acc1) * xs * ws.y; + const float scaled2 = static_cast(acc2) * xs * ws.z; + const float scaled3 = static_cast(acc3) * xs * ws.w; + hip_bfloat16* o = out + idx; + o[0] = hip_bfloat16(scaled0); + o[1] = hip_bfloat16(scaled1); + o[2] = hip_bfloat16(scaled2); + o[3] = hip_bfloat16(scaled3); +} + +// Identity device-to-device copy for the bootstrap pack operation (int8 +// weights). Grid-stride loop: generic fallback for every (K, N). +__global__ void w8a8_pack_identity_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t numel) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < numel; idx += stride) { + dst[idx] = src[idx]; + } +} + +// Identity device-to-device copy for the bootstrap pack operation (fp32 +// per-column weight scales). +__global__ void w8a8_pack_identity_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t numel) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < numel; idx += stride) { + dst[idx] = src[idx]; + } +} + +// N-major transpose pack (iteration 5) for the exact MiniMax TP8 gate_up +// shape (k, n) == (6144, 768): dst[n * k + kk] = src[kk * n + n]. Runs once, +// outside the timed region and outside Graph capture, into the same +// caller-owned packed_weight buffer (byte count unchanged, so captured +// addresses stay valid). Each m16n16k32 B-fragment lane then owns 8 +// contiguous bytes (see w8a8_gemm_dumma_m16_direct_sk_partial_kernel), +// enabling one aligned global_load_dwordx2 per fragment instead of 8 +// strided global_load_ubyte. The identity copy remains the fallback for +// every other (K, N). +__global__ void w8a8_pack_nmajor_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t numel, + int n, + int k) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < numel; idx += stride) { + const int64_t n_idx = idx / k; + const int64_t k_idx = idx - n_idx * k; + dst[idx] = src[k_idx * n + n_idx]; + } +} + +} // namespace + +// gemm_out backend entry point (called by csrc/bindings.cpp on PyTorch's +// current HIP stream). This function must stay allocation/sync/compile-free +// and enqueue only on the caller stream: it is captured into a CUDA/HIP Graph +// and replayed with changed tensor contents. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // workspace (caller-owned, allocated before Graph capture) holds the + // kSplitK int32 partial planes used only by the exact-shape guard below + // (plus the per-tile fused-combine arrival counters in the last 192 B of + // the 16-plane contract workspace); its capacity is validated by the + // Python API layer. + + hip_bfloat16* out_bf16 = static_cast(out); + + // Exact-shape guard for the optimized M=16 MiniMax TP8 gate_up shape + // (16x768x6144): iteration-16 fused single-kernel split-K = 10 DUMMA + // m16n16k32 (480 two-wave partial blocks = 4.0 blocks/CU = 2 + // wavefronts/SIMD, one aligned 8-byte vector load per lane per fragment + // against the n-major packed weight, per-block LDS combine, qkv-proven + // monotonic-arrival fused combine tail replacing the 48-block combine + // kernel), on the caller stream. Every other + // (m, n, k) -- including the paired M=2 API shape with the same (N, K) + // and the paired 16x6144x384 down_proj shape -- falls through to the + // scalar generic fallback below (which decodes the n-major pack for + // (n, k) == (768, 6144)). + if (m == 16 && n == 768 && k == 6144) { + const int n_tiles = n / kDummaTileN; // 48 tiles of 16 columns + constexpr int64_t kSkWorkspaceBytes = + static_cast(kSplitK) * kM16N768Elems * sizeof(int32_t); + if (workspace_bytes >= kSkWorkspaceContractBytes) { + // Iteration-16 fused single-kernel path: the combine runs inside the + // partial kernel (last arrival per tile), so the operator is ONE + // launch and the 48-block combine kernel + its launch gap disappear. + // Grid stays 48 tiles x 10 slices = 480 two-wave blocks (2 + // wavefronts per SIMD). The per-tile monotonic arrival counters + // (n_tiles x int32 = 192 B) live in the last 192 B of the 786432-B + // (16-plane) contract workspace (plane 15's tail; planes 0..9 are + // the only partial planes used, so there is no overlap). The + // counters are zeroed once per workspace with an async memset on the + // caller stream (static workspace-pointer guard): the first launch + // happens during the eager warmup before Graph capture, so the + // memset is never part of the captured graph and every replay is + // exactly one kernel launch. + int32_t* partials = static_cast(workspace); + int32_t* counters = reinterpret_cast( + static_cast(workspace) + kSkWorkspaceContractBytes) - + n_tiles; + static const void* s_fused_counters_ws = nullptr; + if (s_fused_counters_ws != workspace) { + hipMemsetAsync(counters, 0, + static_cast(n_tiles) * sizeof(int32_t), + stream); + s_fused_counters_ws = workspace; + } + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_gemm_dumma_m16_direct_sk_partial_kernel), + dim3(n_tiles * kSplitK), dim3(kPartialBlockThreads), 0, stream, + a, b, partials, x_scale, weight_scale, out_bf16, counters, n, k); + return; + } + if (workspace_bytes >= kSkWorkspaceBytes) { + // Two-kernel path for workspaces below the 16-plane contract (cannot + // happen through the validated API path): the iteration-14 partial + // kernel with the fused tail compiled out (FUSED=false) + the + // iteration-12 vectorized combine kernel, byte-identical to the + // accepted iteration-14 behavior. Grid = 48 tiles x 10 slices = 480 + // two-wave blocks = exactly 4.0 blocks per device CU (120 CUs) = 2 + // wavefronts per SIMD; the FUSED=false instantiation never + // dereferences x_scale / weight_scale / out / counters. + int32_t* partials = static_cast(workspace); + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_gemm_dumma_m16_direct_sk_partial_kernel), + dim3(n_tiles * kSplitK), dim3(kPartialBlockThreads), 0, stream, + a, b, partials, x_scale, weight_scale, out_bf16, nullptr, n, k); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_m16_n768_sk_combine_kernel), + dim3(n_tiles), dim3(kDummaWaveThreads), 0, stream, + partials, x_scale, weight_scale, out_bf16, n); + return; + } + // Workspace too small for the ten partial planes (cannot happen through + // the validated API path): keep the iteration-1 direct kernel under the + // same shape guard. + const int blocks = n / kDummaTileN; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_dumma_m16_direct_kernel), + dim3(blocks), dim3(kDummaWaveThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, n, k); + return; + } + + // Generic scalar fallback for every unmatched (m, n, k), including M=2 + // API shapes that share an assigned (N, K). + const int threads = kScalarBlockThreads; + const int total = m * n; + const int blocks = (total + threads - 1) / threads; + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_kernel), + dim3(blocks), dim3(threads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); +} + +// pack_weight backend entry point (called by csrc/bindings.cpp on PyTorch's +// current HIP stream, outside the timed region and outside Graph capture). +// +// Bootstrap: identity device-to-device copy of raw_weight[K, N] and +// weight_scale[N, 1]. Later Parallel-explore rounds may change only this +// function's HIP implementation (and the matching GEMM interpretation) to +// test packed layouts; the identity copy remains the generic fallback for +// unmatched (K, N). +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + const int64_t weight_numel = static_cast(k) * n; + const int weight_blocks = static_cast( + (weight_numel + kCopyBlockThreads - 1) / kCopyBlockThreads); + const int scale_blocks = + (n + kCopyBlockThreads - 1) / kCopyBlockThreads; + + if (k == 6144 && n == 768) { + // Iteration-5 n-major transpose pack for the exact gate_up shape + // (packed[n*K+kk] = raw[kk*n+n]); the matching GEMM kernels and the + // scalar generic fallback decode this layout. Identity copy for every + // other (k, n). + hipLaunchKernelGGL( + w8a8_pack_nmajor_i8_kernel, + dim3(weight_blocks), dim3(kCopyBlockThreads), 0, stream, + raw_weight, packed_weight, weight_numel, n, k); + } else { + hipLaunchKernelGGL( + w8a8_pack_identity_i8_kernel, + dim3(weight_blocks), dim3(kCopyBlockThreads), 0, stream, + raw_weight, packed_weight, weight_numel); + } + hipLaunchKernelGGL( + w8a8_pack_identity_f32_kernel, + dim3(scale_blocks), dim3(kCopyBlockThreads), 0, stream, + weight_scale, packed_weight_scale, static_cast(n)); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/o_proj.hip new file mode 100644 index 00000000..0235a9c5 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/o_proj.hip @@ -0,0 +1,999 @@ +// @@variant shape=minimax_tp8_o_proj_m4096 commit=bdaa6897e0b608089d609eaa1639ea1ea170f65f added=2026-08-29 +// median_us=397.4 p90_us=398 speedup=50.06 baseline_us=1.99e+04 +// source=minimaxm3-dsh-tp8-m4096-1-b0482833 +// @@variant shape=minimax_tp8_o_proj_m4096 (iteration 2) +// INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. +// +// Worker: worker_2 (physical GPU 2), assigned shape: +// minimax_tp8_o_proj_m4096 : M=4096, N=6144, K=1024 +// +// Bootstrap strategy (iteration 1, correctness-first but a usable profiling +// baseline - a large-Prefill scalar K loop would not be): +// * Large-prefill path (exact (m, n, k) == (4096, 6144, 1024)): +// native INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 +// output tile per block; four wavefronts (256 threads); each wave owns +// a 64x32 quadrant built from eight m16n16k32 accumulator fragments; +// the block cooperatively vector-loads A[128,64] and B[64,64] into one +// single-buffered 64-K LDS stage (15,360 B total: 128*80 + 64*80 with +// 16 B padding per row for bank skew); two __syncthreads per stage +// (one after the cooperative load, one before the next stage +// overwrites LDS - both on all-thread paths); A and B fragments are +// read from LDS with the library du_load_matrix_sync row-major loaders +// (B tile is the raw [K, N] layout because launch_pack_w8a8_weight is +// the bootstrap identity device-to-device copy; later Parallel explore +// rounds may introduce a packed layout together with matching fragment +// loads); fused dot * x_scale[m] * weight_scale[n] epilogue stored +// directly as bf16 from the accumulator fragments (one coalesced +// 8-byte store per lane per fragment). Grid dim3(96, 32) = 3072 +// blocks dwarfs the 120 CUs, so no split-K is needed. +// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), +// including all small-M API cases (M=2, M=16) and any M in (0, 4096] +// with the same (K, N); the fallback reads the identity [K, N] weight +// layout (packing is identity in bootstrap). +// * launch_pack_w8a8_weight: identity device-to-device copy of the raw +// [K, N] weight and of the [N, 1] scales for every (k, n) - bootstrap +// only. Packing never happens inside the timed GEMM and never inside +// Graph capture. +// +// Iteration 2 (operand-reuse / staging round): +// Decision: cooperative A+B LDS staging is KEPT, direct global fragment +// loads are REJECTED. Per 128x64 macro-tile reuse: each A element is used +// by 64 output columns within the block (2 waves share a row-half) and +// re-fetched by 96 N-slice blocks; each B element is used by 128 output +// rows and re-fetched by 32 M-slice blocks. Direct loads would move the +// fragment reads from LDS to global/L2 with 2x fetch amplification per +// block (A 16 KiB vs 8 KiB unique, B 8 KiB vs 4 KiB unique per stage: +// 1.2 GiB vs 603 MB total wavefront traffic) and add per-wave L2-latency +// waits behind the two __syncthreads per stage; the worker-29 TP4 gate_up +// measured the direct-A arm at 7.9x slower (3444.26 vs 220.71 us) - the +// M=16 decode direct-load win does not transfer to M=4096 prefill. The +// binding cost is instead the LDS fragment path: the gfx928 code object +// shows the raw-[K,N] row_major B loader compiling to eight per-byte +// ds_read_u8 (k-strided by ldm) plus a mask/OR reassembly VALU chain on +// both operands (PMC: 9.24M LDS instr, 23.59M bank conflicts = 2.55/instr, +// 12.46M LDS waits). This round applies the worker-29 TP8 o_proj lineage +// iteration-10/12 mechanism for the exact same 128x64x64 geometry: the +// weight pack for (k, n) == (1024, 6144) becomes n-major [N, K] +// (packed[n*K + kk] == raw[kk*N + n]; every B tile of one (n-tile, +// K-stage) is one contiguous 4 KiB vector stream), B is staged n-major in +// LDS (b_tile[n][k], 64 rows x 80 B - the same 5,120 B, same +// 16-byte-aligned bank-skewed stride), and BOTH operand fragments are +// filled by load_fragment8 (one 8-byte LDS read per fragment, the same 8 +// bytes in the same x[0..7] order the library loaders produce => v_mmac +// operand bit pattern and exact int32 accumulation unchanged; mismatch 0 / +// max_abs_error 0.0 expected). Tile geometry, cooperative int4 staging, +// two-barrier single-buffered 64-K stage, per-accumulator MMAC order, the +// coalesced epilogue, the exact (4096, 6144, 1024) guard and the scalar +// fallback (now decoding the n-major layout for (k, n) == (1024, 6144)) +// are otherwise untouched. Falsifiable prediction: lds_instructions and +// lds_bank_conflicts drop sharply (8-byte fragment reads instead of 8 +// strided byte loads), lds_wait_instructions and valu_instructions drop +// (mask/OR chains gone), and the median/P90 improve; the removal of +// reassembly temporaries may also lower arch_vgpr 80 -> ~61-64, which +// would move residency from 3 to 4 blocks/CU (the trusted occupancy probe +// set [2,3,4,5]: 4 blocks x 256 threads x <=64 VGPR = 65,536 exactly +// fits, 4 x 15,360 B LDS = 61,440 B fits 64 KiB). +// +// Iteration 6 (dead-barrier peel round; source restored to the accepted +// iteration-2 digest 92917037c97822f73ae519ba94d92219946b3fe50dafcbf351cd175a38ff2e60 +// after the iteration-5 infrastructure kill): +// The stage loop emits two __syncthreads per stage (32 per block for +// K=1024/64 = 16 stages): barrier A after the cooperative staging writes +// (writes -> reads) and barrier B after the consume (reads -> next stage +// overwrites LDS). Barrier B of the FINAL stage (k0 = 960) is dead work: +// no successor stage overwrites LDS, and the epilogue reads only +// accumulator registers plus global x_scale/weight_scale - it never +// touches LDS. This round guards that barrier with +// `if (k0 + kStageK < k)` so the final stage runs its last MMAC group and +// enters the epilogue one barrier earlier (32 -> 31 barriers per block, +// -3.1%). Everything else is untouched: 128x64 tile, 256 threads, +// single-buffered 64-K stage (15,360 B LDS, 4 blocks/CU, arch_vgpr 64), +// cooperative int4 A+B staging, load_fragment8 on both operands (bit- +// identical int32 order), coalesced 8-byte epilogue, exact (4096,6144,1024) +// guard, n-major [N,K] pack, and the scalar fallback. The guard is +// wavefront-uniform (k0 and k are scalars), so there is no divergence. +// Falsifiable claim: barrier-stall time per block drops by ~1/32 of the +// barrier cost, so median/P90 stay flat-to-improved (down_proj iter-12's +// accepted variant dropped its dead last barrier for the same reason); +// a >2% regression versus 449.53/450.35 us falsifies the mechanism and the +// accepted source is restored. Expected ISA delta: steady-state loop still +// has 2 s_barrier for stages 0..14, the final stage path has only 1. +// +// Iteration 8 (resource round: tune waves per block, VGPR live ranges and LDS +// footprint from measured occupancy; recheck the best tile family with normal +// median/P90 measurements): +// Measured occupancy of the current-best (iteration-6) source (digest +// 8d133a8d90204d4e4e36922249701b2ea78cad8243e89bcd9affbe17d3f667e0): 16 +// resident waves/CU is the hard ceiling - arch_vgpr 64 x 16 waves = 1,024 +// VGPRs/CU exactly (65 VGPRs -> 3 blocks/CU, the measured +6% regression of +// iteration 4) AND LDS 4 x 15,360 B = 61,440 B (93.75% of 64 KiB). 5 +// blocks/CU is doubly unreachable (it would need VGPR <= 51 - 8 accumulator +// chains alone are 32 VGPRs plus 12 fragment VGPRs - AND LDS <= 13,107 B), +// and the probe set [2,3,4] is monotonically worse below 16 waves +// (iteration 3 double-buffer at 2 blocks/CU = 8 waves: 894 us; iteration 4 +// 128x96 at 3 blocks/CU = 12 waves: 477.6 us; accepted 4 blocks/CU = 16 +// waves: 446.4 us). The one resource axis never probed AT the 16-wave +// ceiling is waves PER BLOCK. This round redistributes the same 16 +// waves/CU from 4 blocks x 4 waves to 2 blocks x 8 waves by doubling the +// M-tile to 256x64 (512 threads = 8 waves of 64; each wave still owns a +// 64x32 quadrant of eight m16n16k32 int32 accumulators, so the per-wave +// instruction stream, VGPR live ranges and per-wave LDS fragment footprint +// are IDENTICAL by construction - the probe is a pure block-geometry +// change). Per-block LDS grows 15,360 -> 25,600 B (A[256,80] + B[64,80]); +// per-CU LDS is unchanged at 51,200 B (2 x 25,600 <= 64 KiB) and per-CU +// VGPR is unchanged at 1,024. A staging stays 2 int4s/thread over 512 +// threads (rows [0,128) + [128,256)); B staging (256 int4s) is guarded to +// threads 0..255 via `if (b_n < kBlockN)` so the 64-row b_tile index cannot +// overflow. Barriers: still 2 per stage with the iteration-6 final-stage +// guard (31 per block), now synchronizing 8 waves instead of 4; grid 3072 +// -> 1536 blocks (16 M-slices x 96 N-slices). Bit-identical int32 +// discipline holds (same per-accumulator kk0-then-kk1 order, same v_mmac +// operand bit patterns -> mismatch 0 / max_abs_error 0.0 expected). +// Everything else is untouched: 64-K single-buffered stage, cooperative +// int4 staging, load_fragment8 on both operands, n-major [N,K] pack, +// coalesced 8-byte epilogue, exact (4096,6144,1024) dispatch guard, scalar +// fallback, and the (unchanged) kernel symbol name - an internal +// anonymous-namespace symbol the harness profiler matches by name. +// FALSIFIABLE CLAIM: at the fixed 16-wave ceiling the kernel is +// LDS-issue/latency-bound (PMC of the same geometry: 2.10M lds_wait vs +// 749,568 vmem_read, 5.54 bank conflicts per ds_read2_b64; iteration 7's +// prefetch probe was flat at 447.32/448.24, so LDS latency is already +// hidden by the resident waves), therefore the only changed physical +// variables - barrier-domain size (8 vs 4 waves per __syncthreads) and +// block co-residency (2 vs 4 independent barrier domains per CU) - keep +// median/P90 flat-to-improved vs 446.41/446.99 us; a >2% regression +// (median > 455.34 or p90 > 455.93) falsifies the 8x2 geometry and the +// iteration-6 source is restored before the next experiment. CONFOUND +// GUARD: if the compiled arch_vgpr lands >= 65, residency drops to 1 +// block/CU (8 waves/CU) and the probe is confounded - restore the +// iteration-6 source and report regardless of timing. Expected ISA +// deltas: grid 3072 -> 1536, workgroup 256 -> 512, LDS 15,360 -> 25,600 B, +// arch_vgpr 64; per-wave steady-state consume unchanged (6 ds_read2_b64, +// 16 v_mmac, 5 lgkmcnt waits, 31 s_barrier per block); staging: waves 0..3 +// issue 3 ds_write_b128 per stage (2 A + 1 B), waves 4..7 issue 2 (A only); +// lds_instructions / lds_bank_conflicts / vmem_read flat per output. +// +// Iteration 16 (N-tile expansion 64 -> 128 at the measured winning +// residency; round 15 was killed by the agent infrastructure before a valid +// proposal and does NOT count - the accepted iteration-8 source is restored +// byte-identical and ONE new bounded mechanism is applied): +// Measured state: accepted iteration-8 source (digest +// 88c9a331d5dac906a4e77c882e7fdde0bfd2f51f685730ad8ce0af49b282d66a) at +// 440.21/440.84 us, 1 block/CU x 8 resident waves (arch_vgpr 72 > 64, so +// 2 blocks/CU never fit; PMC: 2.46M lds_instructions vs 16.71M +// lds_wait_instructions, 651K vmem_read, 83.7% L2 hit - LDS-latency/issue +// bound, not HBM- or MMAC-bound). Every co-residency probe on this shape +// was falsified (iter 10 double-buffer 2 blocks/CU 669.35; iter 11 +// launch_bounds diet confounded at vgpr 71; iter 14 N=32 diet 632.94; the +// 128x64 16-wave control 446-449), and the winning direction is +// consistently FEWER, BIGGER blocks per barrier domain (iter 8's 256x64 +// 8-wave 1-block beat every 16-wave layout). DECISION - N-tile expansion +// 64 -> 128 at the UNCHANGED 1 block/CU x 8 waves residency: kBlockN 128, +// per-wave quadrant 64x64 = sixteen m16n16k32 int32 accumulators +// (accumulator VGPR 32 -> 64; arch_vgpr ~72 -> ~95-110, which STRUCTURALLY +// pins 1 block/CU x 8 waves - 2 blocks/CU would need <= 64 VGPR with 64 +// VGPR of accumulators alone, impossible), LDS 25,600 -> 30,720 B +// (A[256,80] + B[128,80]; B row stride stays 80 = kStageK + kBPad, the +// canonical 16-byte-aligned five-bank-phase skew, 2-way conflict floor +// identical to the accepted 64-row B tile), B staging stays one +// int4/thread (128 rows x 64 k = 512 int4s over exactly 512 threads - the +// b_n < kBlockN guard becomes always-true and is elided), grid 1536 -> 768 +// (48 x 16), 2 barriers per stage unchanged (31 per block), kernel symbol, +// n-major pack layout, exact (4096,6144,1024) dispatch guard and scalar +// fallback unchanged. Per-output deltas: LDS fragment reads 3 -> 2 B (A +// read amplification 2 -> 1 at constant B; per wave per stage 12 -> 16 +// ds_read2_b64 for twice the outputs), barrier frequency halves, global +// read traffic 1.25 -> 0.75 B/output (A staging unchanged at 16,384 B per +// stage, B staging doubles to 8,192 B per stage = 24,576 B per block per +// stage). Bit-identical discipline: each of the 16 accumulators keeps the +// same kk0-then-kk1 int32 order and each v_mmac operand reads the SAME 8 +// LDS bytes per fragment (the same addresses receive the same bytes) -> +// mismatch 0 / max_abs_error 0.0 expected; Graph capture + changed-content +// replay must pass. COST ACKNOWLEDGED: grid 768/120 CUs = 6.4 blocks/CU +// -> 7 block-waves vs 13 today (+~9% makespan ceiling from tail +// quantization: 6 full waves + one 48-block partial wave), which must be +// repaid by the per-output LDS-read/barrier amortization; the per-output +// instruction reductions are the only changed physical variables. +// FALSIFIABLE CLAIM: median/P90 flat-to-improved vs the accepted +// 440.21/440.84 us; a >2% regression (median > 449.01 or p90 > 449.66) +// falsifies the 256x128 extension and the iteration-8 source (digest +// 88c9a331...) is restored before the next experiment. CONFOUND GUARD: +// the exact code object of the timed run must show vgpr_count <= 128 with +// scratch == 0 and vgpr_spill_count == 0, grid 768, workgroup 512, +// lds_bytes 30,720, residency 1 block/CU (any spill/scratch or residency +// change confounds the probe - restore and report regardless of timing). +// Expected ISA steady state per wave per stage: 16 ds_read2_b64 (8 per +// kk) + 32 v_mmac (vs 12 + 16), still 3 global_load_dwordx4 -> +// ds_write_b128 x3 per thread (2 A + 1 B, all 8 waves now) -> 2 s_barrier +// (31 per block); epilogue 16 x 8-byte coalesced stores per lane per wave. +// Expected PMC: vmem_read ~651K -> ~326K (halved - fewer blocks), +// vmem_write ~98K flat (same output bytes), lds_instructions ~2.46M -> +// ~1.64M (-33%), lds_bank_conflicts ~13.37M -> ~8.9M (-33%, same 2-way +// floor), lds_wait ~16.71M -> ~11M, valu_instructions ~14.24M flat per +// output, l2_hit_rate flat-to-slightly-lower (A-slice reuse drops 96 -> +// 48 blocks per slice). Rejected alternatives for this round (evidence, +// not re-litigation): raw asm (policy); any occupancy/co-residency change +// (iters 10/11/14 and the 128x64 16-wave control all falsified or +// confounded); one-stage-ahead global prefetch (iter 13 neutral at 442.57 +// - the global round trip is already hidden); N=96 tile (grid 1024 = +// 8.53 blocks/CU -> 9 waves, +5.5% quantization, for only -22% per-output +// LDS reads - dominated by N=128's -33% at the same barrier amortization); +// M-tile 512 (per-wave 128x32 keeps A at 2 B/output - strictly worse than +// N=128's 1 B/output); stage 128 (measured-saturated axis on this +// lineage); split-K (MxN grid dwarfs 120 CUs); 3/4/5 blocks/CU +// (VGPR-infeasible at 512 threads). +// +// Iteration 17 (final conditional round; control plane phase hip_only, +// plateau=false, raw_inline_asm_allowed=false, skill_allowed=false - no prior +// ISA-guided round recorded a compiler limitation, and recent_valid_ +// improvements_percent [-0.53, -30.45, +2.33] does not prove a plateau, so +// raw inline asm stays forbidden and this is ONE HIP-only consolidation +// change): +// Measured state: accepted iteration-16 source (digest +// 45ca0cdc2d38470c7131dae2fa2307c04e7e650e2cec468abf3eae0e0c398d6a) at +// 430.21/430.95 us, 1 block/CU x 8 resident waves, arch_vgpr 112, grid 768. +// The fresh PMC of the exact accepted object (iteration-17 cache) is still +// LDS-latency-dominated: lds_wait_instructions 8.89M vs lds_instructions +// 1.87M (4.75x; the 128x64 16-wave geometry that measured flat for +// iteration 7's read-prefetch probe had lds_wait 2.10M BELOW lds_instr +// 2.56M - latency was hidden there by 4 waves/SIMD, and that flatness does +// NOT transfer to 2 waves/SIMD). The exact-source ISA of the accepted +// 256x128 object (iteration-17 current-best-isa) shows the consume per kk: +// the first group is 6x ds_read2_b64 (4 B + 2 A) with staggered lgkmcnt +// waits interleaved into the first 8 v_mmac, but the rows-32-47 and rows- +// 48-63 A fragments of the wave's 64x64 quadrant are loaded JUST-IN-TIME +// after that first MMAC group (ds_read2_b64 at 0x3754/0x375C, standalone +// s_waitcnt lgkmcnt(2) at 0x3764, first acc20 v_mmac at 0x3768 - ~2 +// instructions of slack = the full LDS latency is exposed twice per stage +// per wave). DECISION - intra-stage A-fragment read-prefetch at the +// 16-accumulator/8-wave geometry: the rows-32-47/48-63 A fragments move +// into dedicated fragment variables (a_frag2/a_frag3) loaded in the FIRST +// read group (all 8 ds_read2_b64: 4 B + 4 A before any v_mmac), giving +// those two reads ~16-32 instructions of MMAC-issue slack per kk. LDS +// traffic is UNCHANGED (still 8 ds_read2_b64 per kk = 16 per stage per +// wave, the same 8 LDS bytes per fragment at the same addresses) and the +// ISA already allocates four distinct A-fragment register groups +// (v[88:91], v[100:103], v[104:107], v[108:111] = 16 VGPR), so arch_vgpr +// stays ~112 (<= 128 -> 1 block/CU x 8 waves unchanged, scratch 0, no +// spills). Bit-identical discipline holds: each accumulator keeps its +// kk0-then-kk1 int32 order, the v_mmac operand bit patterns are unchanged, +// and the only delta is when the same LDS bytes enter the A-fragment +// registers -> mismatch 0 / max_abs_error 0.0 expected; Graph capture + +// changed-content replay must pass. Everything else is byte-for-byte the +// accepted iteration-16 structure: 256x128 tile, 512 threads, per-wave +// 64x64 quadrant of sixteen accumulators, single-buffered 64-K stage +// (A[256,80] + B[128,80] = 30,720 B), 2 barriers per stage with the +// iteration-6 final-stage peel (31 per block), cooperative int4 staging +// with the b_n < kBlockN guard, load_fragment8 on both operands, coalesced +// 8-byte epilogue, exact (4096,6144,1024) dispatch guard, n-major [N,K] +// pack, scalar fallback, unchanged kernel symbol. EXPECTED DELTAS to +// verify the mechanism: ISA steady state - all 8 ds_read2_b64 issue before +// the first v_mmac and the standalone lgkmcnt wait at 0x3764 disappears +// (merged into the staggered waits); PMC - lds_instructions ~1.87M FLAT, +// lds_bank_conflicts ~8.65M FLAT (same 2-way floor), lds_wait 8.89M -> +// ~6.5-7.5M (-15-25%), vmem_read ~454K FLAT, vmem_write ~98K FLAT, +// valu_instructions ~13.25M FLAT; launch record - grid 768, workgroup 512, +// lds 30,720, arch_vgpr ~112, scratch 0, residency 1 block/CU. FALSIFIABLE +// CLAIM: with the two just-in-time A reads per kk covered by MMAC-issue +// slack at the measured 2-waves/SIMD residency, median/P90 stay flat-to- +// improved vs the accepted 430.21/430.95 us; a >2% regression (median > +// 438.81 or p90 > 439.57) falsifies the read-prefetch mechanism and the +// iteration-16 source (digest 45ca0cdc...) is restored before the next +// experiment. CONFOUND GUARD: the exact code object of the timed run must +// show vgpr_count <= 128 with scratch == 0 and vgpr_spill_count == 0, grid +// 768, workgroup 512, lds_bytes 30,720, residency 1 block/CU (any +// spill/scratch or residency change confounds the probe - restore the +// iteration-16 source and report regardless of timing). Rejected +// alternatives for this round (evidence, not re-litigation): raw asm +// (policy); every occupancy/co-residency point in the trusted probe set +// [2,3,4,5] (VGPR-infeasible at 512 threads: 64 accumulator + 16 A + 16 B +// fragment VGPR alone; 2 blocks/CU would need <= 64 VGPR, and every +// measured co-residency probe was falsified - iter 10 double-buffer 669.35, +// iter 11 diet confounded at vgpr 71, iter 14 N=32 diet 632.94, the 128x64 +// 16-wave control 446-449); 128x128 at 256 threads / 2 blocks/CU x 4 waves +// (the domain-splitting direction lost at every measurement - iter 8's +// 4x4-wave control 446.4 vs its 2x8-wave 440.2, then iter 16's 1x8-wave +// 430.2: the trend is consistently FEWER, BIGGER barrier domains, and the +// grid 1536/240 = 6.4 quantization is identical); one-stage-ahead global +// prefetch (iter 13 neutral-worse 442.57 - the global round trip is not +// the binding stall); double-buffer / stage-top publish (iter 10 669.35, +// TP4 o_proj 891.29); stride/bank-conflict probes (flat-to-worse, conflict +// floor not on the critical path); stage 128 (saturated axis at 144-B +// strides); stage 32 (TP4 866.60, and iter 14's stage change was fused +// with the N=32 diet); kBlockN 192/256 (24/32 accumulators -> VGPR +// infeasible); M-tile 512 (VGPR-infeasible); split-K (MxN grid dwarfs 120 +// CUs); epilogue restructures (vmem_write 98K is 0.2% of the stream; scale +// staging lost on down_proj -4.9%); per-wave staging / zero barriers (TP4 +// o_proj 967.17 - duplicated staging traffic swamps barrier savings). +// +// Iteration 21 (final conditional round; control plane phase hip_only, +// plateau=false, raw_inline_asm_allowed=false, skill_allowed=false - no prior +// ISA-guided round recorded a compiler limitation, and recent_valid_ +// improvements_percent [+1.29, -7.18, -0.27] does not prove a plateau, so raw +// inline asm stays forbidden and this is ONE HIP-only consolidation change): +// Measured state: accepted iteration-17 source (digest +// e1b870f66e8cb61f04aac92f5d669405e5a2fd6bca4a01902592b5bab6acff23) at +// 424.71/425.64 us, 1 block/CU x 8 resident waves, arch_vgpr 112, grid 768 +// (iteration 20's loop-carried one-stage-ahead global prefetch measured +// neutral 425.86/426.97 and was NOT accepted; its exact code object shows +// why - the compiler reverted the pipelining: the 'prefetched' loads still +// issue immediately before the visibility barrier with their own +// s_waitcnt vmcnt(0), so the stage-top global latency was never actually +// covered and the +12 loop-carried VGPR bought nothing). DECISION - remove +// the provably-dead b_n < kBlockN staging guard (iteration-8 leftover; +// always true at kBlockN = 128 with 512 threads: b_n = tid >> 2 is in +// [0, 127]) and load all three staging int4s (A rows 0-127, A rows +// 128-255, B) into locals BEFORE the three ds_write_b128 publishes. The +// accepted object's ISA (iteration-18 current-best-isa of digest +// e1b870f6..., vgpr 112) shows the guarded B load issues only after both A +// writes: global_load_dwordx4 x2 -> s_waitcnt vmcnt(1) -> ds_write_b128 -> +// vmcnt(0) -> ds_write_b128 -> s_and_saveexec/cbranch (guard) -> B +// global_load_dwordx4 -> vmcnt(0) -> ds_write_b128 -> s_or_b64, i.e. a +// SECOND serialized L2 round trip on the collective stage-top critical +// path every stage (all 8 waves are barrier-locked between the WAR and the +// visibility barrier, so the sibling wave on the SIMD cannot hide it). A +// stage is ~5,500 cycles at the measured 2-waves/SIMD residency, so one +// L2 round trip per stage is a first-order term. With the guard gone the +// B load issues together with the A loads at the stage top and its latency +// overlaps A's; the exec-mask dance (s_and_saveexec/s_cbranch_execz/ +// s_or_b64, 3-4 instructions per stage per wave) disappears. Same global +// addresses, same LDS addresses, same bytes -> mismatch 0 / max_abs_error +// 0.0 expected; Graph capture + changed-content replay must pass. LDS +// 30,720 B unchanged, 2 barriers per stage with the iteration-6 final-stage +// peel (31 per block), 8 ds_read2_b64 + 32 v_mmac per stage per wave +// unchanged, grid 768, workgroup 512, arch_vgpr ~112 (<= 128 -> residency +// stays 1 block/CU x 8 waves; the trusted occupancy split candidates +// [2,3,4,5] stay structurally impossible: 64 accumulator + 32 fragment +// VGPR alone at 512 threads). EXPECTED DELTAS to verify the mechanism: +// ISA steady state - the three global_load_dwordx4 issue back-to-back at +// the stage top with no branch and no exec-mask, the B ds_write_b128 is +// unconditional (staging still 3 ds_write_b128 per thread per stage), 31 +// s_barrier / 32 v_mmac / 8 ds_read2_b64 unchanged; PMC - vmem_read ~454K +// FLAT, lds_instructions ~1.87M FLAT, lds_bank_conflicts ~8.65M FLAT (same +// LDS layout and reads), lds_wait ~12.32M FLAT-to-lower (no LDS-side +// change), valu_instructions flat minus the exec-mask ops; launch record - +// grid 768, workgroup 512, lds_bytes 30,720, arch_vgpr ~112, scratch 0, +// residency 1 block/CU. FALSIFIABLE CLAIM: with the second serialized +// staging L2 round trip removed from the collective stage-top critical +// path, median/P90 improve or stay flat vs the accepted 424.71/425.64 us; +// a >2% regression (median > 433.20 or p90 > 434.16) falsifies the +// staging-load consolidation and the iteration-17 source (digest +// e1b870f6...) is restored before the next experiment. CONFOUND GUARD: +// the exact code object of the timed run must show vgpr_count <= 128 with +// scratch == 0 and vgpr_spill_count == 0, grid 768, workgroup 512, +// lds_bytes 30,720, residency 1 block/CU (any spill/scratch or residency +// change confounds the probe - restore the iteration-17 source and report +// regardless of timing). Rejected alternatives for this round (evidence, +// not re-litigation): raw asm (policy); loop-carried staging prefetch +// (iter 20 425.86 - neutral, and its code object proves the compiler +// reverted the pipelining); double-buffer / barrier-halving / stage-top +// publish into the idle buffer (iter 18 457.56 - falsified at this +// residency); every occupancy/co-residency point in the trusted probe set +// [2,3,4,5] (VGPR-infeasible at 512 threads: 64 accumulator + 32 fragment +// VGPR alone, and every measured co-residency probe was falsified - iter +// 10 double-buffer 669.35, iter 11 diet confounded at vgpr 71, iter 14 +// N=32 diet 632.94, the 128x64 16-wave control 446-449); 128x128 at 256 +// threads / 2 blocks/CU x 4 waves (the domain-splitting direction lost at +// every measurement); kBlockN 192/256 (24/32 accumulators -> VGPR +// infeasible); M-tile 512 (VGPR-infeasible); stage 128 (saturated axis at +// 144-B strides); stage 32 (TP4 866.60); stride/bank-conflict probes +// (flat-to-worse, conflict floor not on the critical path); split-K (MxN +// grid dwarfs 120 CUs); epilogue restructures (vmem_write 98K is ~0.2% of +// the instruction stream); per-wave staging / zero barriers (TP4 o_proj +// 967.17 - duplicated staging traffic swamps barrier savings). + +#include +#include +#include + +#include + +namespace { + +using namespace du::dumma; + +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 +constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kTargetM = 4096; +constexpr int kTargetN = 6144; +constexpr int kTargetK = 1024; +constexpr int kBlockM = 256; // iteration 8: 512 threads = 8 waves/block +constexpr int kBlockN = 128; // iteration 16: N-tile 64 -> 128 (16 accums/wave) +constexpr int kStageK = 64; +constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) +constexpr int kBStride = kStageK + kBPad; // 80: B rows hold kStageK=64 k + 16 pad +constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride +constexpr int kBlockThreads = 8 * kWaveSize; // iteration 8: 512 = 8 waves/block + +using bf16_t = hip_bfloat16; + +// --------------------------------------------------------------------------- +// Coalesced fragment epilogue: the m16n16k32 accumulator lane mapping +// (row = lane & 15, column group c4 = lane >> 4, frag.x[i] -> column +// c4 + 4*i) gives each lane four elements strided by 4 columns, so the direct +// per-element store is four 2-byte scalar stores per lane at 25% sector +// utilization. This epilogue transposes the 4-element groups within each +// 4-lane column group (lanes r, r+16, r+32, r+48: a 4x4 transpose via two +// 2x2 steps with __shfl_xor 16 then 32; no staging tile round trip), so lane +// (r, c4) ends up holding the four CONTIGUOUS columns 4*c4 .. 4*c4+3, +// converts them to bf16, packs 4 bf16 (8 B), and issues ONE 8-byte store per +// lane (100% store sector efficiency). Only the int32 values are re-routed +// between lanes - the per-element scale multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are +// unchanged, so the stored bits are identical to the direct store. The +// row>=m guard is wavefront-uniform (all 64 lanes of a wave share the same +// 16-row window), the shuffles never mix active and inactive lanes, and +// base_col is a multiple of 16 with n*2 a multiple of 8, so the float4 +// weight_scale load and the 8-byte store are aligned. +// --------------------------------------------------------------------------- +template +__device__ __forceinline__ void store_prefill_fragment_coalesced( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int base_row, + int base_col, + int m, + int n, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M masking: padded rows never write (wavefront-uniform) + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 16, n*2 is a multiple of 8). + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + static_cast(row) * n + col0) = + packed; +} + +// --------------------------------------------------------------------------- +// Iteration 2: direct 8-byte LDS fragment fill (load_fragment8). The library +// du_load_matrix_sync int8 loaders assign lane (r = lane & 15, c4 = lane >> 4) +// eight consecutive bytes x[0..7] and then make the compiler emit per-byte +// LDS reads plus a mask/OR reassembly VALU chain before the v_mmac operand. +// Writing the same 8 bytes straight into the fragment storage (one 8-byte +// LDS read) keeps the operand bit pattern identical (exact int32 +// accumulation unchanged) and removes the dead byte-reassembly VALU and its +// lgkmcnt wait chains. For the m-major A tile and the n-major B tile both +// operands are laid out so each lane's eight bytes are contiguous at +// (row)*stride + c4*8: A lane (r, c4) owns k = kk + c4*8 .. +7 of +// a_tile[(local_row + r)*80 + kk]; B lane (r, c4) owns k = kk + c4*8 .. +7 +// of b_tile[(local_col + r)*80 + kk]. This is the exact pattern validated +// on the worker-29 TP8 o_proj lineage (iterations 10 and 12, 419.05 -> +// 348.27 -> 322.99 us) for the same 128x64x64 geometry. +template +__device__ __forceinline__ void load_fragment8( + Frag& frag, + const signed char* __restrict__ base, + int stride, + int lane) { + const int off = (lane & 15) * stride + ((lane >> 4) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// --------------------------------------------------------------------------- +// Large-M prefill: 256x128 output tile per block, eight wavefronts of 64 lanes +// (iteration 8 resource round: 512 threads = 8 waves/block; iteration 16 +// N-tile round: 64 -> 128 columns per block so each wave owns a 64x64 +// quadrant of sixteen m16n16k32 int32 accumulators). Measured residency of +// the accepted 256x64 source: 1 block/CU x 8 resident waves (arch_vgpr 72 > +// 64, so 2 blocks/CU never fit; every co-resident occupancy probe on this +// shape - iteration 10 double-buffer at 2 blocks/CU, iteration 11 +// launch_bounds diet at vgpr 71, iteration 14 N=32 diet at 632.94 us - was +// falsified, while the 256x64 8-wave single-domain geometry won at +// 440.21/440.84 us). Iteration 16 keeps the winning 1 block/CU x 8 waves +// residency (sixteen int32 accumulators alone are 64 VGPR, so 2 blocks/CU is +// structurally impossible) and doubles per-block work: per-output LDS +// fragment reads drop 3 -> 2 B (A read amplification 2 -> 1 at constant B), +// per-output barrier frequency halves (still 2 barriers per stage, 31 per +// block), and per-output global read traffic drops 1.25 -> 0.75 B (A staging +// unchanged at 2 int4/thread, B staging 1 int4/thread - 128 rows x 64 k = +// 512 int4s over exactly 512 threads). The block cooperatively stages +// A[256,64] from x_q (row-major [M, K], stride k) and B[128,64] from the +// packed n-major [N, K] weight (iteration 2: B row n is contiguous in k, so +// one (n-tile, K-stage) is one contiguous 8 KiB vector stream in both global +// and LDS) into a single-buffered 64-K LDS buffer (A[256,80] + B[128,80] = +// 30,720 B). Two barriers per stage: one after the cooperative load, one +// before the next stage overwrites LDS (iteration 6: the final stage's +// overwrite-guard barrier is skipped because no successor stage exists - 31 +// barriers per block instead of 32). A and B fragments are filled by +// load_fragment8 (one 8-byte LDS read per fragment; iteration 2, replacing +// the library row-major loaders' per-byte ds_read_u8 + mask/OR reassembly +// VALU on both operands). Direct fragment epilogue with fused scales + bf16 +// conversion. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + const int local_row = wave_row * 64; + const int local_col = wave_col * 64; // iteration 16: 2 wave columns x 64 = 128 + + // Single-buffered 64-K stage: A[256, 80] + B[128, 80] = 30,720 B/block + // (iteration 16: N-tile 64 -> 128 at the measured 1 block/CU x 8 waves + // residency - sixteen int32 accumulators alone are 64 VGPR, so 2 blocks/CU + // would need <= 64 VGPR and is structurally impossible; the padded 80-byte + // strides are 16-byte-aligned and break the 64-byte LDS bank periodicity). + // A is m-major [256 rows, 64 k], B is n-major [128 n, 64 k] (iteration 2 + // packed layout; k is the contiguous dimension in the packed [N, K] + // weight). + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kBStride]; + + DUFragment + a_frag0, a_frag1, a_frag2, a_frag3; // iter 17: 4 live A fragments + // (rows 0-15/16-31/32-47/48-63) so + // all 8 ds_read2_b64 issue in the + // first read group + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; + DUFragment + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13, + acc20, acc21, acc22, acc23, acc30, acc31, acc32, acc33; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc22, 0); + du_fill_fragment(acc23, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + du_fill_fragment(acc32, 0); + du_fill_fragment(acc33, 0); + + // Cooperative A staging: each thread owns one int4 in A rows [0,128) and + // one in A rows [128,256): A[256,64] is 1,024 int4s = 512 threads x 2 int4s. + const int vector_byte_offset = tid * static_cast(sizeof(int4)); + const int stage_row = vector_byte_offset / kStageK; + const int stage_col = vector_byte_offset - stage_row * kStageK; + // Cooperative B staging from the packed n-major [N, K] weight (iteration + // 2): the B[128 n, 64 k] tile is 512 int4s = exactly one int4 per thread + // (iteration 16: with kBlockN = 128, b_n = tid >> 2 covers exactly the 128 + // n-rows; the iteration-8 b_n < kBlockN guard became always true then and + // is REMOVED in iteration 21 so the B global load issues with the A loads + // at the stage top - the compiler never elided the guard, and the accepted + // object's ISA shows the guarded B load serialized behind the A writes). + // The four lanes of one n-row read four consecutive 16-byte chunks of one + // contiguous packed row (coalesced 64 B per row) and write them to one + // contiguous 64 B LDS row; the whole (n-tile, K-stage) is one contiguous + // 8 KiB global stream. + const int b_n = tid >> 2; // n row within the tile, 0..127 (512 thr) + const int b_kv = (tid & 3) * 16; // stage-local k offset of the B int4 + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Iteration 21: all three staging int4s are loaded into locals FIRST + // (A rows 0-127, A rows 128-255, B), then all three are published to LDS. + // The b_n < kBlockN guard (iteration-8 leftover, provably always true at + // kBlockN = 128 with 512 threads: b_n = tid >> 2 is in [0, 127]) is + // removed so the B global load issues together with the two A loads at + // the stage top instead of behind the A ds_write_b128s (the accepted + // object's ISA shows the guarded B load starts only after both A writes + // issued, adding a second serialized L2 round trip to the collective + // stage-top critical path). Same global addresses, same LDS addresses, + // same bytes -> bit-identical int32. + const int4 stage_a0 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); + const int4 stage_a1 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + k0 + + stage_col); + const int4 stage_b = *reinterpret_cast( + weight + (static_cast(n0 + b_n) * k + k0 + b_kv)); + *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = + stage_a0; + *reinterpret_cast(a_tile + + (stage_row + kBlockM / 2) * kAStride + + stage_col) = stage_a1; + *reinterpret_cast(b_tile + b_n * kBStride + b_kv) = stage_b; + __syncthreads(); + + // Each wave consumes its 64x64 quadrant: sixteen m16n16k32 MMACs per kk + // (iteration 16: N-tile 64 -> 128; the four B fragments of the wave's 64 + // columns are each reused in registers across all four M-fragments, so + // per-output LDS fragment reads drop 3 -> 2 B). +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + // Iteration 2: B is n-major in LDS, so each lane's eight fragment + // k-values are contiguous; one 8-byte load fills the fragment (the + // same 8 bytes the old row_major loader gathered with eight ds_read_u8 + // plus mask/OR pack VALU), keeping the v_mmac operand bit pattern and + // the int32 accumulation identical. + load_fragment8(b_frag0, b_tile + local_col * kBStride + kk, kBStride, + lane); + load_fragment8(b_frag1, b_tile + (local_col + kTileN) * kBStride + kk, + kBStride, lane); + load_fragment8(b_frag2, b_tile + (local_col + 2 * kTileN) * kBStride + + kk, + kBStride, lane); + load_fragment8(b_frag3, b_tile + (local_col + 3 * kTileN) * kBStride + + kk, + kBStride, lane); + // A gets the same direct 8-byte fill: A is m-major in LDS, so lane + // (r, c4) owns the eight CONSECUTIVE k-values + // a_tile[(local_row + r)*80 + kk + c4*8 .. +7]; the code object shows + // du_load_matrix_sync already vectorized these into ds_read2_b32 but + // still executed the per-dword mask/OR reconstruction before the + // v_mmac A operands. Writing the same 8 bytes straight into the + // fragment storage removes that dead chain; the operand bit pattern + // and the int32 accumulation are unchanged. + load_fragment8(a_frag0, a_tile + local_row * kAStride + kk, kAStride, + lane); + load_fragment8(a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, + kAStride, lane); + // Iteration 17: the rows-32-47 and rows-48-63 A fragments are loaded + // here in the FIRST read group (dedicated a_frag2/a_frag3) instead of + // just-in-time before the acc20/acc30 groups, so their LDS latency is + // covered by the first 16 v_mmac of the kk (the accepted 8-wave ISA + // exposed a standalone lgkmcnt wait immediately before acc20; the + // flat iter-7 probe at the 16-wave geometry does not transfer to 2 + // waves/SIMD where lds_wait is 4.75x lds_instructions). Same 8 LDS + // bytes per fragment at the same addresses -> bit-identical int32. + load_fragment8( + a_frag2, a_tile + (local_row + 2 * kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8( + a_frag3, a_tile + (local_row + 3 * kTileM) * kAStride + kk, + kAStride, lane); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + du_mma_sync(acc20, a_frag2, b_frag0, acc20); + du_mma_sync(acc21, a_frag2, b_frag1, acc21); + du_mma_sync(acc22, a_frag2, b_frag2, acc22); + du_mma_sync(acc23, a_frag2, b_frag3, acc23); + du_mma_sync(acc30, a_frag3, b_frag0, acc30); + du_mma_sync(acc31, a_frag3, b_frag1, acc31); + du_mma_sync(acc32, a_frag3, b_frag2, acc32); + du_mma_sync(acc33, a_frag3, b_frag3, acc33); + } + + // Protect the LDS buffer from the next stage's cooperative overwrite. + // The FINAL stage has no successor stage and the epilogue never touches + // LDS (it reads only accumulator registers and global scales), so this + // guard is dead work on the last iteration: skip it there (iteration 6; + // 32 -> 31 barriers per block). The guard is wavefront-uniform - k0 and + // k are scalars - so all 256 threads take the same path. + if (k0 + kStageK < k) { + __syncthreads(); + } + } + + const int base_row = m0 + local_row; + const int base_col = n0 + local_col; + // Iteration 16: sixteen fragments per wave (4 M x 4 N of the 64x64 + // quadrant); same coalesced 8-byte store per lane per fragment. + store_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, + base_row, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, + base_row, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced(acc02, x_scale, weight_scale, out, + base_row, base_col + 2 * kTileN, m, n, + lane); + store_prefill_fragment_coalesced(acc03, x_scale, weight_scale, out, + base_row, base_col + 3 * kTileN, m, n, + lane); + store_prefill_fragment_coalesced(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, + lane); + store_prefill_fragment_coalesced(acc12, x_scale, weight_scale, out, + base_row + kTileM, base_col + 2 * kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc13, x_scale, weight_scale, out, + base_row + kTileM, base_col + 3 * kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc22, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + 2 * kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc23, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + 3 * kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc31, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc32, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + 2 * kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc33, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + 3 * kTileN, + m, n, lane); +} + +// --------------------------------------------------------------------------- +// Generic scalar fallback: one thread per output element. Exact int32 dot +// over K, then fused x_scale * weight_scale, then bf16 store. Handles every +// unmatched (m, n, k), including all small-M API cases (M=2, M=16) and any +// M in (0, 4096] with the same (K, N). For (k, n) == (1024, 6144) it decodes +// the packed n-major [N, K] layout (raw[kk, col] == packed[col*K + kk]); +// every other (k, n) reads the identity [K, N] layout. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = static_cast(m) * n; + if (linear >= total) { + return; + } + const int row = static_cast(linear / n); + const int col = static_cast(linear - static_cast(row) * n); + int32_t acc = 0; + const int8_t* a_row = x_q + static_cast(row) * k; + if (k == kTargetK && n == kTargetN) { + // Iteration 2: the packed layout for (k, n) == (1024, 6144) is n-major + // [N, K] (packed[col*K + kk] == raw[kk*N + col]). + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast( + weight[static_cast(col) * k + kk]); + } + } else { + const int8_t* b_col = weight + col; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +// --------------------------------------------------------------------------- +// Weight packing (outside the timed region and outside Graph capture; the +// packed buffer keeps the same byte count K*N and the same allocated +// address, so the layout is graph-stable). For the target (k, n) == +// (1024, 6144) the raw [K, N] row-major weight is transposed into the +// n-major [N, K] layout: +// packed[n*K + kk] == raw[kk*N + n] +// so every DUMMA B tile (one 64-column n-tile x one 64-K stage) is a single +// contiguous 4 KiB stream of 16-byte vector loads and each lane's eight +// B-fragment k-values are contiguous in LDS (one 8-byte fragment read). +// The transpose's 16-byte destination vectors run along K, so each thread +// GATHERS its 16 bytes from 16 raw rows (aligned byte loads strided by N) +// instead of copying a contiguous raw vector (a memcpy-style vectorized +// pack would write the wrong axis and leave most of the buffer unwritten - +// the measured correctness failure on the worker-29 TP8 o_proj lineage, +// iteration 10 -> 11 repair). The pack is one-time and out-of-timed-region, +// so the gather cost is free. For every other (k, n) the pack is an +// identity device-to-device copy. Scales are copied identity in both cases. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_o_proj_panels_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int k, + int n) { + // One thread per 16-byte vector of the packed [N, K] layout; the vector is + // a GATHER, not a memcpy. The destination vector + // packed[n*K + kk0 .. kk0+15] holds 16 CONSECUTIVE k-values of raw column n + // (packed[n*K + kk0 + j] == raw[(kk0 + j)*N + n], j = 0..15), so its 16 + // source bytes sit in 16 DIFFERENT raw rows (strided by N). + const int kk_vectors = k / 16; // 16-k vectors per packed row (k % 16 == 0) + const int64_t total4 = static_cast(n) * kk_vectors; + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear >= total4) { + return; + } + const int kk0 = static_cast(linear % kk_vectors) * 16; // multiple of 16 + const int nn = static_cast(linear / kk_vectors); // packed row = raw col + uint32_t d[4] = {0, 0, 0, 0}; +#pragma unroll + for (int j = 0; j < 16; ++j) { + const uint8_t b = raw[static_cast(kk0 + j) * n + nn]; + d[j >> 2] |= static_cast(b) << (8 * (j & 3)); + } + int4 v; + v.x = static_cast(d[0]); + v.y = static_cast(d[1]); + v.z = static_cast(d[2]); + v.w = static_cast(d[3]); + // Destination offset nn*K + kk0 is 16-byte aligned (K % 16 == 0 and + // kk0 % 16 == 0), so the int4 store is aligned. + *reinterpret_cast(packed + static_cast(nn) * k + kk0) = v; +} + +__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( + const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, + int64_t count) { + const int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +__global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( + const float* __restrict__ raw, + float* __restrict__ packed, + int count) { + const int linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear < count) { + packed[linear] = raw[linear]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Stable host launch symbols consumed by csrc/bindings.cpp. +// Both launchers are pure dispatch: no allocation, no packing, no +// synchronization, no default-stream launch; they run on the caller-provided +// PyTorch stream and are CUDA-Graph safe. +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; // no split-K: the GEMM does not use the workspace + (void)workspace_bytes; + auto* out_bf16 = reinterpret_cast(out); + + // Explicit dispatch. The assigned shape (M=4096, N=6144, K=1024) takes the + // DUMMA 256x128 single-buffered path (iteration 8 geometry 256x64/512 + // threads/8 waves extended by iteration 16 to 256x128 with sixteen + // accumulators per wave; measured residency of the accepted 256x64 source + // is 1 block/CU x 8 waves - arch_vgpr 72 - and iteration 16 keeps 1 + // block/CU, pinned structurally by 16 int32 accumulators = 64 VGPR; grid + // (6144/128, 4096/256) = 768 blocks; the kernel symbol name is + // intentionally unchanged - an internal anonymous-namespace symbol matched + // by the harness profiler); every other (m, n, k) - including + // small-M API cases (M=2, M=16) - takes the scalar fallback, which decodes + // the packed n-major [N, K] layout for (k, n) == (1024, 6144) and the + // identity [K, N] layout otherwise. + if (m == kTargetM && n == kTargetN && k == kTargetK) { + const dim3 grid(kTargetN / kBlockN, kTargetM / kBlockM); + const dim3 block(kBlockThreads); + hipLaunchKernelGGL(w8a8_dumma_128x64x64_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else { + constexpr int kBlock = 256; + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + } + (void)hipGetLastError(); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kBlock = 256; + const dim3 block(kBlock); + if (k == kTargetK && n == kTargetN) { + // Iteration 2: n-major [N, K] transpose for the target (k, n) == + // (1024, 6144); one thread per 16-byte destination vector (gather). + const int64_t total4 = static_cast(k) * n / 16; + const dim3 pack_grid( + static_cast((total4 + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_o_proj_panels_kernel, + pack_grid, block, 0, stream, + raw_weight, packed_weight, k, n); + } else { + // Identity device-to-device copy for every other (k, n). The packed + // buffer keeps the same byte count K*N and the same allocated address, + // so the layout is graph-stable. + const int64_t weight_count = static_cast(k) * n; + const dim3 weight_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_identity_kernel, + weight_grid, block, 0, stream, + raw_weight, packed_weight, weight_count); + } + + const int64_t scale_count = n; + const dim3 scale_grid( + static_cast((scale_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, + scale_grid, block, 0, stream, + weight_scale, packed_weight_scale, n); + (void)hipGetLastError(); +} +// @@end +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/qkv_proj.hip new file mode 100644 index 00000000..3a300bca --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/qkv_proj.hip @@ -0,0 +1,941 @@ +// @@variant shape=minimax_tp8_qkv_proj_m4096 commit=76e2b8659429332829a877855ef1178692d453d5 added=2026-08-29 +// median_us=475.5 p90_us=476.5 speedup=86.22 baseline_us=4.1e+04 +// source=minimaxm3-dsh-tp8-m4096-1-b0482833 +// MetaInfer W8A8 INT8 GEMM bootstrap for gfx928 (K500SM_AI). +// +// Worker: worker_0 (physical GPU 0) +// Assigned shape: minimax_tp8_qkv_proj_m4096 (M=4096, N=1280, K=6144) +// +// Operator contract (fixed by the control plane): +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// +// Iteration 2 (2-D macro-tile DUMMA throughput baseline): the exact assigned +// shape routes to the templated packed-B 2-D macro-tile kernel +// w8a8_dumma_prefill_packedb_tiled_kernel, dispatched as the +// <64, 128> instantiation (64x128 output tile per block, 256 threads = +// 4 wavefronts, each wave owns a 32x64 quadrant = eight m16n16k32 int32 +// accumulator fragments). The mandated tile benchmark {64x64, 64x128, +// 128x64} is realized by the same template; the <64, 64> and <128, 64> +// instantiations stay compiled (dead-branch instantiation in +// launch_w8a8_gemm) so a later round can flip the exact-shape dispatch +// without source surgery. Per-tile resources (kStageK = 64, A stride 68 B, +// packed-B n-major stride 80 B, 4 waves/block in all three): +// tile LDS B/block acc frags arch VGPR est. blocks/CU grid +// 64x64 9,472 4 ~80 3-4 1280 +// 64x128 14,592 8 ~96 2 640 +// 128x64 13,824 8 ~96 2 640 +// (<64,128> is the exact-shape timing candidate; the control plane verifies +// VGPR/LDS/occupancy from the code object and measures median/P90/TOPS.) +// +// B is packed once (outside the timed region and Graph capture) to an +// n-major packed[n][k] layout for the exact (K=6144, N=1280); the packed +// kernel stages B into an n-major LDS tile (row stride 80 B = 20 words, five +// bank phases) so every col_major m16n16k32 B fragment load is the explicit +// 8-byte loader load_b_frag8 (one ds_read2_b64 per fragment instead of 8 +// ds_read_u8 + mask/OR reassembly) and every 16-byte staging vector commits +// with one ds_write_b128. A is staged row-major into a_tile[BM][68] (odd +// 17-word stride, five bank phases); global-load latency is fully exposed, +// exactly as in the accepted sibling TP8 qkv pipeline (co-resident blocks +// hide it). +// +// K is staged cooperatively in a SINGLE LDS buffer at 64-K granularity with +// TWO __syncthreads per stage; the final stage is peeled and its provably +// dead all-consumption barrier is dropped (191 barriers over K=6144). All +// fragment loads of both 32-K steps are hoisted before the first v_mmac so +// LDS latency concentrates in one wait and the MMAC stream has no lgkmcnt +// waits. The int32 accumulation order (k0-outer, kk-inner) and the +// element-to-slot fragment mapping are unchanged, so results are +// bit-identical to the scalar reference. +// +// The epilogue is the fused coalesced store: a register 4x4 shuffle +// transpose (one 8-byte store per lane per fragment, 100% store sector +// efficiency), with per-row (xs) and per-column (ws) scale registers hoisted +// once per lane after the final MMAC burst. +// +// Every unmatched (m, n, k) keeps the prior generic path: the 64x64 +// single-buffered 128-K staged DUMMA kernel for large-M compatible shapes +// (identity [K, N] layout), and the scalar int8/int32 fallback, which +// decodes the packed n-major layout when (k, n) == (6144, 1280) (including +// the paired small-M API shapes with the same (N, K)). +// +// launch_pack_w8a8_weight packs (K=6144, N=1280) to packed[n][k] once, +// outside the timed region and outside Graph capture; every other (k, n) +// keeps the identity copy. The packed buffer keeps the same byte count +// (k*n), so allocations and graph-stable addresses are unchanged. +// +// The timed operator (launch_w8a8_gemm) performs no allocation, compilation, +// autotuning, packing, host/device synchronization, or default-stream launch: +// it only dispatches kernels on the caller-provided HIP stream. +// +// Include order is fixed by the control plane: hip_runtime, hip_bfloat16, +// then du_mma (this DTK's du_mma.h is not self-contained otherwise). + +#include +#include +#include + +#include + +namespace { + +constexpr int kWaveSize = 64; +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; +// 4 wavefronts; must remain a multiple of the gfx928 wavefront size (64). +constexpr int kThreadsPerBlock = 256; + +// Generic 64x64 large-M path (unchanged bootstrap, identity [K, N] layout): +// single-buffered 128-K K stage with padded row strides. +constexpr int kBlockM = 64; +constexpr int kBlockN = 64; +constexpr int kStageK = 128; +// Padded LDS row strides (bytes) for the int8 fragment loads. With the +// natural strides (A 128, B 64) every du_load_matrix_sync ds_read is +// 8-way (A) / 16-way (B) bank-conflicted: A lane l reads (l&15)*ldm + +// 8*(l>>4) and the row term 32*(l&15) mod 64 collapses to {0,32} (8 lanes +// per bank); B lane l reads 8*(l>>4)*ldm + (l&15) and the k-chunk term +// 8*64 = 512 B is 0 mod 256 B, stacking all four (l>>4) k-groups on the +// same 16 banks. Padding A rows to 144 B (36 dwords) makes 36*(l&15) mod +// 64 take 16 distinct dword-banks; padding B rows to 72 B moves the four +// k-groups 16 banks apart (16*(l>>4)), leaving at most a 4-way residue on +// the byte-granular (l&15)+i diagonal. Pure layout remap: every staged +// value and every fragment element is unchanged, so the int32 accumulation +// order and results stay bit-identical. LDS grows 16,384 -> 18,432 B/block; +// residency drops 4 -> 3 blocks/CU (3 x 18,432 = 55,296 <= 65,536 B LDS, +// 3 x 56 x 256 = 43,008 <= 65,536 VGPR). +constexpr int kAStride = 144; +constexpr int kBStride = 72; + +// Packed-B 2-D macro-tile family (iteration 2, exact assigned shape): +// 64-K stage, A row stride 68 B (17 words, five bank phases), packed-B +// n-major row stride 80 B (20 words, five bank phases). +constexpr int kStageK128 = 64; +constexpr int kAStride128 = kStageK128 + 4; // 68 bytes per A row +constexpr int kPackedBStride = 80; // 64 data + 16 pad bytes (20 words) +constexpr int kPackedBK = 6144; // exact K of the packed assigned shape +constexpr int kPackedBN = 1280; // exact N of the packed assigned shape +// Default tile of the packed family used by the exact-shape dispatch. +constexpr int kBlockM128 = 64; +constexpr int kBlockN128 = 128; + +using namespace du::dumma; + +// Direct accumulator epilogue for gfx928 INT8 m16n16k32 (verified against +// du_store_matrix_sync): row = lane & 15, col_mod4 = lane >> 4, +// frag.x[i] maps to columns col_mod4 + 4*i. Stores +// bf16(float(dot) * x_scale[row] * weight_scale[col]); out-of-range rows are +// masked (tail-M handling). +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Iteration-2 register-scale coalesced store used by the packed-B macro-tile +// kernel. The m16n16k32 accumulator lane mapping (row = lane & 15, column +// group c4 = lane >> 4, frag.x[i] -> column c4 + 4*i) gives each lane four +// elements strided by 4 columns; a register 4x4 transpose (two 2x2 steps +// with shfl_xor 16 then 32) re-routes the four int32 values so lane (r, c4) +// owns the four CONTIGUOUS columns 4*c4 .. 4*c4+3, then converts them to +// bf16 and writes ONE 8-byte store per lane (64 lanes x 8 B = 512 B per +// fragment per wavefront in 16 fully-used 32-B sectors; vmem_write drops +// 4x). Only the int32 values are re-routed -- the per-element multiply order +// (float(dot) * xs * ws.i) and the bf16 rounding are unchanged, so the +// stored bits are identical to store_prefill_fragment. The row>=m guard is +// wavefront-uniform (lane & 15 cycles the same 16 rows in every 16-lane +// group), so the shuffles never mix active and inactive lanes; base_col is a +// multiple of 64 and n*2 a multiple of 8, so the float4 weight_scale load +// (col0 % 4 == 0) and the 8-byte store are aligned. The per-row (xs) and +// per-column (ws) scales are loaded ONCE per lane after the final MMAC burst +// by the caller and passed in (2 scalar + 4 float4 loads instead of the +// per-fragment reloads). +template +__device__ __forceinline__ void store_prefill_fragment_coalesced_rs( + const AccFragment& frag, + float xs, + const float4 ws, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 64, n*2 is a multiple of 8). + const int col0 = base_col + 4 * c4; + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + row * n + col0) = packed; +} + +// Explicit 8-byte loader for the col_major m16n16k32 B fragment. The +// col_major fragment mapping is n = lane&15, k = 8*(lane>>4) + i; with the +// n-major LDS tile at row stride 80 the lane's 8 elements are contiguous +// (p[row*ldm + col .. +7], 8-byte aligned because ldm=80 and col are +// multiples of 8), so this compiles to ONE ds_read2_b64 instead of 8 +// ds_read_u8 + mask/OR reassembly. The byte placement is identical to +// du_load_matrix_sync, so the v_mmac_i32_16x16x32_i8 fragment +// registers receive the same values. +__device__ __forceinline__ void load_b_frag8( + DUFragment& f, + const int8_t* __restrict__ p, + int ldm) { + const unsigned row = static_cast(__lane_id()) & 0xfu; + const unsigned col = (static_cast(__lane_id()) >> 4) << 3; + const int64_t v = *reinterpret_cast(p + row * ldm + col); + *reinterpret_cast(&f.x[0]) = v; +} + +// Iteration-2 throughput baseline: templated 2-D macro-tile packed-B DUMMA +// prefill kernel for the exact assigned shape (M>=128, N=1280, K=6144). The +// mandated tile set {64x64, 64x128, 128x64} is instantiated from this one +// template; <64,128> is the exact-shape dispatch. 256 threads = 4 +// wavefronts; each wave owns a (BM/2) x (BN/2) quadrant of m16n16k32 int32 +// accumulator fragments (kWaveM16 x kWaveN16 per wave). K is staged +// cooperatively in a SINGLE LDS buffer at 64-K granularity with TWO +// __syncthreads per stage; the final stage is peeled and its provably dead +// all-consumption barrier dropped: +// * A is staged row-major into a_tile[BM][68] (odd 17-word stride, five +// bank phases); global-load latency is fully exposed (co-resident +// blocks hide it). +// * B is staged n-major from packed[n][k] into b_tile[BN][80] (20 words, +// five bank phases): each 16-byte staging vector commits with one +// ds_write_b128 (b_n*80 + b_k16 is 0 mod 16). +// * The 64-K stage is split into two explicit 32-K steps and ALL fragment +// loads (kWaveM16 A + kWaveN16 B per step, two steps) are hoisted before +// the first v_mmac of the stage, so LDS latency concentrates in one wait +// and the MMAC stream has no lgkmcnt waits. +// * The int32 accumulation order (k0-outer over 64-K stages, kk-inner +// kk=0 then kk=32 within a stage) and the element-to-slot fragment +// mapping are unchanged, so the results are bit-identical to the scalar +// reference. +// * The epilogue is the fused coalesced store +// (store_prefill_fragment_coalesced_rs): per-row and per-column scales +// are hoisted into registers loaded once per lane after the peeled final +// stage's MMAC burst; one 8-byte store per lane per fragment. +template +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_packedb_tiled_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + // m16n16k32 accumulator tiles per wave quadrant: rows BM/32, cols BN/32. + constexpr int kWaveM16 = BM / 32; + constexpr int kWaveN16 = BN / 32; + // 16-byte staging vectors per thread (1 for 64-row/64-col sides, 2 for + // 128-row/128-col sides). + constexpr int kAVecsPerThread = + (BM * kStageK128 / static_cast(sizeof(int4))) / kThreadsPerBlock; + constexpr int kBVecsPerThread = + (BN * kStageK128 / static_cast(sizeof(int4))) / kThreadsPerBlock; + static_assert(BM % 32 == 0 && BN % 32 == 0, + "block tile must be a multiple of the wave quadrant"); + static_assert(kAVecsPerThread >= 1 && kBVecsPerThread >= 1, + "staging must fit the block thread count"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * BM; + const int n0 = static_cast(blockIdx.x) * BN; + + __shared__ __align__(16) int8_t a_tile[BM * kAStride128]; + __shared__ __align__(16) int8_t b_tile[BN * kPackedBStride]; + + DUFragment + a_frag[kWaveM16][2]; + DUFragment + b_frag[kWaveN16][2]; + DUFragment + acc[kWaveM16][kWaveN16]; +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_fill_fragment(acc[j][i], 0); + } + } + + const int num_stages = k / kStageK128; + + // Cooperative staging mapping (fixed per thread, reused every stage): + // A[BM,64]: vector v = tid + t*256 owns row v>>2 and 16-B column group + // (v&3)*16; each wavefront covers 16 rows x 64 contiguous B. + // B[BN,64]: n-major; vector v = tid + t*256 owns packed row (n0 + v>>2) + // and the 16-byte k-run (v&3)*16. + int4 vA[kAVecsPerThread]; + int4 vB[kBVecsPerThread]; + + // Prologue: load stage 0 into registers, commit it to the single buffers, + // and make it visible before the first burst. + { +#pragma unroll + for (int t = 0; t < kAVecsPerThread; ++t) { + const int v = tid + t * kThreadsPerBlock; + const int a_row = v >> 2; + const int a_k16 = (v & 3) << 4; + const int g_row = m0 + a_row; + vA[t] = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_k16) + : int4{0, 0, 0, 0}; + // Commit A as int32 stores (68-byte rows are 4 mod 16 -> no b128). + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_k16); + adst[0] = vA[t].x; + adst[1] = vA[t].y; + adst[2] = vA[t].z; + adst[3] = vA[t].w; + } +#pragma unroll + for (int t = 0; t < kBVecsPerThread; ++t) { + const int v = tid + t * kThreadsPerBlock; + const int b_n = v >> 2; + const int b_k16 = (v & 3) << 4; + vB[t] = *reinterpret_cast( + packed_b + (n0 + b_n) * k + b_k16); + // Commit B as 16-byte vector stores: b_n*80 + b_k16 is 0 mod 16. + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = + vB[t]; + } + __syncthreads(); + } + + // Stage consumption shared by the loop body and the peeled final stage: + // hoist all fragment loads of both 32-K steps before the first v_mmac, + // then run the MMAC bursts. + auto consume_stage = [&]() { + const int local_row = wave_row * (BM / 2); + const int local_col = wave_col * (BN / 2); + const int8_t* abase = a_tile + local_row * kAStride128; + const int8_t* bbase = b_tile + local_col * kPackedBStride; +#pragma unroll + for (int st = 0; st < 2; ++st) { + const int kk = st * kTileK; +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { + du_load_matrix_sync( + a_frag[j][st], abase + j * kTileM * kAStride128 + kk, + kAStride128); + } +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + load_b_frag8( + b_frag[i][st], + bbase + i * kTileN * kPackedBStride + kk, kPackedBStride); + } + } + // kk = 0 then kk = 32 bursts (same accumulators, same MMAC order as the + // scalar k-ascending reference). +#pragma unroll + for (int st = 0; st < 2; ++st) { +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_mma_sync(acc[j][i], a_frag[j][st], b_frag[i][st], acc[j][i]); + } + } + } + }; + + for (int s = 0; s < num_stages - 1; ++s) { + consume_stage(); + // All consumption of the single buffer is complete only after every + // wavefront passes this barrier; the stage s+1 stores below overwrite + // that same buffer, so they must wait for it. + __syncthreads(); + + // Load stage s+1 and commit it into the single buffer. The global-load + // latency is fully exposed here (co-resident blocks hide it); the + // compiler's vmcnt wait before the DS stores is on this path. + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; +#pragma unroll + for (int t = 0; t < kAVecsPerThread; ++t) { + const int v = tid + t * kThreadsPerBlock; + const int a_row = v >> 2; + const int a_k16 = (v & 3) << 4; + const int g_row = m0 + a_row; + vA[t] = (g_row < m) + ? *reinterpret_cast( + x_q + g_row * k + s1 + a_k16) + : int4{0, 0, 0, 0}; + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_k16); + adst[0] = vA[t].x; + adst[1] = vA[t].y; + adst[2] = vA[t].z; + adst[3] = vA[t].w; + } +#pragma unroll + for (int t = 0; t < kBVecsPerThread; ++t) { + const int v = tid + t * kThreadsPerBlock; + const int b_n = v >> 2; + const int b_k16 = (v & 3) << 4; + vB[t] = *reinterpret_cast( + packed_b + (n0 + b_n) * k + s1 + b_k16); + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = + vB[t]; + } + // Make the stage s+1 stores visible to every wavefront before the + // next burst. + __syncthreads(); + } + } + + // Peel the final stage. After its MMAC burst no wavefront reads LDS again + // -- the epilogue touches only registers and global memory -- so the final + // all-consumption barrier is provably dead and is dropped. The per-row + // (xs) and per-column (ws) scale registers are prefetched right after the + // last MMACs: 2 scalar + 4 float4 loads per lane instead of the per- + // fragment reloads, with the load latency hidden behind the + // transpose/shuffle chains. The per-element multiply order and bf16 + // rounding are unchanged, so the stored bits are identical. + consume_stage(); + + const int base_row = m0 + wave_row * (BM / 2); + const int base_col = n0 + wave_col * (BN / 2); + const int lane_row = lane & 15; + const int c4 = lane >> 4; + float xs[kWaveM16]; +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { + const int row = base_row + j * kTileM + lane_row; + // Tail rows are clamped to 0.0f (never dereferenced); the store guard + // below skips them, and the guard is uniform across each 4-lane row + // group so the shuffle transpose never mixes active and inactive lanes. + xs[j] = (row < m) ? x_scale[row] : 0.0f; + } + float4 ws[kWaveN16]; +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + // col0 = base_col + i*kTileN + 4*c4 is a multiple of 4 (base_col is a + // multiple of 64), so the float4 load is 16-byte aligned. + ws[i] = *reinterpret_cast( + weight_scale + base_col + i * kTileN + 4 * c4); + } +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + store_prefill_fragment_coalesced_rs( + acc[j][i], xs[j], ws[i], out, m, n, + base_row + j * kTileM, base_col + i * kTileN, lane); + } + } +} + +// Large-M prefill kernel: 64x64 output tile per block, K staged in LDS +// (single buffer at 128-K granularity), four waves each owning a 32x32 +// quadrant = four m16n16k32 int8->int32 DUMMA accumulators. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[64,128] (row-major, zero-filled past M) into the padded + // 144-B-row layout (row stride 144, still 16-B vector stores: 144/16 + // = 9 int4 slots per row). + constexpr int kAVectors = kBlockM * kStageK / static_cast(sizeof(int4)); +#pragma unroll(2) + for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int local_row = byte_offset / kStageK; + const int kk = byte_offset - local_row * kStageK; + const int global_row = m0 + local_row; + const int4 v = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + kk) + : int4{0, 0, 0, 0}; + reinterpret_cast(a_tile)[local_row * (kAStride / 16) + kk / 16] = + v; + } + + // Stage B[128,64] (row-major [K, N] logical weight layout) into the + // padded 72-B-row layout. 72 is not a multiple of 16, so each 16-B + // global vector is stored as two 8-B int2 halves (72*kk + col is + // always 8-B aligned). + constexpr int kBVectors = kStageK * kBlockN / static_cast(sizeof(int4)); +#pragma unroll(2) + for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + const int4 v = *reinterpret_cast( + weight + (k0 + kk) * n + n0 + col); + int2* dst = reinterpret_cast(b_tile + kk * kBStride + col); + dst[0] = *reinterpret_cast(&v); + dst[1] = *reinterpret_cast( + reinterpret_cast(&v) + 8); + } + __syncthreads(); + + // Consume the stage: each wave does four m16n16k32 MMAs per 32-K step. + // Four 32-K tiles per stage: pin the full unroll so the int32 + // accumulation order (k0-outer, kk-inner) is deterministic codegen, not + // compiler-heuristic dependent. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll(4) + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + // Single buffer: make sure every wave finished reading LDS before the + // next stage overwrites it. + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Generic scalar int8/int32 fallback for unmatched shapes and small-M API +// cases (including the paired M=2 shape). One output element per grid-stride +// step; exact int32 accumulation (the assigned K keeps the int8 dot well +// within int32 range), then the fused float scale and bf16 store. The exact +// packed shape (K=6144, N=1280) uses the packed n-major layout packed[n][k] +// (set up once by launch_pack_w8a8_weight); every other (k, n) keeps the +// logical [K, N] row-major layout. The branch is grid-uniform per launch. +__global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const bool packed_b = (k == kPackedBK && n == kPackedBN); + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; + idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = a + static_cast(row) * k; + const int8_t* b_col = + packed_b ? b + static_cast(col) * k : b + col; + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + const int32_t bw = packed_b + ? static_cast(b_col[kk]) + : static_cast( + b_col[static_cast(kk) * n]); + acc += static_cast(a_row[kk]) * bw; + } + const float scaled = static_cast(acc) * + x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// Identity device-to-device byte copy (pack_weight bootstrap, valid for any +// (K, N)). +__global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +// Identity device-to-device float copy for the weight scales. +__global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +// One-time n-major transpose of the logical [K, N] row-major weight into +// packed[n][k] for the exact assigned shape (K=6144, N=1280). Runs only from +// launch_pack_w8a8_weight, outside the timed region and outside Graph +// capture. The packed buffer keeps the same byte count (k*n) as the identity +// pack, so allocations and graph-stable addresses are unchanged. Element +// (k, n) of the logical weight lands at packed[n * K + k]; each thread +// copies one 16-byte k-run (coalesced read side; the strided write side is +// off the critical path). The guard guarantees n % 16 == 0, so every 16-byte +// chunk lies inside one logical row and the transpose is exact. +__global__ __launch_bounds__(256) void w8a8_pack_nmajor_b_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int k, + int n) { + const int64_t chunks = (static_cast(k) * n) >> 4; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t c = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + c < chunks; + c += stride) { + const int64_t off = c << 4; + const int kk = static_cast(off / n); + const int col = static_cast(off - static_cast(kk) * n); + const int4 v = *reinterpret_cast(src + off); + const int8_t* bytes = reinterpret_cast(&v); +#pragma unroll + for (int i = 0; i < 16; ++i) { + // Logical element (k = kk, n = col + i) lands at packed[(col+i)*K+kk]. + dst[static_cast(col + i) * k + kk] = bytes[i]; + } + } +} + +} // namespace + +// Stable host launch symbol consumed by csrc/bindings.cpp. Dispatches on the +// caller-provided stream only; never allocates, synchronizes, or touches the +// default stream. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + auto* out_ptr = static_cast(out); + + // Iteration-2 throughput baseline: the exact assigned shape (M>=128, + // N=1280, K=6144) routes to the packed n-major B variant of the 2-D + // macro-tile kernel, dispatched as the <64,128> instantiation (64x128 + // tile, 4 wavefronts, 8 accumulator fragments per wave, grid + // (N/128)x(M/64) = (10, 64) = 640 blocks). The guard is exact + // (m >= 128, n == 1280, k == 6144) and sits BEFORE the generic large-M + // paths; the packed layout is only valid for this (k, n). + if (m >= 128 && n == kPackedBN && k == kPackedBK) { + const dim3 grid( + static_cast(n / kBlockN128), + static_cast((m + kBlockM128 - 1) / kBlockM128)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_tiled_kernel<64, 128>), + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Mandated 2-D macro-tile benchmark family: keep the <64,64> and <128,64> + // instantiations of the packed kernel compiled (and correct) so later + // rounds can flip the exact-shape dispatch without source surgery. This + // branch can never run (m <= 0 already returned above); it only forces + // template instantiation of the two sibling tiles. + if (m < 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_tiled_kernel<64, 64>), + dim3(1), + dim3(kThreadsPerBlock), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_tiled_kernel<128, 64>), + dim3(1), + dim3(kThreadsPerBlock), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Native INT8 DUMMA m16n16k32 prefill path (2-D macro-tile 64x64, + // single-buffered 128-K K stage, identity [K, N] weight layout) for + // large-M shapes with compatible geometry that are NOT the packed exact + // shape (which returned above), e.g. (4096, 1536, 6144): + // 1536 % 64 == 0, 6144 % 128 == 0. The M >= 128 guard keeps the paired + // small-M API shapes (M=2 / M=16) on the scalar fallback. + if (m >= 128 && (n % kBlockN) == 0 && (k % kStageK) == 0) { + const dim3 grid( + static_cast(n / kBlockN), + static_cast((m + kBlockM - 1) / kBlockM)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x64_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Scalar int8/int32 fallback for unmatched (m, n, k) and small-M API cases. + const int64_t total = static_cast(m) * n; + constexpr int kFallbackThreads = 256; + int64_t blocks = (total + kFallbackThreads - 1) / kFallbackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + const dim3 grid(static_cast(blocks)); + const dim3 block(kFallbackThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); +} + +// Stable host launch symbol consumed by csrc/bindings.cpp. For the exact +// assigned shape (K=6144, N=1280) it performs the one-time n-major B pack +// (logical [K, N] -> packed[n][k]) outside the timed region and outside +// Graph capture; every other (K, N) keeps the identity device-to-device +// copy, valid for every (K, N). The packed buffer size is k*n in both cases, +// so the allocation and graph-stable packed layout are unchanged. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kPackThreads = 256; + const int64_t weight_count = static_cast(k) * n; + if (weight_count > 0) { + if (k == kPackedBK && n == kPackedBN) { + const int64_t chunks = weight_count >> 4; + int64_t blocks = (chunks + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + hipLaunchKernelGGL( + w8a8_pack_nmajor_b_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + int64_t blocks = (weight_count + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + hipLaunchKernelGGL( + w8a8_pack_identity_i8_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count); + } + } + if (n > 0) { + int64_t blocks = (static_cast(n) + kPackThreads - 1) / + kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + hipLaunchKernelGGL( + w8a8_pack_identity_f32_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/qkv_proj_and_indexer_qk.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/qkv_proj_and_indexer_qk.hip new file mode 100644 index 00000000..d1e3660b --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/qkv_proj_and_indexer_qk.hip @@ -0,0 +1,1712 @@ +// @@variant shape=minimax_tp8_qkv_proj_and_indexer_qk_m4096 commit=0f58b3835dbdecf6de8ce1469373ca854898caf2 added=2026-08-29 +// median_us=531.2 p90_us=533.2 speedup=67.02 baseline_us=3.56e+04 +// source=minimaxm3-dsh-tp8-m4096-1-b0482833 +// MetaInfer W8A8 INT8 GEMM for gfx928 (K500SM_AI). +// +// Worker: worker_1 (physical GPU 1) +// Assigned shape: minimax_tp8_qkv_proj_and_indexer_qk_m4096 +// (M=4096, N=1536, K=6144) +// +// Operator contract (fixed by the control plane): +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// * x_scale[m] * weight_scale[n] ) +// +// Iteration 1 (DUMMA 2-D macro-tile baseline sweep, mandatory decision): +// * The bootstrap DUMMA 64x128 baseline (this file, prior state) measured +// 816.66 us median / 94.7 TOPS on the assigned shape (M=4096, N=1536, +// K=6144) with PMC: 21.82M LDS instructions, 7.08M LDS bank conflicts, +// 37.13M LDS waits (1.70x LDS instructions, dominant stall class), 1.00M +// VMEM reads, 91.2% L2 hit, 108.4 GB/s HBM, 3 blocks/CU (arch_vgpr 80). +// Bottleneck hypothesis: the row-major [K, N] B operand is staged into +// LDS but its m16n16k32 B fragments are re-assembled by the generic +// loader as EIGHT byte-granular ds_read_u8 per fragment (4-way bank +// aliasing with the 132-byte stride), which inflates the LDS instruction +// stream and its waits; the B staging also commits as 4 x ds_write_b32 +// (odd 132-byte rows), and the epilogue issues four 2-byte scalar stores +// per lane per fragment at 25% store-sector efficiency. The trusted +// same-operator TP4 lineage (minimax_tp4_qkv_proj_and_indexer_qk, M=4096, +// SAME K=6144; accepted 1082 us / 119 TOPS) fixed exactly this by packing +// the weight ONCE outside the timed region and loading B fragments +// col_major as one conflict-free 8-byte ds_read_b64 per fragment +// (load_b_frag8_swz), hoisting all twelve fragment loads of a 64-K stage +// before the first v_mmac (no lgkmcnt waits inside the 16-MMAC burst), +// staging B with one ds_write_b128 per 16-byte vector, and storing the +// epilogue as one coalesced 8-byte store per lane per fragment (100% +// store sector efficiency). This round ports that accepted lineage to the +// exact assigned TP8 shape (K=6144, N=1536) and installs the complete +// 2-D macro-tile family (64x64, 64x128, 128x64; each 4 wavefronts x 64 +// lanes = 256 threads, one 32x32/32x64/64x32 quadrant per wave) so the +// mandated tile sweep can flip the exact-shape dispatch in the next +// rounds. Round 1 dispatches the exact shape to the 64x128 tile (the +// strongest evidence for this operator: the accepted same-operator TP4 +// lineage and the accepted TP8 hy3 qkv lineage both use 64x128). +// * 64x128 packed-B kernel (dispatched): A staged row-major into 68-byte +// LDS rows (17 words, odd-word bank spread, 4 ds_write_b32 per 16-byte +// vector), B staged from the swizzled pack into an unpadded 8192-B LDS +// tile (one ds_write_b128 per 16-byte vector), B fragments col_major via +// load_b_frag8_swz (one conflict-free ds_read_b64 per fragment), +// single-buffered 64-K K stage with two barriers per stage, all twelve +// fragment loads hoisted before the 16-MMAC burst, fused coalesced +// scale/bf16 epilogue (two 4-fragment row-group calls, one 8-byte store +// per lane per fragment). LDS/block = 64x68 + 64x128 = 12,544 B -> 3 +// resident blocks/CU (VGPR-bound at arch_vgpr <= 85, per the accepted +// lineage). Grid (N/128)x(M/64) = 12x64 = 768 blocks >> 120 CUs, no +// split-K, workspace unused. +// * 64x64 and 128x64 packed-B kernels are compiled family members for the +// mandated tile sweep (grid 24x64 = 1536 and 24x32 = 768 blocks; LDS +// 9,472 B and 15,360 B); they still consume the n-major 80-byte-row tile +// and are NOT dispatched for the exact shape, which uses 64x128 this +// round. +// * Iteration 7 (compute-pipeline round): the dispatched 64x128 kernel's +// per-stage consume is re-grouped into two independent kk halves - +// kk=0 loads -> one wait -> contiguous 8-v_mmac burst -> kk=32 loads +// issued BETWEEN the bursts (LDS latency overlaps the kk=0 MMAs) -> +// contiguous 8-v_mmac burst - replacing the baseline "all twelve +// fragment loads hoisted, seven lgkmcnt waits distributed through the +// 16-MMAC stream" (accepted ISA 0x6B7C..0x6C78). Same loads, same +// operands, same per-accumulator kk=0-then-kk=32 int32 order, same +// single-buffered 2-barrier pipeline, LDS and occupancy -> bit-identical +// output. Rounds 4 (A-stride conflict elimination, 554.00 us regression) +// and 6 (B pair loads, 538.60 us neutral) are NOT accepted; this round +// starts from the accepted iteration-1 digest. +// * Iteration 9 (B-fragment pair round on the iteration-8 split): the +// DISPATCHED 64x128 kernel keeps the iteration-8 two-kk-half structure +// but consumes each kk half's four swizzled B fragments as TWO pair reads +// (load_b_frag8_swz2 -> one ds_read2_b64 per fragment pair instead of two +// ds_read_b64; fragment blocks j and j+1 are kPackedBGroupBytes (512 B) +// apart in the same 32-K region, inside the ds_read2_b64 8-bit-dword +// range, and both sub-reads keep the conflict-free pattern of +// load_b_frag8_swz), cutting the per-wave-stage B read-issue stream from +// 8 to 4 instructions while the split's wait grouping is unchanged. Byte +// placement is identical to two load_b_frag8_swz calls (v0 -> f0.x[0..7], +// v1 -> f1.x[0..7]), so operands, per-accumulator kk=0-then-kk=32 int32 +// order, LDS size (12,544 B), occupancy (3 blocks/CU at 79 VGPR) and +// output bits are unchanged. Round 6 measured this mechanism +0.16% on +// the hoisted baseline (538.60 us); iteration 8 measured the split +// +0.35% (537.60 us). If the two mechanisms are independent (read-path +// instruction count AND wait grouping both on the LDS-issue/latency +// critical path), the combination stacks below 537.60/538.23 us; if the +// kernel is bound elsewhere (MMA throughput / co-residency), it is flat +// or regresses. +// * launch_pack_w8a8_weight packs exactly (k, n) == (6144, 1536) once, +// outside the timed region and Graph capture, into the iteration-5 +// fragment-interleaved swizzle (element (n, k) at +// ((k>>5)*96 + (n>>4))*512 + (((k>>3)&3)*16 + (n&15))*8 + (k&7); same +// byte count k*n, so allocations and graph-stable addresses are +// unchanged); every other (k, n) keeps the identity copy. +// * The scalar int8/int32 fallback decodes the iteration-5 swizzled layout +// when (k, n) == (6144, 1536) (so the paired small-M API shape with the +// same (N, K) stays correct) and the identity layout otherwise. +// * Generic 64x128/64x64 DUMMA kernels (identity-packed B) are preserved +// unchanged for every other large-M shape. +// +// The timed operator (launch_w8a8_gemm) performs no allocation, compilation, +// autotuning, packing, host/device synchronization, or default-stream +// launch: it only dispatches kernels on the caller-provided HIP stream and +// uses only the caller-provided out/workspace tensors (workspace is unused +// in this no-split-K kernel). +// +// Include order is fixed by the control plane: hip_runtime, hip_bfloat16, +// then du_mma (this DTK's du_mma.h is not self-contained otherwise). + +#include +#include +#include + +#include + +namespace { + +constexpr int kWaveSize = 64; +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; +// 4 wavefronts; must remain a multiple of the gfx928 wavefront size (64). +constexpr int kThreadsPerBlock = 256; + +// 64x64 generic kernel geometry: K stage of 128 (unpadded rows). +constexpr int kBlockM = 64; +constexpr int kBlockN = 64; +constexpr int kStageK = 128; + +// 64x128 large-prefill kernel geometry: single-buffered 64-K K stage with +// odd-word padded row strides (A: 68 B = 17 words, B: 132 B = 33 words). +constexpr int kBlockM128 = 64; +constexpr int kBlockN128 = 128; +constexpr int kStageK128 = 64; +constexpr int kAStride128 = kStageK128 + 4; // 68 bytes per A row (17 words) +constexpr int kBStride128 = kBlockN128 + 4; // 132 bytes per B row (33 words) + +// Iteration 1 (packed-B family): the exact assigned shape packs the weight +// once (outside timing) and stages B into an LDS tile. The compiled 64x64 +// and 128x64 family members (not dispatched) still use the round-1 n-major +// layout: packed[n][k] staged into an n-major LDS tile with a 16-byte-aligned +// 80-byte row stride (20 words, five bank phases); every col_major +// m16n16k32 B fragment then reads its lane's 8 elements as one contiguous +// 8-byte chunk (one ds_read_b64 per fragment). 128x64 uses the same 80-byte +// stride for the A rows (16-byte-aligned staging stores, five bank phases). +// The dispatched 64x128 kernel switched to the iteration-5 fragment- +// interleaved swizzle (see kPackedBGroupBytes below) and no longer uses +// kPackedBStride. The 64x128/64x64 packed kernels keep the 68-byte A stride +// (4 ds_write_b32 per staging vector). +constexpr int kPackedBStride = 80; // 64 data + 16 pad bytes (20 words) +constexpr int kPackedAStride128x64 = 80; // 64 data + 16 pad (128x64 A tile) +constexpr int kPackedBK = 6144; // exact K of the packed assigned shape +constexpr int kPackedBN = 1536; // exact N of the packed assigned shape + +// Iteration 5 (packing round): the exact-shape pack is swizzled from the +// plain n-major transpose to a fragment-interleaved layout so the dispatched +// 64x128 packed-B kernel's B tile is vector-loadable AND LDS-bank-safe end +// to end. For each 32-k x 16-n DUMMA B fragment block the pack stores the 4 +// x 8-byte k-groups of the 16 n-rows n-interleaved ([k8][n][8], 128 B per +// k8, 512 B per fragment block, k32-major inside the 64-K stage). Every +// 16-byte staging vector is then 2 n-rows x 8 B of one k8 (int4 global load +// + one ds_write_b128 whose 8-lane LDS cycles each cover all 32 banks once), +// and every col_major B fragment load is one ds_read_b64 whose 16-lane group +// spans exactly one 128-B k8 block (all 32 banks exactly once per cycle -> +// conflict-free; the n-major 80-byte-row tile read the same bytes with 2-way +// conflicts). Same byte count k*n -> graph-stable packed layout. +constexpr int kPackedBGroupBytes = 512; // 16 n x 4 k8 x 8 B per 32-k fragment block +constexpr int kPackedB32Bytes = 8 * kPackedBGroupBytes; // 4096 B per 32-k block of the 64-K stage +constexpr int kPackedBN16 = 96; // 1536 / 16 n-groups per 32-k block (global pack stride) + +using namespace du::dumma; + +// Direct accumulator epilogue for gfx928 INT8 m16n16k32 (verified against +// du_store_matrix_sync): row = lane & 15, col_mod4 = lane >> 4, +// frag.x[i] maps to columns col_mod4 + 4*i. Stores +// bf16(float(dot) * x_scale[row] * weight_scale[col]); out-of-range rows are +// masked (tail-M handling). The scale multiply order +// float(dot) * x_scale[row] * weight_scale[col] matches the reference +// (dot.float() * a_scale * b_scale.T), so bf16 bits are exact. +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Coalesced fragment store (packed-B kernels). The m16n16k32 accumulator +// lane mapping (row = lane & 15, column group c4 = lane >> 4, frag.x[i] -> +// column c4 + 4*i) gives each lane four elements strided by 4 columns, so a +// direct store is four 2-byte scalar stores per lane touching every 32-B +// sector at 25% utilization. This epilogue transposes the 4-element groups +// within each 4-lane column group (lanes r, r+16, r+32, r+48; two 2x2 +// shuffle steps) so lane (r, c4) owns the four CONTIGUOUS columns +// 4*c4 .. 4*c4+3, converts to bf16, packs 4 bf16 (8 B), and issues ONE +// 8-byte store per lane per fragment (100% store sector efficiency; the +// assigned shape's vmem_write_instructions drop 163,840 -> 40,960). Only the +// int32 values are re-routed between lanes; the per-element float scale +// multiply order (float(dot) * x_scale[row] * weight_scale[col]) and bf16 +// rounding are unchanged, so the stored bits are identical. The row >= m +// guard is wavefront-uniform (lane & 15 cycles the same 16 rows in every +// 16-lane group) and base_col is a multiple of 4 (multiple of 32 here), so +// the float4 weight_scale load (col0 % 4 == 0) and the 8-byte store (n even) +// are aligned. This DTK lowers __shfl_xor to ds_bpermute (LDS permute at the +// block tail, where the LDS pipe is otherwise idle); there is no global +// round trip and no staging tile. +template +__device__ __forceinline__ void store_prefill_fragment_coalesced( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte + // aligned: base_col is a multiple of 4, n is even). + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + row * n + col0) = packed; +} + +// Iteration 6 (epilogue round): 4-element 4x4 lane-group transpose for one +// m16n16k32 accumulator fragment. Identical masks (xor 16, xor 32), select +// structure and lane mapping as store_prefill_fragment_coalesced (step 1 +// swaps the 2x2 in-block pairs, step 2 swaps the 2x2 blocks), so the four +// returned ints are the SAME per-lane contiguous-column values; only the +// surrounding scheduling is restructured. Extracted so the fused +// store_prefill_rowgroup_coalesced can issue the transposes of a whole +// 16-row x 64-col group as one software-pipelined block. +template +__device__ __forceinline__ void transpose_frag4( + const AccFragment& frag, + int lane, + int& f0, + int& f1, + int& f2, + int& f3) { + const int c4 = lane >> 4; + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + f0 = hi ? a0 : t2; + f1 = hi ? a1 : t3; + f2 = hi ? t0 : a2; + f3 = hi ? t1 : a3; +} + +// Scale x 4 contiguous bf16 and pack into one 8-byte store word. The +// per-element multiply order float(dot) * x_scale[row] * weight_scale[col] +// (left to right) and the bf16 rounding are identical to +// store_prefill_fragment_coalesced, so the stored bits are unchanged. +__device__ __forceinline__ uint64_t pack_bf16x4( + int a0, + int a1, + int a2, + int a3, + float xs, + float4 ws) { + const float v0 = static_cast(a0) * xs * ws.x; + const float v1 = static_cast(a1) * xs * ws.y; + const float v2 = static_cast(a2) * xs * ws.z; + const float v3 = static_cast(a3) * xs * ws.w; + return static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); +} + +// Iteration 6 (epilogue round): fused 4-fragment epilogue for ONE 16-row +// group of a wave (the four fragments cover the group's 64 contiguous +// columns). This is the dispatched 64x128 packed-B kernel's epilogue, +// restructured from eight independent per-fragment calls into two +// row-group calls so the tail becomes one explicitly-scheduled block: +// (1) x_scale[row] and the FOUR float4 weight_scale groups are loaded into +// registers ONCE per lane BEFORE the first ds_bpermute (one L1 round trip +// for the whole group instead of relying on cross-call CSE/load sinking), +// overlapping the scale-load latency with the transpose chain; (2) the four +// independent 4x4 transposes are issued as one software-pipelined block +// (step-1 xor-16 ds_bpermutes of all fragments precede step-2 xor-32, so +// the tail's dependent LDS-latency chain is ~2 round trips, not up to 8 +// serialized per-fragment chains); (3) the four 8-byte coalesced stores +// (offsets 0/32/64/96 B from obase) are issued back-to-back at the end, +// keeping 100% store-sector efficiency (same 8-byte stores per lane per +// fragment; vmem_write_instructions unchanged at 40,960). The int32 +// transpose routing, per-element float multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and bf16 rounding are +// byte-identical to the iteration-5 epilogue, so output bits are exact. +template +__device__ __forceinline__ void store_prefill_rowgroup_coalesced( + const F0& f0, + const F1& f1, + const F2& f2, + const F3& f3, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int c4 = lane >> 4; + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 w0 = *reinterpret_cast(weight_scale + col0); + const float4 w1 = *reinterpret_cast(weight_scale + col0 + 16); + const float4 w2 = *reinterpret_cast(weight_scale + col0 + 32); + const float4 w3 = *reinterpret_cast(weight_scale + col0 + 48); + hip_bfloat16* const obase = out + row * n + col0; + // Four independent transposes, issued as one pipelined block (the + // compiler sees all 32 ds_bpermute before the first store). + int u0, u1, u2, u3, v0, v1, v2, v3, q0, q1, q2, q3, r0, r1, r2, r3; + transpose_frag4(f0, lane, u0, u1, u2, u3); + transpose_frag4(f1, lane, v0, v1, v2, v3); + transpose_frag4(f2, lane, q0, q1, q2, q3); + transpose_frag4(f3, lane, r0, r1, r2, r3); + // Four 8-byte coalesced stores, back-to-back (16 rows x 32 B per store + // instruction; every 32-B sector covered by 4 lanes -> 100% efficiency). + // Fragment j's slot is 16 columns (32 B) further along the row: obase + + // 16*j elements. + *reinterpret_cast(obase) = pack_bf16x4(u0, u1, u2, u3, xs, w0); + *reinterpret_cast(obase + 16) = + pack_bf16x4(v0, v1, v2, v3, xs, w1); + *reinterpret_cast(obase + 32) = + pack_bf16x4(q0, q1, q2, q3, xs, w2); + *reinterpret_cast(obase + 48) = + pack_bf16x4(r0, r1, r2, r3, xs, w3); +} + +// Explicit 8-byte loader for the col_major m16n16k32 B fragment (accepted +// hy3 lineage, validated on this DTK). The col_major fragment mapping is +// n = lane & 15, k = 8*(lane >> 4) + i; with the n-major LDS tile at row +// stride 80 the lane's 8 elements are contiguous +// (p[row*ldm + col .. +7], 8-byte aligned because ldm=80 and col are +// multiples of 8), so this compiles to ONE ds_read2_b64 instead of 8 +// byte-granular ds_read_u8 + mask/OR reassembly. The byte placement is +// identical to du_load_matrix_sync, so the v_mmac operand +// registers receive the same values. +__device__ __forceinline__ void load_b_frag8( + DUFragment& f, + const int8_t* __restrict__ p, + int ldm) { + const unsigned row = static_cast(__lane_id()) & 0xfu; + const unsigned col = (static_cast(__lane_id()) >> 4) << 3; + const int64_t v = *reinterpret_cast(p + row * ldm + col); + *reinterpret_cast(&f.x[0]) = v; +} + +// Iteration 5 (packing round): 8-byte loader for the col_major B fragment on +// the swizzled exact-shape tile. The fragment block layout is +// [k8 0..3][n 0..15][8 B] with n = lane&15 and k8 = lane>>4, so lane (n, k8) +// reads p + k8*128 + n*8: one ds_read_b64 per fragment, and each 16-lane +// group of the wavefront spans exactly one 128-B k8 block, so every LDS +// cycle touches all 32 banks exactly once (conflict-free; the n-major +// 80-byte-row tile delivered the same bytes with 2-way conflicts). The byte +// placement is identical to du_load_matrix_sync / load_b_frag8, +// so the v_mmac operand registers and the int32 accumulation order are +// unchanged. +__device__ __forceinline__ void load_b_frag8_swz( + DUFragment& f, + const int8_t* __restrict__ p) { + const unsigned n = static_cast(__lane_id()) & 0xfu; + const unsigned k8 = static_cast(__lane_id()) >> 4; + const int64_t v = + *reinterpret_cast(p + (k8 << 7) + (n << 3)); + *reinterpret_cast(&f.x[0]) = v; +} + +// Iteration 9 (B-fragment pair round on the iteration-8 split): two-fragment +// 8-byte loader for the swizzled tile. Lane (n, k8) of fragment j reads its +// 8-byte k-run at p + k8*128 + n*8 and the SAME (n, k8) slot of fragment j+1, +// whose block sits exactly kPackedBGroupBytes (512 B = 128 dwords) further +// along the same 32-K region. Both addresses are 8-byte aligned and the +// second offset is inside the ds_read2_b64 8-bit-dword range (round-6 and the +// accepted worker_0 TP8 qkv lineage ISA both prove this DTK/LLVM merges the +// pair into ONE ds_read2_b64), so the eight single-fragment B reads of a +// wave-stage collapse to four pair reads. Each sub-read keeps the exact +// conflict-free pattern of load_b_frag8_swz (a 16-lane group spans one 128-B +// block and touches all 32 banks exactly once per cycle), so the B side stays +// conflict-free and the LDS cycles are unchanged. The byte placement is +// identical to two load_b_frag8_swz calls (v0 -> f0.x[0..7], v1 -> +// f1.x[0..7]), so the v_mmac operand registers and the int32 accumulation +// order are unchanged. +__device__ __forceinline__ void load_b_frag8_swz2( + DUFragment& f0, + DUFragment& f1, + const int8_t* __restrict__ p) { + const unsigned n = static_cast(__lane_id()) & 0xfu; + const unsigned k8 = static_cast(__lane_id()) >> 4; + const int64_t v0 = + *reinterpret_cast(p + (k8 << 7) + (n << 3)); + const int64_t v1 = *reinterpret_cast( + p + (k8 << 7) + (n << 3) + kPackedBGroupBytes); + *reinterpret_cast(&f0.x[0]) = v0; + *reinterpret_cast(&f1.x[0]) = v1; +} + +// Large-M prefill kernel: 64x128 output tile per block, 256 threads +// (4 wavefronts), single-buffered 64-K K stage, identity-packed B. Each wave +// owns a 32x64 quadrant = eight m16n16k32 int8->int32 DUMMA accumulators. +// Serves large-M shapes with N % 128 == 0 except the exact assigned shape +// (which routes to w8a8_dumma_prefill_64x128_packedb_kernel above). +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x128_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM128; + const int n0 = static_cast(blockIdx.x) * kBlockN128; + + __shared__ __align__(16) int8_t a_tile[kBlockM128 * kAStride128]; + __shared__ __align__(16) int8_t b_tile[kStageK128 * kBStride128]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; + DUFragment + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); + + const int num_stages = k / kStageK128; + + // Cooperative staging mapping (fixed per thread, reused every stage): + // A[64,64] -> 256 int4 vectors; thread `tid` owns vector `tid` + // (row = tid/4, 16-B column group = (tid%4)*16), so each + // wavefront covers 16 rows x 64 B contiguous per row. + // B[64,128] -> 512 int4 vectors; thread `tid` owns vectors `tid` and + // `tid+256` (kk = tid/8, column group = (tid%8)*16), so + // each wavefront covers 8 kk rows x 128 B contiguous. + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + const int b_kk0 = tid >> 3; + const int b_col16 = (tid & 7) << 4; + + // VGPR payload for one full 64-K stage (3 x int4 = 12 VGPR), loaded just + // before it is committed (no one-stage-ahead overlap in this control). + int4 vA, vB0, vB1; + + // Prologue: load stage 0 into registers, commit it to the single buffer, + // and make it visible before the first burst. + { + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + vB0 = *reinterpret_cast( + weight + b_kk0 * n + n0 + b_col16); + vB1 = *reinterpret_cast( + weight + (b_kk0 + kStageK128 / 2) * n + n0 + b_col16); + // Commit as int32 stores (skips the 4-byte pad column; 4 x ds_write_b32 + // per 16-byte vector because the odd strides are 4 B mod 16). + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + int32_t* bdst0 = reinterpret_cast( + b_tile + b_kk0 * kBStride128 + b_col16); + bdst0[0] = vB0.x; + bdst0[1] = vB0.y; + bdst0[2] = vB0.z; + bdst0[3] = vB0.w; + int32_t* bdst1 = bdst0 + (kStageK128 / 2) * (kBStride128 / 4); + bdst1[0] = vB1.x; + bdst1[1] = vB1.y; + bdst1[2] = vB1.z; + bdst1[3] = vB1.w; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + // Consume the single buffer: each wave does eight m16n16k32 MMAs per + // 32-K step over its 32x64 quadrant. + const int local_row = wave_row * 32; + const int local_col = wave_col * 64; +#pragma unroll + for (int kk = 0; kk < kStageK128; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride128 + kk, kAStride128); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride128 + kk, + kAStride128); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride128 + local_col, kBStride128); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride128 + local_col + kTileN, + kBStride128); + du_load_matrix_sync( + b_frag2, b_tile + kk * kBStride128 + local_col + 2 * kTileN, + kBStride128); + du_load_matrix_sync( + b_frag3, b_tile + kk * kBStride128 + local_col + 3 * kTileN, + kBStride128); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + } + // All consumption of the single buffer is complete only after every + // wavefront passes this barrier; the stage s+1 stores below overwrite + // that same buffer, so they must wait for it. + __syncthreads(); + + // Load stage s+1 and commit it into the single buffer. The global-load + // latency is fully exposed here: the loads are issued after the burst + // (nothing to overlap) and the compiler's vmcnt wait before the DS + // stores is on the critical path. `s + 1 < num_stages` is block-uniform, + // so the branch and its barrier are free of divergence. + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vB0 = *reinterpret_cast( + weight + (s1 + b_kk0) * n + n0 + b_col16); + vB1 = *reinterpret_cast( + weight + (s1 + b_kk0 + kStageK128 / 2) * n + n0 + b_col16); + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + int32_t* bdst0 = reinterpret_cast( + b_tile + b_kk0 * kBStride128 + b_col16); + bdst0[0] = vB0.x; + bdst0[1] = vB0.y; + bdst0[2] = vB0.z; + bdst0[3] = vB0.w; + int32_t* bdst1 = bdst0 + (kStageK128 / 2) * (kBStride128 / 4); + bdst1[0] = vB1.x; + bdst1[1] = vB1.y; + bdst1[2] = vB1.z; + bdst1[3] = vB1.w; + // Make the stage s+1 stores visible to every wavefront before the + // next burst. + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 64; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc02, x_scale, weight_scale, out, m, n, + base_row, base_col + 2 * kTileN, lane); + store_prefill_fragment( + acc03, x_scale, weight_scale, out, m, n, + base_row, base_col + 3 * kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); + store_prefill_fragment( + acc12, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + 2 * kTileN, lane); + store_prefill_fragment( + acc13, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + 3 * kTileN, lane); +} + +// --------------------------------------------------------------------------- +// Packed-B family for the exact assigned shape (K=6144, N=1536). +// launch_pack_w8a8_weight stores the logical [K, N] weight once, outside the +// timed region, in the iteration-5 fragment-interleaved swizzle (the 64x64 +// and 128x64 compiled-but-undispatched family members still assume the +// round-1 n-major packed[n][k] / 80-byte-row tile and are never launched by +// launch_w8a8_gemm). The dispatched 64x128 kernel stages B into the swizzled +// [k32][n16][k8][n][8] tile and consumes col_major B fragments with the +// 8-byte loader load_b_frag8_swz (one conflict-free ds_read_b64 per +// fragment). A stays row-major (68-byte LDS rows). All three kernels share +// the accepted pipeline: single-buffered 64-K K stage with two +// __syncthreads per stage and the fused coalesced scale/bf16 epilogue. +// Iteration 7 (compute-pipeline round): the DISPATCHED 64x128 kernel splits +// the per-stage consume into two independent kk halves - kk=0 fragment +// loads, one lgkmcnt wait, a contiguous 8-v_mmac burst, then the kk=32 +// loads issued BETWEEN the bursts (their LDS latency overlaps the kk=0 +// MMAs) and a second contiguous 8-v_mmac burst - replacing the baseline's +// "all twelve loads hoisted, seven lgkmcnt waits distributed through the +// 16-MMAC stream". The 64x64/128x64 compiled family members keep the +// hoisted form. The int32 accumulation order is unchanged from the generic +// kernels (every accumulator still accumulates kk=0 before kk=32), so +// results are bit-identical to the logical [K, N] layout. +// Iteration 9 (B-fragment pair round): each kk half's four swizzled B +// fragments load as TWO load_b_frag8_swz2 pair reads (one ds_read2_b64 per +// fragment pair; blocks j and j+1 are 512 B apart in the same 32-K region, +// inside the ds_read2_b64 8-bit-dword range), cutting the per-wave-stage B +// read-issue stream 8 -> 4 while the iteration-8 wait grouping is kept. +// Byte placement equals two load_b_frag8_swz calls, so operands, the +// kk=0-then-kk=32 int32 order, LDS size and occupancy are unchanged. +// Iteration 12 (staging-prefetch consolidation; rounds 10/11 both moved +// kk=32 LDS fragment loads earlier inside the consume and regressed +// 537.84/538.41 and 537.81/538.24 us, so the LDS read-issue stream is at +// its local optimum and the remaining exposed latency is on the OTHER side +// of the pipeline): the exact iteration-9 gfx928 ISA shows the three stage +// s+1 global payload loads (global_load_dwordx4 vA/vB0/vB1) issuing AFTER +// the consume s_barrier with only ~6-20 instructions of cover before the +// s_waitcnt vmcnt(2)/(1)/(0) waits in front of the LDS commit stores - one +// L2-hit round trip (~150-250 cycles) of exposed VMEM stall per wave-stage +// at the loop top, never touched by rounds 4-11 (which all modified the LDS +// consume path only). This round issues the three payload loads at the TOP +// of the consume (the payload registers are dead after the commit and the +// consume never touches them, so the same 12 VGPR are reused), giving the +// VMEM latency the whole 16-MMA consume as cover; the backend's pre-barrier +// s_waitcnt vmcnt drain retires as a no-op and the commit stores after the +// barrier issue with no VMEM stall. Same bytes into the same LDS addresses, +// same two barriers per stage with the stores still between them, same +// consume, LDS (12,544 B), VGPR (79) and occupancy (3 blocks/CU) -> output +// bits identical. +// --------------------------------------------------------------------------- + +// 64x128 tile: 4 waves, 32x64 quadrant per wave, eight accumulators. This is +// the exact-shape dispatch (strongest evidence for this operator: the accepted +// hy3 qkv_proj lineage at 113.5 TOPS, K=4096). Grid (12, 64) = 768 blocks. +// LDS/block = 64x68 + 64x128 = 12,544 B (iteration 5: the swizzled B tile +// needs no padding) -> 3 blocks/CU (VGPR-bound at arch_vgpr <= 85). +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x128_packedb_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM128; + const int n0 = static_cast(blockIdx.x) * kBlockN128; + + __shared__ __align__(16) int8_t a_tile[kBlockM128 * kAStride128]; + // Swizzled B tile (iteration 5): the exact-shape pack stores the weight + // fragment-interleaved [k32 0..1][n16 0..7][k8 0..3][n 0..15][8 B] + // (64-K stage = 2 x 4096 B), so staging is int4 -> ds_write_b128 and the + // col_major B fragment loads are one conflict-free ds_read_b64 each. + __shared__ __align__(16) int8_t b_tile[kStageK128 * kBlockN128]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; + DUFragment + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); + + const int num_stages = k / kStageK128; + + // A staging (row-major x_q -> 68-byte LDS rows): thread `tid` owns the + // 16-byte k-run (a_col16..a_col16+15) of row a_row. + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + // B staging on the swizzled pack: thread `tid` owns the 16-byte vectors + // (k32=0, n16, k8, n_pair) = v = tid and (k32=1, n16, k8, n_pair) = v = + // tid+256, i.e. bytes + // ((k32*96 + n0/16 + n16)*512 + k8*128 + n_pair*16) of packed_b, which + // are the 8-byte k-runs of n-rows 2*n_pair and 2*n_pair+1 in one k8 group. + const int b_n16 = tid >> 5; // 0..7 (n16 group inside the 128-n tile) + const int b_k8 = (tid >> 3) & 3; // 0..3 (8-byte k8 group inside the 32-k block) + const int b_np = tid & 7; // 0..7 (n pair inside the 16-n group) + // Tile-local offset of the 16-byte chunk (k32 block 0, n16 group b_n16, + // k8 group, n pair) inside b_tile. + const int b_off = + b_n16 * kPackedBGroupBytes + (b_k8 << 7) + (b_np << 4); + // Global packed_b offset of the same 16-byte chunk inside ONE 32-k block + // (the n16-group term is already carried by the (k32g*kPackedBN16 + + // n0_16 + b_n16) index, so it must NOT be repeated here). + const int b_pack_off = (b_k8 << 7) + (b_np << 4); + + // VGPR payload for one full 64-K stage (3 x int4 = 12 VGPR). + int4 vA, vB0, vB1; + + // Prologue: load stage 0 into registers, commit it to the single buffers, + // and make it visible before the first burst. + { + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + const int n0_16 = n0 >> 4; + vB0 = *reinterpret_cast( + packed_b + (n0_16 + b_n16) * kPackedBGroupBytes + b_pack_off); + vB1 = *reinterpret_cast( + packed_b + + (kPackedBN16 + n0_16 + b_n16) * kPackedBGroupBytes + b_pack_off); + // Commit A as int32 stores (68-byte rows are 4 B mod 16 -> no b128). + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + // Commit B as 16-byte vector stores (all offsets are 0 mod 16): each + // 8-lane LDS cycle covers all 32 banks exactly once (conflict-free). + *reinterpret_cast(b_tile + b_off) = vB0; + *reinterpret_cast(b_tile + kPackedB32Bytes + b_off) = vB1; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + // Consume the single buffer: each wave does eight m16n16k32 MMAs per + // 32-K step over its 32x64 quadrant. Iteration 7 (compute-pipeline + // round): the stage body is split into two independent kk halves. The + // kk=0 loads issue first and their single lgkmcnt wait retires before + // the kk=0 burst, so the first eight v_mmac issue as ONE contiguous, + // wait-free issue group (the accepted baseline hoisted all twelve + // loads and distributed seven lgkmcnt waits through the 16-MMAC + // stream, stalling the issue pipeline whenever a wait's count had not + // yet dropped). The kk=32 fragment loads are issued BETWEEN the two + // bursts: their LDS latency overlaps the kk=0 MMAC burst (prefetch + // distance = one 8-MMAC burst), and the second eight v_mmac issue as + // another contiguous group. Each burst's MMAs write eight independent + // accumulator chains (no intra-burst dependency), and every + // accumulator still accumulates kk=0 before kk=32 in the exact int32 + // order -> output bits identical. + const int local_row = wave_row * 32; + const int8_t* abase = a_tile + local_row * kAStride128; + // B fragment j covers n16 group wave_col*4 + j of the swizzled tile + // (512 B per 16-n fragment block); the kk=32 burst lives in k32 block 1. + const int8_t* bbase = b_tile + (wave_col << 2) * kPackedBGroupBytes; + // Iteration 12 (staging-prefetch consolidation): issue the stage s+1 + // global payload loads (vA/vB0/vB1) at the TOP of the consume. The + // payload registers are dead after this stage's commit, and the consume + // never touches them, so the loads sit as early as possible in the loop + // body and their VMEM (L2-hit ~150-250 cycles) latency is covered by the + // entire 16-MMA consume below; the backend's pre-barrier s_waitcnt + // vmcnt drain then retires as a no-op and the LDS commit stores after + // the barrier issue with no VMEM stall. (The exact iteration-9 ISA shows + // the same three loads issuing AFTER the consume barrier with only ~6-20 + // instructions of cover before exposed s_waitcnt vmcnt waits in front of + // the stores - one L2 round trip of stall per wave-stage at the loop + // top; rounds 4-11 never touched this staging side, they only modified + // the LDS consume stream, which rounds 10/11 proved is at its local + // optimum.) Same bytes into the same LDS addresses, same two barriers + // per stage with the stores still between them, same 12 payload VGPR, + // same consume and int32 accumulation order -> output bits identical. + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int k32g = s1 >> 5; // global 32-k block of the stage s+1 payload + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + const int n0_16 = n0 >> 4; + vB0 = *reinterpret_cast( + packed_b + + (k32g * kPackedBN16 + n0_16 + b_n16) * kPackedBGroupBytes + + b_pack_off); + vB1 = *reinterpret_cast( + packed_b + + ((k32g + 1) * kPackedBN16 + n0_16 + b_n16) * kPackedBGroupBytes + + b_pack_off); + } + du_load_matrix_sync(a_frag0, abase, kAStride128); + du_load_matrix_sync( + a_frag1, abase + kTileM * kAStride128, kAStride128); + // Iteration 9: the four kk=0 B fragments load as TWO pair reads (one + // ds_read2_b64 per fragment pair instead of two ds_read_b64; fragment + // blocks j and j+1 are kPackedBGroupBytes apart in the same 32-K region, + // inside the ds_read2_b64 offset range, and both sub-reads keep the exact + // conflict-free pattern of load_b_frag8_swz). + load_b_frag8_swz2(b_frag0, b_frag1, bbase); + load_b_frag8_swz2(b_frag2, b_frag3, bbase + 2 * kPackedBGroupBytes); + // kk = 0 burst (eight m16n16k32 MMAs; all six operands are ready after + // the single lgkmcnt wait, so the group issues contiguously). + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + // kk = 32 fragment loads (prefetch distance: issued after the kk=0 + // burst so their LDS latency overlaps the eight MMAs above). + DUFragment + a2_frag0, a2_frag1; + DUFragment + b2_frag0, b2_frag1, b2_frag2, b2_frag3; + du_load_matrix_sync(a2_frag0, abase + kTileK, kAStride128); + du_load_matrix_sync( + a2_frag1, abase + kTileM * kAStride128 + kTileK, kAStride128); + // Iteration 9: the four kk=32 B fragments also load as TWO pair reads + // (fragment blocks j and j+1 of the k32=1 region are kPackedBGroupBytes + // apart; the k32 block itself is kPackedB32Bytes from the kk=0 region). + load_b_frag8_swz2(b2_frag0, b2_frag1, bbase + kPackedB32Bytes); + load_b_frag8_swz2( + b2_frag2, b2_frag3, bbase + kPackedB32Bytes + 2 * kPackedBGroupBytes); + // kk = 32 burst (same accumulators, same int32 accumulation order). + du_mma_sync(acc00, a2_frag0, b2_frag0, acc00); + du_mma_sync(acc01, a2_frag0, b2_frag1, acc01); + du_mma_sync(acc02, a2_frag0, b2_frag2, acc02); + du_mma_sync(acc03, a2_frag0, b2_frag3, acc03); + du_mma_sync(acc10, a2_frag1, b2_frag0, acc10); + du_mma_sync(acc11, a2_frag1, b2_frag1, acc11); + du_mma_sync(acc12, a2_frag1, b2_frag2, acc12); + du_mma_sync(acc13, a2_frag1, b2_frag3, acc13); + // All consumption of the single buffer is complete only after every + // wavefront passes this barrier; the stage s+1 stores below overwrite + // that same buffer. + __syncthreads(); + + // Commit the stage s+1 payload (already in vA/vB0/vB1, loaded at the + // top of this iteration) into the single buffer. The VMEM data landed + // during the 16-MMA consume, so the stores issue with no VMEM stall. + if (s + 1 < num_stages) { + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + *reinterpret_cast(b_tile + b_off) = vB0; + *reinterpret_cast(b_tile + kPackedB32Bytes + b_off) = vB1; + // Make the stage s+1 stores visible to every wavefront before the + // next burst. + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 64; + // Iteration 6 (epilogue round): fused 4-fragment row-group epilogues. + // acc00..acc03 cover rows base_row..base_row+15, columns base_col..+63; + // acc10..acc13 cover the same columns one 16-row group lower. Each call + // loads x_scale[row] and the four float4 weight_scale groups once per + // lane before any ds_bpermute, issues the four transposes as one + // software-pipelined block, then stores back-to-back. + store_prefill_rowgroup_coalesced( + acc00, acc01, acc02, acc03, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_rowgroup_coalesced( + acc10, acc11, acc12, acc13, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); +} + +// 64x64 packed-B tile: 4 waves, 32x32 quadrant per wave, four accumulators. +// Compiled family member for the mandated tile sweep (grid 24x64 = 1536 +// blocks for the assigned shape; LDS/block = 64x68 + 64x80 = 9,472 B). Not +// dispatched this round; the exact shape uses the 64x128 packed-B kernel. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x64_packedb_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride128]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kPackedBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + const int num_stages = k / kStageK128; + + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + const int b_n = tid >> 2; + const int b_k16 = (tid & 3) << 4; + + int4 vA, vB; + + // Prologue: stage 0. + { + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + vB = *reinterpret_cast( + packed_b + (n0 + b_n) * k + b_k16); + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + const int8_t* abase = a_tile + local_row * kAStride128; + const int8_t* bbase = b_tile + local_col * kPackedBStride; + du_load_matrix_sync(a_frag0, abase, kAStride128); + du_load_matrix_sync( + a_frag1, abase + kTileM * kAStride128, kAStride128); + load_b_frag8(b_frag0, bbase, kPackedBStride); + load_b_frag8( + b_frag1, bbase + kTileN * kPackedBStride, kPackedBStride); + DUFragment + a2_frag0, a2_frag1; + DUFragment + b2_frag0, b2_frag1; + du_load_matrix_sync(a2_frag0, abase + kTileK, kAStride128); + du_load_matrix_sync( + a2_frag1, abase + kTileM * kAStride128 + kTileK, kAStride128); + load_b_frag8(b2_frag0, bbase + kTileK, kPackedBStride); + load_b_frag8( + b2_frag1, bbase + kTileN * kPackedBStride + kTileK, kPackedBStride); + // kk = 0 burst. + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + // kk = 32 burst. + du_mma_sync(acc00, a2_frag0, b2_frag0, acc00); + du_mma_sync(acc01, a2_frag0, b2_frag1, acc01); + du_mma_sync(acc10, a2_frag1, b2_frag0, acc10); + du_mma_sync(acc11, a2_frag1, b2_frag1, acc11); + __syncthreads(); + + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int g_row = m0 + a_row; + vA = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vB = *reinterpret_cast( + packed_b + (n0 + b_n) * k + s1 + b_k16); + int32_t* adst = reinterpret_cast( + a_tile + a_row * kAStride128 + a_col16); + adst[0] = vA.x; + adst[1] = vA.y; + adst[2] = vA.z; + adst[3] = vA.w; + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB; + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment_coalesced( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment_coalesced( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment_coalesced( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment_coalesced( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// 128x64 packed-B tile: 4 waves, 64x32 quadrant per wave, eight +// accumulators (the o_proj lineage geometry). Compiled family member for the +// mandated tile sweep (grid 24x32 = 768 blocks for the assigned shape; +// LDS/block = 128x80 + 64x80 = 15,360 B). Not dispatched this round. +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_128x64_packedb_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ packed_b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM128 * 2; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM128 * 2 * kPackedAStride128x64]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kPackedBStride]; + + DUFragment + a_frag0, a_frag1, a_frag2, a_frag3; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + du_fill_fragment(acc20, 0); + du_fill_fragment(acc21, 0); + du_fill_fragment(acc30, 0); + du_fill_fragment(acc31, 0); + + const int num_stages = k / kStageK128; + + // A staging: thread owns the 16-byte k-run (a_col16..+15) of A rows + // a_row and a_row + 64 (A[128,64] = 512 int4 = 2 per thread). + const int a_row = tid >> 2; + const int a_col16 = (tid & 3) << 4; + // B staging: thread owns the 16-byte k-run (b_k16..+15) of n-row b_n. + const int b_n = tid >> 2; + const int b_k16 = (tid & 3) << 4; + + int4 vA0, vA1, vB; + + // Prologue: stage 0. + { + const int g_row = m0 + a_row; + vA0 = (g_row < m) + ? *reinterpret_cast(x_q + g_row * k + a_col16) + : int4{0, 0, 0, 0}; + vA1 = (g_row + kBlockM128 < m) + ? *reinterpret_cast( + x_q + (g_row + kBlockM128) * k + a_col16) + : int4{0, 0, 0, 0}; + vB = *reinterpret_cast( + packed_b + (n0 + b_n) * k + b_k16); + // 80-byte A rows are 0 mod 16: one b128 store per vector. + *reinterpret_cast( + a_tile + a_row * kPackedAStride128x64 + a_col16) = vA0; + *reinterpret_cast( + a_tile + (a_row + kBlockM128) * kPackedAStride128x64 + a_col16) = vA1; + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB; + __syncthreads(); + } + + for (int s = 0; s < num_stages; ++s) { + // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMAs per 32-K + // step (4 A fragments x 2 B fragments), all fragment loads hoisted. + const int local_row = wave_row * 64; + const int local_col = wave_col * 32; + const int8_t* abase = a_tile + local_row * kPackedAStride128x64; + const int8_t* bbase = b_tile + local_col * kPackedBStride; + du_load_matrix_sync(a_frag0, abase, kPackedAStride128x64); + du_load_matrix_sync( + a_frag1, abase + kTileM * kPackedAStride128x64, + kPackedAStride128x64); + du_load_matrix_sync( + a_frag2, abase + 2 * kTileM * kPackedAStride128x64, + kPackedAStride128x64); + du_load_matrix_sync( + a_frag3, abase + 3 * kTileM * kPackedAStride128x64, + kPackedAStride128x64); + load_b_frag8(b_frag0, bbase, kPackedBStride); + load_b_frag8( + b_frag1, bbase + kTileN * kPackedBStride, kPackedBStride); + DUFragment + a2_frag0, a2_frag1, a2_frag2, a2_frag3; + DUFragment + b2_frag0, b2_frag1; + du_load_matrix_sync(a2_frag0, abase + kTileK, kPackedAStride128x64); + du_load_matrix_sync( + a2_frag1, abase + kTileM * kPackedAStride128x64 + kTileK, + kPackedAStride128x64); + du_load_matrix_sync( + a2_frag2, abase + 2 * kTileM * kPackedAStride128x64 + kTileK, + kPackedAStride128x64); + du_load_matrix_sync( + a2_frag3, abase + 3 * kTileM * kPackedAStride128x64 + kTileK, + kPackedAStride128x64); + load_b_frag8(b2_frag0, bbase + kTileK, kPackedBStride); + load_b_frag8( + b2_frag1, bbase + kTileN * kPackedBStride + kTileK, kPackedBStride); + // kk = 0 burst. + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc20, a_frag2, b_frag0, acc20); + du_mma_sync(acc21, a_frag2, b_frag1, acc21); + du_mma_sync(acc30, a_frag3, b_frag0, acc30); + du_mma_sync(acc31, a_frag3, b_frag1, acc31); + // kk = 32 burst. + du_mma_sync(acc00, a2_frag0, b2_frag0, acc00); + du_mma_sync(acc01, a2_frag0, b2_frag1, acc01); + du_mma_sync(acc10, a2_frag1, b2_frag0, acc10); + du_mma_sync(acc11, a2_frag1, b2_frag1, acc11); + du_mma_sync(acc20, a2_frag2, b2_frag0, acc20); + du_mma_sync(acc21, a2_frag2, b2_frag1, acc21); + du_mma_sync(acc30, a2_frag3, b2_frag0, acc30); + du_mma_sync(acc31, a2_frag3, b2_frag1, acc31); + __syncthreads(); + + if (s + 1 < num_stages) { + const int s1 = (s + 1) * kStageK128; + const int g_row = m0 + a_row; + vA0 = (g_row < m) + ? *reinterpret_cast( + x_q + g_row * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vA1 = (g_row + kBlockM128 < m) + ? *reinterpret_cast( + x_q + (g_row + kBlockM128) * k + s1 + a_col16) + : int4{0, 0, 0, 0}; + vB = *reinterpret_cast( + packed_b + (n0 + b_n) * k + s1 + b_k16); + *reinterpret_cast( + a_tile + a_row * kPackedAStride128x64 + a_col16) = vA0; + *reinterpret_cast( + a_tile + (a_row + kBlockM128) * kPackedAStride128x64 + a_col16) = + vA1; + *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB; + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 64; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment_coalesced( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment_coalesced( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment_coalesced( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment_coalesced( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); + store_prefill_fragment_coalesced( + acc20, x_scale, weight_scale, out, m, n, + base_row + 2 * kTileM, base_col, lane); + store_prefill_fragment_coalesced( + acc21, x_scale, weight_scale, out, m, n, + base_row + 2 * kTileM, base_col + kTileN, lane); + store_prefill_fragment_coalesced( + acc30, x_scale, weight_scale, out, m, n, + base_row + 3 * kTileM, base_col, lane); + store_prefill_fragment_coalesced( + acc31, x_scale, weight_scale, out, m, n, + base_row + 3 * kTileM, base_col + kTileN, lane); +} + +// Generic large-M prefill kernel: 64x64 output tile per block, K staged in +// LDS (single buffer), four waves each owning a 32x32 quadrant = four +// m16n16k32 int8->int32 DUMMA accumulators. Covers large-M shapes whose N is +// a multiple of 64 but not 128 (and K a multiple of 128). +__global__ __launch_bounds__(kThreadsPerBlock) void +w8a8_dumma_prefill_64x64_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM * kStageK]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBlockN]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[64,128] (row-major, zero-filled past M). + constexpr int kAVectors = kBlockM * kStageK / static_cast(sizeof(int4)); + for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int local_row = byte_offset / kStageK; + const int kk = byte_offset - local_row * kStageK; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile)[vec] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + kk) + : int4{0, 0, 0, 0}; + } + + // Stage B[128,64] (row-major [K, N] logical weight layout). + constexpr int kBVectors = kStageK * kBlockN / static_cast(sizeof(int4)); + for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { + const int byte_offset = vec * static_cast(sizeof(int4)); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + reinterpret_cast(b_tile)[vec] = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + col); + } + __syncthreads(); + + // Consume the stage: each wave does four m16n16k32 MMAs per 32-K step. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kStageK + kk, kStageK); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kStageK + kk, kStageK); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBlockN + local_col, kBlockN); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBlockN + local_col + kTileN, kBlockN); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + // Single buffer: make sure every wave finished reading LDS before the + // next stage overwrites it. + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Generic scalar int8/int32 fallback for unmatched shapes and small-M API +// cases. One output element per grid-stride step; exact int32 accumulation +// (K <= 8192 keeps the int8 dot well within int32 range: 8192*128*128 = +// 134,217,728 < 2^31), then the fused float scale and bf16 store. The exact +// assigned shape (K=6144, N=1536) uses the iteration-5 swizzled pack (set up +// once by launch_pack_w8a8_weight; decoded per element below); every other +// (k, n) keeps the logical [K, N] row-major layout. The branch is +// grid-uniform per launch, so the paired M=2 API shape with the same (N, K) +// stays correct. +__global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int64_t total = static_cast(m) * n; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + const bool packed_b = (k == kPackedBK && n == kPackedBN); + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; + idx += stride) { + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = a + static_cast(row) * k; + int32_t acc = 0; + for (int kk = 0; kk < k; ++kk) { + int32_t bw; + if (packed_b) { + // Iteration-5 swizzled exact-shape layout: element (n=col, k=kk) at + // ((kk>>5)*96 + (col>>4))*512 + (((kk>>3)&3)*16 + (col&15))*8 + + // (kk&7) (fragment-interleaved [k32][n16][k8][n][8]). + const int64_t off = + ((static_cast(kk >> 5) * kPackedBN16 + (col >> 4)) * + kPackedBGroupBytes) + + ((((kk >> 3) & 3) * 16 + (col & 15)) << 3) + (kk & 7); + bw = static_cast(b[off]); + } else { + bw = static_cast( + b[static_cast(kk) * n + col]); + } + acc += static_cast(a_row[kk]) * bw; + } + const float scaled = static_cast(acc) * + x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); + } +} + +// Identity device-to-device byte copy (pack_weight fallback, valid for any +// (K, N)). +__global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +// One-time swizzled pack of the logical [K, N] weight for the exact assigned +// shape (K=6144, N=1536), iteration 5 (packing round). Runs only from +// launch_pack_w8a8_weight, outside the timed region and outside Graph +// capture. The packed buffer keeps the same byte count (k*n) as the identity +// pack, so allocations and graph-stable addresses are unchanged. Layout: +// element (n, k) lands at +// ((k>>5)*n16g + (n>>4))*512 + (((k>>3)&3)*16 + (n&15))*8 + (k&7) +// with n16g = n>>4 (96 for the exact shape): for each 32-k x 16-n DUMMA B +// fragment block the 4 k8 groups x 16 n-rows x 8 B are stored n-interleaved +// (k8-major within the block). Each thread writes one contiguous 16-byte +// chunk (two n-rows of one k8) with a strided 2x8 gather on the read side +// (packing is outside timing, so the scalar gather is off the critical +// path). The guard guarantees n % 16 == 0. +__global__ __launch_bounds__(256) void w8a8_pack_b_swz_kernel( + const int8_t* __restrict__ src, + int8_t* __restrict__ dst, + int k, + int n) { + const int64_t chunks = (static_cast(k) * n) >> 4; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + const int n16g = n >> 4; + for (int64_t c = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + c < chunks; + c += stride) { + const int k32 = static_cast(c / (static_cast(n16g) * 32)); + const int rem = + static_cast(c % (static_cast(n16g) * 32)); + const int n16 = rem >> 5; + const int w = rem & 31; + const int k8 = w >> 3; + const int np = w & 7; + const int n0 = (n16 << 4) + (np << 1); + const int k0 = (k32 << 5) + (k8 << 3); + int8_t* dst16 = dst + (c << 4); +#pragma unroll + for (int i = 0; i < 8; ++i) { + // 16-byte chunk = [n0 row: k0..k0+7][n0+1 row: k0..k0+7] (n-major + // within the k8 group, matching the [k8][n][8] tile layout). + dst16[i] = src[static_cast(k0 + i) * n + n0]; + dst16[8 + i] = src[static_cast(k0 + i) * n + n0 + 1]; + } + } +} + +// Identity device-to-device float copy for the weight scales. +__global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; + i += stride) { + dst[i] = src[i]; + } +} + +} // namespace + +// Stable host launch symbol consumed by csrc/bindings.cpp. Dispatches on the +// caller-provided stream only; never allocates, synchronizes, or touches the +// default stream. +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + (void)workspace; + (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + auto* out_ptr = static_cast(out); + + // Packed-B family: the exact assigned shape (K=6144, N=1536) routes to + // the swizzled-packed-B 64x128 kernel (iteration 5 layout). The guard is + // exact (m >= 128, n == 1536, k == 6144) and sits BEFORE the generic + // 64x128 path; every other shape keeps its existing path (generic 64x128 + // for n % 128 == 0, 64x64 for n % 64 == 0, scalar fallback otherwise) and + // the identity-packed layout. Grid (N/128)x(M/64) = 12x64 = 768 blocks. + if (m >= 128 && n == kPackedBN && k == kPackedBK) { + const dim3 grid( + static_cast(n / kBlockN128), + static_cast((m + kBlockM128 - 1) / kBlockM128)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x128_packedb_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Native INT8 DUMMA m16n16k32 prefill path (2-D macro-tile 64x128, + // single-buffered 64-K K stage, identity-packed B) for every other + // large-M shape with compatible geometry. + if (m >= 128 && (n % kBlockN128) == 0 && (k % kStageK128) == 0) { + const dim3 grid( + static_cast(n / kBlockN128), + static_cast((m + kBlockM128 - 1) / kBlockM128)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x128_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Native INT8 DUMMA m16n16k32 prefill path (64x64 tile) for large-M shapes + // whose N is not a multiple of 128 but is a multiple of 64. + if (m >= 128 && (n % kBlockN) == 0 && (k % kStageK) == 0) { + const dim3 grid( + static_cast(n / kBlockN), + static_cast((m + kBlockM - 1) / kBlockM)); + const dim3 block(static_cast(kThreadsPerBlock)); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x64_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + return; + } + + // Scalar int8/int32 fallback for unmatched (m, n, k) and small-M API cases + // (M < 128, including the paired M=2 API shape with the same (N,K)). + const int64_t total = static_cast(m) * n; + constexpr int kFallbackThreads = 256; + int64_t blocks = (total + kFallbackThreads - 1) / kFallbackThreads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + const dim3 grid(static_cast(blocks)); + const dim3 block(kFallbackThreads); + hipLaunchKernelGGL( + w8a8_scalar_gemm_kernel, + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); +} + +// Stable host launch symbol consumed by csrc/bindings.cpp. For the exact +// assigned shape (K=6144, N=1536) it performs the one-time swizzled B pack +// (iteration 5: logical [K, N] -> fragment-interleaved +// [k32][n16][k8][n][8]) outside the timed region and outside Graph capture; +// every other (K, N) keeps the identity device-to-device copy, valid for +// every (K, N). The packed buffer size is k*n in both cases, so the +// allocation and graph-stable packed layout are unchanged. +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + constexpr int kPackThreads = 256; + const int64_t weight_count = static_cast(k) * n; + if (weight_count > 0) { + if (k == kPackedBK && n == kPackedBN) { + const int64_t chunks = weight_count >> 4; + int64_t blocks = (chunks + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_b_swz_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + k, + n); + } else { + int64_t blocks = (weight_count + kPackThreads - 1) / kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_identity_i8_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + raw_weight, + packed_weight, + weight_count); + } + } + if (n > 0) { + int64_t blocks = (static_cast(n) + kPackThreads - 1) / + kPackThreads; + if (blocks > 4096) { + blocks = 4096; + } + hipLaunchKernelGGL( + w8a8_pack_identity_f32_kernel, + dim3(static_cast(blocks)), + dim3(kPackThreads), + 0, + stream, + weight_scale, + packed_weight_scale, + n); + } +} +// @@end +// @@end +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/shared_down_proj.hip new file mode 100644 index 00000000..1cdb65ed --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/shared_down_proj.hip @@ -0,0 +1,1514 @@ +// @@variant shape=minimax_tp8_shared_down_proj_m4096 commit=d7ed2409b78292728a0b5bbd5056f01cd38775be added=2026-08-29 +// median_us=312.5 p90_us=313.1 speedup=71.65 baseline_us=2.239e+04 +// source=minimaxm3-dsh-tp8-m4096-1-b0482833 +// MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). +// worker_3 native INT8 DUMMA m16n16k32 prefill kernels for the assigned +// MiniMax TP8 M=4096 shapes: +// * minimax_tp8_shared_gate_up_proj_m4096: (M=4096, N=768, K=6144) +// * minimax_tp8_shared_down_proj_m4096: (M=4096, N=6144, K=384) +// +// This file provides the two stable host launch symbols consumed by +// csrc/bindings.cpp: +// * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's +// HIP stream; no allocation, packing, +// autotuning, or synchronization. Uses +// only the caller-provided out and +// workspace (workspace untouched by this +// no-split-K path). +// * launch_pack_w8a8_weight(...) - one-time out-of-timed-region packing; +// the exact (K, N) = (6144, 768) and +// (384, 6144) pairs are stored transposed +// [N, K], every other pair keeps the +// identity copy. +// +// Strategy (round 1: establish the 2-D macro-tile DUMMA baseline for the +// exact gate_up shape by porting the sibling-validated gate_up architecture +// family; the bootstrap's single-buffered 64x128 tile was measured at +// 595 us / 64.95 TOPS): +// * Large-M prefill (M >= 128): native INT8 DUMMA m16n16k32 2-D macro-tiles +// (64x64 / 64x128 / 128x64), each wave owning a 32x32 quadrant (four +// 16x16 int32 accumulators) resident across the whole K loop. A and B are +// cooperatively staged into bank-skewed LDS with 16-byte vectorized +// coalesced global loads. +// * Exact gate_up shape (M=4096, N=768, K=6144): dispatches to +// w8a8_gemm_prefill_tiled_kernel_w4<64,128,64> -- the 64x128 macro-tile, +// 256 threads = 4 wavefronts, each wave owning a 32x64 quadrant (eight +// m16n16k32 accumulators), K stage 64 DOUBLE-buffered with one +// __syncthreads per stage (stage s+1's global loads issue into registers +// before the stage-s MMAC burst and commit after it), B packed [N, K] +// and loaded with col_major fragments as ONE 8-byte LDS read per lane +// (load_fragment8: identical element-to-slot mapping and byte order, no +// per-dword mask/OR reassembly VALU), fused fragment -> scale -> bf16 +// epilogue. LDS 30,720 B/block -> 2 blocks/CU; grid (N/128) x (M/64) = +// 6 x 64 = 384 blocks; 96 K stages. No split-K: the M x N output-tile +// grid dwarfs the 120 CUs. +// * Other large-M shapes sharing the packed (k, n) = (6144, 768) pair use +// the generic 8-wave double-buffered <64,128,64,2> kernel (bit-identical +// path); the generic single-buffer 64x128/128x64/64x64 arms and the +// identity packing for unmatched (k, n) pairs are unchanged. +// * The scalar fallback is transposed-aware for the exact (k, n) pair so +// every consumer of the packed (6144, 768) buffer (including the paired +// M=2/M=16 decode shapes) stays correct. +// +// Mathematical contract (exact int32 dot before float scaling): +// out[m, n] = bf16(int32_dot(a[m, :], b[:, n]) * x_scale[m] * weight_scale[n]) +// +// Header order is fixed by the control plane for this DTK: +// hip_runtime.h -> hip_bfloat16.h -> du_mma.h + +#include +#include +#include + +#include + +namespace { + +// gfx928 INT8 DUMMA primitive: m16n16k32, int8 x int8 -> int32. +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; + +// gfx928 wavefront is 64 lanes; every blockDim must be a multiple of 64. +constexpr int kWaveSize = 64; + +// Large-prefill dispatch geometry. +constexpr int kPrefillMinM = 128; +constexpr int kPrefillStageK = 128; +constexpr int kDefaultBlockN = 64; + +// LDS padding in bytes added to each staged row. A keeps 16 so every A row +// start stays 16-byte aligned for int4 staging while shifting fragment rows +// off the same 32-bank phase (power-of-two strides alias every row onto +// identical banks). B deliberately uses only 4 bytes of pad: with the gfx928 +// INT8 DUMMA B-fragment layout (lane (c, g) reads bytes at k-rows 8g..8g+7, +// one byte per row, ldm apart), any 16-byte-aligned row stride makes +// 2*stride*g == 0 (mod 32), so all four k-groups of every ds_read_u8 land on +// the same bank phase -> 4 banks x 16 lanes (16-way). Stride 132 (== 4 mod +// 16) moves the k-groups to bank phases {0,8,16,24} -> 16 banks x 4 lanes. +constexpr int kLdsPad = 16; +constexpr int kBLdsPad = 4; + +constexpr int kScalarThreads = 256; + +using namespace du::dumma; + +// Verified gfx928 INT8 DUMMA accumulator ownership: +// row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> columns +// col_mod4 + 4*i. Scale by x_scale[row] and weight_scale[col] and store +// bf16 directly from the fragment (no accumulator LDS round trip). +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Direct 8-byte LDS fragment load for the double-buffered exact-shape path. +// du_load_matrix_sync's int8 loaders assign x[0..7] = 8 consecutive bytes at +// (lane & 15) * ldm + ((lane >> 4) << 3) for both matrix_a row_major and +// matrix_b col_major, and du_mma_sync passes the fragment to the v_mmac +// builtin as one packed 8-byte operand, so the compiler emits a per-byte +// mask/OR reassembly chain for every loaded dword. Writing the same 8 bytes +// directly into the fragment storage keeps the operand bit pattern identical +// (exact int32 accumulation unchanged) and lets the compiler feed the +// ds_read2_b32 pair straight to the v_mmac (one ds_read2_b32 per fragment, +// ~4-way bank floor instead of the 16-way aliasing of strided byte reads). +__device__ __forceinline__ void load_fragment8( + DUFragmentBase& frag, + const signed char* __restrict__ base, + int stride, + int lane) { + const int off = (lane & 15) * stride + ((lane >> 4) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// Large-M prefill path: one block computes a kBlockM x kBlockN output tile. +// kWaveRows x kWaveCols wavefronts each own a 32x32 quadrant (four 16x16 +// DUMMA accumulators), while the block cooperatively stages +// A[kBlockM, kStageK] and B[kStageK, kBlockN] in bank-skewed LDS. The LDS is +// single-buffered with two barriers per stage (stage in, compute, protect). +// Tail-M rows are zero-filled on load and masked on store, so any M >= 128 +// is supported. The packed weight is the bootstrap identity [K, N] layout, +// so B rows are staged n-major contiguous chunks. +template +__global__ __launch_bounds__((kBlockM / 32) * (kBlockN / 32) * kWaveSize) void +w8a8_dumma_prefill_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWaveRows = kBlockM / 32; + constexpr int kWaveCols = kBlockN / 32; + constexpr int kThreads = kWaveRows * kWaveCols * kWaveSize; + constexpr int kAStride = kStageK + kLdsPad; + constexpr int kBStride = kBlockN + kBLdsPad; + static_assert(kAStride % sizeof(int4) == 0, + "A LDS row stride must stay 16-byte aligned"); + static_assert(kBStride % sizeof(int32_t) == 0, + "B LDS row stride must stay 4-byte aligned (int32 stores)"); + static_assert(kStageK % kTileK == 0, + "K stage must be a multiple of the DUMMA K unit"); + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kBlockN / sizeof(int4); + constexpr int kBVectors = kStageK * kBVectorsPerRow; + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWaveCols; + const int wave_col = wave - wave_row * kWaveCols; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Cooperatively stage A[kBlockM, kStageK] rows into LDS. Out-of-range + // (tail-M) rows are zero-filled so the DUMMA math stays defined. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + // Cooperatively stage B[kStageK, kBlockN] rows into LDS. The packed + // weight is the bootstrap identity [K, N] layout, so each B row is a + // contiguous chunk of kBlockN columns at row stride n. Global loads stay + // 16-byte vectorized and fully coalesced (one int4 per element); the LDS + // row stride kBStride = kBlockN + 4 is not 16-byte aligned by design (see + // kBLdsPad), so ds_write_b128 is unavailable and each int4 is committed + // as four int32 stores (identical total bank traffic to one b128). + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int kk = vec / kBVectorsPerRow; + const int v = vec - kk * kBVectorsPerRow; + const int4 b = *reinterpret_cast( + weight + (k0 + kk) * n + n0 + + v * static_cast(sizeof(int4))); + int32_t* b_row = reinterpret_cast( + b_tile + kk * kBStride + v * static_cast(sizeof(int4))); + b_row[0] = b.x; + b_row[1] = b.y; + b_row[2] = b.z; + b_row[3] = b.w; + } + __syncthreads(); + + // Compute the kTileK steps of this stage. Each wave owns a 32x32 + // quadrant: two A fragments (rows 0-15 / 16-31) and two B fragments + // (cols 0-15 / 16-31) per step, four du_mma_sync per step. The kk + // ordering (k0-outer, kk-inner step kTileK) keeps the int32 accumulation + // in ascending k order, bit-identical to the scalar fallback. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + + // Protect the single LDS buffer from being overwritten by the next + // stage's staging while this stage's fragment loads are still in flight. + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Double-buffered variant of the large-prefill path above (ported from the +// sibling-validated gate_up lineage; only instantiated as <64,128,64,2>). +// +// kBuffers == 1: single-buffered K loop (K stage 128, two barriers per stage): +// stage s+1's global loads cannot start until the stage-s barrier, so each +// stage serializes one global-load round trip in front of the compute. +// kBuffers == 2: software-pipelined K loop (K stage 64, one barrier per +// stage): stage s+1's loads are issued into registers before the stage-s +// MMA compute and committed to the other LDS buffer after it, overlapping +// the global round trip with the LDS-wait-bound compute. Requires at most +// one int4 per thread per operand (static_asserted), which holds for the +// exact shape. The (K, N) = (6144, 768) packed weight is transposed +// [N, K], so B is staged [N, K] and loaded with the col_major fragment +// loader (8 consecutive k-values per lane -> one ds_read2_b32 per fragment, +// ~4-way bank floor). +template +__global__ __launch_bounds__((kBlockM / 32) * (kBlockN / 32) * kWaveSize) void +w8a8_gemm_prefill_tiled_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWaveRows = kBlockM / 32; + constexpr int kWaveCols = kBlockN / 32; + constexpr int kThreads = kWaveRows * kWaveCols * kWaveSize; + // 16-byte bank skew: keeps int4 staging stores aligned while spreading the + // fragment rows across distinct LDS banks (~4-way floor for 8-byte/lane + // reads instead of the 16-way aliasing of power-of-two strides). + constexpr int kAStride = kStageK + sizeof(int4); + // Double-buffered path (exact shape) stages B transposed [N, K] so the + // DUMMA B fragments read 8 consecutive k-values per lane (one ds_read2_b32, + // ~4-way) instead of 8 strided k-rows (eight ds_read_u8, 16-way: with any + // 16-byte-aligned row stride S, B k-rows 8 apart alias onto one bank group + // because (S/4)*8*g == 0 mod 32 for all g). The generic single-buffer path + // keeps [K, N] B and the row-major B loader. + constexpr int kBStride = (kBuffers == 2) ? (kStageK + sizeof(int4)) + : (kBlockN + sizeof(int4)); + constexpr int kBRows = (kBuffers == 2) ? kBlockN : kStageK; + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = + (kBuffers == 2) ? (kStageK / sizeof(int4)) : (kBlockN / sizeof(int4)); + constexpr int kBVectors = kBRows * kBVectorsPerRow; + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWaveCols; + const int wave_col = wave - wave_row * kWaveCols; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBuffers][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBuffers][kBRows * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + if constexpr (kBuffers == 2) { + static_assert(kAVectors <= kThreads && kBVectors <= kThreads, + "double-buffered path stages at most one int4 per thread"); + // The exact shape's packed weight is [N, K] (transposed by + // launch_pack_w8a8_weight), so B fragments load 8 consecutive k-values + // per lane from the [N, K] LDS tile via the col_major loader; the + // element-to-slot mapping (slot i = B[k = kk + 8*g + i][n = col + row]) + // is identical to the row-major loader on a [K, N] tile, so the v_mmac + // operand values and the exact int32 accumulation are unchanged. + DUFragment + b_frag_t0, b_frag_t1; + // Prologue: stage K tile 0 into buffer 0, then one barrier. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + // B is packed [N, K] for the exact shape: each thread stages 16 + // consecutive k-values of one n row (n-stride in global is k). + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + n_local * kBStride)[v] = + *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + const int n_stages = k / kStageK; + for (int s = 0; s < n_stages; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + const bool has_next = (s + 1) < n_stages; + // Issue the stage s+1 global loads now; the data is consumed by the + // ds_write after the compute, so the vmcnt wait lands after the MMA + // loop instead of stalling the front of the stage. + int4 a_reg{0, 0, 0, 0}; + int4 b_reg{0, 0, 0, 0}; + if (has_next) { + const int k1 = (s + 1) * kStageK; + if (tid < kAVectors) { + const int local_row = tid / kAVectorsPerRow; + const int v = tid - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + a_reg = global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k1 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + const int n_local = tid / kBVectorsPerRow; + const int v = tid - n_local * kBVectorsPerRow; + b_reg = *reinterpret_cast( + weight + (n0 + n_local) * k + k1 + + v * static_cast(sizeof(int4))); + } + + // Compute stage s from the buffer staged last iteration. + // Load the int8 fragments as raw 8-byte LDS reads into the fragment + // storage (same bytes, same element-to-slot mapping), so the compiler + // feeds the ds_read2_b32 pair straight to the v_mmac and the per-dword + // mask/OR reassembly VALU disappears from the steady state. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + load_fragment8( + a_frag0, a_tile[cur] + local_row * kAStride + kk, kAStride, + lane); + load_fragment8( + a_frag1, a_tile[cur] + (local_row + kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8( + b_frag_t0, b_tile[cur] + local_col * kBStride + kk, kBStride, + lane); + load_fragment8( + b_frag_t1, b_tile[cur] + (local_col + kTileN) * kBStride + kk, + kBStride, lane); + du_mma_sync(acc00, a_frag0, b_frag_t0, acc00); + du_mma_sync(acc01, a_frag0, b_frag_t1, acc01); + du_mma_sync(acc10, a_frag1, b_frag_t0, acc10); + du_mma_sync(acc11, a_frag1, b_frag_t1, acc11); + } + + // Commit the prefetched stage s+1 into the other LDS buffer, then one + // barrier: it orders both this iteration's compute reads of buffer cur + // (against the next-next prefetch, which reuses cur) and the prefetch + // writes of buffer nxt (against the next iteration's compute reads). + if (has_next) { + if (tid < kAVectors) { + const int local_row = tid / kAVectorsPerRow; + const int v = tid - local_row * kAVectorsPerRow; + reinterpret_cast(a_tile[nxt] + local_row * kAStride)[v] = + a_reg; + } + const int n_local = tid / kBVectorsPerRow; + const int v = tid - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[nxt] + n_local * kBStride)[v] = b_reg; + } + __syncthreads(); + } + } else { + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[kBlockM, kStageK] into LDS (zero-filled tail M rows). + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + // Stage B[kStageK, kBlockN] into LDS from the packed weight. The + // generic path keeps the identity [K, N] layout (pack_weight is an + // identity copy for every (k, n) pair other than the exact shape). + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int kk = vec / kBVectorsPerRow; + const int v = vec - kk * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + kk * kBStride)[v] = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + du_load_matrix_sync( + a_frag0, a_tile[0] + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile[0] + (local_row + kTileM) * kAStride + kk, + kAStride); + du_load_matrix_sync( + b_frag0, b_tile[0] + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile[0] + kk * kBStride + local_col + kTileN, + kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// 8-wave / 4-accumulator variant of the exact-shape double-buffered 64x128 +// kernel (w8a8_gemm_prefill_tiled_kernel_w8, instantiated as <64,128,64>). +// The macro-tile, K stage 64 double buffer (one barrier/stage), [N,K] packed +// B, bank-skewed LDS layout (30,720 B/block), staging bytes, element-to-slot +// fragment mapping, byte order, and the k0-outer/kk-inner int32 accumulation +// order are IDENTICAL to the sibling-accepted 4-wave kernel, so the result is +// bit-identical; the only deltas are the wave partition and the per-thread +// staging split: +// * 512 threads (8 wavefronts); each wave owns a 32x32 quadrant (two +// 16-row bands x two 16-col bands) = FOUR independent m16n16k32 +// accumulators (4-deep MMAC ILP per kk slice). Per-wavefront MMACs per +// stage = 8 (512 MMACs/wavefront over a K=4096 loop). +// * A staging is one int4/thread for the first 256 threads, B staging one +// int4/thread for all 512 threads (same 12,288 B/stage/block). +// Grid (N/128) x (M/64) = 6 x 64 = 384 blocks; LDS 30,720 B x 2 = 61,440 B +// <= 64 KiB; target <= 64 VGPR x 512 x 2 <= 512 KiB -> 2 blocks/CU -> 4 +// waves/SIMD (must be confirmed in the code object; > 64 VGPR degenerates to +// 1 block/CU = 2 waves/SIMD). Loop-invariant staging/fragment addresses are +// hoisted before the K loop, and the eight fragment loads of each stage are +// pinned ahead of the eight-MMAC burst. +template +__global__ __launch_bounds__(8 * kWaveSize) void +w8a8_gemm_prefill_tiled_kernel_w8( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kThreads = 8 * kWaveSize; + constexpr int kAStride = kStageK + sizeof(int4); + constexpr int kBStride = kStageK + sizeof(int4); + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kStageK / sizeof(int4); + constexpr int kBVectors = kBlockN * kBVectorsPerRow; + static_assert(kAVectors <= kThreads, + "w8 path stages A with at most one int4 per thread"); + static_assert(kBVectors == kThreads, + "w8 path stages exactly one B int4 per thread"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / 4; + const int wave_col = wave - wave_row * 4; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBStride]; + + DUFragment acc[2][2]; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 2; ++c) { + du_fill_fragment(acc[r][c], 0); + } + } + + // Prologue: stage K tile 0 into buffer 0 (A: first 256 threads, B: all 512 + // threads; same vectors and addresses as the generic kernel distributed + // over 512 threads), then one barrier. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + n_local * kBStride)[v] = + *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + const int n_stages = k / kStageK; + // Hoist every loop-invariant address of the steady state out of the K + // loop. The A staging pointer is only dereferenced by threads tid < + // kAVectors. The three global staging pointers (A/B) walk +kStageK per + // stage; the LDS commit slots add only the buffer offset nxt * kBufInt4; + // the fragment reads use the per-wave buffer-0 bases and add only the + // buffer offset cur * kBufBytes and the unrolled kk slice. + const bool a_active = tid < kAVectors; + const int local_row_a = tid / kAVectorsPerRow; + const int v_a = tid - local_row_a * kAVectorsPerRow; + const int global_row_a = m0 + local_row_a; + const bool a_row_valid = global_row_a < m; + const int8_t* __restrict__ a_g = + x_q + static_cast(global_row_a) * k + + v_a * static_cast(sizeof(int4)) + kStageK; + const int n_local = tid / kBVectorsPerRow; + const int v_b = tid - n_local * kBVectorsPerRow; + const int8_t* __restrict__ b_g = + weight + static_cast(n0 + n_local) * k + + v_b * static_cast(sizeof(int4)) + kStageK; + int4* const a_wslot = + reinterpret_cast(a_tile[0] + local_row_a * kAStride) + v_a; + int4* const b_wslot = + reinterpret_cast(b_tile[0] + n_local * kBStride) + v_b; + constexpr int kABufInt4 = + (kBlockM * kAStride) / static_cast(sizeof(int4)); + constexpr int kBBufInt4 = + (kBlockN * kBStride) / static_cast(sizeof(int4)); + // Per-wave fragment-read bases in buffer 0 (32x32 quadrant per wave); the + // loop adds the buffer offset cur * kBufBytes and the kk slice. + const int8_t* const a_rd = a_tile[0] + (wave_row * 32) * kAStride; + const int8_t* const b_rd = b_tile[0] + (wave_col * 32) * kBStride; + constexpr int kABufBytes = kBlockM * kAStride; + constexpr int kBBufBytes = kBlockN * kBStride; + + // Prologue prefetch: issue the tile-1 global loads into the loop-carried + // payload now (a_g/b_g already address tile 1). The prologue barrier above + // already ordered the tile-0 staging; the publish of this payload happens + // at the top of stage 0, one full stage later. + int4 a_pld{0, 0, 0, 0}; + int4 b_pld{0, 0, 0, 0}; + if (a_active) { + a_pld = a_row_valid ? *reinterpret_cast(a_g) + : int4{0, 0, 0, 0}; + } + b_pld = *reinterpret_cast(b_g); + a_g += kStageK; + b_g += kStageK; + + // Steady state (stages 0..n_stages-2): publish-at-top double buffering. + // The tile-(s+1) payload prefetched one iteration earlier is written into + // the idle LDS buffer at the TOP of stage s, BEFORE the MMAC burst, so the + // ds_write latency (and the already-satisfied vmcnt wait) overlaps the + // compute instead of sitting between the MMACs and the barrier (the exact + // code object shows both ds_write_b128 after the v_mmac group, immediately + // before s_barrier). The tile-(s+2) prefetch then issues from the walking + // pointers; its vmcnt wait lands at the next stage's publish. One barrier + // per stage orders both this stage's fragment reads of buffer cur (against + // the next publish into it) and this stage's publish into buffer nxt + // (against the next stage's reads of it) -- identical semantics to the old + // commit-after-compute loop, so every global load, LDS write, fragment + // read, and v_mmac is unchanged and the int32 accumulation stays + // bit-identical. + DUFragment + a_frag00, a_frag01, a_frag10, a_frag11; + DUFragment + b_frag00, b_frag01, b_frag10, b_frag11; + static_assert(kStageK == 2 * kTileK, + "w8 burst schedule expects exactly two kk slices per stage"); + for (int s = 0; s < n_stages - 1; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + // Publish tile s+1 into buffer nxt at the top of the stage. + if (a_active) { + *(a_wslot + nxt * kABufInt4) = a_pld; + } + *(b_wslot + nxt * kBBufInt4) = b_pld; + // Prefetch tile s+2 (skipped on the last loop iteration; the peeled + // final stage below needs no payload). + if (s + 2 < n_stages) { + if (a_active) { + a_pld = a_row_valid ? *reinterpret_cast(a_g) + : int4{0, 0, 0, 0}; + } + b_pld = *reinterpret_cast(b_g); + a_g += kStageK; + b_g += kStageK; + } + + // Compute stage s: each wave owns a 32x32 quadrant (two 16-row bands x + // two 16-col bands), four independent m16n16k32 accumulators. The + // k0-outer/kk-inner order per output element is unchanged, so the int32 + // accumulation is bit-identical to the generic kernel. + const int8_t* const a_cur = a_rd + cur * kABufBytes; + const int8_t* const b_cur = b_rd + cur * kBBufBytes; + // kk = 0 slice loads (row bands x0/x1, col bands 0..1). + load_fragment8(a_frag00, a_cur, kAStride, lane); + load_fragment8(a_frag01, a_cur + kTileM * kAStride, kAStride, lane); + load_fragment8(b_frag00, b_cur, kBStride, lane); + load_fragment8(b_frag01, b_cur + kTileN * kBStride, kBStride, lane); + // kk = kTileK slice loads (same rows/cols, k offset kTileK). + load_fragment8(a_frag10, a_cur + kTileK, kAStride, lane); + load_fragment8( + a_frag11, a_cur + kTileM * kAStride + kTileK, kAStride, lane); + load_fragment8(b_frag10, b_cur + kTileK, kBStride, lane); + load_fragment8( + b_frag11, b_cur + kTileN * kBStride + kTileK, kBStride, lane); + // kk = 0 MMAC group. + du_mma_sync(acc[0][0], a_frag00, b_frag00, acc[0][0]); + du_mma_sync(acc[0][1], a_frag00, b_frag01, acc[0][1]); + du_mma_sync(acc[1][0], a_frag01, b_frag00, acc[1][0]); + du_mma_sync(acc[1][1], a_frag01, b_frag01, acc[1][1]); + // kk = kTileK MMAC group (same accumulators, same order). + du_mma_sync(acc[0][0], a_frag10, b_frag10, acc[0][0]); + du_mma_sync(acc[0][1], a_frag10, b_frag11, acc[0][1]); + du_mma_sync(acc[1][0], a_frag11, b_frag10, acc[1][0]); + du_mma_sync(acc[1][1], a_frag11, b_frag11, acc[1][1]); + + __syncthreads(); + } + + // Peeled final stage: compute tile n_stages-1 from buffer (n_stages-1)&1, + // which the last loop iteration published above its barrier. No prefetch, + // no publish, and no dead barrier after the last stage: each wave's + // epilogue writes its own disjoint 32x32 output quadrant, so no + // cross-wave ordering is required. + { + const int cur = (n_stages - 1) & 1; + const int8_t* const a_cur = a_rd + cur * kABufBytes; + const int8_t* const b_cur = b_rd + cur * kBBufBytes; + // kk = 0 slice loads (row bands x0/x1, col bands 0..1). + load_fragment8(a_frag00, a_cur, kAStride, lane); + load_fragment8(a_frag01, a_cur + kTileM * kAStride, kAStride, lane); + load_fragment8(b_frag00, b_cur, kBStride, lane); + load_fragment8(b_frag01, b_cur + kTileN * kBStride, kBStride, lane); + // kk = kTileK slice loads (same rows/cols, k offset kTileK). + load_fragment8(a_frag10, a_cur + kTileK, kAStride, lane); + load_fragment8( + a_frag11, a_cur + kTileM * kAStride + kTileK, kAStride, lane); + load_fragment8(b_frag10, b_cur + kTileK, kBStride, lane); + load_fragment8( + b_frag11, b_cur + kTileN * kBStride + kTileK, kBStride, lane); + // kk = 0 MMAC group. + du_mma_sync(acc[0][0], a_frag00, b_frag00, acc[0][0]); + du_mma_sync(acc[0][1], a_frag00, b_frag01, acc[0][1]); + du_mma_sync(acc[1][0], a_frag01, b_frag00, acc[1][0]); + du_mma_sync(acc[1][1], a_frag01, b_frag01, acc[1][1]); + // kk = kTileK MMAC group (same accumulators, same order). + du_mma_sync(acc[0][0], a_frag10, b_frag10, acc[0][0]); + du_mma_sync(acc[0][1], a_frag10, b_frag11, acc[0][1]); + du_mma_sync(acc[1][0], a_frag11, b_frag10, acc[1][0]); + du_mma_sync(acc[1][1], a_frag11, b_frag11, acc[1][1]); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 2; ++c) { + store_prefill_fragment( + acc[r][c], x_scale, weight_scale, out, m, n, + base_row + r * kTileM, base_col + c * kTileN, lane); + } + } +} + +// 4-wave / 8-accumulator variant of the exact-shape double-buffered 64x128 +// kernel (w8a8_gemm_prefill_tiled_kernel_w4, instantiated as <64,128,64>). +// The macro-tile, K stage 64 double buffer (one barrier/stage, publish-at-top, +// peeled final stage), [N,K] packed B, bank-skewed LDS layout (30,720 +// B/block), staging bytes, element-to-slot fragment mapping, byte order, and +// the k0-outer/kk-inner int32 accumulation order are IDENTICAL to the +// accepted 8-wave kernel, so the result is bit-identical; the only deltas are +// the wave partition and the per-thread staging split: +// * 256 threads (4 wavefronts); each wave owns a 32x64 quadrant (two +// 16-row bands x four 16-col bands) = EIGHT independent m16n16k32 int32 +// accumulators (8-deep MMAC ILP per kk slice). Per block-stage the LDS +// fragment reads drop from 8 waves x 8 = 64 to 4 waves x 12 = 48 +// load_fragment8 (six ds_read2_b64 per wave) and the LDS read BYTES drop +// from 32 KB to 24 KB: the 32x32 quadrant layout re-reads A 4x and B 2x, +// the 32x64 quadrant re-reads both 2x (per-wave per stage: 4 A fragments +// + 8 B fragments), so per block-stage the LDS read instruction count +// drops 64 -> 48 and the family lds_instructions 1.62M -> ~1.33M. +// * A staging is one int4/thread for all 256 threads, B staging two +// int4/thread for all 256 threads (same 12,288 B/stage/block). +// Grid (N/128) x (M/64) = 6 x 64 = 384 blocks; LDS 30,720 B x 2 = 61,440 B +// <= 64 KiB; 256 threads x 2 blocks = 8 waves/CU = 2 waves/SIMD at +// arch_vgpr <= 128 (8 accumulators = 32 VGPR + 12 loop-carried payload + +// 12 in-flight fragment VGPR + hoisted addresses ~= 80; 2 x 256 x 128 = +// 65,536 exactly fits; > 128 VGPR degenerates to 1 block/CU and predicts a +// regression -- falsifiable). Loop-invariant staging/fragment addresses are +// hoisted before the K loop, and the twelve fragment loads of each stage are +// pinned ahead of the sixteen-MMAC burst. +template +__global__ __launch_bounds__(4 * kWaveSize) void +w8a8_gemm_prefill_tiled_kernel_w4( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kThreads = 4 * kWaveSize; + constexpr int kAStride = kStageK + sizeof(int4); + constexpr int kBStride = kStageK + sizeof(int4); + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kStageK / sizeof(int4); + constexpr int kBVectors = kBlockN * kBVectorsPerRow; + static_assert(kAVectors == kThreads, + "w4 path stages A with exactly one int4 per thread"); + static_assert(kBVectors == 2 * kThreads, + "w4 path stages B with exactly two int4 per thread"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBStride]; + + DUFragment acc[2][4]; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 4; ++c) { + du_fill_fragment(acc[r][c], 0); + } + } + + // Prologue: stage K tile 0 into buffer 0 (A: one int4/thread, B: two + // int4/thread; the same vectors and addresses as the 8-wave kernel + // distributed over 256 threads), then one barrier. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + n_local * kBStride)[v] = + *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + const int n_stages = k / kStageK; + // Hoist every loop-invariant address of the steady state out of the K + // loop (same scheme as the 8-wave kernel; B staging is two int4/thread: + // vectors tid for rows tid/4 (0..63) and tid+256 for rows 64 + tid/4). + const int local_row_a = tid / kAVectorsPerRow; + const int v_a = tid - local_row_a * kAVectorsPerRow; + const int global_row_a = m0 + local_row_a; + const bool a_row_valid = global_row_a < m; + const int8_t* __restrict__ a_g = + x_q + static_cast(global_row_a) * k + + v_a * static_cast(sizeof(int4)) + kStageK; + const int n_local0 = tid / kBVectorsPerRow; + const int v_b0 = tid - n_local0 * kBVectorsPerRow; + const int8_t* __restrict__ b_g0 = + weight + static_cast(n0 + n_local0) * k + + v_b0 * static_cast(sizeof(int4)) + kStageK; + const int vec1 = tid + kThreads; + const int n_local1 = vec1 / kBVectorsPerRow; + const int v_b1 = vec1 - n_local1 * kBVectorsPerRow; + const int8_t* __restrict__ b_g1 = + weight + static_cast(n0 + n_local1) * k + + v_b1 * static_cast(sizeof(int4)) + kStageK; + int4* const a_wslot = + reinterpret_cast(a_tile[0] + local_row_a * kAStride) + v_a; + int4* const b_wslot0 = + reinterpret_cast(b_tile[0] + n_local0 * kBStride) + v_b0; + int4* const b_wslot1 = + reinterpret_cast(b_tile[0] + n_local1 * kBStride) + v_b1; + constexpr int kABufInt4 = + (kBlockM * kAStride) / static_cast(sizeof(int4)); + constexpr int kBBufInt4 = + (kBlockN * kBStride) / static_cast(sizeof(int4)); + // Per-wave fragment-read bases in buffer 0 (32x64 quadrant per wave); the + // loop adds the buffer offset cur * kBufBytes and the kk slice. + const int8_t* const a_rd = a_tile[0] + (wave_row * 32) * kAStride; + const int8_t* const b_rd = b_tile[0] + (wave_col * 64) * kBStride; + constexpr int kABufBytes = kBlockM * kAStride; + constexpr int kBBufBytes = kBlockN * kBStride; + + // Prologue prefetch: issue the tile-1 global loads into the loop-carried + // payload now (a_g/b_g0/b_g1 already address tile 1). The publish of this + // payload happens at the top of stage 0, one full stage later. + int4 a_pld{0, 0, 0, 0}; + int4 b_pld0{0, 0, 0, 0}; + int4 b_pld1{0, 0, 0, 0}; + a_pld = a_row_valid ? *reinterpret_cast(a_g) + : int4{0, 0, 0, 0}; + b_pld0 = *reinterpret_cast(b_g0); + b_pld1 = *reinterpret_cast(b_g1); + a_g += kStageK; + b_g0 += kStageK; + b_g1 += kStageK; + + // Steady state (stages 0..n_stages-2): publish-at-top double buffering + // with identical semantics to the 8-wave kernel (one barrier per stage; + // 1 prologue + 95 loop = 96 dynamic barriers/block for K=6144, stage 64). + DUFragment + a_frag00, a_frag01, a_frag10, a_frag11; + DUFragment + b_frag00, b_frag01, b_frag02, b_frag03; + DUFragment + b_frag10, b_frag11, b_frag12, b_frag13; + static_assert(kStageK == 2 * kTileK, + "w4 burst schedule expects exactly two kk slices per stage"); + for (int s = 0; s < n_stages - 1; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + // Publish tile s+1 into buffer nxt at the top of the stage. + *(a_wslot + nxt * kABufInt4) = a_pld; + *(b_wslot0 + nxt * kBBufInt4) = b_pld0; + *(b_wslot1 + nxt * kBBufInt4) = b_pld1; + // Prefetch tile s+2 (skipped on the last loop iteration; the peeled + // final stage below needs no payload). + if (s + 2 < n_stages) { + a_pld = a_row_valid ? *reinterpret_cast(a_g) + : int4{0, 0, 0, 0}; + b_pld0 = *reinterpret_cast(b_g0); + b_pld1 = *reinterpret_cast(b_g1); + a_g += kStageK; + b_g0 += kStageK; + b_g1 += kStageK; + } + + // Compute stage s: each wave owns a 32x64 quadrant (two 16-row bands x + // four 16-col bands), eight independent m16n16k32 accumulators. The + // k0-outer/kk-inner order per output element is unchanged, so the int32 + // accumulation is bit-identical to the 8-wave/generic kernels. + const int8_t* const a_cur = a_rd + cur * kABufBytes; + const int8_t* const b_cur = b_rd + cur * kBBufBytes; + // kk = 0 slice loads (row bands x0/x1, col bands 0..3). + load_fragment8(a_frag00, a_cur, kAStride, lane); + load_fragment8(a_frag01, a_cur + kTileM * kAStride, kAStride, lane); + load_fragment8(b_frag00, b_cur, kBStride, lane); + load_fragment8(b_frag01, b_cur + kTileN * kBStride, kBStride, lane); + load_fragment8(b_frag02, b_cur + 2 * kTileN * kBStride, kBStride, lane); + load_fragment8(b_frag03, b_cur + 3 * kTileN * kBStride, kBStride, lane); + // kk = kTileK slice loads (same rows/cols, k offset kTileK). + load_fragment8(a_frag10, a_cur + kTileK, kAStride, lane); + load_fragment8( + a_frag11, a_cur + kTileM * kAStride + kTileK, kAStride, lane); + load_fragment8(b_frag10, b_cur + kTileK, kBStride, lane); + load_fragment8( + b_frag11, b_cur + kTileN * kBStride + kTileK, kBStride, lane); + load_fragment8( + b_frag12, b_cur + 2 * kTileN * kBStride + kTileK, kBStride, lane); + load_fragment8( + b_frag13, b_cur + 3 * kTileN * kBStride + kTileK, kBStride, lane); + // kk = 0 MMAC group. + du_mma_sync(acc[0][0], a_frag00, b_frag00, acc[0][0]); + du_mma_sync(acc[0][1], a_frag00, b_frag01, acc[0][1]); + du_mma_sync(acc[0][2], a_frag00, b_frag02, acc[0][2]); + du_mma_sync(acc[0][3], a_frag00, b_frag03, acc[0][3]); + du_mma_sync(acc[1][0], a_frag01, b_frag00, acc[1][0]); + du_mma_sync(acc[1][1], a_frag01, b_frag01, acc[1][1]); + du_mma_sync(acc[1][2], a_frag01, b_frag02, acc[1][2]); + du_mma_sync(acc[1][3], a_frag01, b_frag03, acc[1][3]); + // kk = kTileK MMAC group (same accumulators, same order). + du_mma_sync(acc[0][0], a_frag10, b_frag10, acc[0][0]); + du_mma_sync(acc[0][1], a_frag10, b_frag11, acc[0][1]); + du_mma_sync(acc[0][2], a_frag10, b_frag12, acc[0][2]); + du_mma_sync(acc[0][3], a_frag10, b_frag13, acc[0][3]); + du_mma_sync(acc[1][0], a_frag11, b_frag10, acc[1][0]); + du_mma_sync(acc[1][1], a_frag11, b_frag11, acc[1][1]); + du_mma_sync(acc[1][2], a_frag11, b_frag12, acc[1][2]); + du_mma_sync(acc[1][3], a_frag11, b_frag13, acc[1][3]); + + __syncthreads(); + } + + // Peeled final stage: compute tile n_stages-1 from buffer (n_stages-1)&1, + // which the last loop iteration published above its barrier. No prefetch, + // no publish, and no dead barrier after the last stage: each wave's + // epilogue writes its own disjoint 32x64 output quadrant, so no + // cross-wave ordering is required. + { + const int cur = (n_stages - 1) & 1; + const int8_t* const a_cur = a_rd + cur * kABufBytes; + const int8_t* const b_cur = b_rd + cur * kBBufBytes; + // kk = 0 slice loads (row bands x0/x1, col bands 0..3). + load_fragment8(a_frag00, a_cur, kAStride, lane); + load_fragment8(a_frag01, a_cur + kTileM * kAStride, kAStride, lane); + load_fragment8(b_frag00, b_cur, kBStride, lane); + load_fragment8(b_frag01, b_cur + kTileN * kBStride, kBStride, lane); + load_fragment8(b_frag02, b_cur + 2 * kTileN * kBStride, kBStride, lane); + load_fragment8(b_frag03, b_cur + 3 * kTileN * kBStride, kBStride, lane); + // kk = kTileK slice loads (same rows/cols, k offset kTileK). + load_fragment8(a_frag10, a_cur + kTileK, kAStride, lane); + load_fragment8( + a_frag11, a_cur + kTileM * kAStride + kTileK, kAStride, lane); + load_fragment8(b_frag10, b_cur + kTileK, kBStride, lane); + load_fragment8( + b_frag11, b_cur + kTileN * kBStride + kTileK, kBStride, lane); + load_fragment8( + b_frag12, b_cur + 2 * kTileN * kBStride + kTileK, kBStride, lane); + load_fragment8( + b_frag13, b_cur + 3 * kTileN * kBStride + kTileK, kBStride, lane); + // kk = 0 MMAC group. + du_mma_sync(acc[0][0], a_frag00, b_frag00, acc[0][0]); + du_mma_sync(acc[0][1], a_frag00, b_frag01, acc[0][1]); + du_mma_sync(acc[0][2], a_frag00, b_frag02, acc[0][2]); + du_mma_sync(acc[0][3], a_frag00, b_frag03, acc[0][3]); + du_mma_sync(acc[1][0], a_frag01, b_frag00, acc[1][0]); + du_mma_sync(acc[1][1], a_frag01, b_frag01, acc[1][1]); + du_mma_sync(acc[1][2], a_frag01, b_frag02, acc[1][2]); + du_mma_sync(acc[1][3], a_frag01, b_frag03, acc[1][3]); + // kk = kTileK MMAC group (same accumulators, same order). + du_mma_sync(acc[0][0], a_frag10, b_frag10, acc[0][0]); + du_mma_sync(acc[0][1], a_frag10, b_frag11, acc[0][1]); + du_mma_sync(acc[0][2], a_frag10, b_frag12, acc[0][2]); + du_mma_sync(acc[0][3], a_frag10, b_frag13, acc[0][3]); + du_mma_sync(acc[1][0], a_frag11, b_frag10, acc[1][0]); + du_mma_sync(acc[1][1], a_frag11, b_frag11, acc[1][1]); + du_mma_sync(acc[1][2], a_frag11, b_frag12, acc[1][2]); + du_mma_sync(acc[1][3], a_frag11, b_frag13, acc[1][3]); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 64; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 4; ++c) { + store_prefill_fragment( + acc[r][c], x_scale, weight_scale, out, m, n, + base_row + r * kTileM, base_col + c * kTileN, lane); + } + } +} + +// Exact-shape 2-D macro-tile DUMMA baseline for the TP8 shared down_proj +// shape (M=4096, N=6144, K=384): double-buffered K stage 64 with one +// __syncthreads per stage (prologue + n_stages = 7 barriers/block), A +// staged row-major with the 16-byte bank-skew stride kStageK+16 = 80, B +// staged n-major [N, K] (the packed layout for (k, n) = (384, 6144)) with +// the 8-byte skew stride kStageK+8 = 72 (conflict-free ds_read2_b64 B +// fragments: 16 lanes x 8 bytes land on 16 distinct bank phases), and +// direct 8-byte fragment loads (load_fragment8) for BOTH operands so the +// compiler feeds the ds_read2_b64 pair straight to the v_mmac without the +// per-dword mask/OR reassembly of du_load_matrix_sync. Each wave owns a +// 32x32 quadrant (four m16n16k32 int32 accumulators resident over K); the +// int32 accumulation order (k0-outer stage 64, kk-inner kTileK, ascending +// k) and the element-to-slot fragment mapping are identical to the generic +// family, so the result is bit-identical. +// +// Template family over (kBlockM, kBlockN) with kStageK = 64: +// * <64, 128, 64>: 8 waves / 512 threads, grid (N/128) x ceil(M/64), +// LDS 28,672 B/block -> 2 blocks/CU. Selected for the exact shape +// (iteration-4 tile-aspect round: same per-block global bytes and +// identical fragment-read mix as <128,64,64>, 2 blocks/CU retained). +// * <128, 64, 64>: 8 waves / 512 threads, grid (N/64) x ceil(M/128), +// LDS 29,696 B/block -> 2 blocks/CU (TP4 K=384 lineage winner; the +// sibling arm if the iteration-4 flip regresses). +// * <64, 64, 64>: 4 waves / 256 threads, LDS 19,456 B/block. +// B is staged [N, K] from the packed weight: each thread commits one int4 +// as two int64 halves (kBStride % 16 == 8, so no ds_write_b128), the same +// staging form the TP4 down_proj lineage validated at exactly K=384. +template +__global__ __launch_bounds__((kBlockM / 32) * (kBlockN / 32) * kWaveSize) void +w8a8_dumma_prefill_packedb_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWaveRows = kBlockM / 32; + constexpr int kWaveCols = kBlockN / 32; + constexpr int kThreads = kWaveRows * kWaveCols * kWaveSize; + constexpr int kAStride = kStageK + sizeof(int4); + constexpr int kBStride = kStageK + sizeof(int64_t); + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kStageK / sizeof(int4); + constexpr int kBVectors = kBlockN * kBVectorsPerRow; + static_assert(kStageK % kTileK == 0, + "K stage must be a multiple of the DUMMA K unit"); + static_assert(kAStride % sizeof(int4) == 0, + "A LDS row stride must stay 16-byte aligned"); + static_assert(kBStride % sizeof(int64_t) == 0, + "B LDS row stride must stay 8-byte aligned"); + static_assert(kAVectors <= kThreads && kBVectors <= kThreads, + "packed-B path stages at most one int4 per thread per operand"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWaveCols; + const int wave_col = wave - wave_row * kWaveCols; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + // Prologue: stage K tile 0 into buffer 0, then one barrier. A is + // [kBlockM, kStageK] row-major (zero-filled tail-M rows); B is + // [kBlockN, kStageK] n-major from the packed weight, one int4/thread + // committed as two int64 halves. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + const int4 b = *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + uint64_t* const b_row = reinterpret_cast( + b_tile[0] + n_local * kBStride + v * static_cast(sizeof(int4))); + b_row[0] = static_cast(static_cast(b.x)) | + (static_cast(static_cast(b.y)) << 32); + b_row[1] = static_cast(static_cast(b.z)) | + (static_cast(static_cast(b.w)) << 32); + } + __syncthreads(); + + const int n_stages = k / kStageK; + for (int s = 0; s < n_stages; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + const bool has_next = (s + 1) < n_stages; + // Issue the stage s+1 global loads now; the commit after the MMAC + // burst consumes them, so the vmcnt wait lands after the compute + // instead of stalling the front of the stage. + int4 a_reg{0, 0, 0, 0}; + int4 b_reg{0, 0, 0, 0}; + if (has_next) { + const int k1 = (s + 1) * kStageK; + if (tid < kAVectors) { + const int local_row = tid / kAVectorsPerRow; + const int v = tid - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + a_reg = global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k1 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + if (tid < kBVectors) { + const int n_local = tid / kBVectorsPerRow; + const int v = tid - n_local * kBVectorsPerRow; + b_reg = *reinterpret_cast( + weight + (n0 + n_local) * k + k1 + + v * static_cast(sizeof(int4))); + } + } + + // Compute stage s: each wave owns a 32x32 quadrant; the direct 8-byte + // LDS fragment reads keep the exact element-to-slot mapping of + // du_load_matrix_sync (row_major A, col_major B on the [N, K] tile), + // so the v_mmac operand values and the int32 accumulation are + // bit-identical to the generic family. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + load_fragment8( + a_frag0, a_tile[cur] + local_row * kAStride + kk, kAStride, lane); + load_fragment8( + a_frag1, a_tile[cur] + (local_row + kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8( + b_frag0, b_tile[cur] + local_col * kBStride + kk, kBStride, lane); + load_fragment8( + b_frag1, b_tile[cur] + (local_col + kTileN) * kBStride + kk, + kBStride, lane); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + + // Commit the prefetched stage s+1 into the other LDS buffer, then one + // barrier ordering both this stage's compute reads of buffer cur and + // the prefetch writes of buffer nxt. + if (has_next) { + if (tid < kAVectors) { + const int local_row = tid / kAVectorsPerRow; + const int v = tid - local_row * kAVectorsPerRow; + reinterpret_cast(a_tile[nxt] + local_row * kAStride)[v] = + a_reg; + } + if (tid < kBVectors) { + const int n_local = tid / kBVectorsPerRow; + const int v = tid - n_local * kBVectorsPerRow; + uint64_t* const b_row = reinterpret_cast( + b_tile[nxt] + n_local * kBStride + + v * static_cast(sizeof(int4))); + b_row[0] = static_cast(static_cast(b_reg.x)) | + (static_cast(static_cast(b_reg.y)) + << 32); + b_row[1] = static_cast(static_cast(b_reg.z)) | + (static_cast(static_cast(b_reg.w)) + << 32); + } + } + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Generic scalar fallback: one output element per thread with exact int32 +// accumulation in k order. Correct for any (m, n, k), including the small-M +// API shapes (M=2/M=16) and any unmatched geometry. The exact (k, n) = +// (6144, 768) and (384, 6144) pairs read the transposed [N, K] packed +// weight (the pack specialization is keyed on (k, n), so every consumer of +// those packed pairs, including the paired M=2/M=16 decode shapes that +// reach this fallback, must interpret them transposed). b_col_stride +// carries the identity [K, N] column stride (n) for all other pairs. +__global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k, + int b_col_stride) { + const bool packed_transposed = + (k == 6144 && n == 768) || (k == 384 && n == 6144); + const int64_t total = static_cast(m) * n; + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = x_q + static_cast(row) * k; + // packed[n * k + kk] = raw[kk * n + n]: the column's k values are + // contiguous at packed + col * k. + const int8_t* b_col = packed_transposed + ? weight + static_cast(col) * k + : weight + static_cast(col) * b_col_stride; + int32_t acc = 0; + if (packed_transposed) { + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[kk]); + } + } else { + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast( + b_col[static_cast(kk) * b_col_stride]); + } + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); +} + +// pack_weight. The packed layout is opaque per the API contract; for the +// exact (K=6144, N=768) pair the weight is stored transposed [N, K] +// (n-major) so the double-buffered GEMM stages B in [N, K] LDS order and the +// DUMMA B fragments read 8 consecutive k-values per lane. Every other (k, n) +// pair keeps the identity [K, N] copy (generic fallback unchanged). Runs +// outside the timed/CUDA-Graph region. +__global__ __launch_bounds__(kScalarThreads) void +w8a8_pack_kernel( + const int8_t* __restrict__ raw_weight, + const float* __restrict__ weight_scale, + int8_t* __restrict__ packed_weight, + float* __restrict__ packed_weight_scale, + int64_t weight_elems, + int n, + int k) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if ((k == 6144 && n == 768) || (k == 384 && n == 6144)) { + // Transpose [K, N] -> [N, K]: packed[n * k + kk] = raw[kk * n + n]. + if (idx < weight_elems) { + const int64_t n_idx = idx / static_cast(k); + const int64_t k_idx = idx - n_idx * static_cast(k); + packed_weight[idx] = raw_weight[k_idx * n + n_idx]; + } + } else if (idx < weight_elems) { + packed_weight[idx] = raw_weight[idx]; + } + if (idx < n) { + packed_weight_scale[idx] = weight_scale[idx]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Host launch symbols (extern "C", consumed by csrc/bindings.cpp). +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // No split-K: the workspace is not touched. The timed operator performs no + // allocation, synchronization, packing, or default-stream launch; it + // writes only the caller-provided out. + (void)workspace; + (void)workspace_bytes; + auto* out_bf16 = static_cast(out); + + // Native INT8 DUMMA m16n16k32 tiled path for large-M prefill. All three + // macro-tile instantiations share the exact guard; the launch geometry is + // selected explicitly so every assigned shape is covered (no + // single-shape specialization). The guard divisibility (k % 128 == 0, + // n % 64 == 0) holds for both assigned shapes and for every supported + // catalog shape. + if (m >= kPrefillMinM && (n % kDefaultBlockN == 0) && + (k % kPrefillStageK == 0)) { + if (n == 768 && k == 6144) { + // Packed (k, n) = (6144, 768) weight is transposed [N, K], so every + // large-M shape with this pair uses the double-buffered transposed-B + // kernel (the generic [K, N] arms below would misinterpret it). Exact + // assigned shape (M=4096, N=768, K=6144): 64x128 macro-tile, grid + // 6x64 = 384 blocks, K stage 64 double-buffered, 30,720 B LDS/block + // -> 2 blocks/CU. + const dim3 grid( + static_cast(n / 128), + static_cast((m + 63) / 64)); + if (m == 4096) { + // 4-wave variant w4<64,128,64> (256 threads, 32x64 quadrant per + // wave, eight accumulators, 96 K stages of 64; 25% fewer LDS + // fragment reads per block-stage than the 8-wave arm). + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel_w4<64, 128, 64>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + // Other large-M shapes with the packed (k, n) pair keep the generic + // 8-wave double-buffered <64,128,64,2> kernel (bit-identical path). + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel<64, 128, 64, 2>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (n == 6144 && k == 384) { + // Packed (k, n) = (384, 6144) weight is n-major [N, K] (transposed by + // launch_pack_w8a8_weight), so every large-M shape with this pair uses + // the double-buffered packed-B kernel (the generic [K, N] arms below + // would misinterpret it). Exact assigned shape (M=4096, N=6144, + // K=384): 64x128 macro-tile (iteration-4 tile-aspect round), grid + // 48x64 = 3072 blocks, K stage 64 double-buffered, 28,672 B LDS/block + // -> 2 blocks/CU, 6 K stages, one barrier per stage (7 dynamic + // barriers/block). Aspect accounting vs the 128x64 arm: the B tile + // (64x128) is now shared by M/64 = 64 blocks (was 32) while the A band + // is shared by N/128 = 48 blocks (was 96); per-block global bytes + // (72 KB) and the per-stage fragment-read mix (32 A + 32 B loads) are + // invariant, so only the LDS write mix (A 256 ds_write_b128 + B 512 + // ds_write2_b64 vs the reverse) and the A:B L2 mix swap. + const dim3 grid( + static_cast(n / 128), + static_cast((m + 63) / 64)); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_kernel<64, 128, 64>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (n % 128 == 0) { + // 64x128 tile, 8 wavefronts / 512 threads, single-buffered 128-K stage + // (generic [K, N] packed-B arm for unmatched (k, n) pairs; the packed + // (6144, 768) and (384, 6144) pairs never reach this arm): + const dim3 grid( + static_cast(n / 128), + static_cast((m + 63) / 64)); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 128, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (m % 128 == 0) { + // 128x64 tile, 8 wavefronts / 512 threads. + const dim3 grid( + static_cast(n / 64), + static_cast(m / 128)); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<128, 64, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + { + // 64x64 tile, 4 wavefronts / 256 threads (generic default). + const dim3 grid( + static_cast(n / 64), + static_cast((m + 63) / 64)); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 64, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + } + + // Generic scalar fallback for every unmatched (m, n, k) and M < 128 + // (including the paired M=2/M=16 API shapes). The kernel decodes the + // transposed [N, K] pack for the exact (k, n) = (6144, 768) and + // (384, 6144) pairs and uses the identity [K, N] column stride n for + // every other pair. + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast( + (total + kScalarThreads - 1) / kScalarThreads)); + const dim3 block(static_cast(kScalarThreads)); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k, n); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Exact (K, N) = (6144, 768) and (384, 6144) pairs are stored transposed + // [N, K]; every other pair keeps the identity [K, N] copy. The scale copy + // is N-length and order-independent. Runs outside the timed/CUDA-Graph + // region. + const int64_t weight_elems = static_cast(k) * n; + const int64_t total = weight_elems > n ? weight_elems : n; + const dim3 grid(static_cast( + (total + kScalarThreads - 1) / kScalarThreads)); + hipLaunchKernelGGL( + w8a8_pack_kernel, + grid, dim3(kScalarThreads), 0, stream, + raw_weight, weight_scale, packed_weight, packed_weight_scale, + weight_elems, n, k); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/shared_gate_up_proj.hip new file mode 100644 index 00000000..e93fb813 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M4096/shared_gate_up_proj.hip @@ -0,0 +1,1281 @@ +// @@variant shape=minimax_tp8_shared_gate_up_proj_m4096 commit=7c4c216bf06688139ccf5d4bbc70df27a420f661 added=2026-08-29 +// median_us=319.2 p90_us=320.6 speedup=99.89 baseline_us=3.188e+04 +// source=minimaxm3-dsh-tp8-m4096-1-b0482833 +// MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). +// worker_3 native INT8 DUMMA m16n16k32 prefill kernels for the assigned +// MiniMax TP8 M=4096 shapes: +// * minimax_tp8_shared_gate_up_proj_m4096: (M=4096, N=768, K=6144) +// * minimax_tp8_shared_down_proj_m4096: (M=4096, N=6144, K=384) +// +// This file provides the two stable host launch symbols consumed by +// csrc/bindings.cpp: +// * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's +// HIP stream; no allocation, packing, +// autotuning, or synchronization. Uses +// only the caller-provided out and +// workspace (workspace untouched by this +// no-split-K path). +// * launch_pack_w8a8_weight(...) - one-time out-of-timed-region packing; +// the exact (K, N) = (6144, 768) pair is +// stored transposed [N, K], every other +// pair keeps the identity copy. +// +// Strategy (round 1: establish the 2-D macro-tile DUMMA baseline for the +// exact gate_up shape by porting the sibling-validated gate_up architecture +// family; the bootstrap's single-buffered 64x128 tile was measured at +// 595 us / 64.95 TOPS): +// * Large-M prefill (M >= 128): native INT8 DUMMA m16n16k32 2-D macro-tiles +// (64x64 / 64x128 / 128x64), each wave owning a 32x32 quadrant (four +// 16x16 int32 accumulators) resident across the whole K loop. A and B are +// cooperatively staged into bank-skewed LDS with 16-byte vectorized +// coalesced global loads. +// * Exact gate_up shape (M=4096, N=768, K=6144): dispatches to +// w8a8_gemm_prefill_tiled_kernel_w4<64,128,64> -- the 64x128 macro-tile, +// 256 threads = 4 wavefronts, each wave owning a 32x64 quadrant (eight +// m16n16k32 accumulators), K stage 64 DOUBLE-buffered with one +// __syncthreads per stage (stage s+1's global loads issue into registers +// before the stage-s MMAC burst and commit after it), B packed [N, K] +// and loaded with col_major fragments as ONE 8-byte LDS read per lane +// (load_fragment8: identical element-to-slot mapping and byte order, no +// per-dword mask/OR reassembly VALU), fused fragment -> scale -> bf16 +// epilogue. LDS 30,720 B/block -> 2 blocks/CU; grid (N/128) x (M/64) = +// 6 x 64 = 384 blocks; 96 K stages. No split-K: the M x N output-tile +// grid dwarfs the 120 CUs. +// * Other large-M shapes sharing the packed (k, n) = (6144, 768) pair use +// the generic 8-wave double-buffered <64,128,64,2> kernel (bit-identical +// path); the generic single-buffer 64x128/128x64/64x64 arms and the +// identity packing for unmatched (k, n) pairs are unchanged. +// * The scalar fallback is transposed-aware for the exact (k, n) pair so +// every consumer of the packed (6144, 768) buffer (including the paired +// M=2/M=16 decode shapes) stays correct. +// +// Mathematical contract (exact int32 dot before float scaling): +// out[m, n] = bf16(int32_dot(a[m, :], b[:, n]) * x_scale[m] * weight_scale[n]) +// +// Header order is fixed by the control plane for this DTK: +// hip_runtime.h -> hip_bfloat16.h -> du_mma.h + +#include +#include +#include + +#include + +namespace { + +// gfx928 INT8 DUMMA primitive: m16n16k32, int8 x int8 -> int32. +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; + +// gfx928 wavefront is 64 lanes; every blockDim must be a multiple of 64. +constexpr int kWaveSize = 64; + +// Large-prefill dispatch geometry. +constexpr int kPrefillMinM = 128; +constexpr int kPrefillStageK = 128; +constexpr int kDefaultBlockN = 64; + +// LDS padding in bytes added to each staged row. A keeps 16 so every A row +// start stays 16-byte aligned for int4 staging while shifting fragment rows +// off the same 32-bank phase (power-of-two strides alias every row onto +// identical banks). B deliberately uses only 4 bytes of pad: with the gfx928 +// INT8 DUMMA B-fragment layout (lane (c, g) reads bytes at k-rows 8g..8g+7, +// one byte per row, ldm apart), any 16-byte-aligned row stride makes +// 2*stride*g == 0 (mod 32), so all four k-groups of every ds_read_u8 land on +// the same bank phase -> 4 banks x 16 lanes (16-way). Stride 132 (== 4 mod +// 16) moves the k-groups to bank phases {0,8,16,24} -> 16 banks x 4 lanes. +constexpr int kLdsPad = 16; +constexpr int kBLdsPad = 4; + +constexpr int kScalarThreads = 256; + +using namespace du::dumma; + +// Verified gfx928 INT8 DUMMA accumulator ownership: +// row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> columns +// col_mod4 + 4*i. Scale by x_scale[row] and weight_scale[col] and store +// bf16 directly from the fragment (no accumulator LDS round trip). +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Direct 8-byte LDS fragment load for the double-buffered exact-shape path. +// du_load_matrix_sync's int8 loaders assign x[0..7] = 8 consecutive bytes at +// (lane & 15) * ldm + ((lane >> 4) << 3) for both matrix_a row_major and +// matrix_b col_major, and du_mma_sync passes the fragment to the v_mmac +// builtin as one packed 8-byte operand, so the compiler emits a per-byte +// mask/OR reassembly chain for every loaded dword. Writing the same 8 bytes +// directly into the fragment storage keeps the operand bit pattern identical +// (exact int32 accumulation unchanged) and lets the compiler feed the +// ds_read2_b32 pair straight to the v_mmac (one ds_read2_b32 per fragment, +// ~4-way bank floor instead of the 16-way aliasing of strided byte reads). +__device__ __forceinline__ void load_fragment8( + DUFragmentBase& frag, + const signed char* __restrict__ base, + int stride, + int lane) { + const int off = (lane & 15) * stride + ((lane >> 4) << 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(base + off); +} + +// Large-M prefill path: one block computes a kBlockM x kBlockN output tile. +// kWaveRows x kWaveCols wavefronts each own a 32x32 quadrant (four 16x16 +// DUMMA accumulators), while the block cooperatively stages +// A[kBlockM, kStageK] and B[kStageK, kBlockN] in bank-skewed LDS. The LDS is +// single-buffered with two barriers per stage (stage in, compute, protect). +// Tail-M rows are zero-filled on load and masked on store, so any M >= 128 +// is supported. The packed weight is the bootstrap identity [K, N] layout, +// so B rows are staged n-major contiguous chunks. +template +__global__ __launch_bounds__((kBlockM / 32) * (kBlockN / 32) * kWaveSize) void +w8a8_dumma_prefill_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWaveRows = kBlockM / 32; + constexpr int kWaveCols = kBlockN / 32; + constexpr int kThreads = kWaveRows * kWaveCols * kWaveSize; + constexpr int kAStride = kStageK + kLdsPad; + constexpr int kBStride = kBlockN + kBLdsPad; + static_assert(kAStride % sizeof(int4) == 0, + "A LDS row stride must stay 16-byte aligned"); + static_assert(kBStride % sizeof(int32_t) == 0, + "B LDS row stride must stay 4-byte aligned (int32 stores)"); + static_assert(kStageK % kTileK == 0, + "K stage must be a multiple of the DUMMA K unit"); + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kBlockN / sizeof(int4); + constexpr int kBVectors = kStageK * kBVectorsPerRow; + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWaveCols; + const int wave_col = wave - wave_row * kWaveCols; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Cooperatively stage A[kBlockM, kStageK] rows into LDS. Out-of-range + // (tail-M) rows are zero-filled so the DUMMA math stays defined. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + // Cooperatively stage B[kStageK, kBlockN] rows into LDS. The packed + // weight is the bootstrap identity [K, N] layout, so each B row is a + // contiguous chunk of kBlockN columns at row stride n. Global loads stay + // 16-byte vectorized and fully coalesced (one int4 per element); the LDS + // row stride kBStride = kBlockN + 4 is not 16-byte aligned by design (see + // kBLdsPad), so ds_write_b128 is unavailable and each int4 is committed + // as four int32 stores (identical total bank traffic to one b128). + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int kk = vec / kBVectorsPerRow; + const int v = vec - kk * kBVectorsPerRow; + const int4 b = *reinterpret_cast( + weight + (k0 + kk) * n + n0 + + v * static_cast(sizeof(int4))); + int32_t* b_row = reinterpret_cast( + b_tile + kk * kBStride + v * static_cast(sizeof(int4))); + b_row[0] = b.x; + b_row[1] = b.y; + b_row[2] = b.z; + b_row[3] = b.w; + } + __syncthreads(); + + // Compute the kTileK steps of this stage. Each wave owns a 32x32 + // quadrant: two A fragments (rows 0-15 / 16-31) and two B fragments + // (cols 0-15 / 16-31) per step, four du_mma_sync per step. The kk + // ordering (k0-outer, kk-inner step kTileK) keeps the int32 accumulation + // in ascending k order, bit-identical to the scalar fallback. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + + // Protect the single LDS buffer from being overwritten by the next + // stage's staging while this stage's fragment loads are still in flight. + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// Double-buffered variant of the large-prefill path above (ported from the +// sibling-validated gate_up lineage; only instantiated as <64,128,64,2>). +// +// kBuffers == 1: single-buffered K loop (K stage 128, two barriers per stage): +// stage s+1's global loads cannot start until the stage-s barrier, so each +// stage serializes one global-load round trip in front of the compute. +// kBuffers == 2: software-pipelined K loop (K stage 64, one barrier per +// stage): stage s+1's loads are issued into registers before the stage-s +// MMA compute and committed to the other LDS buffer after it, overlapping +// the global round trip with the LDS-wait-bound compute. Requires at most +// one int4 per thread per operand (static_asserted), which holds for the +// exact shape. The (K, N) = (6144, 768) packed weight is transposed +// [N, K], so B is staged [N, K] and loaded with the col_major fragment +// loader (8 consecutive k-values per lane -> one ds_read2_b32 per fragment, +// ~4-way bank floor). +template +__global__ __launch_bounds__((kBlockM / 32) * (kBlockN / 32) * kWaveSize) void +w8a8_gemm_prefill_tiled_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWaveRows = kBlockM / 32; + constexpr int kWaveCols = kBlockN / 32; + constexpr int kThreads = kWaveRows * kWaveCols * kWaveSize; + // 16-byte bank skew: keeps int4 staging stores aligned while spreading the + // fragment rows across distinct LDS banks (~4-way floor for 8-byte/lane + // reads instead of the 16-way aliasing of power-of-two strides). + constexpr int kAStride = kStageK + sizeof(int4); + // Double-buffered path (exact shape) stages B transposed [N, K] so the + // DUMMA B fragments read 8 consecutive k-values per lane (one ds_read2_b32, + // ~4-way) instead of 8 strided k-rows (eight ds_read_u8, 16-way: with any + // 16-byte-aligned row stride S, B k-rows 8 apart alias onto one bank group + // because (S/4)*8*g == 0 mod 32 for all g). The generic single-buffer path + // keeps [K, N] B and the row-major B loader. + constexpr int kBStride = (kBuffers == 2) ? (kStageK + sizeof(int4)) + : (kBlockN + sizeof(int4)); + constexpr int kBRows = (kBuffers == 2) ? kBlockN : kStageK; + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = + (kBuffers == 2) ? (kStageK / sizeof(int4)) : (kBlockN / sizeof(int4)); + constexpr int kBVectors = kBRows * kBVectorsPerRow; + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / kWaveCols; + const int wave_col = wave - wave_row * kWaveCols; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kBuffers][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBuffers][kBRows * kBStride]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + if constexpr (kBuffers == 2) { + static_assert(kAVectors <= kThreads && kBVectors <= kThreads, + "double-buffered path stages at most one int4 per thread"); + // The exact shape's packed weight is [N, K] (transposed by + // launch_pack_w8a8_weight), so B fragments load 8 consecutive k-values + // per lane from the [N, K] LDS tile via the col_major loader; the + // element-to-slot mapping (slot i = B[k = kk + 8*g + i][n = col + row]) + // is identical to the row-major loader on a [K, N] tile, so the v_mmac + // operand values and the exact int32 accumulation are unchanged. + DUFragment + b_frag_t0, b_frag_t1; + // Prologue: stage K tile 0 into buffer 0, then one barrier. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + // B is packed [N, K] for the exact shape: each thread stages 16 + // consecutive k-values of one n row (n-stride in global is k). + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + n_local * kBStride)[v] = + *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + const int n_stages = k / kStageK; + for (int s = 0; s < n_stages; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + const bool has_next = (s + 1) < n_stages; + // Issue the stage s+1 global loads now; the data is consumed by the + // ds_write after the compute, so the vmcnt wait lands after the MMA + // loop instead of stalling the front of the stage. + int4 a_reg{0, 0, 0, 0}; + int4 b_reg{0, 0, 0, 0}; + if (has_next) { + const int k1 = (s + 1) * kStageK; + if (tid < kAVectors) { + const int local_row = tid / kAVectorsPerRow; + const int v = tid - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + a_reg = global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k1 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + const int n_local = tid / kBVectorsPerRow; + const int v = tid - n_local * kBVectorsPerRow; + b_reg = *reinterpret_cast( + weight + (n0 + n_local) * k + k1 + + v * static_cast(sizeof(int4))); + } + + // Compute stage s from the buffer staged last iteration. + // Load the int8 fragments as raw 8-byte LDS reads into the fragment + // storage (same bytes, same element-to-slot mapping), so the compiler + // feeds the ds_read2_b32 pair straight to the v_mmac and the per-dword + // mask/OR reassembly VALU disappears from the steady state. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + load_fragment8( + a_frag0, a_tile[cur] + local_row * kAStride + kk, kAStride, + lane); + load_fragment8( + a_frag1, a_tile[cur] + (local_row + kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8( + b_frag_t0, b_tile[cur] + local_col * kBStride + kk, kBStride, + lane); + load_fragment8( + b_frag_t1, b_tile[cur] + (local_col + kTileN) * kBStride + kk, + kBStride, lane); + du_mma_sync(acc00, a_frag0, b_frag_t0, acc00); + du_mma_sync(acc01, a_frag0, b_frag_t1, acc01); + du_mma_sync(acc10, a_frag1, b_frag_t0, acc10); + du_mma_sync(acc11, a_frag1, b_frag_t1, acc11); + } + + // Commit the prefetched stage s+1 into the other LDS buffer, then one + // barrier: it orders both this iteration's compute reads of buffer cur + // (against the next-next prefetch, which reuses cur) and the prefetch + // writes of buffer nxt (against the next iteration's compute reads). + if (has_next) { + if (tid < kAVectors) { + const int local_row = tid / kAVectorsPerRow; + const int v = tid - local_row * kAVectorsPerRow; + reinterpret_cast(a_tile[nxt] + local_row * kAStride)[v] = + a_reg; + } + const int n_local = tid / kBVectorsPerRow; + const int v = tid - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[nxt] + n_local * kBStride)[v] = b_reg; + } + __syncthreads(); + } + } else { + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[kBlockM, kStageK] into LDS (zero-filled tail M rows). + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + // Stage B[kStageK, kBlockN] into LDS from the packed weight. The + // generic path keeps the identity [K, N] layout (pack_weight is an + // identity copy for every (k, n) pair other than the exact shape). + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int kk = vec / kBVectorsPerRow; + const int v = vec - kk * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + kk * kBStride)[v] = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + du_load_matrix_sync( + a_frag0, a_tile[0] + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile[0] + (local_row + kTileM) * kAStride + kk, + kAStride); + du_load_matrix_sync( + b_frag0, b_tile[0] + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile[0] + kk * kBStride + local_col + kTileN, + kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + __syncthreads(); + } + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + kTileM, base_col + kTileN, lane); +} + +// 8-wave / 4-accumulator variant of the exact-shape double-buffered 64x128 +// kernel (w8a8_gemm_prefill_tiled_kernel_w8, instantiated as <64,128,64>). +// The macro-tile, K stage 64 double buffer (one barrier/stage), [N,K] packed +// B, bank-skewed LDS layout (30,720 B/block), staging bytes, element-to-slot +// fragment mapping, byte order, and the k0-outer/kk-inner int32 accumulation +// order are IDENTICAL to the sibling-accepted 4-wave kernel, so the result is +// bit-identical; the only deltas are the wave partition and the per-thread +// staging split: +// * 512 threads (8 wavefronts); each wave owns a 32x32 quadrant (two +// 16-row bands x two 16-col bands) = FOUR independent m16n16k32 +// accumulators (4-deep MMAC ILP per kk slice). Per-wavefront MMACs per +// stage = 8 (512 MMACs/wavefront over a K=4096 loop). +// * A staging is one int4/thread for the first 256 threads, B staging one +// int4/thread for all 512 threads (same 12,288 B/stage/block). +// Grid (N/128) x (M/64) = 6 x 64 = 384 blocks; LDS 30,720 B x 2 = 61,440 B +// <= 64 KiB; target <= 64 VGPR x 512 x 2 <= 512 KiB -> 2 blocks/CU -> 4 +// waves/SIMD (must be confirmed in the code object; > 64 VGPR degenerates to +// 1 block/CU = 2 waves/SIMD). Loop-invariant staging/fragment addresses are +// hoisted before the K loop, and the eight fragment loads of each stage are +// pinned ahead of the eight-MMAC burst. +template +__global__ __launch_bounds__(8 * kWaveSize) void +w8a8_gemm_prefill_tiled_kernel_w8( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kThreads = 8 * kWaveSize; + constexpr int kAStride = kStageK + sizeof(int4); + constexpr int kBStride = kStageK + sizeof(int4); + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kStageK / sizeof(int4); + constexpr int kBVectors = kBlockN * kBVectorsPerRow; + static_assert(kAVectors <= kThreads, + "w8 path stages A with at most one int4 per thread"); + static_assert(kBVectors == kThreads, + "w8 path stages exactly one B int4 per thread"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave / 4; + const int wave_col = wave - wave_row * 4; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBStride]; + + DUFragment acc[2][2]; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 2; ++c) { + du_fill_fragment(acc[r][c], 0); + } + } + + // Prologue: stage K tile 0 into buffer 0 (A: first 256 threads, B: all 512 + // threads; same vectors and addresses as the generic kernel distributed + // over 512 threads), then one barrier. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + n_local * kBStride)[v] = + *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + const int n_stages = k / kStageK; + // Hoist every loop-invariant address of the steady state out of the K + // loop. The A staging pointer is only dereferenced by threads tid < + // kAVectors. The three global staging pointers (A/B) walk +kStageK per + // stage; the LDS commit slots add only the buffer offset nxt * kBufInt4; + // the fragment reads use the per-wave buffer-0 bases and add only the + // buffer offset cur * kBufBytes and the unrolled kk slice. + const bool a_active = tid < kAVectors; + const int local_row_a = tid / kAVectorsPerRow; + const int v_a = tid - local_row_a * kAVectorsPerRow; + const int global_row_a = m0 + local_row_a; + const bool a_row_valid = global_row_a < m; + const int8_t* __restrict__ a_g = + x_q + static_cast(global_row_a) * k + + v_a * static_cast(sizeof(int4)) + kStageK; + const int n_local = tid / kBVectorsPerRow; + const int v_b = tid - n_local * kBVectorsPerRow; + const int8_t* __restrict__ b_g = + weight + static_cast(n0 + n_local) * k + + v_b * static_cast(sizeof(int4)) + kStageK; + int4* const a_wslot = + reinterpret_cast(a_tile[0] + local_row_a * kAStride) + v_a; + int4* const b_wslot = + reinterpret_cast(b_tile[0] + n_local * kBStride) + v_b; + constexpr int kABufInt4 = + (kBlockM * kAStride) / static_cast(sizeof(int4)); + constexpr int kBBufInt4 = + (kBlockN * kBStride) / static_cast(sizeof(int4)); + // Per-wave fragment-read bases in buffer 0 (32x32 quadrant per wave); the + // loop adds the buffer offset cur * kBufBytes and the kk slice. + const int8_t* const a_rd = a_tile[0] + (wave_row * 32) * kAStride; + const int8_t* const b_rd = b_tile[0] + (wave_col * 32) * kBStride; + constexpr int kABufBytes = kBlockM * kAStride; + constexpr int kBBufBytes = kBlockN * kBStride; + + // Prologue prefetch: issue the tile-1 global loads into the loop-carried + // payload now (a_g/b_g already address tile 1). The prologue barrier above + // already ordered the tile-0 staging; the publish of this payload happens + // at the top of stage 0, one full stage later. + int4 a_pld{0, 0, 0, 0}; + int4 b_pld{0, 0, 0, 0}; + if (a_active) { + a_pld = a_row_valid ? *reinterpret_cast(a_g) + : int4{0, 0, 0, 0}; + } + b_pld = *reinterpret_cast(b_g); + a_g += kStageK; + b_g += kStageK; + + // Steady state (stages 0..n_stages-2): publish-at-top double buffering. + // The tile-(s+1) payload prefetched one iteration earlier is written into + // the idle LDS buffer at the TOP of stage s, BEFORE the MMAC burst, so the + // ds_write latency (and the already-satisfied vmcnt wait) overlaps the + // compute instead of sitting between the MMACs and the barrier (the exact + // code object shows both ds_write_b128 after the v_mmac group, immediately + // before s_barrier). The tile-(s+2) prefetch then issues from the walking + // pointers; its vmcnt wait lands at the next stage's publish. One barrier + // per stage orders both this stage's fragment reads of buffer cur (against + // the next publish into it) and this stage's publish into buffer nxt + // (against the next stage's reads of it) -- identical semantics to the old + // commit-after-compute loop, so every global load, LDS write, fragment + // read, and v_mmac is unchanged and the int32 accumulation stays + // bit-identical. + DUFragment + a_frag00, a_frag01, a_frag10, a_frag11; + DUFragment + b_frag00, b_frag01, b_frag10, b_frag11; + static_assert(kStageK == 2 * kTileK, + "w8 burst schedule expects exactly two kk slices per stage"); + for (int s = 0; s < n_stages - 1; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + // Publish tile s+1 into buffer nxt at the top of the stage. + if (a_active) { + *(a_wslot + nxt * kABufInt4) = a_pld; + } + *(b_wslot + nxt * kBBufInt4) = b_pld; + // Prefetch tile s+2 (skipped on the last loop iteration; the peeled + // final stage below needs no payload). + if (s + 2 < n_stages) { + if (a_active) { + a_pld = a_row_valid ? *reinterpret_cast(a_g) + : int4{0, 0, 0, 0}; + } + b_pld = *reinterpret_cast(b_g); + a_g += kStageK; + b_g += kStageK; + } + + // Compute stage s: each wave owns a 32x32 quadrant (two 16-row bands x + // two 16-col bands), four independent m16n16k32 accumulators. The + // k0-outer/kk-inner order per output element is unchanged, so the int32 + // accumulation is bit-identical to the generic kernel. + const int8_t* const a_cur = a_rd + cur * kABufBytes; + const int8_t* const b_cur = b_rd + cur * kBBufBytes; + // kk = 0 slice loads (row bands x0/x1, col bands 0..1). + load_fragment8(a_frag00, a_cur, kAStride, lane); + load_fragment8(a_frag01, a_cur + kTileM * kAStride, kAStride, lane); + load_fragment8(b_frag00, b_cur, kBStride, lane); + load_fragment8(b_frag01, b_cur + kTileN * kBStride, kBStride, lane); + // kk = kTileK slice loads (same rows/cols, k offset kTileK). + load_fragment8(a_frag10, a_cur + kTileK, kAStride, lane); + load_fragment8( + a_frag11, a_cur + kTileM * kAStride + kTileK, kAStride, lane); + load_fragment8(b_frag10, b_cur + kTileK, kBStride, lane); + load_fragment8( + b_frag11, b_cur + kTileN * kBStride + kTileK, kBStride, lane); + // kk = 0 MMAC group. + du_mma_sync(acc[0][0], a_frag00, b_frag00, acc[0][0]); + du_mma_sync(acc[0][1], a_frag00, b_frag01, acc[0][1]); + du_mma_sync(acc[1][0], a_frag01, b_frag00, acc[1][0]); + du_mma_sync(acc[1][1], a_frag01, b_frag01, acc[1][1]); + // kk = kTileK MMAC group (same accumulators, same order). + du_mma_sync(acc[0][0], a_frag10, b_frag10, acc[0][0]); + du_mma_sync(acc[0][1], a_frag10, b_frag11, acc[0][1]); + du_mma_sync(acc[1][0], a_frag11, b_frag10, acc[1][0]); + du_mma_sync(acc[1][1], a_frag11, b_frag11, acc[1][1]); + + __syncthreads(); + } + + // Peeled final stage: compute tile n_stages-1 from buffer (n_stages-1)&1, + // which the last loop iteration published above its barrier. No prefetch, + // no publish, and no dead barrier after the last stage: each wave's + // epilogue writes its own disjoint 32x32 output quadrant, so no + // cross-wave ordering is required. + { + const int cur = (n_stages - 1) & 1; + const int8_t* const a_cur = a_rd + cur * kABufBytes; + const int8_t* const b_cur = b_rd + cur * kBBufBytes; + // kk = 0 slice loads (row bands x0/x1, col bands 0..1). + load_fragment8(a_frag00, a_cur, kAStride, lane); + load_fragment8(a_frag01, a_cur + kTileM * kAStride, kAStride, lane); + load_fragment8(b_frag00, b_cur, kBStride, lane); + load_fragment8(b_frag01, b_cur + kTileN * kBStride, kBStride, lane); + // kk = kTileK slice loads (same rows/cols, k offset kTileK). + load_fragment8(a_frag10, a_cur + kTileK, kAStride, lane); + load_fragment8( + a_frag11, a_cur + kTileM * kAStride + kTileK, kAStride, lane); + load_fragment8(b_frag10, b_cur + kTileK, kBStride, lane); + load_fragment8( + b_frag11, b_cur + kTileN * kBStride + kTileK, kBStride, lane); + // kk = 0 MMAC group. + du_mma_sync(acc[0][0], a_frag00, b_frag00, acc[0][0]); + du_mma_sync(acc[0][1], a_frag00, b_frag01, acc[0][1]); + du_mma_sync(acc[1][0], a_frag01, b_frag00, acc[1][0]); + du_mma_sync(acc[1][1], a_frag01, b_frag01, acc[1][1]); + // kk = kTileK MMAC group (same accumulators, same order). + du_mma_sync(acc[0][0], a_frag10, b_frag10, acc[0][0]); + du_mma_sync(acc[0][1], a_frag10, b_frag11, acc[0][1]); + du_mma_sync(acc[1][0], a_frag11, b_frag10, acc[1][0]); + du_mma_sync(acc[1][1], a_frag11, b_frag11, acc[1][1]); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 2; ++c) { + store_prefill_fragment( + acc[r][c], x_scale, weight_scale, out, m, n, + base_row + r * kTileM, base_col + c * kTileN, lane); + } + } +} + +// 4-wave / 8-accumulator variant of the exact-shape double-buffered 64x128 +// kernel (w8a8_gemm_prefill_tiled_kernel_w4, instantiated as <64,128,64>). +// The macro-tile, K stage 64 double buffer (one barrier/stage, publish-at-top, +// peeled final stage), [N,K] packed B, bank-skewed LDS layout (30,720 +// B/block), staging bytes, element-to-slot fragment mapping, byte order, and +// the k0-outer/kk-inner int32 accumulation order are IDENTICAL to the +// accepted 8-wave kernel, so the result is bit-identical; the only deltas are +// the wave partition and the per-thread staging split: +// * 256 threads (4 wavefronts); each wave owns a 32x64 quadrant (two +// 16-row bands x four 16-col bands) = EIGHT independent m16n16k32 int32 +// accumulators (8-deep MMAC ILP per kk slice). Per block-stage the LDS +// fragment reads drop from 8 waves x 8 = 64 to 4 waves x 12 = 48 +// load_fragment8 (six ds_read2_b64 per wave) and the LDS read BYTES drop +// from 32 KB to 24 KB: the 32x32 quadrant layout re-reads A 4x and B 2x, +// the 32x64 quadrant re-reads both 2x (per-wave per stage: 4 A fragments +// + 8 B fragments), so per block-stage the LDS read instruction count +// drops 64 -> 48 and the family lds_instructions 1.62M -> ~1.33M. +// * A staging is one int4/thread for all 256 threads, B staging two +// int4/thread for all 256 threads (same 12,288 B/stage/block). +// Grid (N/128) x (M/64) = 6 x 64 = 384 blocks; LDS 30,720 B x 2 = 61,440 B +// <= 64 KiB; 256 threads x 2 blocks = 8 waves/CU = 2 waves/SIMD at +// arch_vgpr <= 128 (8 accumulators = 32 VGPR + 12 loop-carried payload + +// 12 in-flight fragment VGPR + hoisted addresses ~= 80; 2 x 256 x 128 = +// 65,536 exactly fits; > 128 VGPR degenerates to 1 block/CU and predicts a +// regression -- falsifiable). Loop-invariant staging/fragment addresses are +// hoisted before the K loop, and the twelve fragment loads of each stage are +// pinned ahead of the sixteen-MMAC burst. +template +__global__ __launch_bounds__(4 * kWaveSize) void +w8a8_gemm_prefill_tiled_kernel_w4( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kThreads = 4 * kWaveSize; + constexpr int kAStride = kStageK + sizeof(int4); + constexpr int kBStride = kStageK + sizeof(int4); + constexpr int kAVectorsPerRow = kStageK / sizeof(int4); + constexpr int kAVectors = kBlockM * kAVectorsPerRow; + constexpr int kBVectorsPerRow = kStageK / sizeof(int4); + constexpr int kBVectors = kBlockN * kBVectorsPerRow; + static_assert(kAVectors == kThreads, + "w4 path stages A with exactly one int4 per thread"); + static_assert(kBVectors == 2 * kThreads, + "w4 path stages B with exactly two int4 per thread"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBStride]; + + DUFragment acc[2][4]; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 4; ++c) { + du_fill_fragment(acc[r][c], 0); + } + } + + // Prologue: stage K tile 0 into buffer 0 (A: one int4/thread, B: two + // int4/thread; the same vectors and addresses as the 8-wave kernel + // distributed over 256 threads), then one barrier. + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int local_row = vec / kAVectorsPerRow; + const int v = vec - local_row * kAVectorsPerRow; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile[0] + local_row * kAStride)[v] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + v * static_cast(sizeof(int4))) + : int4{0, 0, 0, 0}; + } + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int n_local = vec / kBVectorsPerRow; + const int v = vec - n_local * kBVectorsPerRow; + reinterpret_cast(b_tile[0] + n_local * kBStride)[v] = + *reinterpret_cast( + weight + (n0 + n_local) * k + + v * static_cast(sizeof(int4))); + } + __syncthreads(); + + const int n_stages = k / kStageK; + // Hoist every loop-invariant address of the steady state out of the K + // loop (same scheme as the 8-wave kernel; B staging is two int4/thread: + // vectors tid for rows tid/4 (0..63) and tid+256 for rows 64 + tid/4). + const int local_row_a = tid / kAVectorsPerRow; + const int v_a = tid - local_row_a * kAVectorsPerRow; + const int global_row_a = m0 + local_row_a; + const bool a_row_valid = global_row_a < m; + const int8_t* __restrict__ a_g = + x_q + static_cast(global_row_a) * k + + v_a * static_cast(sizeof(int4)) + kStageK; + const int n_local0 = tid / kBVectorsPerRow; + const int v_b0 = tid - n_local0 * kBVectorsPerRow; + const int8_t* __restrict__ b_g0 = + weight + static_cast(n0 + n_local0) * k + + v_b0 * static_cast(sizeof(int4)) + kStageK; + const int vec1 = tid + kThreads; + const int n_local1 = vec1 / kBVectorsPerRow; + const int v_b1 = vec1 - n_local1 * kBVectorsPerRow; + const int8_t* __restrict__ b_g1 = + weight + static_cast(n0 + n_local1) * k + + v_b1 * static_cast(sizeof(int4)) + kStageK; + int4* const a_wslot = + reinterpret_cast(a_tile[0] + local_row_a * kAStride) + v_a; + int4* const b_wslot0 = + reinterpret_cast(b_tile[0] + n_local0 * kBStride) + v_b0; + int4* const b_wslot1 = + reinterpret_cast(b_tile[0] + n_local1 * kBStride) + v_b1; + constexpr int kABufInt4 = + (kBlockM * kAStride) / static_cast(sizeof(int4)); + constexpr int kBBufInt4 = + (kBlockN * kBStride) / static_cast(sizeof(int4)); + // Per-wave fragment-read bases in buffer 0 (32x64 quadrant per wave); the + // loop adds the buffer offset cur * kBufBytes and the kk slice. + const int8_t* const a_rd = a_tile[0] + (wave_row * 32) * kAStride; + const int8_t* const b_rd = b_tile[0] + (wave_col * 64) * kBStride; + constexpr int kABufBytes = kBlockM * kAStride; + constexpr int kBBufBytes = kBlockN * kBStride; + + // Prologue prefetch: issue the tile-1 global loads into the loop-carried + // payload now (a_g/b_g0/b_g1 already address tile 1). The publish of this + // payload happens at the top of stage 0, one full stage later. + int4 a_pld{0, 0, 0, 0}; + int4 b_pld0{0, 0, 0, 0}; + int4 b_pld1{0, 0, 0, 0}; + a_pld = a_row_valid ? *reinterpret_cast(a_g) + : int4{0, 0, 0, 0}; + b_pld0 = *reinterpret_cast(b_g0); + b_pld1 = *reinterpret_cast(b_g1); + a_g += kStageK; + b_g0 += kStageK; + b_g1 += kStageK; + + // Steady state (stages 0..n_stages-2): publish-at-top double buffering + // with identical semantics to the 8-wave kernel (one barrier per stage; + // 1 prologue + 95 loop = 96 dynamic barriers/block for K=6144, stage 64). + DUFragment + a_frag00, a_frag01, a_frag10, a_frag11; + DUFragment + b_frag00, b_frag01, b_frag02, b_frag03; + DUFragment + b_frag10, b_frag11, b_frag12, b_frag13; + static_assert(kStageK == 2 * kTileK, + "w4 burst schedule expects exactly two kk slices per stage"); + for (int s = 0; s < n_stages - 1; ++s) { + const int cur = s & 1; + const int nxt = cur ^ 1; + // Publish tile s+1 into buffer nxt at the top of the stage. + *(a_wslot + nxt * kABufInt4) = a_pld; + *(b_wslot0 + nxt * kBBufInt4) = b_pld0; + *(b_wslot1 + nxt * kBBufInt4) = b_pld1; + // Prefetch tile s+2 (skipped on the last loop iteration; the peeled + // final stage below needs no payload). + if (s + 2 < n_stages) { + a_pld = a_row_valid ? *reinterpret_cast(a_g) + : int4{0, 0, 0, 0}; + b_pld0 = *reinterpret_cast(b_g0); + b_pld1 = *reinterpret_cast(b_g1); + a_g += kStageK; + b_g0 += kStageK; + b_g1 += kStageK; + } + + // Compute stage s: each wave owns a 32x64 quadrant (two 16-row bands x + // four 16-col bands), eight independent m16n16k32 accumulators. The + // k0-outer/kk-inner order per output element is unchanged, so the int32 + // accumulation is bit-identical to the 8-wave/generic kernels. + const int8_t* const a_cur = a_rd + cur * kABufBytes; + const int8_t* const b_cur = b_rd + cur * kBBufBytes; + // kk = 0 slice loads (row bands x0/x1, col bands 0..3). + load_fragment8(a_frag00, a_cur, kAStride, lane); + load_fragment8(a_frag01, a_cur + kTileM * kAStride, kAStride, lane); + load_fragment8(b_frag00, b_cur, kBStride, lane); + load_fragment8(b_frag01, b_cur + kTileN * kBStride, kBStride, lane); + load_fragment8(b_frag02, b_cur + 2 * kTileN * kBStride, kBStride, lane); + load_fragment8(b_frag03, b_cur + 3 * kTileN * kBStride, kBStride, lane); + // kk = kTileK slice loads (same rows/cols, k offset kTileK). + load_fragment8(a_frag10, a_cur + kTileK, kAStride, lane); + load_fragment8( + a_frag11, a_cur + kTileM * kAStride + kTileK, kAStride, lane); + load_fragment8(b_frag10, b_cur + kTileK, kBStride, lane); + load_fragment8( + b_frag11, b_cur + kTileN * kBStride + kTileK, kBStride, lane); + load_fragment8( + b_frag12, b_cur + 2 * kTileN * kBStride + kTileK, kBStride, lane); + load_fragment8( + b_frag13, b_cur + 3 * kTileN * kBStride + kTileK, kBStride, lane); + // kk = 0 MMAC group. + du_mma_sync(acc[0][0], a_frag00, b_frag00, acc[0][0]); + du_mma_sync(acc[0][1], a_frag00, b_frag01, acc[0][1]); + du_mma_sync(acc[0][2], a_frag00, b_frag02, acc[0][2]); + du_mma_sync(acc[0][3], a_frag00, b_frag03, acc[0][3]); + du_mma_sync(acc[1][0], a_frag01, b_frag00, acc[1][0]); + du_mma_sync(acc[1][1], a_frag01, b_frag01, acc[1][1]); + du_mma_sync(acc[1][2], a_frag01, b_frag02, acc[1][2]); + du_mma_sync(acc[1][3], a_frag01, b_frag03, acc[1][3]); + // kk = kTileK MMAC group (same accumulators, same order). + du_mma_sync(acc[0][0], a_frag10, b_frag10, acc[0][0]); + du_mma_sync(acc[0][1], a_frag10, b_frag11, acc[0][1]); + du_mma_sync(acc[0][2], a_frag10, b_frag12, acc[0][2]); + du_mma_sync(acc[0][3], a_frag10, b_frag13, acc[0][3]); + du_mma_sync(acc[1][0], a_frag11, b_frag10, acc[1][0]); + du_mma_sync(acc[1][1], a_frag11, b_frag11, acc[1][1]); + du_mma_sync(acc[1][2], a_frag11, b_frag12, acc[1][2]); + du_mma_sync(acc[1][3], a_frag11, b_frag13, acc[1][3]); + + __syncthreads(); + } + + // Peeled final stage: compute tile n_stages-1 from buffer (n_stages-1)&1, + // which the last loop iteration published above its barrier. No prefetch, + // no publish, and no dead barrier after the last stage: each wave's + // epilogue writes its own disjoint 32x64 output quadrant, so no + // cross-wave ordering is required. + { + const int cur = (n_stages - 1) & 1; + const int8_t* const a_cur = a_rd + cur * kABufBytes; + const int8_t* const b_cur = b_rd + cur * kBBufBytes; + // kk = 0 slice loads (row bands x0/x1, col bands 0..3). + load_fragment8(a_frag00, a_cur, kAStride, lane); + load_fragment8(a_frag01, a_cur + kTileM * kAStride, kAStride, lane); + load_fragment8(b_frag00, b_cur, kBStride, lane); + load_fragment8(b_frag01, b_cur + kTileN * kBStride, kBStride, lane); + load_fragment8(b_frag02, b_cur + 2 * kTileN * kBStride, kBStride, lane); + load_fragment8(b_frag03, b_cur + 3 * kTileN * kBStride, kBStride, lane); + // kk = kTileK slice loads (same rows/cols, k offset kTileK). + load_fragment8(a_frag10, a_cur + kTileK, kAStride, lane); + load_fragment8( + a_frag11, a_cur + kTileM * kAStride + kTileK, kAStride, lane); + load_fragment8(b_frag10, b_cur + kTileK, kBStride, lane); + load_fragment8( + b_frag11, b_cur + kTileN * kBStride + kTileK, kBStride, lane); + load_fragment8( + b_frag12, b_cur + 2 * kTileN * kBStride + kTileK, kBStride, lane); + load_fragment8( + b_frag13, b_cur + 3 * kTileN * kBStride + kTileK, kBStride, lane); + // kk = 0 MMAC group. + du_mma_sync(acc[0][0], a_frag00, b_frag00, acc[0][0]); + du_mma_sync(acc[0][1], a_frag00, b_frag01, acc[0][1]); + du_mma_sync(acc[0][2], a_frag00, b_frag02, acc[0][2]); + du_mma_sync(acc[0][3], a_frag00, b_frag03, acc[0][3]); + du_mma_sync(acc[1][0], a_frag01, b_frag00, acc[1][0]); + du_mma_sync(acc[1][1], a_frag01, b_frag01, acc[1][1]); + du_mma_sync(acc[1][2], a_frag01, b_frag02, acc[1][2]); + du_mma_sync(acc[1][3], a_frag01, b_frag03, acc[1][3]); + // kk = kTileK MMAC group (same accumulators, same order). + du_mma_sync(acc[0][0], a_frag10, b_frag10, acc[0][0]); + du_mma_sync(acc[0][1], a_frag10, b_frag11, acc[0][1]); + du_mma_sync(acc[0][2], a_frag10, b_frag12, acc[0][2]); + du_mma_sync(acc[0][3], a_frag10, b_frag13, acc[0][3]); + du_mma_sync(acc[1][0], a_frag11, b_frag10, acc[1][0]); + du_mma_sync(acc[1][1], a_frag11, b_frag11, acc[1][1]); + du_mma_sync(acc[1][2], a_frag11, b_frag12, acc[1][2]); + du_mma_sync(acc[1][3], a_frag11, b_frag13, acc[1][3]); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 64; +#pragma unroll + for (int r = 0; r < 2; ++r) { +#pragma unroll + for (int c = 0; c < 4; ++c) { + store_prefill_fragment( + acc[r][c], x_scale, weight_scale, out, m, n, + base_row + r * kTileM, base_col + c * kTileN, lane); + } + } +} + +// Generic scalar fallback: one output element per thread with exact int32 +// accumulation in k order. Correct for any (m, n, k), including the small-M +// API shapes (M=2/M=16) and any unmatched geometry. The exact (k, n) = +// (6144, 768) pair reads the transposed [N, K] packed weight (the pack +// specialization is keyed on (k, n), so every consumer of that packed pair, +// including the paired M=2/M=16 decode shapes that reach this fallback, must +// interpret it transposed). b_col_stride carries the identity [K, N] column +// stride (n) for all other pairs. +__global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k, + int b_col_stride) { + const bool packed_transposed = (k == 6144 && n == 768); + const int64_t total = static_cast(m) * n; + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + const int8_t* a_row = x_q + static_cast(row) * k; + // packed[n * k + kk] = raw[kk * n + n]: the column's k values are + // contiguous at packed + col * k. + const int8_t* b_col = packed_transposed + ? weight + static_cast(col) * k + : weight + static_cast(col) * b_col_stride; + int32_t acc = 0; + if (packed_transposed) { + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[kk]); + } + } else { + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast( + b_col[static_cast(kk) * b_col_stride]); + } + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[idx] = __float2bfloat16(scaled); +} + +// pack_weight. The packed layout is opaque per the API contract; for the +// exact (K=6144, N=768) pair the weight is stored transposed [N, K] +// (n-major) so the double-buffered GEMM stages B in [N, K] LDS order and the +// DUMMA B fragments read 8 consecutive k-values per lane. Every other (k, n) +// pair keeps the identity [K, N] copy (generic fallback unchanged). Runs +// outside the timed/CUDA-Graph region. +__global__ __launch_bounds__(kScalarThreads) void +w8a8_pack_kernel( + const int8_t* __restrict__ raw_weight, + const float* __restrict__ weight_scale, + int8_t* __restrict__ packed_weight, + float* __restrict__ packed_weight_scale, + int64_t weight_elems, + int n, + int k) { + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (k == 6144 && n == 768) { + // Transpose [K, N] -> [N, K]: packed[n * k + kk] = raw[kk * n + n]. + if (idx < weight_elems) { + const int64_t n_idx = idx / static_cast(k); + const int64_t k_idx = idx - n_idx * static_cast(k); + packed_weight[idx] = raw_weight[k_idx * n + n_idx]; + } + } else if (idx < weight_elems) { + packed_weight[idx] = raw_weight[idx]; + } + if (idx < n) { + packed_weight_scale[idx] = weight_scale[idx]; + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Host launch symbols (extern "C", consumed by csrc/bindings.cpp). +// --------------------------------------------------------------------------- +extern "C" void launch_w8a8_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // No split-K: the workspace is not touched. The timed operator performs no + // allocation, synchronization, packing, or default-stream launch; it + // writes only the caller-provided out. + (void)workspace; + (void)workspace_bytes; + auto* out_bf16 = static_cast(out); + + // Native INT8 DUMMA m16n16k32 tiled path for large-M prefill. All three + // macro-tile instantiations share the exact guard; the launch geometry is + // selected explicitly so every assigned shape is covered (no + // single-shape specialization). The guard divisibility (k % 128 == 0, + // n % 64 == 0) holds for both assigned shapes and for every supported + // catalog shape. + if (m >= kPrefillMinM && (n % kDefaultBlockN == 0) && + (k % kPrefillStageK == 0)) { + if (n == 768 && k == 6144) { + // Packed (k, n) = (6144, 768) weight is transposed [N, K], so every + // large-M shape with this pair uses the double-buffered transposed-B + // kernel (the generic [K, N] arms below would misinterpret it). Exact + // assigned shape (M=4096, N=768, K=6144): 64x128 macro-tile, grid + // 6x64 = 384 blocks, K stage 64 double-buffered, 30,720 B LDS/block + // -> 2 blocks/CU. + const dim3 grid( + static_cast(n / 128), + static_cast((m + 63) / 64)); + if (m == 4096) { + // 4-wave variant w4<64,128,64> (256 threads, 32x64 quadrant per + // wave, eight accumulators, 96 K stages of 64; 25% fewer LDS + // fragment reads per block-stage than the 8-wave arm). + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel_w4<64, 128, 64>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + // Other large-M shapes with the packed (k, n) pair keep the generic + // 8-wave double-buffered <64,128,64,2> kernel (bit-identical path). + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_prefill_tiled_kernel<64, 128, 64, 2>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (n % 128 == 0) { + // 64x128 tile, 8 wavefronts / 512 threads, single-buffered 128-K stage + // (generic [K, N] packed-B arm for unmatched (k, n) pairs; the packed + // (6144, 768) pair never reaches this arm): + // down_proj: grid (48, 64) = 3072 blocks x 3 stages + const dim3 grid( + static_cast(n / 128), + static_cast((m + 63) / 64)); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 128, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (m % 128 == 0) { + // 128x64 tile, 8 wavefronts / 512 threads. + const dim3 grid( + static_cast(n / 64), + static_cast(m / 128)); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<128, 64, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + { + // 64x64 tile, 4 wavefronts / 256 threads (generic default). + const dim3 grid( + static_cast(n / 64), + static_cast((m + 63) / 64)); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 64, kPrefillStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + } + + // Generic scalar fallback for every unmatched (m, n, k) and M < 128 + // (including the paired M=2/M=16 API shapes). The kernel decodes the + // transposed [N, K] pack for the exact (k, n) = (6144, 768) pair and uses + // the identity [K, N] column stride n for every other pair. + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast( + (total + kScalarThreads - 1) / kScalarThreads)); + const dim3 block(static_cast(kScalarThreads)); + hipLaunchKernelGGL( + w8a8_gemm_scalar_kernel, + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k, n); +} + +extern "C" void launch_pack_w8a8_weight( + const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Exact (K, N) = (6144, 768) is stored transposed [N, K]; every other pair + // keeps the identity [K, N] copy. The scale copy is N-length and + // order-independent. Runs outside the timed/CUDA-Graph region. + const int64_t weight_elems = static_cast(k) * n; + const int64_t total = weight_elems > n ? weight_elems : n; + const dim3 grid(static_cast( + (total + kScalarThreads - 1) / kScalarThreads)); + hipLaunchKernelGGL( + w8a8_pack_kernel, + grid, dim3(kScalarThreads), 0, stream, + raw_weight, weight_scale, packed_weight, packed_weight_scale, + weight_elems, n, k); +} +// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/w8a8_gemm_variants.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/w8a8_gemm_variants.hip new file mode 100644 index 00000000..ab7eae58 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/w8a8_gemm_variants.hip @@ -0,0 +1,1658 @@ +// MetaInfer optional W8A8 implementation variants for gfx928. +// +// Provenance: ISA_test_codex/wqkv_a_compute_tuning. This file contains +// measured kernels and experiment launchers, including the M=16,K=4096, +// N=1536 split10 / 4-wave / stage64 family. It is staged read-only under +// references/ in newly generated kernel repositories and is NOT included by +// setup.py or the default extension build. +// +// This is evidence, not policy: optimization workers may adapt, combine, or +// ignore any implementation here and must explore other legal strategies. +// Every reused idea requires current-shape exact correctness, Graph replay, +// median/P90, workspace, resource, and ISA validation before acceptance. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace { + +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; +constexpr int kWaveSize = 64; +constexpr int kStageK = 256; +constexpr int kPrefillBlockM = 64; +constexpr int kPrefillBlockN = 64; +constexpr int kPrefillStageK = 128; + +using namespace du::dumma; + +// Fast path for M=16. Both DUMMA operands are loaded directly from their +// row-major global tensors, removing the per-K-stage LDS copy and barriers in +// the generic padded path below. +template +__global__ __launch_bounds__(kWaveSize * kWaves) void +w8a8_dumma_m16_direct_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n, + int k) { + const int wave = static_cast(threadIdx.x) / kWaveSize; + const int lane = static_cast(threadIdx.x) % kWaveSize; + const int n0 = + (static_cast(blockIdx.x) * kWaves + wave) * kTileN; + + __shared__ __align__(16) int32_t + acc_tiles[kWaves * kTileM * kTileN]; + int32_t* acc_tile = acc_tiles + wave * kTileM * kTileN; + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + for (int k0 = 0; k0 < k; k0 += kTileK) { + du_load_matrix_sync(a_frag, x_q + k0, k); + du_load_matrix_sync(b_frag, weight + k0 * n + n0, n); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + du_store_matrix_sync(acc_tile, acc_frag, kTileN, mem_row_major); + __syncthreads(); + +#pragma unroll + for (int linear = lane; linear < kTileM * kTileN; + linear += kWaveSize) { + const int row = linear / kTileN; + const int col = linear - row * kTileN; + const float scaled = static_cast(acc_tile[linear]) * + x_scale[row] * weight_scale[n0 + col]; + out[row * n + n0 + col] = __float2bfloat16(scaled); + } + return; +} + +// One or two waves cooperatively stage A[16,stage_k] and +// B[stage_k,16*waves], then consume DUMMA K=32 fragments per barrier. +template +__global__ __launch_bounds__(kWaves * kWaveSize) void +w8a8_dumma_m16_staged_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n, + int k) { + constexpr int kBlockN = kWaves * kTileN; + constexpr int kBStride = kBlockN + kBPad; + constexpr int kThreads = kWaves * kWaveSize; + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + __shared__ __align__(16) int8_t a_tile[kTileM * kStageKValue]; + __shared__ __align__(16) int8_t b_tile[kStageKValue * kBStride]; + // The legacy epilogue materializes the fragment in LDS. The direct + // epilogue needs only a dummy word here; if constexpr removes all accesses. + __shared__ __align__(16) int32_t + acc_tiles[kDirectEpilogue ? 1 : kWaves * kTileM * kTileN]; + int32_t* acc_tile = acc_tiles + wave * kTileM * kTileN; + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + for (int k0 = 0; k0 < k; k0 += kStageKValue) { + // Vectorized 16-byte cooperative copies. All target K/N values and tile + // origins are 16-byte aligned. + constexpr int kAVectors = kTileM * kStageKValue / sizeof(int4); + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int byte_offset = vec * sizeof(int4); + const int row = byte_offset / kStageKValue; + const int kk = byte_offset - row * kStageKValue; + reinterpret_cast(a_tile)[vec] = + *reinterpret_cast(x_q + row * k + k0 + kk); + } + + if constexpr (kBPad == 0) { + constexpr int kBVectorsPerRow = kBlockN / sizeof(int4); + constexpr int kBVectors = kStageKValue * kBVectorsPerRow; + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int kk = vec / kBVectorsPerRow; + const int segment = vec - kk * kBVectorsPerRow; + const int col = segment * sizeof(int4); + *reinterpret_cast(b_tile + kk * kBStride + col) = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + col); + } + } else if constexpr (kBPad % sizeof(int32_t) == 0) { + constexpr int kBWordsPerRow = kBlockN / sizeof(int32_t); + constexpr int kBWords = kStageKValue * kBWordsPerRow; + for (int word = tid; word < kBWords; word += kThreads) { + const int kk = word / kBWordsPerRow; + const int segment = word - kk * kBWordsPerRow; + const int col = segment * sizeof(int32_t); + *reinterpret_cast(b_tile + kk * kBStride + col) = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + col); + } + } else { + constexpr int kBElements = kStageKValue * kBlockN; + for (int elem = tid; elem < kBElements; elem += kThreads) { + const int kk = elem / kBlockN; + const int col = elem - kk * kBlockN; + b_tile[kk * kBStride + col] = weight[(k0 + kk) * n + n0 + col]; + } + } + __syncthreads(); + +#pragma unroll + for (int kk = 0; kk < kStageKValue; kk += kTileK) { + du_load_matrix_sync(a_frag, a_tile + kk, kStageKValue); + du_load_matrix_sync( + b_frag, b_tile + kk * kBStride + wave * kTileN, kBStride); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + __syncthreads(); + } + + if constexpr (kDirectEpilogue) { + // gfx928 int8 m16n16k32 accumulator ownership, established against + // du_store_matrix_sync: lane%16 selects the row, lane/16 selects col%4, + // and x[i] selects columns separated by four. + const int row = lane & 15; + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int out_col = n0 + wave * kTileN + col_mod4 + 4 * i; + const float scaled = static_cast(acc_frag.x[i]) * + x_scale[row] * weight_scale[out_col]; + out[row * n + out_col] = __float2bfloat16(scaled); + } + } else { + du_store_matrix_sync(acc_tile, acc_frag, kTileN, mem_row_major); + __syncthreads(); + +#pragma unroll + for (int linear = lane; linear < kTileM * kTileN; + linear += kWaveSize) { + const int row = linear / kTileN; + const int col = linear - row * kTileN; + const int out_col = n0 + wave * kTileN + col; + const float scaled = static_cast(acc_tile[linear]) * + x_scale[row] * weight_scale[out_col]; + out[row * n + out_col] = __float2bfloat16(scaled); + } + } + return; +} + +// Split-K partial kernel. Each block computes one 16x32 tile for one K slice +// and stores int32 accumulators into caller-owned workspace. No atomics or +// workspace clearing are required. +template +__global__ __launch_bounds__(kWaves * kWaveSize) void +w8a8_dumma_m16_staged2_splitk_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int32_t* __restrict__ workspace, + int n, + int k, + int split_k) { + constexpr int kBlockN = kWaves * kTileN; + constexpr int kThreads = kWaves * kWaveSize; + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int num_n_tiles = n / kBlockN; + const int split_id = static_cast(blockIdx.x) / num_n_tiles; + const int tile_n = static_cast(blockIdx.x) - split_id * num_n_tiles; + const int n0 = tile_n * kBlockN; + const int k_per_split = k / split_k; + const int k_begin = split_id * k_per_split; + const int k_end = k_begin + k_per_split; + + __shared__ __align__(16) int8_t a_tile[kTileM * kStageKValue]; + __shared__ __align__(16) int8_t b_tile[kStageKValue * kBlockN]; + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + for (int k0 = k_begin; k0 < k_end; k0 += kStageKValue) { + constexpr int kAVectors = kTileM * kStageKValue / sizeof(int4); + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int byte_offset = vec * sizeof(int4); + const int row = byte_offset / kStageKValue; + const int kk = byte_offset - row * kStageKValue; + reinterpret_cast(a_tile)[vec] = + *reinterpret_cast(x_q + row * k + k0 + kk); + } + + constexpr int kBVectors = kStageKValue * kBlockN / sizeof(int4); + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int byte_offset = vec * sizeof(int4); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + reinterpret_cast(b_tile)[vec] = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + col); + } + __syncthreads(); + +#pragma unroll + for (int kk = 0; kk < kStageKValue; kk += kTileK) { + du_load_matrix_sync(a_frag, a_tile + kk, kStageKValue); + du_load_matrix_sync( + b_frag, b_tile + kk * kBlockN + wave * kTileN, kBlockN); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + __syncthreads(); + } + + int32_t* partial = workspace + + split_id * kTileM * n + n0 + wave * kTileN; + du_store_matrix_sync(partial, acc_frag, n, mem_row_major); + return; +} + +// Fixed-shape software-prefetch family. The current winner is 4-wave/stage64 +// with non-uniform split10. Global int4 loads for stage i+1 are issued before +// the two MMACs of +// stage i, then committed to the alternate LDS buffer after all waves finish +// consuming the current buffer. Raw memory asm is intentionally avoided; the +// only inline ISA is the conservative, previously validated VMEM wait. +template +__global__ __launch_bounds__(kWaves * kWaveSize) void +w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int32_t* __restrict__ workspace, + int n, + int k, + int split_k) { + constexpr int kStage = 64; + constexpr int kBlockN = kWaves * kTileN; + constexpr int kThreads = kWaves * kWaveSize; + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int num_n_tiles = n / kBlockN; + const int split_id = static_cast(blockIdx.x) / num_n_tiles; + const int tile_n = static_cast(blockIdx.x) - split_id * num_n_tiles; + const int n0 = tile_n * kBlockN; + // split10 is a fixed-shape occupancy experiment for the 120-CU K500SM_AI: + // 24 N tiles x 10 gives exactly two blocks/CU. Distribute the 64 K/64 + // stages as 4x7 + 6x6 so every boundary remains stage aligned. + const int base_stages = (k / kStage) / split_k; + const int extra_stages = (k / kStage) - base_stages * split_k; + const int begin_stage = split_id * base_stages + + (split_id < extra_stages ? split_id : extra_stages); + const int split_stages = base_stages + (split_id < extra_stages ? 1 : 0); + const int k_begin = begin_stage * kStage; + const int k_end = k_begin + split_stages * kStage; + + __shared__ __align__(16) int8_t a_tile[2][kTileM * kStage]; + __shared__ __align__(16) int8_t b_tile[2][kStage * kBlockN]; + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + int4 a_reg{}; + int4 b_reg{}; + const bool owns_a = kSpreadA ? ((tid & 3) == 0) + : (tid < kTileM * kStage / + static_cast(sizeof(int4))); + const int a_vec = kSpreadA ? (tid >> 2) : tid; + if (owns_a) { + const int byte_offset = a_vec * sizeof(int4); + const int row = byte_offset / kStage; + const int kk = byte_offset - row * kStage; + a_reg = *reinterpret_cast(x_q + row * k + k_begin + kk); + } + { + const int byte_offset = tid * sizeof(int4); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + b_reg = *reinterpret_cast( + weight + (k_begin + kk) * n + n0 + col); + } + if constexpr (kSyncMode != 1) { + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + } + if (owns_a) { + reinterpret_cast(a_tile[0])[a_vec] = a_reg; + } + reinterpret_cast(b_tile[0])[tid] = b_reg; + if constexpr (kSyncMode == 3) { + asm volatile("s_waitcnt lgkmcnt(0)\n\ts_barrier" ::: "memory"); + } else { + __syncthreads(); + } + + int current = 0; + for (int k0 = k_begin; k0 < k_end; k0 += kStage) { + const int next_k = k0 + kStage; + int4 next_a{}; + int4 next_b{}; + if (next_k < k_end) { + if (owns_a) { + const int byte_offset = a_vec * sizeof(int4); + const int row = byte_offset / kStage; + const int kk = byte_offset - row * kStage; + next_a = *reinterpret_cast( + x_q + row * k + next_k + kk); + } + const int byte_offset = tid * sizeof(int4); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + next_b = *reinterpret_cast( + weight + (next_k + kk) * n + n0 + col); + } + +#pragma unroll + for (int kk = 0; kk < kStage; kk += kTileK) { + du_load_matrix_sync(a_frag, a_tile[current] + kk, kStage); + du_load_matrix_sync( + b_frag, b_tile[current] + kk * kBlockN + wave * kTileN, + kBlockN); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + if (next_k < k_end) { + if constexpr (kSyncMode >= 2) { + asm volatile("s_barrier" ::: "memory"); + } else { + __syncthreads(); + } + if constexpr (kSyncMode != 1) { + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + } + const int next = current ^ 1; + if (owns_a) { + reinterpret_cast(a_tile[next])[a_vec] = next_a; + } + reinterpret_cast(b_tile[next])[tid] = next_b; + current = next; + if constexpr (kSyncMode == 3) { + asm volatile("s_waitcnt lgkmcnt(0)\n\ts_barrier" ::: "memory"); + } else { + __syncthreads(); + } + } + } + + int32_t* partial = workspace + + split_id * kTileM * n + n0 + wave * kTileN; + du_store_matrix_sync(partial, acc_frag, n, mem_row_major); + return; +} + +// One-kernel split-K finalization for the 4-wave/stage64 winner. Each N=64 +// tile owns one counter placed after the eight partial planes. The last split +// block to arrive combines its tile and resets the counter for the next +// stream-ordered Graph replay. +__global__ __launch_bounds__(4 * kWaveSize) void +w8a8_dumma_m16_splitk8_w4_s64_fused_finalize_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int32_t* __restrict__ workspace, + int n, + int k) { + constexpr int kWaves = 4; + constexpr int kStage = 64; + constexpr int kBlockN = kWaves * kTileN; + constexpr int kThreads = kWaves * kWaveSize; + constexpr int kSplitK = 8; + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int num_n_tiles = n / kBlockN; + const int split_id = static_cast(blockIdx.x) / num_n_tiles; + const int tile_n = static_cast(blockIdx.x) - split_id * num_n_tiles; + const int n0 = tile_n * kBlockN; + constexpr int kPerSplit = 4096 / kSplitK; + const int k_begin = split_id * kPerSplit; + const int k_end = k_begin + kPerSplit; + + __shared__ __align__(16) int8_t a_tile[2][kTileM * kStage]; + __shared__ __align__(16) int8_t b_tile[2][kStage * kBlockN]; + __shared__ int arrival_ticket; + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + int4 a_reg{}; + int4 b_reg{}; + if (tid < 64) { + const int byte_offset = tid * sizeof(int4); + const int row = byte_offset / kStage; + const int kk = byte_offset - row * kStage; + a_reg = *reinterpret_cast(x_q + row * k + k_begin + kk); + } + { + const int byte_offset = tid * sizeof(int4); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + b_reg = *reinterpret_cast( + weight + (k_begin + kk) * n + n0 + col); + } + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + if (tid < 64) reinterpret_cast(a_tile[0])[tid] = a_reg; + reinterpret_cast(b_tile[0])[tid] = b_reg; + __syncthreads(); + + int current = 0; + for (int k0 = k_begin; k0 < k_end; k0 += kStage) { + const int next_k = k0 + kStage; + int4 next_a{}; + int4 next_b{}; + if (next_k < k_end) { + if (tid < 64) { + const int byte_offset = tid * sizeof(int4); + const int row = byte_offset / kStage; + const int kk = byte_offset - row * kStage; + next_a = *reinterpret_cast( + x_q + row * k + next_k + kk); + } + const int byte_offset = tid * sizeof(int4); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + next_b = *reinterpret_cast( + weight + (next_k + kk) * n + n0 + col); + } + +#pragma unroll + for (int kk = 0; kk < kStage; kk += kTileK) { + du_load_matrix_sync(a_frag, a_tile[current] + kk, kStage); + du_load_matrix_sync( + b_frag, b_tile[current] + kk * kBlockN + wave * kTileN, + kBlockN); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + if (next_k < k_end) { + asm volatile("s_barrier" ::: "memory"); + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + const int next = current ^ 1; + if (tid < 64) reinterpret_cast(a_tile[next])[tid] = next_a; + reinterpret_cast(b_tile[next])[tid] = next_b; + current = next; + __syncthreads(); + } + } + + int32_t* partial = workspace + + split_id * kTileM * n + n0 + wave * kTileN; + du_store_matrix_sync(partial, acc_frag, n, mem_row_major); + __syncthreads(); + + int32_t* counters = workspace + kSplitK * kTileM * n; + if (tid == 0) { + __threadfence(); + arrival_ticket = atomicAdd(counters + tile_n, 1); + } + __syncthreads(); + + if (arrival_ticket == kSplitK - 1) { +#pragma unroll + for (int linear = tid; linear < kTileM * kBlockN; + linear += kThreads) { + const int row = linear / kBlockN; + const int col = linear - row * kBlockN; + const int global_linear = row * n + n0 + col; + int32_t acc = 0; +#pragma unroll + for (int split = 0; split < kSplitK; ++split) { + acc += workspace[split * kTileM * n + global_linear]; + } + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[n0 + col]; + out[global_linear] = __float2bfloat16(scaled); + } + __syncthreads(); + if (tid == 0) { + __threadfence(); + atomicExch(counters + tile_n, 0); + } + } + return; +} + +// Cooperative alternative to the atomic last-arriver experiment. The full +// 192-block grid is resident together on K500SM_AI (120 CUs, two blocks/CU +// required). A device-wide barrier makes every partial visible, after which +// all grid threads share the 16x1536 epilogue without a second Graph node. +__global__ __launch_bounds__(4 * kWaveSize) void +w8a8_dumma_m16_splitk8_w4_s64_cooperative_finalize_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int32_t* __restrict__ workspace, + int n, + int k) { + constexpr int kWaves = 4; + constexpr int kStage = 64; + constexpr int kBlockN = kWaves * kTileN; + constexpr int kThreads = kWaves * kWaveSize; + constexpr int kSplitK = 8; + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int num_n_tiles = n / kBlockN; + const int split_id = static_cast(blockIdx.x) / num_n_tiles; + const int tile_n = static_cast(blockIdx.x) - split_id * num_n_tiles; + const int n0 = tile_n * kBlockN; + constexpr int kPerSplit = 4096 / kSplitK; + const int k_begin = split_id * kPerSplit; + const int k_end = k_begin + kPerSplit; + + __shared__ __align__(16) int8_t a_tile[2][kTileM * kStage]; + __shared__ __align__(16) int8_t b_tile[2][kStage * kBlockN]; + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + int4 a_reg{}; + int4 b_reg{}; + if (tid < 64) { + const int byte_offset = tid * sizeof(int4); + const int row = byte_offset / kStage; + const int kk = byte_offset - row * kStage; + a_reg = *reinterpret_cast(x_q + row * k + k_begin + kk); + } + { + const int byte_offset = tid * sizeof(int4); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + b_reg = *reinterpret_cast( + weight + (k_begin + kk) * n + n0 + col); + } + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + if (tid < 64) reinterpret_cast(a_tile[0])[tid] = a_reg; + reinterpret_cast(b_tile[0])[tid] = b_reg; + __syncthreads(); + + int current = 0; + for (int k0 = k_begin; k0 < k_end; k0 += kStage) { + const int next_k = k0 + kStage; + int4 next_a{}; + int4 next_b{}; + if (next_k < k_end) { + if (tid < 64) { + const int byte_offset = tid * sizeof(int4); + const int row = byte_offset / kStage; + const int kk = byte_offset - row * kStage; + next_a = *reinterpret_cast( + x_q + row * k + next_k + kk); + } + const int byte_offset = tid * sizeof(int4); + const int kk = byte_offset / kBlockN; + const int col = byte_offset - kk * kBlockN; + next_b = *reinterpret_cast( + weight + (next_k + kk) * n + n0 + col); + } + +#pragma unroll + for (int kk = 0; kk < kStage; kk += kTileK) { + du_load_matrix_sync(a_frag, a_tile[current] + kk, kStage); + du_load_matrix_sync( + b_frag, b_tile[current] + kk * kBlockN + wave * kTileN, + kBlockN); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + if (next_k < k_end) { + asm volatile("s_barrier" ::: "memory"); + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + const int next = current ^ 1; + if (tid < 64) reinterpret_cast(a_tile[next])[tid] = next_a; + reinterpret_cast(b_tile[next])[tid] = next_b; + current = next; + __syncthreads(); + } + } + + int32_t* partial = workspace + + split_id * kTileM * n + n0 + wave * kTileN; + du_store_matrix_sync(partial, acc_frag, n, mem_row_major); + + cooperative_groups::this_grid().sync(); + const int linear = static_cast(blockIdx.x) * kThreads + tid; + const int elements = kTileM * n; + if (linear < elements) { + int32_t acc = 0; +#pragma unroll + for (int split = 0; split < kSplitK; ++split) { + acc += workspace[split * elements + linear]; + } + const int row = linear / n; + const int col = linear - row * n; + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[linear] = __float2bfloat16(scaled); + } + return; +} + +// Deeper prefetch candidate: four MMACs per global stage. This intentionally +// remains a separate symbol so the stage64 winner is preserved byte-for-byte. +__global__ __launch_bounds__(3 * kWaveSize) void +w8a8_dumma_m16_splitk8_w3_s128_prefetch_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + int32_t* __restrict__ workspace, + int n, + int k, + int split_k) { + constexpr int kWaves = 3; + constexpr int kStage = 128; + constexpr int kBlockN = kWaves * kTileN; + constexpr int kThreads = kWaves * kWaveSize; + constexpr int kAVectors = kTileM * kStage / sizeof(int4); + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int num_n_tiles = n / kBlockN; + const int split_id = static_cast(blockIdx.x) / num_n_tiles; + const int tile_n = static_cast(blockIdx.x) - split_id * num_n_tiles; + const int n0 = tile_n * kBlockN; + const int k_per_split = k / split_k; + const int k_begin = split_id * k_per_split; + const int k_end = k_begin + k_per_split; + + __shared__ __align__(16) int8_t a_tile[2][kTileM * kStage]; + __shared__ __align__(16) int8_t b_tile[2][kStage * kBlockN]; + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + int4 a_reg{}; + int4 b_reg0{}; + int4 b_reg1{}; + if (tid < kAVectors) { + const int byte_offset = tid * sizeof(int4); + const int row = byte_offset / kStage; + const int kk = byte_offset - row * kStage; + a_reg = *reinterpret_cast(x_q + row * k + k_begin + kk); + } + { + const int byte_offset0 = tid * sizeof(int4); + const int kk0 = byte_offset0 / kBlockN; + const int col0 = byte_offset0 - kk0 * kBlockN; + b_reg0 = *reinterpret_cast( + weight + (k_begin + kk0) * n + n0 + col0); + const int vec1 = tid + kThreads; + const int byte_offset1 = vec1 * sizeof(int4); + const int kk1 = byte_offset1 / kBlockN; + const int col1 = byte_offset1 - kk1 * kBlockN; + b_reg1 = *reinterpret_cast( + weight + (k_begin + kk1) * n + n0 + col1); + } + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + if (tid < kAVectors) reinterpret_cast(a_tile[0])[tid] = a_reg; + reinterpret_cast(b_tile[0])[tid] = b_reg0; + reinterpret_cast(b_tile[0])[tid + kThreads] = b_reg1; + __syncthreads(); + + int current = 0; + for (int k0 = k_begin; k0 < k_end; k0 += kStage) { + const int next_k = k0 + kStage; + int4 next_a{}; + int4 next_b0{}; + int4 next_b1{}; + if (next_k < k_end) { + if (tid < kAVectors) { + const int byte_offset = tid * sizeof(int4); + const int row = byte_offset / kStage; + const int kk = byte_offset - row * kStage; + next_a = *reinterpret_cast( + x_q + row * k + next_k + kk); + } + const int byte_offset0 = tid * sizeof(int4); + const int kk0 = byte_offset0 / kBlockN; + const int col0 = byte_offset0 - kk0 * kBlockN; + next_b0 = *reinterpret_cast( + weight + (next_k + kk0) * n + n0 + col0); + const int vec1 = tid + kThreads; + const int byte_offset1 = vec1 * sizeof(int4); + const int kk1 = byte_offset1 / kBlockN; + const int col1 = byte_offset1 - kk1 * kBlockN; + next_b1 = *reinterpret_cast( + weight + (next_k + kk1) * n + n0 + col1); + } + +#pragma unroll + for (int kk = 0; kk < kStage; kk += kTileK) { + du_load_matrix_sync(a_frag, a_tile[current] + kk, kStage); + du_load_matrix_sync( + b_frag, b_tile[current] + kk * kBlockN + wave * kTileN, + kBlockN); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + if (next_k < k_end) { + __syncthreads(); + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + const int next = current ^ 1; + if (tid < kAVectors) reinterpret_cast(a_tile[next])[tid] = next_a; + reinterpret_cast(b_tile[next])[tid] = next_b0; + reinterpret_cast(b_tile[next])[tid + kThreads] = next_b1; + current = next; + __syncthreads(); + } + } + + int32_t* partial = workspace + + split_id * kTileM * n + n0 + wave * kTileN; + du_store_matrix_sync(partial, acc_frag, n, mem_row_major); + return; +} + +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int col = base_col + col_mod4 + 4 * i; + const float scaled = static_cast(frag.x[i]) * + x_scale[row] * weight_scale[col]; + out[row * n + col] = __float2bfloat16(scaled); + } +} + +// Large-M prefill path. Four waves compute a 64x64 output tile. Each wave +// owns a 32x32 quadrant (four 16x16 DUMMA accumulators), while the block +// cooperatively stages A[64,128] and B[128,64] in 16KB LDS. +__global__ __launch_bounds__(4 * kWaveSize) void +w8a8_dumma_prefill_64x64x128_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kThreads = 4 * kWaveSize; + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid % kWaveSize; + const int wave_row = wave >> 1; + const int wave_col = wave & 1; + const int m0 = static_cast(blockIdx.y) * kPrefillBlockM; + const int n0 = static_cast(blockIdx.x) * kPrefillBlockN; + + __shared__ __align__(16) int8_t + a_tile[kPrefillBlockM * kPrefillStageK]; + __shared__ __align__(16) int8_t + b_tile[kPrefillStageK * kPrefillBlockN]; + + DUFragment + a_frag0, a_frag1; + DUFragment + b_frag0, b_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + for (int k0 = 0; k0 < k; k0 += kPrefillStageK) { + constexpr int kAVectors = + kPrefillBlockM * kPrefillStageK / sizeof(int4); + for (int vec = tid; vec < kAVectors; vec += kThreads) { + const int byte_offset = vec * sizeof(int4); + const int local_row = byte_offset / kPrefillStageK; + const int kk = byte_offset - local_row * kPrefillStageK; + const int global_row = m0 + local_row; + reinterpret_cast(a_tile)[vec] = + global_row < m + ? *reinterpret_cast( + x_q + global_row * k + k0 + kk) + : int4{0, 0, 0, 0}; + } + + constexpr int kBVectors = + kPrefillStageK * kPrefillBlockN / sizeof(int4); + for (int vec = tid; vec < kBVectors; vec += kThreads) { + const int byte_offset = vec * sizeof(int4); + const int kk = byte_offset / kPrefillBlockN; + const int col = byte_offset - kk * kPrefillBlockN; + reinterpret_cast(b_tile)[vec] = + *reinterpret_cast( + weight + (k0 + kk) * n + n0 + col); + } + __syncthreads(); + +#pragma unroll + for (int kk = 0; kk < kPrefillStageK; kk += kTileK) { + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + du_load_matrix_sync( + a_frag0, a_tile + local_row * kPrefillStageK + kk, + kPrefillStageK); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + 16) * kPrefillStageK + kk, + kPrefillStageK); + du_load_matrix_sync( + b_frag0, b_tile + kk * kPrefillBlockN + local_col, + kPrefillBlockN); + du_load_matrix_sync( + b_frag1, b_tile + kk * kPrefillBlockN + local_col + 16, + kPrefillBlockN); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } + __syncthreads(); + } + + const int base_row = m0 + wave_row * 32; + const int base_col = n0 + wave_col * 32; + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, + base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, + base_row, base_col + 16, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, + base_row + 16, base_col, lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, + base_row + 16, base_col + 16, lane); + return; +} + +__global__ __launch_bounds__(256) void w8a8_splitk_combine_kernel( + const int32_t* __restrict__ workspace, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n, + int split_k) { + const int linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int elements = kTileM * n; + if (linear >= elements) { + return; + } + int32_t acc = 0; +#pragma unroll + for (int split = 0; split < split_k; ++split) { + acc += workspace[split * elements + linear]; + } + const int row = linear / n; + const int col = linear - row * n; + const float scaled = static_cast(acc) * + x_scale[row] * weight_scale[col]; + out[linear] = __float2bfloat16(scaled); +} + +struct __align__(4) BFloat16x2 { + __hip_bfloat16 x; + __hip_bfloat16 y; +}; + +struct __align__(8) BFloat16x4 { + __hip_bfloat16 x; + __hip_bfloat16 y; + __hip_bfloat16 z; + __hip_bfloat16 w; +}; + +__global__ __launch_bounds__(256) void w8a8_splitk_combine_vec2_kernel( + const int32_t* __restrict__ workspace, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n, + int split_k) { + const int pair = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + constexpr int kVec = 2; + const int pairs = kTileM * 1536 / kVec; + if (pair >= pairs) return; + const int base = pair * kVec; + int2 acc{0, 0}; +#pragma unroll + for (int split = 0; split < split_k; ++split) { + const int2 value = *reinterpret_cast( + workspace + split * kTileM * n + base); + acc.x += value.x; + acc.y += value.y; + } + const int row = base / n; + const int col = base - row * n; + const float xs = x_scale[row]; + const float2 ws = *reinterpret_cast(weight_scale + col); + const BFloat16x2 result{ + __float2bfloat16(static_cast(acc.x) * xs * ws.x), + __float2bfloat16(static_cast(acc.y) * xs * ws.y)}; + *reinterpret_cast(out + base) = result; +} + +__global__ __launch_bounds__(256) void w8a8_splitk_combine_vec4_kernel( + const int32_t* __restrict__ workspace, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int n, + int split_k) { + const int quad = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + constexpr int kVec = 4; + const int quads = kTileM * 1536 / kVec; + if (quad >= quads) return; + const int base = quad * kVec; + int4 acc{0, 0, 0, 0}; +#pragma unroll + for (int split = 0; split < split_k; ++split) { + const int4 value = *reinterpret_cast( + workspace + split * kTileM * n + base); + acc.x += value.x; + acc.y += value.y; + acc.z += value.z; + acc.w += value.w; + } + const int row = base / n; + const int col = base - row * n; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col); + const BFloat16x4 result{ + __float2bfloat16(static_cast(acc.x) * xs * ws.x), + __float2bfloat16(static_cast(acc.y) * xs * ws.y), + __float2bfloat16(static_cast(acc.z) * xs * ws.z), + __float2bfloat16(static_cast(acc.w) * xs * ws.w)}; + *reinterpret_cast(out + base) = result; +} + +// Generic 1 <= M < 16 path. DUMMA always consumes 16 rows, so absent rows are +// zero-padded in LDS before each 16x16x32 matrix operation. +__global__ __launch_bounds__(kWaveSize) void w8a8_dumma_m_lt16_padded_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + __hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + const int lane = static_cast(threadIdx.x); + const int n0 = static_cast(blockIdx.x) * kTileN; + + __shared__ __align__(16) int8_t a_tile[kTileM * kTileK]; + __shared__ __align__(16) int32_t acc_tile[kTileM * kTileN]; + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + for (int k0 = 0; k0 < k; k0 += kTileK) { + // Stage A because DUMMA always consumes 16 rows, while decode M may be + // smaller. Each lane copies eight int8 values for a full 16x32 tile. +#pragma unroll + for (int linear = lane; linear < kTileM * kTileK; + linear += kWaveSize) { + const int row = linear / kTileK; + const int kk = linear - row * kTileK; + a_tile[linear] = row < m ? x_q[row * k + k0 + kk] : int8_t{0}; + } + __syncthreads(); + + du_load_matrix_sync(a_frag, a_tile, kTileK); + du_load_matrix_sync(b_frag, weight + k0 * n + n0, n); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + + // Ensure every lane has consumed a_tile before it is overwritten. + __syncthreads(); + } + + du_store_matrix_sync(acc_tile, acc_frag, kTileN, mem_row_major); + __syncthreads(); + + // A 16x16 tile has 256 values, so each of the 64 lanes writes four. +#pragma unroll + for (int linear = lane; linear < kTileM * kTileN; + linear += kWaveSize) { + const int row = linear / kTileN; + const int col = linear - row * kTileN; + if (row < m) { + const float scaled = static_cast(acc_tile[linear]) * + x_scale[row] * weight_scale[n0 + col]; + out[row * n + n0 + col] = __float2bfloat16(scaled); + } + } + return; +} + +void check_inputs( + const torch::Tensor& x_q, + const torch::Tensor& weight, + const torch::Tensor& x_scale, + const torch::Tensor& weight_scale, + const torch::Tensor& out) { + TORCH_CHECK(x_q.is_cuda(), "x_q must be a CUDA/HIP tensor"); + TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA/HIP tensor"); + TORCH_CHECK(x_scale.is_cuda(), "x_scale must be a CUDA/HIP tensor"); + TORCH_CHECK(weight_scale.is_cuda(), + "weight_scale must be a CUDA/HIP tensor"); + TORCH_CHECK(out.is_cuda(), "out must be a CUDA/HIP tensor"); + TORCH_CHECK(x_q.device() == weight.device() && + x_q.device() == x_scale.device() && + x_q.device() == weight_scale.device() && + x_q.device() == out.device(), + "all tensors must be on the same device"); + + TORCH_CHECK(x_q.scalar_type() == at::kChar, "x_q must be int8"); + TORCH_CHECK(weight.scalar_type() == at::kChar, "weight must be int8"); + TORCH_CHECK(x_scale.scalar_type() == at::kFloat, + "x_scale must be float32"); + TORCH_CHECK(weight_scale.scalar_type() == at::kFloat, + "weight_scale must be float32"); + TORCH_CHECK(out.scalar_type() == at::kBFloat16, "out must be bfloat16"); + + TORCH_CHECK(x_q.dim() == 2, "x_q must have shape [M, K]"); + TORCH_CHECK(weight.dim() == 2, "weight must have shape [K, N]"); + TORCH_CHECK(x_scale.dim() == 2 && x_scale.size(1) == 1, + "x_scale must have shape [M, 1]"); + TORCH_CHECK(weight_scale.dim() == 2 && weight_scale.size(1) == 1, + "weight_scale must have shape [N, 1]"); + TORCH_CHECK(out.dim() == 2, "out must have shape [M, N]"); + + const auto m = x_q.size(0); + const auto k = x_q.size(1); + TORCH_CHECK(m >= 1, "M must be positive"); + TORCH_CHECK(weight.size(0) == k, "K mismatch between x_q and weight"); + const auto n = weight.size(1); + TORCH_CHECK(k % kTileK == 0, "K must be divisible by 32"); + TORCH_CHECK(n % kTileN == 0, "N must be divisible by 16"); + TORCH_CHECK(x_scale.size(0) == m, "x_scale shape mismatch"); + TORCH_CHECK(weight_scale.size(0) == n, "weight_scale shape mismatch"); + TORCH_CHECK(out.size(0) == m && out.size(1) == n, + "out shape must be [M, N]"); + + TORCH_CHECK(x_q.is_contiguous(), "x_q must be contiguous"); + TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous"); + TORCH_CHECK(x_scale.is_contiguous(), "x_scale must be contiguous"); + TORCH_CHECK(weight_scale.is_contiguous(), + "weight_scale must be contiguous"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); +} + +void launch_m16_variant( + int waves, + const int8_t* x_q, + const int8_t* weight, + const float* x_scale, + const float* weight_scale, + __hip_bfloat16* out, + int n, + int k, + hipStream_t stream) { + TORCH_CHECK( + waves == 1 || waves == 2 || waves == 4 || waves == 16 || + waves == 32 || waves == 33 || waves == 36 || waves == 40 || + waves == 132 || waves == 1128 || waves == 1512, + "variant must be direct 1/2/4, staged1=16, staged2 stride " + "32/33/36/40, or staged2 direct-epilogue K-stage 256/128/512=" + "132/1128/1512"); + if (waves == 16 || waves == 32 || waves == 33 || waves == 36 || + waves == 40 || waves == 132 || waves == 1128 || waves == 1512) { + const int stage_waves = waves == 16 ? 1 : 2; + TORCH_CHECK(n % (stage_waves * kTileN) == 0, + "N is not divisible by the staged block N"); + const int requested_stage_k = + waves == 1128 ? 128 : (waves == 1512 ? 512 : kStageK); + TORCH_CHECK(k % requested_stage_k == 0, + "K must be divisible by the staged kernel K depth"); + const dim3 grid(static_cast(n / (stage_waves * kTileN))); + const dim3 block(static_cast(stage_waves * kWaveSize)); + if (stage_waves == 1) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_staged_kernel<1, 0>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } else { + if (waves == 1128) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_staged_kernel<2, 0, true, 128>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } else if (waves == 1512) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_staged_kernel<2, 0, true, 512>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } else if (waves == 132) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_staged_kernel<2, 0, true>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } else if (waves == 32) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_staged_kernel<2, 0>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } else if (waves == 33) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_staged_kernel<2, 1>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } else if (waves == 36) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_staged_kernel<2, 4>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_staged_kernel<2, 8>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } + } + return; + } + TORCH_CHECK(n % (kTileN * waves) == 0, + "N must be divisible by 16 * waves"); + const dim3 grid(static_cast(n / (kTileN * waves))); + const dim3 block(static_cast(kWaveSize * waves)); + + if (waves == 1) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_direct_kernel<1>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } else if (waves == 2) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_direct_kernel<2>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_direct_kernel<4>), + grid, block, 0, stream, + x_q, weight, x_scale, weight_scale, out, n, k); + } +} + +} // namespace + +torch::Tensor w8a8_gemm_out_hip( + const torch::Tensor& x_q, + const torch::Tensor& weight, + const torch::Tensor& x_scale, + const torch::Tensor& weight_scale, + torch::Tensor out) { + check_inputs(x_q, weight, x_scale, weight_scale, out); + TORCH_CHECK(x_q.size(0) <= kTileM, "gemm_out requires M in [1, 16]"); + + const c10::cuda::CUDAGuard device_guard(x_q.device()); + const int m = static_cast(x_q.size(0)); + const int k = static_cast(x_q.size(1)); + const int n = static_cast(weight.size(1)); + const dim3 grid(static_cast(n / kTileN)); + const dim3 block(kWaveSize); + const auto stream = at::cuda::getCurrentCUDAStream(x_q.get_device()); + + if (m == kTileM) { + // Offline TP4/M=16 tuning: shallow K stages improve occupancy for the + // K<=2048 decode GEMMs, while K=4096 retains the 256-deep stage. + const int variant = k <= 2048 ? 1128 : 132; + launch_m16_variant( + variant, + x_q.data_ptr(), + weight.data_ptr(), + x_scale.data_ptr(), + weight_scale.data_ptr(), + reinterpret_cast<__hip_bfloat16*>(out.data_ptr()), + n, + k, + stream); + } else { + hipLaunchKernelGGL( + w8a8_dumma_m_lt16_padded_kernel, + grid, + block, + 0, + stream, + x_q.data_ptr(), + weight.data_ptr(), + x_scale.data_ptr(), + weight_scale.data_ptr(), + reinterpret_cast<__hip_bfloat16*>(out.data_ptr()), + m, + n, + k); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +torch::Tensor w8a8_gemm_out_variant_hip( + const torch::Tensor& x_q, + const torch::Tensor& weight, + const torch::Tensor& x_scale, + const torch::Tensor& weight_scale, + torch::Tensor out, + int64_t waves) { + check_inputs(x_q, weight, x_scale, weight_scale, out); + TORCH_CHECK(x_q.size(0) == kTileM, + "gemm_out_variant currently requires M=16"); + + const c10::cuda::CUDAGuard device_guard(x_q.device()); + const int k = static_cast(x_q.size(1)); + const int n = static_cast(weight.size(1)); + const auto stream = at::cuda::getCurrentCUDAStream(x_q.get_device()); + launch_m16_variant( + static_cast(waves), + x_q.data_ptr(), + weight.data_ptr(), + x_scale.data_ptr(), + weight_scale.data_ptr(), + reinterpret_cast<__hip_bfloat16*>(out.data_ptr()), + n, + k, + stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +torch::Tensor w8a8_gemm_out_prefill_hip( + const torch::Tensor& x_q, + const torch::Tensor& weight, + const torch::Tensor& x_scale, + const torch::Tensor& weight_scale, + torch::Tensor out) { + check_inputs(x_q, weight, x_scale, weight_scale, out); + TORCH_CHECK(x_q.size(0) > kTileM, "prefill kernel requires M > 16"); + TORCH_CHECK(x_q.size(1) % kPrefillStageK == 0, + "K must be divisible by 128"); + TORCH_CHECK(weight.size(1) % kPrefillBlockN == 0, + "N must be divisible by 64"); + + const c10::cuda::CUDAGuard device_guard(x_q.device()); + const int m = static_cast(x_q.size(0)); + const int k = static_cast(x_q.size(1)); + const int n = static_cast(weight.size(1)); + const auto stream = at::cuda::getCurrentCUDAStream(x_q.get_device()); + const dim3 grid( + static_cast(n / kPrefillBlockN), + static_cast((m + kPrefillBlockM - 1) / kPrefillBlockM)); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + w8a8_dumma_prefill_64x64x128_kernel, + grid, + block, + 0, + stream, + x_q.data_ptr(), + weight.data_ptr(), + x_scale.data_ptr(), + weight_scale.data_ptr(), + reinterpret_cast<__hip_bfloat16*>(out.data_ptr()), + m, + n, + k); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +torch::Tensor w8a8_gemm_out_splitk_hip( + const torch::Tensor& x_q, + const torch::Tensor& weight, + const torch::Tensor& x_scale, + const torch::Tensor& weight_scale, + torch::Tensor out, + torch::Tensor workspace, + int64_t split_k) { + check_inputs(x_q, weight, x_scale, weight_scale, out); + TORCH_CHECK(x_q.size(0) == kTileM, "split-K currently requires M=16"); + TORCH_CHECK(split_k == 2 || split_k == 4 || split_k == 8 || split_k == 16, + "split_k must be 2, 4, 8, or 16"); + TORCH_CHECK(workspace.is_cuda() && workspace.device() == x_q.device(), + "workspace must be on the same CUDA/HIP device"); + TORCH_CHECK(workspace.scalar_type() == at::kInt, + "workspace must be int32"); + TORCH_CHECK(workspace.is_contiguous(), "workspace must be contiguous"); + + const int k = static_cast(x_q.size(1)); + const int n = static_cast(weight.size(1)); + TORCH_CHECK((k / split_k) % kStageK == 0, + "each K split must be divisible by 256"); + TORCH_CHECK(workspace.numel() >= split_k * kTileM * n, + "workspace needs at least split_k * 16 * N int32 elements"); + + const c10::cuda::CUDAGuard device_guard(x_q.device()); + const auto stream = at::cuda::getCurrentCUDAStream(x_q.get_device()); + const dim3 partial_grid( + static_cast((n / (2 * kTileN)) * split_k)); + const dim3 partial_block(2 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_staged2_splitk_kernel<2, kStageK>), + partial_grid, + partial_block, + 0, + stream, + x_q.data_ptr(), + weight.data_ptr(), + workspace.data_ptr(), + n, + k, + static_cast(split_k)); + + const int elements = kTileM * n; + const dim3 combine_block(256); + const dim3 combine_grid(static_cast((elements + 255) / 256)); + hipLaunchKernelGGL( + w8a8_splitk_combine_kernel, + combine_grid, + combine_block, + 0, + stream, + workspace.data_ptr(), + x_scale.data_ptr(), + weight_scale.data_ptr(), + reinterpret_cast<__hip_bfloat16*>(out.data_ptr()), + n, + static_cast(split_k)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +torch::Tensor w8a8_gemm_out_splitk_config_hip( + const torch::Tensor& x_q, + const torch::Tensor& weight, + const torch::Tensor& x_scale, + const torch::Tensor& weight_scale, + torch::Tensor out, + torch::Tensor workspace, + int64_t split_k, + int64_t waves, + int64_t stage_k, + int64_t combine_block, + bool prefetch, + int64_t combine_mode, + int64_t phase, + int64_t sync_mode, + bool spread_a, + bool fused_finalize, + bool cooperative_finalize) { + check_inputs(x_q, weight, x_scale, weight_scale, out); + TORCH_CHECK(x_q.size(0) == kTileM, "configured split-K requires M=16"); + TORCH_CHECK(split_k == 2 || split_k == 4 || split_k == 8 || split_k == 10 || + split_k == 16, + "split_k must be 2, 4, 8, 10, or 16"); + TORCH_CHECK(waves == 1 || waves == 2 || waves == 3 || waves == 4 || + waves == 6 || waves == 8, + "waves must be 1, 2, 3, 4, 6, or 8"); + TORCH_CHECK(stage_k == 32 || stage_k == 64 || stage_k == 128 || stage_k == 256 || + stage_k == 512, + "stage_k must be 32, 64, 128, 256, or 512"); + TORCH_CHECK(combine_block == 64 || combine_block == 128 || + combine_block == 256, + "combine_block must be 64, 128, or 256"); + TORCH_CHECK(combine_mode >= 0 && combine_mode <= 2, + "combine_mode must be 0 scalar, 1 vec2, or 2 vec4"); + TORCH_CHECK(phase >= 0 && phase <= 2, + "phase must be 0 both, 1 partial, or 2 combine"); + TORCH_CHECK(sync_mode >= 0 && sync_mode <= 3, + "sync_mode must be 0 baseline, 1 compiler wait, 2 raw first " + "barrier, or 3 full raw synchronization"); + TORCH_CHECK(!spread_a || + (prefetch && split_k == 8 && waves == 4 && stage_k == 64 && + sync_mode == 2), + "spread_a is specialized for split8/waves4/stage64/sync2"); + TORCH_CHECK(!fused_finalize || + (phase == 0 && prefetch && split_k == 8 && waves == 4 && + stage_k == 64 && sync_mode == 2 && combine_mode == 0), + "fused_finalize requires the split8/waves4/stage64/sync2 path"); + TORCH_CHECK(!cooperative_finalize || + (phase == 0 && prefetch && split_k == 8 && waves == 4 && + stage_k == 64 && sync_mode == 2 && combine_mode == 0 && + !fused_finalize), + "cooperative_finalize requires the split8/waves4/stage64/sync2 path"); + TORCH_CHECK(workspace.is_cuda() && workspace.device() == x_q.device(), + "workspace must be on the same CUDA/HIP device"); + TORCH_CHECK(workspace.scalar_type() == at::kInt && workspace.is_contiguous(), + "workspace must be contiguous int32"); + + const int k = static_cast(x_q.size(1)); + const int n = static_cast(weight.size(1)); + TORCH_CHECK(k == 4096 && n == 1536, + "configured kernel is specialized for K=4096,N=1536"); + TORCH_CHECK(split_k == 10 || (k / split_k) % stage_k == 0, + "each uniform K split must be divisible by stage_k"); + TORCH_CHECK(workspace.numel() >= split_k * kTileM * n, + "workspace needs split_k * 16 * N int32 elements"); + + const c10::cuda::CUDAGuard device_guard(x_q.device()); + const auto stream = at::cuda::getCurrentCUDAStream(x_q.get_device()); + const unsigned grid_x = static_cast((n / (waves * kTileN)) * split_k); + +#define LAUNCH_SPLITK_CONFIG(W, S) \ + hipLaunchKernelGGL( \ + HIP_KERNEL_NAME(w8a8_dumma_m16_staged2_splitk_kernel), \ + dim3(grid_x), dim3(W * kWaveSize), 0, stream, \ + x_q.data_ptr(), weight.data_ptr(), \ + workspace.data_ptr(), n, k, static_cast(split_k)) + + if (cooperative_finalize) { + const int8_t* x_ptr = x_q.data_ptr(); + const int8_t* w_ptr = weight.data_ptr(); + const float* xs_ptr = x_scale.data_ptr(); + const float* ws_ptr = weight_scale.data_ptr(); + auto* out_ptr = reinterpret_cast<__hip_bfloat16*>( + out.data_ptr()); + int32_t* workspace_ptr = workspace.data_ptr(); + int kernel_n = n; + int kernel_k = k; + void* kernel_args[] = {&x_ptr, &w_ptr, &xs_ptr, &ws_ptr, &out_ptr, + &workspace_ptr, &kernel_n, &kernel_k}; + C10_CUDA_CHECK(hipLaunchCooperativeKernel( + reinterpret_cast( + w8a8_dumma_m16_splitk8_w4_s64_cooperative_finalize_kernel), + dim3(grid_x), dim3(4 * kWaveSize), kernel_args, 0, stream)); + } else if (fused_finalize) { + hipLaunchKernelGGL( + w8a8_dumma_m16_splitk8_w4_s64_fused_finalize_kernel, + dim3(grid_x), dim3(4 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + x_scale.data_ptr(), weight_scale.data_ptr(), + reinterpret_cast<__hip_bfloat16*>(out.data_ptr()), + workspace.data_ptr(), n, k); + } else if (phase != 2) { + if (prefetch) { + TORCH_CHECK((split_k == 4 || split_k == 8 || split_k == 10 || + split_k == 16) && + ((stage_k == 64 && + (waves == 1 || waves == 2 || waves == 3 || + waves == 4 || waves == 6 || waves == 8)) || + (stage_k == 128 && waves == 3)), + "prefetch supports split4|8|10|16/stage64 waves1|2|3|4|6|8, or " + "split8/stage128 waves3"); + if (stage_k == 64) { + if (waves != 3) { + TORCH_CHECK(sync_mode == 2, + "non-3-wave prefetch candidates require sync_mode=2"); + if (waves == 1) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel<1, 2>), + dim3(grid_x), dim3(kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } else if (waves == 2) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel<2, 2>), + dim3(grid_x), dim3(2 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } else if (waves == 4) { + if (spread_a) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel<4, 2, true>), + dim3(grid_x), dim3(4 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel<4, 2, false>), + dim3(grid_x), dim3(4 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } + } else if (waves == 6) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel<6, 2>), + dim3(grid_x), dim3(6 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel<8, 2>), + dim3(grid_x), dim3(8 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } + } else if (sync_mode == 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel<3, 0>), + dim3(grid_x), dim3(3 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } else if (sync_mode == 1) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel<3, 1>), + dim3(grid_x), dim3(3 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } else if (sync_mode == 2) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel<3, 2>), + dim3(grid_x), dim3(3 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16_splitk8_w3_s64_prefetch_kernel<3, 3>), + dim3(grid_x), dim3(3 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } + } else { + hipLaunchKernelGGL( + w8a8_dumma_m16_splitk8_w3_s128_prefetch_kernel, + dim3(grid_x), dim3(3 * kWaveSize), 0, stream, + x_q.data_ptr(), weight.data_ptr(), + workspace.data_ptr(), n, k, static_cast(split_k)); + } + } else if (waves == 1 && stage_k == 32) LAUNCH_SPLITK_CONFIG(1, 32); + else if (waves == 1 && stage_k == 64) LAUNCH_SPLITK_CONFIG(1, 64); + else if (waves == 1 && stage_k == 128) LAUNCH_SPLITK_CONFIG(1, 128); + else if (waves == 1 && stage_k == 256) LAUNCH_SPLITK_CONFIG(1, 256); + else if (waves == 1 && stage_k == 512) LAUNCH_SPLITK_CONFIG(1, 512); + else if (waves == 2 && stage_k == 32) LAUNCH_SPLITK_CONFIG(2, 32); + else if (waves == 2 && stage_k == 64) LAUNCH_SPLITK_CONFIG(2, 64); + else if (waves == 2 && stage_k == 128) LAUNCH_SPLITK_CONFIG(2, 128); + else if (waves == 2 && stage_k == 256) LAUNCH_SPLITK_CONFIG(2, 256); + else if (waves == 2 && stage_k == 512) LAUNCH_SPLITK_CONFIG(2, 512); + else if (waves == 3 && stage_k == 32) LAUNCH_SPLITK_CONFIG(3, 32); + else if (waves == 3 && stage_k == 64) LAUNCH_SPLITK_CONFIG(3, 64); + else if (waves == 3 && stage_k == 128) LAUNCH_SPLITK_CONFIG(3, 128); + else if (waves == 3 && stage_k == 256) LAUNCH_SPLITK_CONFIG(3, 256); + else if (waves == 3 && stage_k == 512) LAUNCH_SPLITK_CONFIG(3, 512); + else if (waves == 4 && stage_k == 32) LAUNCH_SPLITK_CONFIG(4, 32); + else if (waves == 4 && stage_k == 64) LAUNCH_SPLITK_CONFIG(4, 64); + else if (waves == 4 && stage_k == 128) LAUNCH_SPLITK_CONFIG(4, 128); + else if (waves == 4 && stage_k == 256) LAUNCH_SPLITK_CONFIG(4, 256); + else if (waves == 4 && stage_k == 512) LAUNCH_SPLITK_CONFIG(4, 512); + else if (waves == 6 && stage_k == 32) LAUNCH_SPLITK_CONFIG(6, 32); + else if (waves == 6 && stage_k == 64) LAUNCH_SPLITK_CONFIG(6, 64); + else if (waves == 6 && stage_k == 128) LAUNCH_SPLITK_CONFIG(6, 128); + else if (waves == 6 && stage_k == 256) LAUNCH_SPLITK_CONFIG(6, 256); + else LAUNCH_SPLITK_CONFIG(6, 512); + } +#undef LAUNCH_SPLITK_CONFIG + + if (!fused_finalize && !cooperative_finalize && phase != 1) { + const int elements = kTileM * n; + const int vector_width = combine_mode == 0 ? 1 : combine_mode == 1 ? 2 : 4; + const dim3 combine_grid(static_cast( + ((elements / vector_width) + combine_block - 1) / combine_block)); + const dim3 combine_threads(static_cast(combine_block)); + if (combine_mode == 0) { + hipLaunchKernelGGL( + w8a8_splitk_combine_kernel, combine_grid, combine_threads, 0, + stream, workspace.data_ptr(), x_scale.data_ptr(), + weight_scale.data_ptr(), + reinterpret_cast<__hip_bfloat16*>(out.data_ptr()), n, + static_cast(split_k)); + } else if (combine_mode == 1) { + hipLaunchKernelGGL( + w8a8_splitk_combine_vec2_kernel, combine_grid, combine_threads, 0, + stream, workspace.data_ptr(), x_scale.data_ptr(), + weight_scale.data_ptr(), + reinterpret_cast<__hip_bfloat16*>(out.data_ptr()), n, + static_cast(split_k)); + } else { + hipLaunchKernelGGL( + w8a8_splitk_combine_vec4_kernel, combine_grid, combine_threads, 0, + stream, workspace.data_ptr(), x_scale.data_ptr(), + weight_scale.data_ptr(), + reinterpret_cast<__hip_bfloat16*>(out.data_ptr()), n, + static_cast(split_k)); + } + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +}